relay_knowledge/api/
agent.rs

1use std::collections::{HashMap, HashSet};
2
3use serde::{Deserialize, Serialize};
4
5use crate::domain::{
6    ContextPackItem, FreshnessPolicy, FusionDiagnostics, IndexStatus, RetrievalBackendStatus,
7    RetrievalHit, RetrievalMode, RetrievedContextPack,
8};
9use crate::project::{ACP_LOCAL_ADAPTER_NAME, MCP_ADAPTER_NAME};
10use crate::storage::{IndexCursor, IndexRefreshDiagnostics};
11
12use super::{ApiMetadata, RequestContext};
13
14/// Agent protocol family used by external resident-process adapters.
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
16#[serde(rename_all = "kebab-case")]
17pub enum AgentProtocolKind {
18    Mcp,
19    Acp,
20}
21
22/// Runtime identity captured from an agent protocol request.
23#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
24pub struct RuntimeIdentity {
25    pub protocol: AgentProtocolKind,
26    pub adapter_name: String,
27    #[serde(skip_serializing_if = "Option::is_none")]
28    pub adapter_version: Option<String>,
29    #[serde(skip_serializing_if = "Option::is_none")]
30    pub client_name: Option<String>,
31    #[serde(skip_serializing_if = "Option::is_none")]
32    pub client_version: Option<String>,
33    #[serde(skip_serializing_if = "Option::is_none")]
34    pub host_name: Option<String>,
35    #[serde(skip_serializing_if = "Option::is_none")]
36    pub actor_id: Option<String>,
37    #[serde(skip_serializing_if = "Option::is_none")]
38    pub session_id: Option<String>,
39    #[serde(skip_serializing_if = "Option::is_none")]
40    pub tool_call_id: Option<String>,
41}
42
43impl RuntimeIdentity {
44    /// Creates the resident MCP adapter identity for a single request.
45    pub fn mcp(tool_call_id: Option<String>) -> Self {
46        Self {
47            protocol: AgentProtocolKind::Mcp,
48            adapter_name: MCP_ADAPTER_NAME.to_owned(),
49            adapter_version: Some(env!("CARGO_PKG_VERSION").to_owned()),
50            client_name: None,
51            client_version: None,
52            host_name: None,
53            actor_id: None,
54            session_id: None,
55            tool_call_id,
56        }
57    }
58
59    /// Creates the local ACP adapter identity for one session request.
60    pub fn acp(
61        client_name: Option<String>,
62        client_version: Option<String>,
63        actor_id: Option<String>,
64        session_id: String,
65        request_id: Option<String>,
66    ) -> Self {
67        Self {
68            protocol: AgentProtocolKind::Acp,
69            adapter_name: ACP_LOCAL_ADAPTER_NAME.to_owned(),
70            adapter_version: Some(env!("CARGO_PKG_VERSION").to_owned()),
71            client_name,
72            client_version,
73            host_name: None,
74            actor_id,
75            session_id: Some(session_id),
76            tool_call_id: request_id,
77        }
78    }
79}
80
81/// Unified API context plus agent protocol identity and policy provenance.
82#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
83pub struct AgentRequestContext {
84    pub request: RequestContext,
85    pub runtime_identity: RuntimeIdentity,
86    pub policy_id: String,
87}
88
89/// Local access policy applied before agent protocol requests reach services.
90#[derive(Debug, Clone, PartialEq, Eq)]
91pub struct AgentAccessPolicy {
92    pub allowed_scopes: Vec<String>,
93    pub allow_unspecified_scope: bool,
94    pub max_limit: usize,
95    pub max_context_bytes: usize,
96    pub max_runtime_ms: u64,
97    pub allow_remote_clients: bool,
98}
99
100impl AgentAccessPolicy {
101    pub const DEFAULT_MAX_LIMIT: usize = 10;
102    pub const DEFAULT_MAX_CONTEXT_BYTES: usize = 65_536;
103
104    /// Creates a validated access policy for agent protocol adapters.
105    pub fn new(
106        allowed_scopes: Vec<String>,
107        allow_unspecified_scope: bool,
108        max_limit: usize,
109        max_context_bytes: usize,
110        max_runtime_ms: u64,
111        allow_remote_clients: bool,
112    ) -> Result<Self, AgentPolicyError> {
113        if max_limit == 0 {
114            return Err(AgentPolicyError::ZeroMaxLimit);
115        }
116        if max_context_bytes == 0 {
117            return Err(AgentPolicyError::ZeroMaxContextBytes);
118        }
119        if max_runtime_ms == 0 {
120            return Err(AgentPolicyError::ZeroMaxRuntime);
121        }
122
123        Ok(Self {
124            allowed_scopes,
125            allow_unspecified_scope,
126            max_limit,
127            max_context_bytes,
128            max_runtime_ms,
129            allow_remote_clients,
130        })
131    }
132
133    /// Summarizes policy without exposing scope names or secrets.
134    pub fn summary(&self) -> AgentAccessPolicySummary {
135        AgentAccessPolicySummary {
136            allowed_scope_count: self.allowed_scopes.len(),
137            allow_unspecified_scope: self.allow_unspecified_scope,
138            max_limit: self.max_limit,
139            max_context_bytes: self.max_context_bytes,
140            max_runtime_ms: self.max_runtime_ms,
141            allow_remote_clients: self.allow_remote_clients,
142        }
143    }
144}
145
146/// Stable policy validation error.
147#[derive(Debug, Clone, PartialEq, Eq)]
148pub enum AgentPolicyError {
149    ZeroMaxLimit,
150    ZeroMaxContextBytes,
151    ZeroMaxRuntime,
152}
153
154impl std::fmt::Display for AgentPolicyError {
155    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
156        match self {
157            Self::ZeroMaxLimit => write!(formatter, "MCP max limit must be greater than zero"),
158            Self::ZeroMaxContextBytes => {
159                write!(formatter, "MCP max context bytes must be greater than zero")
160            }
161            Self::ZeroMaxRuntime => write!(formatter, "MCP max runtime must be greater than zero"),
162        }
163    }
164}
165
166impl std::error::Error for AgentPolicyError {}
167
168/// Redacted policy status for service diagnostics.
169#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
170pub struct AgentAccessPolicySummary {
171    pub allowed_scope_count: usize,
172    pub allow_unspecified_scope: bool,
173    pub max_limit: usize,
174    pub max_context_bytes: usize,
175    pub max_runtime_ms: u64,
176    pub allow_remote_clients: bool,
177}
178
179/// Service status projection for resident agent protocols.
180#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
181pub struct AgentProtocolStatus {
182    pub mcp_streamable_http_enabled: bool,
183    pub mcp_endpoint: String,
184    pub mcp_resources_enabled: bool,
185    pub mcp_prompts_enabled: bool,
186    pub metrics_endpoint: String,
187    pub http_bind: String,
188    pub allowed_origin_count: usize,
189    pub mcp_allowed_origins: Vec<String>,
190    pub policy: AgentAccessPolicySummary,
191    pub audit_sink_enabled: bool,
192    pub audit_log_path: String,
193    pub audit_queue_depth: usize,
194}
195
196/// Canonical retrieval result shared by MCP and future agent protocols.
197#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
198pub struct AgentRetrievalResult {
199    pub metadata: ApiMetadata,
200    pub runtime_identity: RuntimeIdentity,
201    #[serde(skip_serializing_if = "Option::is_none")]
202    pub source_scope: Option<String>,
203    pub freshness: String,
204    pub retrieval_mode: RetrievalMode,
205    pub context_pack: RetrievedContextPack,
206    pub results: Vec<RetrievalHit>,
207    pub fusion: FusionDiagnostics,
208    pub rerank: crate::domain::RerankDiagnostics,
209    #[serde(default, skip_serializing_if = "Vec::is_empty")]
210    pub backend_statuses: Vec<RetrievalBackendStatus>,
211    pub indexes: Vec<IndexStatus>,
212    #[serde(default)]
213    pub index_cursors: Vec<IndexCursor>,
214    #[serde(default)]
215    pub index_refresh: IndexRefreshDiagnostics,
216    #[serde(skip_serializing_if = "Option::is_none")]
217    pub degraded_reason: Option<String>,
218    pub truncated: bool,
219    pub budget_used: AgentBudgetUsed,
220}
221
222#[derive(Debug, Clone, PartialEq, Eq, Hash)]
223struct AgentResultKey {
224    result_id: String,
225    source_scope: String,
226    source_path: Option<String>,
227}
228
229impl AgentResultKey {
230    fn from_hit(hit: &RetrievalHit) -> Self {
231        Self {
232            result_id: hit.evidence_id.clone(),
233            source_scope: hit.source_scope.clone(),
234            source_path: agent_hit_source_path(hit),
235        }
236    }
237
238    fn from_item(item: &ContextPackItem) -> Self {
239        Self {
240            result_id: item.result_id.clone(),
241            source_scope: item.source_scope.clone(),
242            source_path: item
243                .source_path
244                .clone()
245                .or_else(|| item.code_artifact.as_ref().and_then(agent_artifact_path)),
246        }
247    }
248}
249
250fn agent_hit_source_path(hit: &RetrievalHit) -> Option<String> {
251    hit.source_path
252        .clone()
253        .or_else(|| hit.code_artifact.as_ref().and_then(agent_artifact_path))
254}
255
256fn agent_artifact_path(artifact: &crate::domain::CodeGraphArtifact) -> Option<String> {
257    (!artifact.path.is_empty()).then(|| artifact.path.clone())
258}
259
260impl AgentRetrievalResult {
261    /// Builds the canonical agent result and applies the context byte budget.
262    pub fn from_retrieval(
263        response: crate::api::HybridRetrievalResponse,
264        identity: RuntimeIdentity,
265        max_context_bytes: usize,
266        elapsed_ms: u64,
267    ) -> Self {
268        let crate::api::HybridRetrievalResponse {
269            metadata,
270            mut context_pack,
271            retrieval_mode,
272            source_scope,
273            freshness,
274            results: response_results,
275            fusion,
276            mut rerank,
277            mut backend_statuses,
278            truncated: response_truncated,
279            budget_used,
280            degraded_reason,
281            indexes,
282            index_cursors,
283            index_refresh,
284        } = response;
285        let item_bytes = context_pack
286            .items
287            .iter()
288            .map(|item| {
289                (
290                    AgentResultKey::from_item(item),
291                    serialized_context_bytes(item),
292                )
293            })
294            .collect::<HashMap<_, _>>();
295        let mut context_bytes = serialized_context_bytes(&context_pack.backend_statuses)
296            .saturating_add(serialized_context_bytes(&backend_statuses));
297        let mut truncated = response_truncated;
298        if context_bytes > max_context_bytes {
299            context_pack.backend_statuses.clear();
300            backend_statuses.clear();
301            context_bytes = 0;
302            truncated = true;
303        }
304        let mut results = Vec::new();
305
306        for hit in response_results {
307            let hit_key = AgentResultKey::from_hit(&hit);
308            let hit_bytes = serialized_context_bytes(&hit)
309                .saturating_add(item_bytes.get(&hit_key).copied().unwrap_or_default());
310            if context_bytes.saturating_add(hit_bytes) > max_context_bytes {
311                truncated = true;
312                continue;
313            }
314            context_bytes += hit_bytes;
315            results.push(hit);
316        }
317        let returned_count = results.len();
318        rerank.returned_count = returned_count;
319        let retained_result_keys = results
320            .iter()
321            .map(AgentResultKey::from_hit)
322            .collect::<HashSet<_>>();
323        context_pack.truncated = truncated;
324        context_pack
325            .items
326            .retain(|item| retained_result_keys.contains(&AgentResultKey::from_item(item)));
327        if let Some(trace) = &mut context_pack.provenance_trace {
328            trace.retain_hits(results.iter());
329            trace.mark_citations_for_hits(results.iter());
330            trace.truncated |= truncated;
331            trace.apply_budget(
332                returned_count
333                    .saturating_mul(4)
334                    .max(returned_count + 8)
335                    .min(64),
336            );
337            if trace.truncated {
338                truncated = true;
339                context_pack.truncated = true;
340            }
341        }
342        if let Some(trace) = &mut context_pack.provenance_trace {
343            let mut trace_bytes = serialized_context_bytes(trace);
344            if context_bytes.saturating_add(trace_bytes) > max_context_bytes {
345                trace.apply_budget(returned_count.max(1));
346                trace.truncated = true;
347                truncated = true;
348                context_pack.truncated = true;
349                trace_bytes = serialized_context_bytes(trace);
350            }
351            if context_bytes.saturating_add(trace_bytes) > max_context_bytes {
352                context_pack.provenance_trace = None;
353                truncated = true;
354                context_pack.truncated = true;
355            } else {
356                context_bytes += trace_bytes;
357            }
358        }
359
360        Self {
361            metadata,
362            runtime_identity: identity,
363            source_scope,
364            freshness: freshness_label(freshness).to_owned(),
365            retrieval_mode,
366            context_pack,
367            results,
368            fusion,
369            rerank,
370            backend_statuses,
371            indexes,
372            index_cursors,
373            index_refresh,
374            degraded_reason,
375            truncated,
376            budget_used: AgentBudgetUsed {
377                limit: budget_used.limit,
378                candidate_count: budget_used.candidate_count,
379                returned_count,
380                context_bytes,
381                elapsed_ms,
382            },
383        }
384    }
385}
386
387fn serialized_context_bytes<T: Serialize>(value: &T) -> usize {
388    serde_json::to_vec(value)
389        .map(|bytes| bytes.len())
390        .unwrap_or(usize::MAX / 4)
391}
392
393/// Runtime budget consumed by a completed agent retrieval.
394#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
395pub struct AgentBudgetUsed {
396    pub limit: usize,
397    pub candidate_count: usize,
398    pub returned_count: usize,
399    pub context_bytes: usize,
400    pub elapsed_ms: u64,
401}
402
403pub fn freshness_label(freshness: FreshnessPolicy) -> &'static str {
404    match freshness {
405        FreshnessPolicy::AllowStale => "allow-stale",
406        FreshnessPolicy::WaitUntilFresh => "wait-until-fresh",
407        FreshnessPolicy::GraphOnly => "graph-only",
408    }
409}
410
411#[cfg(test)]
412mod tests {
413    use super::*;
414    use crate::{
415        api::InterfaceKind,
416        domain::{
417            ConfidenceScore, ContextGraphFact, ContextGraphFactKind, ContextPackItem, FactStatus,
418            FusionDiagnostics, GraphVersion, GraphVersionRange, RerankDiagnostics, RerankMode,
419            RetrievalBackendState, RetrievalBackendStatus, RetrievalBudgetUsed, RetrievalHit,
420            RetrievedContextPack, RetrieverSource, TraversalProvenanceTrace,
421        },
422    };
423
424    #[test]
425    fn truncates_retrieval_results_to_context_byte_budget() {
426        let items = vec![pack_item("ev-1"), pack_item("ev-2"), pack_item("ev-3")];
427        let results = vec![
428            hit("ev-1", "abcd"),
429            hit("ev-2", "efgh"),
430            hit("ev-3", "ijkl"),
431        ];
432        let max_context_bytes = serialized_context_bytes(&Vec::<RetrievalBackendStatus>::new())
433            + serialized_context_bytes(&Vec::<RetrievalBackendStatus>::new())
434            + serialized_context_bytes(&results[0])
435            + serialized_context_bytes(&items[0])
436            + serialized_context_bytes(&results[1])
437            + serialized_context_bytes(&items[1]);
438        let response = crate::api::HybridRetrievalResponse {
439            metadata: ApiMetadata {
440                trace_id: "trace".to_owned(),
441                request_id: "req".to_owned(),
442                graph_version: 1,
443                index_version: None,
444                indexed_graph_version: None,
445                stale: false,
446            },
447            context_pack: RetrievedContextPack {
448                graph_version: GraphVersion::new(1),
449                source_scope: Some("docs".to_owned()),
450                freshness: FreshnessPolicy::AllowStale,
451                truncated: false,
452                backend_statuses: Vec::new(),
453                provenance_trace: None,
454                items,
455            },
456            retrieval_mode: RetrievalMode::Hybrid,
457            source_scope: Some("docs".to_owned()),
458            freshness: FreshnessPolicy::AllowStale,
459            results,
460            fusion: FusionDiagnostics {
461                algorithm: "reciprocal_rank_fusion".to_owned(),
462                k: 60.0,
463                candidate_count: 3,
464            },
465            rerank: rerank_diagnostics(3, 3),
466            backend_statuses: Vec::new(),
467            truncated: false,
468            budget_used: RetrievalBudgetUsed {
469                limit: 3,
470                candidate_count: 3,
471                returned_count: 3,
472                context_bytes: 12,
473            },
474            degraded_reason: None,
475            indexes: Vec::new(),
476            index_cursors: Vec::new(),
477            index_refresh: IndexRefreshDiagnostics {
478                queue_depth: 2,
479                ..IndexRefreshDiagnostics::default()
480            },
481        };
482
483        let result = AgentRetrievalResult::from_retrieval(
484            response,
485            RuntimeIdentity::mcp(Some("call-1".to_owned())),
486            max_context_bytes,
487            4,
488        );
489
490        assert!(result.truncated);
491        assert_eq!(result.results.len(), 2);
492        assert_eq!(result.context_pack.items.len(), 2);
493        assert_eq!(result.budget_used.returned_count, 2);
494        assert_eq!(result.rerank.returned_count, 2);
495        assert_eq!(result.budget_used.context_bytes, max_context_bytes);
496        assert_eq!(result.freshness, "allow-stale");
497        assert_eq!(result.index_refresh.queue_depth, 2);
498    }
499
500    #[test]
501    fn omits_backend_metadata_when_it_exceeds_agent_context_budget() {
502        let backend_statuses = vec![RetrievalBackendStatus {
503            source: RetrieverSource::Semantic,
504            state: RetrievalBackendState::Unavailable,
505            scope_post_filter: true,
506            indexed_graph_version: Some(GraphVersion::new(1)),
507            reason: Some("semantic backend disabled by local policy".repeat(8)),
508        }];
509        let response = crate::api::HybridRetrievalResponse {
510            metadata: ApiMetadata {
511                trace_id: "trace".to_owned(),
512                request_id: "req".to_owned(),
513                graph_version: 1,
514                index_version: None,
515                indexed_graph_version: None,
516                stale: false,
517            },
518            context_pack: RetrievedContextPack {
519                graph_version: GraphVersion::new(1),
520                source_scope: Some("docs".to_owned()),
521                freshness: FreshnessPolicy::AllowStale,
522                truncated: false,
523                backend_statuses: backend_statuses.clone(),
524                provenance_trace: None,
525                items: Vec::new(),
526            },
527            retrieval_mode: RetrievalMode::Hybrid,
528            source_scope: Some("docs".to_owned()),
529            freshness: FreshnessPolicy::AllowStale,
530            results: Vec::new(),
531            fusion: FusionDiagnostics {
532                algorithm: "reciprocal_rank_fusion".to_owned(),
533                k: 60.0,
534                candidate_count: 0,
535            },
536            rerank: rerank_diagnostics(0, 0),
537            backend_statuses,
538            truncated: false,
539            budget_used: RetrievalBudgetUsed {
540                limit: 3,
541                candidate_count: 0,
542                returned_count: 0,
543                context_bytes: 0,
544            },
545            degraded_reason: None,
546            indexes: Vec::new(),
547            index_cursors: Vec::new(),
548            index_refresh: IndexRefreshDiagnostics::default(),
549        };
550
551        let result = AgentRetrievalResult::from_retrieval(
552            response,
553            RuntimeIdentity::mcp(Some("call-1".to_owned())),
554            8,
555            4,
556        );
557
558        assert!(result.truncated);
559        assert!(result.backend_statuses.is_empty());
560        assert!(result.context_pack.backend_statuses.is_empty());
561        assert!(result.budget_used.context_bytes <= 8);
562    }
563
564    #[test]
565    fn omits_trace_before_dropping_cited_results_when_context_budget_is_tight() {
566        let results = vec![hit("ev-1", "grounded answer content")];
567        let items = vec![pack_item("ev-1")];
568        let mut trace = TraversalProvenanceTrace::from_hits(
569            GraphVersion::new(1),
570            Some("docs".to_owned()),
571            "direct_context_lookup".to_owned(),
572            &results,
573        );
574        trace.mark_citations(["ev-1"]);
575        let max_context_bytes = serialized_context_bytes(&Vec::<RetrievalBackendStatus>::new())
576            + serialized_context_bytes(&Vec::<RetrievalBackendStatus>::new())
577            + serialized_context_bytes(&results[0])
578            + serialized_context_bytes(&items[0]);
579        let response = crate::api::HybridRetrievalResponse {
580            metadata: ApiMetadata {
581                trace_id: "trace".to_owned(),
582                request_id: "req".to_owned(),
583                graph_version: 1,
584                index_version: None,
585                indexed_graph_version: None,
586                stale: false,
587            },
588            context_pack: RetrievedContextPack {
589                graph_version: GraphVersion::new(1),
590                source_scope: Some("docs".to_owned()),
591                freshness: FreshnessPolicy::AllowStale,
592                truncated: false,
593                backend_statuses: Vec::new(),
594                provenance_trace: Some(trace),
595                items,
596            },
597            retrieval_mode: RetrievalMode::Hybrid,
598            source_scope: Some("docs".to_owned()),
599            freshness: FreshnessPolicy::AllowStale,
600            results,
601            fusion: FusionDiagnostics {
602                algorithm: "reciprocal_rank_fusion".to_owned(),
603                k: 60.0,
604                candidate_count: 1,
605            },
606            rerank: rerank_diagnostics(1, 1),
607            backend_statuses: Vec::new(),
608            truncated: false,
609            budget_used: RetrievalBudgetUsed {
610                limit: 1,
611                candidate_count: 1,
612                returned_count: 1,
613                context_bytes: 0,
614            },
615            degraded_reason: None,
616            indexes: Vec::new(),
617            index_cursors: Vec::new(),
618            index_refresh: IndexRefreshDiagnostics::default(),
619        };
620
621        let result = AgentRetrievalResult::from_retrieval(
622            response,
623            RuntimeIdentity::mcp(Some("call-1".to_owned())),
624            max_context_bytes,
625            4,
626        );
627
628        assert!(result.truncated);
629        assert_eq!(result.results.len(), 1);
630        assert!(result.context_pack.provenance_trace.is_none());
631    }
632
633    #[test]
634    fn reports_truncated_agent_result_when_trace_is_budgeted_but_retained() {
635        let mut result_hit = hit("ev-1", "grounded answer content");
636        result_hit.graph_facts = (0..16)
637            .map(|index| graph_fact(index, "ev-1"))
638            .collect::<Vec<_>>();
639        result_hit.retriever_sources = vec![RetrieverSource::GraphPath];
640        let results = vec![result_hit];
641        let items = vec![pack_item("ev-1")];
642        let mut trace = TraversalProvenanceTrace::from_hits(
643            GraphVersion::new(1),
644            Some("docs".to_owned()),
645            "direct_context_lookup".to_owned(),
646            &results,
647        );
648        trace.mark_citations(["ev-1"]);
649        let mut budgeted_trace = trace.clone();
650        budgeted_trace.apply_budget(9);
651        budgeted_trace.apply_budget(1);
652        budgeted_trace.truncated = true;
653        let max_context_bytes = serialized_context_bytes(&Vec::<RetrievalBackendStatus>::new())
654            + serialized_context_bytes(&Vec::<RetrievalBackendStatus>::new())
655            + serialized_context_bytes(&results[0])
656            + serialized_context_bytes(&items[0])
657            + serialized_context_bytes(&budgeted_trace);
658        let response = crate::api::HybridRetrievalResponse {
659            metadata: ApiMetadata {
660                trace_id: "trace".to_owned(),
661                request_id: "req".to_owned(),
662                graph_version: 1,
663                index_version: None,
664                indexed_graph_version: None,
665                stale: false,
666            },
667            context_pack: RetrievedContextPack {
668                graph_version: GraphVersion::new(1),
669                source_scope: Some("docs".to_owned()),
670                freshness: FreshnessPolicy::AllowStale,
671                truncated: false,
672                backend_statuses: Vec::new(),
673                provenance_trace: Some(trace),
674                items,
675            },
676            retrieval_mode: RetrievalMode::Hybrid,
677            source_scope: Some("docs".to_owned()),
678            freshness: FreshnessPolicy::AllowStale,
679            results,
680            fusion: FusionDiagnostics {
681                algorithm: "reciprocal_rank_fusion".to_owned(),
682                k: 60.0,
683                candidate_count: 1,
684            },
685            rerank: rerank_diagnostics(1, 1),
686            backend_statuses: Vec::new(),
687            truncated: false,
688            budget_used: RetrievalBudgetUsed {
689                limit: 1,
690                candidate_count: 1,
691                returned_count: 1,
692                context_bytes: 0,
693            },
694            degraded_reason: None,
695            indexes: Vec::new(),
696            index_cursors: Vec::new(),
697            index_refresh: IndexRefreshDiagnostics::default(),
698        };
699
700        let result = AgentRetrievalResult::from_retrieval(
701            response,
702            RuntimeIdentity::mcp(Some("call-1".to_owned())),
703            max_context_bytes,
704            4,
705        );
706
707        assert!(result.truncated);
708        assert!(result.context_pack.truncated);
709        assert!(
710            result
711                .context_pack
712                .provenance_trace
713                .as_ref()
714                .is_some_and(|trace| trace.truncated)
715        );
716    }
717
718    #[test]
719    fn reports_truncated_agent_result_when_trace_items_are_budgeted() {
720        let mut result_hit = hit("ev-1", "grounded answer content");
721        result_hit.graph_facts = (0..16)
722            .map(|index| graph_fact(index, "ev-1"))
723            .collect::<Vec<_>>();
724        result_hit.retriever_sources = vec![RetrieverSource::GraphPath];
725        let results = vec![result_hit];
726        let items = vec![pack_item("ev-1")];
727        let mut trace = TraversalProvenanceTrace::from_hits(
728            GraphVersion::new(1),
729            Some("docs".to_owned()),
730            "direct_context_lookup".to_owned(),
731            &results,
732        );
733        trace.mark_citations(["ev-1"]);
734        let response = crate::api::HybridRetrievalResponse {
735            metadata: ApiMetadata {
736                trace_id: "trace".to_owned(),
737                request_id: "req".to_owned(),
738                graph_version: 1,
739                index_version: None,
740                indexed_graph_version: None,
741                stale: false,
742            },
743            context_pack: RetrievedContextPack {
744                graph_version: GraphVersion::new(1),
745                source_scope: Some("docs".to_owned()),
746                freshness: FreshnessPolicy::AllowStale,
747                truncated: false,
748                backend_statuses: Vec::new(),
749                provenance_trace: Some(trace),
750                items,
751            },
752            retrieval_mode: RetrievalMode::Hybrid,
753            source_scope: Some("docs".to_owned()),
754            freshness: FreshnessPolicy::AllowStale,
755            results,
756            fusion: FusionDiagnostics {
757                algorithm: "reciprocal_rank_fusion".to_owned(),
758                k: 60.0,
759                candidate_count: 1,
760            },
761            rerank: rerank_diagnostics(1, 1),
762            backend_statuses: Vec::new(),
763            truncated: false,
764            budget_used: RetrievalBudgetUsed {
765                limit: 1,
766                candidate_count: 1,
767                returned_count: 1,
768                context_bytes: 0,
769            },
770            degraded_reason: None,
771            indexes: Vec::new(),
772            index_cursors: Vec::new(),
773            index_refresh: IndexRefreshDiagnostics::default(),
774        };
775
776        let result = AgentRetrievalResult::from_retrieval(
777            response,
778            RuntimeIdentity::mcp(Some("call-1".to_owned())),
779            usize::MAX,
780            4,
781        );
782
783        assert!(result.truncated);
784        assert!(result.context_pack.truncated);
785        assert!(
786            result
787                .context_pack
788                .provenance_trace
789                .as_ref()
790                .is_some_and(|trace| trace.truncated)
791        );
792    }
793
794    #[test]
795    fn filters_dropped_hits_from_agent_trace_before_byte_budget() {
796        let retained_hit = hit("ev-1", "grounded answer content");
797        let mut dropped_hit = hit("ev-2", "omitted answer content");
798        dropped_hit.graph_facts = (0..32)
799            .map(|index| graph_fact(index, "ev-2"))
800            .collect::<Vec<_>>();
801        dropped_hit.retriever_sources = vec![RetrieverSource::GraphPath];
802        let results = vec![retained_hit, dropped_hit];
803        let items = vec![pack_item("ev-1"), pack_item("ev-2")];
804        let mut trace = TraversalProvenanceTrace::from_hits(
805            GraphVersion::new(1),
806            Some("docs".to_owned()),
807            "direct_context_lookup".to_owned(),
808            &results,
809        );
810        trace.mark_citations(["ev-1", "ev-2"]);
811        let mut retained_trace = trace.clone();
812        retained_trace.retain_hits([&results[0]]);
813        retained_trace.mark_citations_for_hits([&results[0]]);
814        retained_trace.truncated = true;
815        retained_trace.apply_budget(9);
816        let max_context_bytes = serialized_context_bytes(&Vec::<RetrievalBackendStatus>::new())
817            + serialized_context_bytes(&Vec::<RetrievalBackendStatus>::new())
818            + serialized_context_bytes(&results[0])
819            + serialized_context_bytes(&items[0])
820            + serialized_context_bytes(&retained_trace);
821        let response = crate::api::HybridRetrievalResponse {
822            metadata: ApiMetadata {
823                trace_id: "trace".to_owned(),
824                request_id: "req".to_owned(),
825                graph_version: 1,
826                index_version: None,
827                indexed_graph_version: None,
828                stale: false,
829            },
830            context_pack: RetrievedContextPack {
831                graph_version: GraphVersion::new(1),
832                source_scope: Some("docs".to_owned()),
833                freshness: FreshnessPolicy::AllowStale,
834                truncated: false,
835                backend_statuses: Vec::new(),
836                provenance_trace: Some(trace),
837                items,
838            },
839            retrieval_mode: RetrievalMode::Hybrid,
840            source_scope: Some("docs".to_owned()),
841            freshness: FreshnessPolicy::AllowStale,
842            results,
843            fusion: FusionDiagnostics {
844                algorithm: "reciprocal_rank_fusion".to_owned(),
845                k: 60.0,
846                candidate_count: 2,
847            },
848            rerank: rerank_diagnostics(2, 2),
849            backend_statuses: Vec::new(),
850            truncated: false,
851            budget_used: RetrievalBudgetUsed {
852                limit: 2,
853                candidate_count: 2,
854                returned_count: 2,
855                context_bytes: 0,
856            },
857            degraded_reason: None,
858            indexes: Vec::new(),
859            index_cursors: Vec::new(),
860            index_refresh: IndexRefreshDiagnostics::default(),
861        };
862
863        let result = AgentRetrievalResult::from_retrieval(
864            response,
865            RuntimeIdentity::mcp(Some("call-1".to_owned())),
866            max_context_bytes,
867            4,
868        );
869
870        assert!(result.truncated);
871        assert_eq!(result.results.len(), 1);
872        assert_eq!(result.results[0].evidence_id, "ev-1");
873        let trace = result
874            .context_pack
875            .provenance_trace
876            .as_ref()
877            .expect("retained-only trace should fit");
878        assert!(
879            trace
880                .cited_evidence
881                .iter()
882                .all(|evidence| evidence.evidence_id == "ev-1")
883        );
884        assert!(
885            trace
886                .ranking_contributions
887                .iter()
888                .all(|contribution| contribution.result_id == "ev-1")
889        );
890    }
891
892    #[test]
893    fn rejects_zero_policy_budgets() {
894        let error = AgentAccessPolicy::new(Vec::new(), false, 0, 1, 1, false).expect_err("zero");
895
896        assert_eq!(error, AgentPolicyError::ZeroMaxLimit);
897    }
898
899    fn hit(evidence_id: &str, content: &str) -> RetrievalHit {
900        RetrievalHit {
901            evidence_id: evidence_id.to_owned(),
902            source_scope: "docs".to_owned(),
903            source_path: None,
904            source_span: None,
905            content: content.to_owned(),
906            entity_labels: Vec::new(),
907            entities: Vec::new(),
908            graph_facts: Vec::new(),
909            code_artifact: None,
910            retriever_sources: Vec::new(),
911            ranking: Vec::new(),
912            rerank: None,
913            score: 1.0,
914        }
915    }
916
917    fn pack_item(result_id: &str) -> ContextPackItem {
918        ContextPackItem {
919            result_id: result_id.to_owned(),
920            source_scope: "docs".to_owned(),
921            source_path: None,
922            source_span: None,
923            entities: Vec::new(),
924            graph_facts: Vec::new(),
925            graph_paths: Vec::new(),
926            code_artifact: None,
927            retriever_sources: Vec::new(),
928            ranking: Vec::new(),
929            rerank: None,
930        }
931    }
932
933    fn graph_fact(index: usize, evidence_id: &str) -> ContextGraphFact {
934        ContextGraphFact {
935            fact_id: format!("fact-{index}"),
936            kind: ContextGraphFactKind::Relation,
937            subject: format!("source-{index}"),
938            predicate: "supports".to_owned(),
939            object: Some(format!("target-{index}")),
940            evidence_ids: vec![evidence_id.to_owned()],
941            confidence: ConfidenceScore { basis_points: 9000 },
942            status: FactStatus::Accepted,
943            version_range: GraphVersionRange::open_from(GraphVersion::new(1)),
944        }
945    }
946
947    fn rerank_diagnostics(candidate_count: usize, returned_count: usize) -> RerankDiagnostics {
948        RerankDiagnostics {
949            requested_mode: RerankMode::Local,
950            effective_mode: RerankMode::Local,
951            algorithm: "deterministic_feature_rerank".to_owned(),
952            candidate_count,
953            returned_count,
954            degraded: false,
955            reason: None,
956        }
957    }
958
959    #[test]
960    fn carries_agent_context_without_domain_identity_leakage() {
961        let context = AgentRequestContext {
962            request: RequestContext::with_ids(InterfaceKind::Mcp, "req", "trace"),
963            runtime_identity: RuntimeIdentity::mcp(Some("tool".to_owned())),
964            policy_id: "default".to_owned(),
965        };
966
967        assert_eq!(context.request.interface, InterfaceKind::Mcp);
968        assert_eq!(context.runtime_identity.protocol, AgentProtocolKind::Mcp);
969    }
970}
relay_knowledge/api/agent.rs

relay_knowledge/api/
agent.rs
