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#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
16#[serde(rename_all = "kebab-case")]
17pub enum AgentProtocolKind {
18 Mcp,
19 Acp,
20}
21
22#[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 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 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#[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#[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 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 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#[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#[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#[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#[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 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#[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}