Skip to main content

relay_knowledge/api/
agent.rs

1use std::collections::{HashMap, HashSet};
2
3use serde::{Deserialize, Serialize};
4
5use crate::domain::{
6    FreshnessPolicy, FusionDiagnostics, IndexStatus, RetrievalBackendStatus, RetrievalHit,
7    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
222impl AgentRetrievalResult {
223    /// Builds the canonical agent result and applies the context byte budget.
224    pub fn from_retrieval(
225        response: crate::api::HybridRetrievalResponse,
226        identity: RuntimeIdentity,
227        max_context_bytes: usize,
228        elapsed_ms: u64,
229    ) -> Self {
230        let crate::api::HybridRetrievalResponse {
231            metadata,
232            mut context_pack,
233            retrieval_mode,
234            source_scope,
235            freshness,
236            results: response_results,
237            fusion,
238            mut rerank,
239            mut backend_statuses,
240            truncated: response_truncated,
241            budget_used,
242            degraded_reason,
243            indexes,
244            index_cursors,
245            index_refresh,
246        } = response;
247        let item_bytes = context_pack
248            .items
249            .iter()
250            .map(|item| (item.result_id.clone(), serialized_context_bytes(item)))
251            .collect::<HashMap<_, _>>();
252        let mut context_bytes = serialized_context_bytes(&context_pack.backend_statuses)
253            .saturating_add(serialized_context_bytes(&backend_statuses));
254        let mut truncated = response_truncated;
255        if context_bytes > max_context_bytes {
256            context_pack.backend_statuses.clear();
257            backend_statuses.clear();
258            context_bytes = 0;
259            truncated = true;
260        }
261        let mut results = Vec::new();
262
263        for hit in response_results {
264            let hit_bytes = serialized_context_bytes(&hit).saturating_add(
265                item_bytes
266                    .get(hit.evidence_id.as_str())
267                    .copied()
268                    .unwrap_or_default(),
269            );
270            if context_bytes.saturating_add(hit_bytes) > max_context_bytes {
271                truncated = true;
272                continue;
273            }
274            context_bytes += hit_bytes;
275            results.push(hit);
276        }
277        let returned_count = results.len();
278        rerank.returned_count = returned_count;
279        let retained_result_ids = results
280            .iter()
281            .map(|hit| hit.evidence_id.as_str())
282            .collect::<HashSet<_>>();
283        context_pack.truncated = truncated;
284        context_pack
285            .items
286            .retain(|item| retained_result_ids.contains(item.result_id.as_str()));
287
288        Self {
289            metadata,
290            runtime_identity: identity,
291            source_scope,
292            freshness: freshness_label(freshness).to_owned(),
293            retrieval_mode,
294            context_pack,
295            results,
296            fusion,
297            rerank,
298            backend_statuses,
299            indexes,
300            index_cursors,
301            index_refresh,
302            degraded_reason,
303            truncated,
304            budget_used: AgentBudgetUsed {
305                limit: budget_used.limit,
306                candidate_count: budget_used.candidate_count,
307                returned_count,
308                context_bytes,
309                elapsed_ms,
310            },
311        }
312    }
313}
314
315fn serialized_context_bytes<T: Serialize>(value: &T) -> usize {
316    serde_json::to_vec(value)
317        .map(|bytes| bytes.len())
318        .unwrap_or(usize::MAX / 4)
319}
320
321/// Runtime budget consumed by a completed agent retrieval.
322#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
323pub struct AgentBudgetUsed {
324    pub limit: usize,
325    pub candidate_count: usize,
326    pub returned_count: usize,
327    pub context_bytes: usize,
328    pub elapsed_ms: u64,
329}
330
331pub fn freshness_label(freshness: FreshnessPolicy) -> &'static str {
332    match freshness {
333        FreshnessPolicy::AllowStale => "allow-stale",
334        FreshnessPolicy::WaitUntilFresh => "wait-until-fresh",
335        FreshnessPolicy::GraphOnly => "graph-only",
336    }
337}
338
339#[cfg(test)]
340mod tests {
341    use super::*;
342    use crate::{
343        api::InterfaceKind,
344        domain::{
345            ContextPackItem, FusionDiagnostics, GraphVersion, RerankDiagnostics, RerankMode,
346            RetrievalBackendState, RetrievalBackendStatus, RetrievalBudgetUsed, RetrievalHit,
347            RetrievedContextPack, RetrieverSource,
348        },
349    };
350
351    #[test]
352    fn truncates_retrieval_results_to_context_byte_budget() {
353        let items = vec![pack_item("ev-1"), pack_item("ev-2"), pack_item("ev-3")];
354        let results = vec![
355            hit("ev-1", "abcd"),
356            hit("ev-2", "efgh"),
357            hit("ev-3", "ijkl"),
358        ];
359        let max_context_bytes = serialized_context_bytes(&Vec::<RetrievalBackendStatus>::new())
360            + serialized_context_bytes(&Vec::<RetrievalBackendStatus>::new())
361            + serialized_context_bytes(&results[0])
362            + serialized_context_bytes(&items[0])
363            + serialized_context_bytes(&results[1])
364            + serialized_context_bytes(&items[1]);
365        let response = crate::api::HybridRetrievalResponse {
366            metadata: ApiMetadata {
367                trace_id: "trace".to_owned(),
368                request_id: "req".to_owned(),
369                graph_version: 1,
370                index_version: None,
371                indexed_graph_version: None,
372                stale: false,
373            },
374            context_pack: RetrievedContextPack {
375                graph_version: GraphVersion::new(1),
376                source_scope: Some("docs".to_owned()),
377                freshness: FreshnessPolicy::AllowStale,
378                truncated: false,
379                backend_statuses: Vec::new(),
380                items,
381            },
382            retrieval_mode: RetrievalMode::Hybrid,
383            source_scope: Some("docs".to_owned()),
384            freshness: FreshnessPolicy::AllowStale,
385            results,
386            fusion: FusionDiagnostics {
387                algorithm: "reciprocal_rank_fusion".to_owned(),
388                k: 60.0,
389                candidate_count: 3,
390            },
391            rerank: rerank_diagnostics(3, 3),
392            backend_statuses: Vec::new(),
393            truncated: false,
394            budget_used: RetrievalBudgetUsed {
395                limit: 3,
396                candidate_count: 3,
397                returned_count: 3,
398                context_bytes: 12,
399            },
400            degraded_reason: None,
401            indexes: Vec::new(),
402            index_cursors: Vec::new(),
403            index_refresh: IndexRefreshDiagnostics {
404                queue_depth: 2,
405                ..IndexRefreshDiagnostics::default()
406            },
407        };
408
409        let result = AgentRetrievalResult::from_retrieval(
410            response,
411            RuntimeIdentity::mcp(Some("call-1".to_owned())),
412            max_context_bytes,
413            4,
414        );
415
416        assert!(result.truncated);
417        assert_eq!(result.results.len(), 2);
418        assert_eq!(result.context_pack.items.len(), 2);
419        assert_eq!(result.budget_used.returned_count, 2);
420        assert_eq!(result.rerank.returned_count, 2);
421        assert_eq!(result.budget_used.context_bytes, max_context_bytes);
422        assert_eq!(result.freshness, "allow-stale");
423        assert_eq!(result.index_refresh.queue_depth, 2);
424    }
425
426    #[test]
427    fn omits_backend_metadata_when_it_exceeds_agent_context_budget() {
428        let backend_statuses = vec![RetrievalBackendStatus {
429            source: RetrieverSource::Semantic,
430            state: RetrievalBackendState::Unavailable,
431            scope_post_filter: true,
432            indexed_graph_version: Some(GraphVersion::new(1)),
433            reason: Some("semantic backend disabled by local policy".repeat(8)),
434        }];
435        let response = crate::api::HybridRetrievalResponse {
436            metadata: ApiMetadata {
437                trace_id: "trace".to_owned(),
438                request_id: "req".to_owned(),
439                graph_version: 1,
440                index_version: None,
441                indexed_graph_version: None,
442                stale: false,
443            },
444            context_pack: RetrievedContextPack {
445                graph_version: GraphVersion::new(1),
446                source_scope: Some("docs".to_owned()),
447                freshness: FreshnessPolicy::AllowStale,
448                truncated: false,
449                backend_statuses: backend_statuses.clone(),
450                items: Vec::new(),
451            },
452            retrieval_mode: RetrievalMode::Hybrid,
453            source_scope: Some("docs".to_owned()),
454            freshness: FreshnessPolicy::AllowStale,
455            results: Vec::new(),
456            fusion: FusionDiagnostics {
457                algorithm: "reciprocal_rank_fusion".to_owned(),
458                k: 60.0,
459                candidate_count: 0,
460            },
461            rerank: rerank_diagnostics(0, 0),
462            backend_statuses,
463            truncated: false,
464            budget_used: RetrievalBudgetUsed {
465                limit: 3,
466                candidate_count: 0,
467                returned_count: 0,
468                context_bytes: 0,
469            },
470            degraded_reason: None,
471            indexes: Vec::new(),
472            index_cursors: Vec::new(),
473            index_refresh: IndexRefreshDiagnostics::default(),
474        };
475
476        let result = AgentRetrievalResult::from_retrieval(
477            response,
478            RuntimeIdentity::mcp(Some("call-1".to_owned())),
479            8,
480            4,
481        );
482
483        assert!(result.truncated);
484        assert!(result.backend_statuses.is_empty());
485        assert!(result.context_pack.backend_statuses.is_empty());
486        assert!(result.budget_used.context_bytes <= 8);
487    }
488
489    #[test]
490    fn rejects_zero_policy_budgets() {
491        let error = AgentAccessPolicy::new(Vec::new(), false, 0, 1, 1, false).expect_err("zero");
492
493        assert_eq!(error, AgentPolicyError::ZeroMaxLimit);
494    }
495
496    fn hit(evidence_id: &str, content: &str) -> RetrievalHit {
497        RetrievalHit {
498            evidence_id: evidence_id.to_owned(),
499            source_scope: "docs".to_owned(),
500            source_path: None,
501            source_span: None,
502            content: content.to_owned(),
503            entity_labels: Vec::new(),
504            entities: Vec::new(),
505            graph_facts: Vec::new(),
506            code_artifact: None,
507            retriever_sources: Vec::new(),
508            ranking: Vec::new(),
509            rerank: None,
510            score: 1.0,
511        }
512    }
513
514    fn pack_item(result_id: &str) -> ContextPackItem {
515        ContextPackItem {
516            result_id: result_id.to_owned(),
517            source_scope: "docs".to_owned(),
518            source_path: None,
519            source_span: None,
520            entities: Vec::new(),
521            graph_facts: Vec::new(),
522            graph_paths: Vec::new(),
523            code_artifact: None,
524            retriever_sources: Vec::new(),
525            ranking: Vec::new(),
526            rerank: None,
527        }
528    }
529
530    fn rerank_diagnostics(candidate_count: usize, returned_count: usize) -> RerankDiagnostics {
531        RerankDiagnostics {
532            requested_mode: RerankMode::Local,
533            effective_mode: RerankMode::Local,
534            algorithm: "deterministic_feature_rerank".to_owned(),
535            candidate_count,
536            returned_count,
537            degraded: false,
538            reason: None,
539        }
540    }
541
542    #[test]
543    fn carries_agent_context_without_domain_identity_leakage() {
544        let context = AgentRequestContext {
545            request: RequestContext::with_ids(InterfaceKind::Mcp, "req", "trace"),
546            runtime_identity: RuntimeIdentity::mcp(Some("tool".to_owned())),
547            policy_id: "default".to_owned(),
548        };
549
550        assert_eq!(context.request.interface, InterfaceKind::Mcp);
551        assert_eq!(context.runtime_identity.protocol, AgentProtocolKind::Mcp);
552    }
553}