1use std::collections::HashMap;
17use std::time::Duration;
18
19use contextgraph_types::{
20 ConsentReceipt, ContextFrame, ContextQuery, ContextQueryResult, DataFlow, EgressScope, FrameId,
21 ProviderUsage, ServedFrame, UsageReport, Verdict, VerifyRequest,
22};
23
24use crate::consent::{ConsentDecision, ConsentRecord, ConsentStore};
25use crate::error::HostError;
26use crate::provider::{ContextProvider, capability_matches};
27use crate::stdio::StdioProvider;
28use crate::trust::{AttestationLedger, FrameAttestationOutcome, TrustStore, TrustedKey};
29
30const DEFAULT_PROVIDER_TIMEOUT: Duration = Duration::from_secs(30);
34
35pub struct Host {
38 providers: Vec<Box<dyn ContextProvider>>,
39 consent: ConsentStore,
40 trust: TrustStore,
41 per_provider_timeout: Duration,
42}
43
44impl Default for Host {
45 fn default() -> Self {
46 Self::new()
47 }
48}
49
50impl Host {
51 pub fn new() -> Self {
53 Self {
54 providers: Vec::new(),
55 consent: ConsentStore::new(),
56 trust: TrustStore::new(),
57 per_provider_timeout: DEFAULT_PROVIDER_TIMEOUT,
58 }
59 }
60
61 pub fn with_timeout(per_provider_timeout: Duration) -> Self {
63 Self {
64 per_provider_timeout,
65 ..Self::new()
66 }
67 }
68
69 pub fn register(&mut self, provider: Box<dyn ContextProvider>) {
71 self.providers.push(provider);
72 }
73
74 pub async fn add_stdio(
77 &mut self,
78 id: impl Into<String>,
79 program: &str,
80 args: &[String],
81 ) -> Result<(), HostError> {
82 let provider = StdioProvider::spawn(id, program, args).await?;
83 self.providers.push(Box::new(provider));
84 Ok(())
85 }
86
87 pub async fn add_http(
95 &mut self,
96 id: impl Into<String>,
97 url: impl Into<String>,
98 credential: Option<crate::http::Credential>,
99 ) -> Result<(), HostError> {
100 let provider = crate::http::HttpProvider::connect_with_auth(id, url, credential).await?;
101 self.providers.push(Box::new(provider));
102 Ok(())
103 }
104
105 pub fn record_consent(&mut self, record: ConsentRecord) {
108 self.consent.record(record);
109 }
110
111 pub fn record_receipt(&mut self, receipt: ConsentReceipt) {
116 self.consent.record_receipt(receipt);
117 }
118
119 pub fn consent(&self) -> &ConsentStore {
121 &self.consent
122 }
123
124 pub fn trust_key(&mut self, provider_id: impl Into<String>, key: TrustedKey) {
138 self.trust.trust(provider_id, key);
139 }
140
141 pub fn trust(&self) -> &TrustStore {
144 &self.trust
145 }
146
147 pub fn set_trust_store(&mut self, trust: TrustStore) {
149 self.trust = trust;
150 }
151
152 pub fn provider_ids(&self) -> Vec<&str> {
154 self.providers.iter().map(|p| p.id()).collect()
155 }
156
157 pub fn provider(&self, id: &str) -> Option<&dyn ContextProvider> {
160 self.providers
161 .iter()
162 .find(|p| p.id() == id)
163 .map(|p| p.as_ref())
164 }
165
166 pub fn len(&self) -> usize {
168 self.providers.len()
169 }
170
171 pub fn is_empty(&self) -> bool {
172 self.providers.is_empty()
173 }
174
175 pub async fn verify_frames(&self, held: &[FrameId]) -> VerifyOutcome {
192 use futures_util::future::join_all;
193
194 let mut order: Vec<&str> = Vec::new();
197 let mut grouped: HashMap<&str, Vec<FrameId>> = HashMap::new();
198 for frame in held {
199 let id = frame.provider_id.as_str();
200 if !grouped.contains_key(id) {
201 order.push(id);
202 }
203 grouped.entry(id).or_default().push(frame.clone());
204 }
205
206 let legs = order.into_iter().map(|provider_id| {
207 let frames = grouped.remove(provider_id).unwrap_or_default();
208 self.verify_one_provider(provider_id, frames)
209 });
210
211 let mut outcome = VerifyOutcome::default();
212 for leg in join_all(legs).await {
213 outcome.retained.extend(leg.retained);
214 outcome.dropped.extend(leg.dropped);
215 }
216 outcome
217 }
218
219 async fn verify_one_provider(&self, provider_id: &str, frames: Vec<FrameId>) -> VerifyOutcome {
223 let mut outcome = VerifyOutcome::default();
224
225 let Some(provider) = self.provider(provider_id) else {
226 outcome.drop_all(frames, DropReason::UnknownProvider);
227 return outcome;
228 };
229 if !provider.capabilities().verify {
230 outcome.drop_all(frames, DropReason::VerifyUnsupported);
233 return outcome;
234 }
235
236 let (verifiable, undigested): (Vec<FrameId>, Vec<FrameId>) =
240 frames.into_iter().partition(FrameId::is_verifiable);
241 outcome.drop_all(undigested, DropReason::NoDigest);
242 if verifiable.is_empty() {
243 return outcome;
244 }
245
246 let request = VerifyRequest::new(verifiable.clone());
247 let response = match tokio::time::timeout(
248 self.per_provider_timeout,
249 provider.verify(&request),
250 )
251 .await
252 {
253 Ok(Ok(response)) => response,
254 Ok(Err(error)) => {
255 outcome.drop_all(verifiable, DropReason::VerifyFailed(error.to_string()));
256 return outcome;
257 }
258 Err(_) => {
259 let error = HostError::Timeout {
260 id: provider_id.to_string(),
261 timeout_ms: self.per_provider_timeout.as_millis() as u64,
262 };
263 outcome.drop_all(verifiable, DropReason::VerifyFailed(error.to_string()));
264 return outcome;
265 }
266 };
267
268 for frame in verifiable {
269 match response.verdict_for(&frame) {
272 Some(Verdict::Valid) => outcome.retained.push(frame),
273 Some(Verdict::Stale { replacement_digest }) => outcome.drop_one(
274 frame,
275 DropReason::Stale {
276 replacement_digest: replacement_digest.clone(),
277 },
278 ),
279 Some(Verdict::Gone) => outcome.drop_one(frame, DropReason::Gone),
280 Some(Verdict::Unknown) | None => outcome.drop_one(frame, DropReason::Unknown),
281 }
282 }
283 outcome
284 }
285
286 pub async fn query_provider(
292 &self,
293 id: &str,
294 query: &ContextQuery,
295 ) -> Result<ContextQueryResult, HostError> {
296 Ok(self.query_provider_attested(id, query).await?.0)
297 }
298
299 pub async fn query_provider_attested(
310 &self,
311 id: &str,
312 query: &ContextQuery,
313 ) -> Result<(ContextQueryResult, Vec<FrameAttestationOutcome>), HostError> {
314 let provider = self
315 .providers
316 .iter()
317 .find(|p| p.id() == id)
318 .ok_or_else(|| HostError::UnknownProvider(id.to_string()))?;
319
320 match self.consent.evaluate(provider.id(), provider.info()) {
321 ConsentDecision::Permitted => {}
322 ConsentDecision::NeedsConsent => {
323 return Err(HostError::ConsentRequired {
324 id: id.to_string(),
325 data_flow: provider.info().data_flow.clone(),
326 });
327 }
328 ConsentDecision::NeedsReceipts(scopes) => {
329 return Err(HostError::ConsentScopeRequired {
330 id: id.to_string(),
331 scopes,
332 });
333 }
334 }
335
336 let result =
337 match tokio::time::timeout(self.per_provider_timeout, provider.query(query)).await {
338 Ok(result) => result?,
339 Err(_) => {
340 return Err(HostError::Timeout {
341 id: id.to_string(),
342 timeout_ms: self.per_provider_timeout.as_millis() as u64,
343 });
344 }
345 };
346
347 let signing_id = provider.info().name.clone();
356 let outcomes = self.trust.check_result_signed_as(id, &signing_id, &result);
357 Ok((result, outcomes))
358 }
359
360 pub async fn query_all(&self, query: &ContextQuery) -> FanOut {
366 use futures_util::future::join_all;
367
368 let futures: Vec<_> = self
369 .providers
370 .iter()
371 .filter(|p| capability_matches(p.capabilities(), query))
372 .map(|p| self.query_one_isolated(p.as_ref(), query))
373 .collect();
374
375 FanOut {
376 outcomes: join_all(futures).await,
377 }
378 }
379
380 pub async fn query_all_budgeted(&self, template: &ContextQuery, global_budget: u32) -> FanOut {
400 use futures_util::future::join_all;
401
402 let matching: Vec<&dyn ContextProvider> = self
405 .providers
406 .iter()
407 .map(|provider| provider.as_ref())
408 .filter(|provider| capability_matches(provider.capabilities(), template))
409 .collect();
410
411 let shares = crate::compose::budget_split(global_budget, matching.len());
414
415 let queries: Vec<ContextQuery> = shares
418 .iter()
419 .map(|&share| ContextQuery {
420 max_tokens: share,
421 ..template.clone()
422 })
423 .collect();
424
425 let futures: Vec<_> = matching
426 .iter()
427 .zip(queries.iter())
428 .map(|(provider, query)| self.query_one_isolated(*provider, query))
429 .collect();
430
431 FanOut {
432 outcomes: join_all(futures).await,
433 }
434 }
435
436 async fn query_one_isolated(
439 &self,
440 provider: &dyn ContextProvider,
441 query: &ContextQuery,
442 ) -> ProviderOutcome {
443 let id = provider.id().to_string();
444
445 match self.consent.evaluate(provider.id(), provider.info()) {
450 ConsentDecision::Permitted => {}
451 ConsentDecision::NeedsConsent => {
452 return ProviderOutcome::unattested(
453 id,
454 ProviderResult::ConsentRequired(provider.info().data_flow.clone()),
455 );
456 }
457 ConsentDecision::NeedsReceipts(scopes) => {
458 return ProviderOutcome::unattested(
459 id,
460 ProviderResult::ConsentScopeRequired {
461 data_flow: provider.info().data_flow.clone(),
462 missing: scopes,
463 },
464 );
465 }
466 }
467
468 let result =
469 match tokio::time::timeout(self.per_provider_timeout, provider.query(query)).await {
470 Ok(Ok(result)) => result,
471 Ok(Err(error)) => {
472 return ProviderOutcome::unattested(id, ProviderResult::Failed(error));
473 }
474 Err(_) => {
475 let error = HostError::Timeout {
476 id: id.clone(),
477 timeout_ms: self.per_provider_timeout.as_millis() as u64,
478 };
479 return ProviderOutcome::unattested(id, ProviderResult::Failed(error));
480 }
481 };
482
483 if !result.respects_budget(query.max_tokens) {
486 return ProviderOutcome::unattested(
487 id,
488 ProviderResult::BudgetLie {
489 claimed_tokens: result.total_token_cost(),
490 max_tokens: query.max_tokens,
491 dropped_frames: result.frames.len(),
492 },
493 );
494 }
495
496 if !result.respects_frame_limit(query.max_frames) {
501 return ProviderOutcome::unattested(
502 id,
503 ProviderResult::FrameFlood {
504 returned_frames: result.frames.len(),
505 max_frames: query.max_frames,
506 },
507 );
508 }
509
510 let signing_id = provider.info().name.clone();
522 let attestations = self.trust.check_result_signed_as(&id, &signing_id, &result);
523
524 ProviderOutcome {
525 provider_id: id,
526 result: ProviderResult::Frames(result),
527 attestations,
528 }
529 }
530
531 pub async fn shutdown(self) -> Vec<(String, Result<(), HostError>)> {
535 let mut results = Vec::with_capacity(self.providers.len());
536 for provider in &self.providers {
537 results.push((provider.id().to_string(), provider.shutdown().await));
538 }
539 results
540 }
541}
542
543#[derive(Debug)]
546pub struct FanOut {
547 pub outcomes: Vec<ProviderOutcome>,
550}
551
552impl FanOut {
553 pub fn accepted_frames(&self) -> impl Iterator<Item = &ContextFrame> {
557 self.outcomes
558 .iter()
559 .filter_map(|outcome| match &outcome.result {
560 ProviderResult::Frames(result) => Some(result.frames.iter()),
561 _ => None,
562 })
563 .flatten()
564 }
565
566 pub fn total_accepted_tokens(&self) -> u64 {
568 self.accepted_frames().map(|f| f.token_cost as u64).sum()
569 }
570
571 pub fn accepted_with_provider(&self) -> impl Iterator<Item = (&str, &ContextFrame)> {
574 self.outcomes
575 .iter()
576 .filter_map(|outcome| match &outcome.result {
577 ProviderResult::Frames(result) => Some(
578 result
579 .frames
580 .iter()
581 .map(move |frame| (outcome.provider_id.as_str(), frame)),
582 ),
583 _ => None,
584 })
585 .flatten()
586 }
587
588 pub fn compose(&self) -> String {
594 crate::compose::compose_context(self.accepted_with_provider())
595 }
596
597 pub fn attestation_ledger(&self) -> AttestationLedger {
607 self.outcomes
608 .iter()
609 .flat_map(|outcome| outcome.attestations.iter().cloned())
610 .collect()
611 }
612
613 pub fn any_attested(&self) -> bool {
617 self.outcomes
618 .iter()
619 .flat_map(|outcome| outcome.attestations.iter())
620 .any(|outcome| outcome.state.is_attested())
621 }
622
623 pub fn compose_for_prompt(&self, global_budget: u32) -> crate::compose::ComposedPrompt {
643 self.compose_for_prompt_with(global_budget, &crate::compose::ranking::ScoreDescending)
644 }
645
646 pub fn compose_for_prompt_with<S>(
655 &self,
656 global_budget: u32,
657 strategy: &S,
658 ) -> crate::compose::ComposedPrompt
659 where
660 S: crate::compose::ranking::RankingStrategy + ?Sized,
661 {
662 crate::compose::compose_for_prompt_attested(
663 self.accepted_with_provider(),
664 global_budget,
665 strategy,
666 &self.attestation_ledger(),
667 )
668 }
669
670 pub fn usage_report(&self, query: &ContextQuery, as_of: impl Into<String>) -> UsageReport {
683 let providers: Vec<ProviderUsage> = self
684 .outcomes
685 .iter()
686 .map(|outcome| {
687 let provider_id = outcome.provider_id.clone();
688 match &outcome.result {
689 ProviderResult::Frames(result) => {
690 let served_frames: Vec<ServedFrame> = result
691 .frames
692 .iter()
693 .map(|frame| ServedFrame {
694 frame: frame.identity(&provider_id),
695 token_cost: frame.token_cost,
696 })
697 .collect();
698 let token_cost = served_frames.iter().map(|s| s.token_cost as u64).sum();
699 ProviderUsage {
700 provider_id,
701 frames_served: served_frames.len() as u32,
702 frames_rejected: 0,
703 token_cost,
704 served_frames,
705 }
706 }
707 ProviderResult::BudgetLie { dropped_frames, .. } => ProviderUsage {
710 provider_id,
711 frames_served: 0,
712 frames_rejected: *dropped_frames as u32,
713 token_cost: 0,
714 served_frames: vec![],
715 },
716 ProviderResult::FrameFlood {
720 returned_frames, ..
721 } => ProviderUsage {
722 provider_id,
723 frames_served: 0,
724 frames_rejected: *returned_frames as u32,
725 token_cost: 0,
726 served_frames: vec![],
727 },
728 ProviderResult::ConsentRequired(_)
731 | ProviderResult::ConsentScopeRequired { .. }
732 | ProviderResult::Failed(_) => ProviderUsage {
733 provider_id,
734 frames_served: 0,
735 frames_rejected: 0,
736 token_cost: 0,
737 served_frames: vec![],
738 },
739 }
740 })
741 .collect();
742
743 let budget_consumed = providers.iter().map(|p| p.token_cost).sum();
744 UsageReport {
745 budget_requested: query.max_tokens,
746 budget_consumed,
747 as_of: as_of.into(),
748 providers,
749 }
750 }
751
752 pub fn failures(&self) -> impl Iterator<Item = (&str, &HostError)> {
754 self.outcomes
755 .iter()
756 .filter_map(|outcome| match &outcome.result {
757 ProviderResult::Failed(error) => Some((outcome.provider_id.as_str(), error)),
758 _ => None,
759 })
760 }
761
762 pub fn budget_liars(&self) -> impl Iterator<Item = &ProviderOutcome> {
765 self.outcomes
766 .iter()
767 .filter(|outcome| matches!(outcome.result, ProviderResult::BudgetLie { .. }))
768 }
769
770 pub fn frame_floods(&self) -> impl Iterator<Item = &ProviderOutcome> {
774 self.outcomes
775 .iter()
776 .filter(|outcome| matches!(outcome.result, ProviderResult::FrameFlood { .. }))
777 }
778}
779
780#[derive(Debug)]
782pub struct ProviderOutcome {
783 pub provider_id: String,
784 pub result: ProviderResult,
785 pub attestations: Vec<FrameAttestationOutcome>,
794}
795
796impl ProviderOutcome {
797 fn unattested(provider_id: String, result: ProviderResult) -> Self {
800 Self {
801 provider_id,
802 result,
803 attestations: Vec::new(),
804 }
805 }
806}
807
808#[derive(Debug)]
811pub enum ProviderResult {
812 Frames(ContextQueryResult),
814 BudgetLie {
817 claimed_tokens: u64,
818 max_tokens: u32,
819 dropped_frames: usize,
820 },
821 FrameFlood {
825 returned_frames: usize,
826 max_frames: u32,
827 },
828 ConsentRequired(DataFlow),
831 ConsentScopeRequired {
835 data_flow: DataFlow,
836 missing: Vec<EgressScope>,
837 },
838 Failed(HostError),
840}
841
842#[derive(Debug, Clone, PartialEq, Eq)]
850pub enum DropReason {
851 Stale { replacement_digest: Option<String> },
854 Gone,
856 Unknown,
859 NoDigest,
862 VerifyUnsupported,
865 UnknownProvider,
867 VerifyFailed(String),
869}
870
871impl DropReason {
872 pub fn warrants_requery(&self) -> bool {
876 !matches!(self, Self::Gone)
877 }
878}
879
880#[derive(Debug, Clone, PartialEq, Eq)]
882pub struct DroppedFrame {
883 pub frame: FrameId,
885 pub reason: DropReason,
887}
888
889#[derive(Debug, Clone, Default, PartialEq, Eq)]
896pub struct VerifyOutcome {
897 pub retained: Vec<FrameId>,
900 pub dropped: Vec<DroppedFrame>,
902}
903
904impl VerifyOutcome {
905 pub fn requery(&self) -> impl Iterator<Item = &FrameId> {
908 self.dropped
909 .iter()
910 .filter(|dropped| dropped.reason.warrants_requery())
911 .map(|dropped| &dropped.frame)
912 }
913
914 pub fn was_dropped(&self, frame: &FrameId) -> bool {
916 self.dropped.iter().any(|dropped| &dropped.frame == frame)
917 }
918
919 pub fn drop_reason(&self, frame: &FrameId) -> Option<&DropReason> {
921 self.dropped
922 .iter()
923 .find(|dropped| &dropped.frame == frame)
924 .map(|dropped| &dropped.reason)
925 }
926
927 fn drop_one(&mut self, frame: FrameId, reason: DropReason) {
928 self.dropped.push(DroppedFrame { frame, reason });
929 }
930
931 fn drop_all(&mut self, frames: impl IntoIterator<Item = FrameId>, reason: DropReason) {
932 for frame in frames {
933 self.drop_one(frame, reason.clone());
934 }
935 }
936}
937
938#[cfg(test)]
939mod tests {
940 use super::*;
941 use async_trait::async_trait;
942 use contextgraph_types::Grantor;
943 use contextgraph_types::capability::QueryCapability;
944 use contextgraph_types::{Capabilities, ContextFrame, FrameKind, ProviderInfo};
945 use std::sync::Arc;
946 use std::sync::atomic::{AtomicBool, Ordering};
947
948 struct FakeProvider {
950 id: String,
951 info: ProviderInfo,
952 capabilities: Capabilities,
953 behavior: Behavior,
954 queried: Arc<AtomicBool>,
955 }
956
957 enum Behavior {
958 Frames(Vec<ContextFrame>),
959 Fail(String),
960 Slow(Duration),
961 }
962
963 impl FakeProvider {
964 fn new(id: &str, egress: bool, behavior: Behavior) -> Self {
965 Self::with_data_flow(
966 id,
967 DataFlow {
968 reads: true,
969 writes: false,
970 egress,
971 egress_scopes: vec![],
972 },
973 behavior,
974 )
975 }
976
977 fn scoped(id: &str, scopes: Vec<EgressScope>, behavior: Behavior) -> Self {
979 Self::with_data_flow(
980 id,
981 DataFlow {
982 reads: true,
983 writes: false,
984 egress: true,
985 egress_scopes: scopes,
986 },
987 behavior,
988 )
989 }
990
991 fn with_data_flow(id: &str, data_flow: DataFlow, behavior: Behavior) -> Self {
992 Self {
993 id: id.into(),
994 info: ProviderInfo {
995 name: id.into(),
996 version: "0.0.1".into(),
997 data_flow,
998 },
999 capabilities: Capabilities {
1000 query: QueryCapability {
1001 kinds: vec!["doc".into()],
1002 },
1003 ..Capabilities::default()
1004 },
1005 behavior,
1006 queried: Arc::new(AtomicBool::new(false)),
1007 }
1008 }
1009 }
1010
1011 #[async_trait]
1012 impl ContextProvider for FakeProvider {
1013 fn id(&self) -> &str {
1014 &self.id
1015 }
1016 fn info(&self) -> &ProviderInfo {
1017 &self.info
1018 }
1019 fn capabilities(&self) -> &Capabilities {
1020 &self.capabilities
1021 }
1022 async fn query(&self, _query: &ContextQuery) -> Result<ContextQueryResult, HostError> {
1023 self.queried.store(true, Ordering::SeqCst);
1024 match &self.behavior {
1025 Behavior::Frames(frames) => Ok(ContextQueryResult {
1026 frames: frames.clone(),
1027 truncated: false,
1028 dropped_estimate: None,
1029 ..Default::default()
1030 }),
1031 Behavior::Fail(message) => Err(HostError::Provider {
1032 id: self.id.clone(),
1033 code: None,
1034 message: message.clone(),
1035 }),
1036 Behavior::Slow(duration) => {
1037 tokio::time::sleep(*duration).await;
1038 Ok(ContextQueryResult {
1039 frames: vec![],
1040 truncated: false,
1041 dropped_estimate: None,
1042 ..Default::default()
1043 })
1044 }
1045 }
1046 }
1047 }
1048
1049 fn frame(id: &str, cost: u32) -> ContextFrame {
1050 ContextFrame {
1051 id: id.into(),
1052 kind: FrameKind::Doc,
1053 title: id.into(),
1054 content: Some("c".into()),
1055 content_digest: None,
1056 uri: None,
1057 representation: Default::default(),
1058 content_fidelity: None,
1059 canonical_content_hash: None,
1060 content_ref: None,
1061 transform: None,
1062 minimum_content_fidelity: None,
1063 inline_content_requirement: None,
1064 score: 0.5,
1065 token_cost: cost,
1066 canonical_token_cost: None,
1067 tokenizer_ref: None,
1068 valid_from: None,
1069 valid_to: None,
1070 recorded_at: None,
1071 provenance: vec![],
1072 citation_label: Some(id.into()),
1073 embedding: None,
1074 relations: vec![],
1075 }
1076 }
1077
1078 fn query() -> ContextQuery {
1079 ContextQuery {
1080 goal: "g".into(),
1081 query_text: None,
1082 embedding: None,
1083 kinds: vec![],
1084 anchors: vec![],
1085 max_frames: 10,
1086 max_tokens: 1000,
1087 as_of: None,
1088 representation_preferences: vec![],
1089 }
1090 }
1091
1092 #[tokio::test]
1093 async fn query_all_collects_frames_from_healthy_providers() {
1094 let mut host = Host::new();
1095 host.register(Box::new(FakeProvider::new(
1096 "a",
1097 false,
1098 Behavior::Frames(vec![frame("f1", 100), frame("f2", 100)]),
1099 )));
1100 host.register(Box::new(FakeProvider::new(
1101 "b",
1102 false,
1103 Behavior::Frames(vec![frame("f3", 50)]),
1104 )));
1105
1106 let fanout = host.query_all(&query()).await;
1107 assert_eq!(fanout.outcomes.len(), 2);
1108 assert_eq!(fanout.accepted_frames().count(), 3);
1109 assert_eq!(fanout.total_accepted_tokens(), 250);
1110 }
1111
1112 #[tokio::test]
1113 async fn a_budgeted_fan_out_keeps_honest_legs_under_the_global_budget() {
1114 let mut host = Host::new();
1120 for id in ["a", "b", "c", "d"] {
1121 host.register(Box::new(FakeProvider::new(
1122 id,
1123 false,
1124 Behavior::Frames(vec![frame(&format!("{id}1"), 200)]),
1125 )));
1126 }
1127 let template = query(); let fanout = host.query_all_budgeted(&template, 1000).await;
1129 assert_eq!(
1130 fanout.accepted_frames().count(),
1131 4,
1132 "each share fits its leg"
1133 );
1134 assert!(
1135 fanout.total_accepted_tokens() <= 1000,
1136 "honest legs must sum to <= the global budget, got {}",
1137 fanout.total_accepted_tokens()
1138 );
1139 assert_eq!(fanout.budget_liars().count(), 0);
1140 }
1141
1142 #[tokio::test]
1143 async fn the_budget_split_enforces_a_per_leg_ceiling_the_flat_fan_out_does_not() {
1144 let mut host = Host::new();
1150 host.register(Box::new(FakeProvider::new(
1151 "greedy",
1152 false,
1153 Behavior::Frames(vec![frame("g", 300)]),
1154 )));
1155 for id in ["b", "c", "d"] {
1156 host.register(Box::new(FakeProvider::new(
1157 id,
1158 false,
1159 Behavior::Frames(vec![frame(&format!("{id}1"), 100)]),
1160 )));
1161 }
1162
1163 let budgeted = host.query_all_budgeted(&query(), 1000).await;
1164 assert!(
1165 budgeted
1166 .budget_liars()
1167 .any(|outcome| outcome.provider_id == "greedy"),
1168 "a leg overspending its share is dropped by the existing B2 audit"
1169 );
1170 assert!(
1171 budgeted
1172 .accepted_with_provider()
1173 .all(|(id, _)| id != "greedy"),
1174 "the greedy leg contributes nothing under the split"
1175 );
1176
1177 let flat = host.query_all(&query()).await;
1180 assert!(flat.accepted_with_provider().any(|(id, _)| id == "greedy"));
1181 }
1182
1183 #[tokio::test]
1184 async fn the_host_composes_the_same_frame_set_to_identical_bytes_across_turns() {
1185 let mut host = Host::new();
1190 host.register(Box::new(FakeProvider::new(
1191 "prov-b",
1192 false,
1193 Behavior::Frames(vec![frame("f2", 100), frame("f1", 100)]),
1194 )));
1195 host.register(Box::new(FakeProvider::new(
1196 "prov-a",
1197 false,
1198 Behavior::Frames(vec![frame("f3", 50)]),
1199 )));
1200
1201 let first = host.query_all(&query()).await.compose();
1202 let second = host.query_all(&query()).await.compose();
1203 assert_eq!(
1204 first, second,
1205 "an unchanged frame set must compose to identical bytes"
1206 );
1207 assert_eq!(first.matches("<frame ").count(), 3);
1211 assert!(first.find("prov-a").unwrap() < first.find("prov-b").unwrap());
1212 }
1213
1214 #[tokio::test]
1215 async fn a_fan_out_rolls_up_into_a_self_consistent_usage_report() {
1216 let mut host = Host::new();
1217 host.register(Box::new(FakeProvider::new(
1218 "a",
1219 false,
1220 Behavior::Frames(vec![frame("f1", 100), frame("f2", 100)]),
1221 )));
1222 host.register(Box::new(FakeProvider::new(
1224 "liar",
1225 false,
1226 Behavior::Frames(vec![frame("big", 1200)]),
1227 )));
1228
1229 let query = query();
1230 let fanout = host.query_all(&query).await;
1231 let report = fanout.usage_report(&query, "2026-07-21T00:00:00Z");
1232
1233 assert_eq!(report.budget_requested, 1000);
1234 assert_eq!(report.budget_consumed, 200);
1235 assert_eq!(report.as_of, "2026-07-21T00:00:00Z");
1236 assert!(report.is_consistent());
1238 assert!(report.within_budget());
1239 let independent: u64 = fanout.accepted_frames().map(|f| f.token_cost as u64).sum();
1242 assert_eq!(report.budget_consumed, independent);
1243 assert_eq!(report.budget_consumed, fanout.total_accepted_tokens());
1244
1245 let a = report
1246 .providers
1247 .iter()
1248 .find(|p| p.provider_id == "a")
1249 .expect("provider a is in the report");
1250 assert_eq!(a.frames_served, 2);
1251 assert_eq!(a.frames_rejected, 0);
1252 assert_eq!(a.token_cost, 200);
1253 let ids: Vec<&str> = a
1255 .served_frames
1256 .iter()
1257 .map(|s| s.frame.frame_id.as_str())
1258 .collect();
1259 assert!(ids.contains(&"f1") && ids.contains(&"f2"));
1260 assert!(a.served_frames.iter().all(|s| s.frame.provider_id == "a"));
1261
1262 let liar = report
1263 .providers
1264 .iter()
1265 .find(|p| p.provider_id == "liar")
1266 .expect("the liar is still accounted for");
1267 assert_eq!(liar.frames_served, 0);
1268 assert_eq!(liar.frames_rejected, 1);
1269 assert_eq!(liar.token_cost, 0);
1270 assert!(liar.served_frames.is_empty());
1271 }
1272
1273 #[tokio::test]
1274 async fn a_provider_lying_about_token_cost_has_its_frames_dropped_loudly() {
1275 let mut host = Host::new();
1276 host.register(Box::new(FakeProvider::new(
1278 "liar",
1279 false,
1280 Behavior::Frames(vec![frame("big", 1200)]),
1281 )));
1282 host.register(Box::new(FakeProvider::new(
1283 "honest",
1284 false,
1285 Behavior::Frames(vec![frame("ok", 200)]),
1286 )));
1287
1288 let fanout = host.query_all(&query()).await;
1289 assert_eq!(fanout.accepted_frames().count(), 1);
1291 assert_eq!(fanout.total_accepted_tokens(), 200);
1292 let liars: Vec<_> = fanout.budget_liars().collect();
1294 assert_eq!(liars.len(), 1);
1295 assert_eq!(liars[0].provider_id, "liar");
1296 match liars[0].result {
1297 ProviderResult::BudgetLie {
1298 claimed_tokens,
1299 max_tokens,
1300 dropped_frames,
1301 } => {
1302 assert_eq!(claimed_tokens, 1200);
1303 assert_eq!(max_tokens, 1000);
1304 assert_eq!(dropped_frames, 1);
1305 }
1306 _ => unreachable!(),
1307 }
1308 }
1309
1310 #[tokio::test]
1311 async fn a_provider_returning_more_than_max_frames_has_them_dropped_loudly() {
1312 let query = query();
1317 let flood: Vec<ContextFrame> = (0..12).map(|i| frame(&format!("f{i}"), 1)).collect();
1318 let as_result = ContextQueryResult {
1321 frames: flood.clone(),
1322 truncated: false,
1323 dropped_estimate: None,
1324 ..Default::default()
1325 };
1326 assert!(as_result.respects_budget(query.max_tokens));
1327 assert!(!as_result.respects_frame_limit(query.max_frames));
1328
1329 let mut host = Host::new();
1330 host.register(Box::new(FakeProvider::new(
1331 "flood",
1332 false,
1333 Behavior::Frames(flood),
1334 )));
1335 host.register(Box::new(FakeProvider::new(
1336 "honest",
1337 false,
1338 Behavior::Frames(vec![frame("ok", 200)]),
1339 )));
1340
1341 let fanout = host.query_all(&query).await;
1342 assert_eq!(fanout.accepted_frames().count(), 1);
1344 assert_eq!(fanout.total_accepted_tokens(), 200);
1345
1346 let floods: Vec<_> = fanout.frame_floods().collect();
1348 assert_eq!(floods.len(), 1);
1349 assert_eq!(floods[0].provider_id, "flood");
1350 match floods[0].result {
1351 ProviderResult::FrameFlood {
1352 returned_frames,
1353 max_frames,
1354 } => {
1355 assert_eq!(returned_frames, 12);
1356 assert_eq!(max_frames, 10);
1357 }
1358 _ => unreachable!(),
1359 }
1360 assert_eq!(fanout.budget_liars().count(), 0);
1362
1363 let report = fanout.usage_report(&query, "2026-07-21T00:00:00Z");
1365 let flooder = report
1366 .providers
1367 .iter()
1368 .find(|p| p.provider_id == "flood")
1369 .expect("the flooder is still accounted for");
1370 assert_eq!(flooder.frames_served, 0);
1371 assert_eq!(flooder.frames_rejected, 12);
1372 assert_eq!(flooder.token_cost, 0);
1373 assert!(flooder.served_frames.is_empty());
1374 assert!(report.is_consistent());
1375 }
1376
1377 #[tokio::test]
1378 async fn one_failing_provider_never_poisons_the_others() {
1379 let mut host = Host::new();
1380 host.register(Box::new(FakeProvider::new(
1381 "healthy",
1382 false,
1383 Behavior::Frames(vec![frame("f", 10)]),
1384 )));
1385 host.register(Box::new(FakeProvider::new(
1386 "broken",
1387 false,
1388 Behavior::Fail("kaboom".into()),
1389 )));
1390
1391 let fanout = host.query_all(&query()).await;
1392 assert_eq!(fanout.accepted_frames().count(), 1);
1393 let failures: Vec<_> = fanout.failures().collect();
1394 assert_eq!(failures.len(), 1);
1395 assert_eq!(failures[0].0, "broken");
1396 }
1397
1398 #[tokio::test]
1399 async fn a_slow_provider_is_timed_out_without_stalling_the_fan_out() {
1400 let mut host = Host::with_timeout(Duration::from_millis(50));
1401 host.register(Box::new(FakeProvider::new(
1402 "fast",
1403 false,
1404 Behavior::Frames(vec![frame("f", 10)]),
1405 )));
1406 host.register(Box::new(FakeProvider::new(
1407 "slow",
1408 false,
1409 Behavior::Slow(Duration::from_secs(30)),
1410 )));
1411
1412 let fanout = host.query_all(&query()).await;
1413 assert_eq!(fanout.accepted_frames().count(), 1);
1414 let failures: Vec<_> = fanout.failures().collect();
1415 assert_eq!(failures.len(), 1);
1416 assert_eq!(failures[0].0, "slow");
1417 assert!(matches!(failures[0].1, HostError::Timeout { .. }));
1418 }
1419
1420 #[tokio::test]
1421 async fn an_egress_provider_is_not_queried_until_consent_is_recorded() {
1422 let mut host = Host::new();
1423 let provider = FakeProvider::new("github", true, Behavior::Frames(vec![frame("f", 10)]));
1424 let queried = provider.queried.clone();
1425 host.register(Box::new(provider));
1426
1427 let fanout = host.query_all(&query()).await;
1430 assert_eq!(fanout.accepted_frames().count(), 0);
1431 assert!(!queried.load(Ordering::SeqCst), "payload must not be sent");
1432 assert!(matches!(
1433 fanout.outcomes[0].result,
1434 ProviderResult::ConsentRequired(_)
1435 ));
1436 assert!(matches!(
1438 host.query_provider("github", &query()).await,
1439 Err(HostError::ConsentRequired { .. })
1440 ));
1441
1442 host.record_consent(ConsentRecord::new(
1444 "github",
1445 DataFlow {
1446 reads: true,
1447 writes: false,
1448 egress: true,
1449 egress_scopes: vec![],
1450 },
1451 "issue titles leave to github.com",
1452 ));
1453 let fanout = host.query_all(&query()).await;
1454 assert!(queried.load(Ordering::SeqCst));
1455 assert_eq!(fanout.accepted_frames().count(), 1);
1456 }
1457
1458 #[tokio::test]
1459 async fn a_scoped_egress_provider_is_not_queried_until_a_receipt_is_recorded() {
1460 let mut host = Host::new();
1461 let provider = FakeProvider::scoped(
1462 "cloud",
1463 vec![EgressScope::ThirdPartyModel],
1464 Behavior::Frames(vec![frame("f", 10)]),
1465 );
1466 let queried = provider.queried.clone();
1467 let info = provider.info().clone();
1468 host.register(Box::new(provider));
1469
1470 let fanout = host.query_all(&query()).await;
1473 assert_eq!(fanout.accepted_frames().count(), 0);
1474 assert!(!queried.load(Ordering::SeqCst), "payload must not be sent");
1475 match &fanout.outcomes[0].result {
1476 ProviderResult::ConsentScopeRequired { missing, .. } => {
1477 assert_eq!(missing, &vec![EgressScope::ThirdPartyModel]);
1478 }
1479 other => panic!("expected ConsentScopeRequired, got {other:?}"),
1480 }
1481 match host.query_provider("cloud", &query()).await {
1483 Err(HostError::ConsentScopeRequired { scopes, .. }) => {
1484 assert_eq!(scopes, vec![EgressScope::ThirdPartyModel]);
1485 }
1486 other => panic!("expected ConsentScopeRequired error, got {other:?}"),
1487 }
1488
1489 host.record_receipt(ConsentReceipt::new(
1491 "cloud",
1492 &info,
1493 EgressScope::ThirdPartyModel,
1494 Grantor::Human("ops@oxagen.sh".into()),
1495 "2026-07-21T00:00:00Z",
1496 ));
1497 let fanout = host.query_all(&query()).await;
1498 assert!(queried.load(Ordering::SeqCst));
1499 assert_eq!(fanout.accepted_frames().count(), 1);
1500 }
1501
1502 #[tokio::test]
1503 async fn query_provider_reports_unknown_ids() {
1504 let host = Host::new();
1505 assert!(matches!(
1506 host.query_provider("nope", &query()).await,
1507 Err(HostError::UnknownProvider(_))
1508 ));
1509 }
1510
1511 use contextgraph_types::{FrameVerdict, VerifyResponse};
1514 use std::collections::HashMap as StdHashMap;
1515
1516 struct VerifyingProvider {
1518 id: String,
1519 capabilities: Capabilities,
1520 verdicts: StdHashMap<String, Verdict>,
1523 verify_error: Option<String>,
1525 asked: Arc<std::sync::Mutex<Vec<FrameId>>>,
1527 }
1528
1529 impl VerifyingProvider {
1530 fn new(id: &str, supports_verify: bool, verdicts: &[(&str, Verdict)]) -> Self {
1531 Self {
1532 id: id.into(),
1533 capabilities: Capabilities {
1534 query: QueryCapability {
1535 kinds: vec!["doc".into()],
1536 },
1537 verify: supports_verify,
1538 ..Capabilities::default()
1539 },
1540 verdicts: verdicts
1541 .iter()
1542 .map(|(f, v)| ((*f).to_string(), v.clone()))
1543 .collect(),
1544 verify_error: None,
1545 asked: Arc::new(std::sync::Mutex::new(Vec::new())),
1546 }
1547 }
1548
1549 fn failing(id: &str) -> Self {
1550 let mut provider = Self::new(id, true, &[]);
1551 provider.verify_error = Some("index unavailable".into());
1552 provider
1553 }
1554 }
1555
1556 #[async_trait]
1557 impl ContextProvider for VerifyingProvider {
1558 fn id(&self) -> &str {
1559 &self.id
1560 }
1561 fn info(&self) -> &ProviderInfo {
1562 static_info()
1564 }
1565 fn capabilities(&self) -> &Capabilities {
1566 &self.capabilities
1567 }
1568 async fn query(&self, _query: &ContextQuery) -> Result<ContextQueryResult, HostError> {
1569 Ok(ContextQueryResult {
1570 frames: vec![],
1571 truncated: false,
1572 dropped_estimate: None,
1573 ..Default::default()
1574 })
1575 }
1576 async fn verify(&self, request: &VerifyRequest) -> Result<VerifyResponse, HostError> {
1577 self.asked
1578 .lock()
1579 .unwrap()
1580 .extend(request.frames.iter().cloned());
1581 if let Some(message) = &self.verify_error {
1582 return Err(HostError::Provider {
1583 id: self.id.clone(),
1584 code: None,
1585 message: message.clone(),
1586 });
1587 }
1588 Ok(VerifyResponse::new(
1589 request
1590 .frames
1591 .iter()
1592 .filter_map(|frame| {
1593 self.verdicts
1594 .get(&frame.frame_id)
1595 .map(|verdict| FrameVerdict::new(frame.clone(), verdict.clone()))
1596 })
1597 .collect(),
1598 ))
1599 }
1600 }
1601
1602 fn static_info() -> &'static ProviderInfo {
1603 use std::sync::OnceLock;
1604 static INFO: OnceLock<ProviderInfo> = OnceLock::new();
1605 INFO.get_or_init(|| ProviderInfo {
1606 name: "verifier".into(),
1607 version: "0.0.1".into(),
1608 data_flow: DataFlow {
1609 reads: true,
1610 writes: false,
1611 egress: false,
1612 egress_scopes: vec![],
1613 },
1614 })
1615 }
1616
1617 fn held(provider: &str, frame: &str, digest: Option<&str>) -> FrameId {
1618 FrameId::new(provider, frame, digest.map(String::from))
1619 }
1620
1621 #[tokio::test]
1622 async fn a_stale_frame_is_dropped_and_a_valid_one_is_retained() {
1623 let mut host = Host::new();
1626 host.register(Box::new(VerifyingProvider::new(
1627 "docs",
1628 true,
1629 &[
1630 ("fresh", Verdict::Valid),
1631 (
1632 "changed",
1633 Verdict::Stale {
1634 replacement_digest: Some("sha256:new".into()),
1635 },
1636 ),
1637 ],
1638 )));
1639
1640 let fresh = held("docs", "fresh", Some("sha256:a"));
1641 let changed = held("docs", "changed", Some("sha256:b"));
1642 let outcome = host.verify_frames(&[fresh.clone(), changed.clone()]).await;
1643
1644 assert_eq!(outcome.retained, vec![fresh]);
1645 assert!(outcome.was_dropped(&changed));
1646 assert_eq!(
1647 outcome.drop_reason(&changed),
1648 Some(&DropReason::Stale {
1649 replacement_digest: Some("sha256:new".into())
1650 })
1651 );
1652 assert_eq!(outcome.requery().collect::<Vec<_>>(), vec![&changed]);
1655 }
1656
1657 #[tokio::test]
1658 async fn a_gone_frame_is_dropped_and_not_worth_re_querying() {
1659 let mut host = Host::new();
1660 host.register(Box::new(VerifyingProvider::new(
1661 "docs",
1662 true,
1663 &[("deleted", Verdict::Gone)],
1664 )));
1665 let deleted = held("docs", "deleted", Some("sha256:a"));
1666 let outcome = host.verify_frames(std::slice::from_ref(&deleted)).await;
1667
1668 assert!(outcome.retained.is_empty());
1669 assert_eq!(outcome.drop_reason(&deleted), Some(&DropReason::Gone));
1670 assert_eq!(outcome.requery().count(), 0);
1672 }
1673
1674 #[tokio::test]
1675 async fn an_unknown_verdict_and_a_missing_verdict_both_drop_the_frame() {
1676 let mut host = Host::new();
1679 host.register(Box::new(VerifyingProvider::new(
1680 "docs",
1681 true,
1682 &[("shrugged", Verdict::Unknown)],
1683 )));
1684 let shrugged = held("docs", "shrugged", Some("sha256:a"));
1685 let unanswered = held("docs", "never-mentioned", Some("sha256:b"));
1686 let outcome = host
1687 .verify_frames(&[shrugged.clone(), unanswered.clone()])
1688 .await;
1689
1690 assert!(outcome.retained.is_empty());
1691 assert_eq!(outcome.drop_reason(&shrugged), Some(&DropReason::Unknown));
1692 assert_eq!(outcome.drop_reason(&unanswered), Some(&DropReason::Unknown));
1693 assert_eq!(outcome.requery().count(), 2);
1694 }
1695
1696 #[tokio::test]
1697 async fn a_provider_without_verify_support_is_never_asked_and_falls_back_to_requery() {
1698 let mut host = Host::new();
1701 let provider = VerifyingProvider::new("docs", false, &[("anything", Verdict::Valid)]);
1702 let asked = provider.asked.clone();
1703 host.register(Box::new(provider));
1704
1705 let frame = held("docs", "anything", Some("sha256:a"));
1706 let outcome = host.verify_frames(std::slice::from_ref(&frame)).await;
1707
1708 assert!(asked.lock().unwrap().is_empty(), "must not be asked");
1709 assert!(outcome.retained.is_empty());
1710 assert_eq!(
1711 outcome.drop_reason(&frame),
1712 Some(&DropReason::VerifyUnsupported)
1713 );
1714 assert_eq!(outcome.requery().count(), 1);
1715 }
1716
1717 #[tokio::test]
1718 async fn a_frame_without_a_digest_is_unverifiable_and_never_sent() {
1719 let mut host = Host::new();
1722 let provider = VerifyingProvider::new("docs", true, &[("bare", Verdict::Valid)]);
1723 let asked = provider.asked.clone();
1724 host.register(Box::new(provider));
1725
1726 let bare = held("docs", "bare", None);
1727 let digested = held("docs", "digested", Some("sha256:a"));
1728 let outcome = host.verify_frames(&[bare.clone(), digested.clone()]).await;
1729
1730 assert_eq!(outcome.drop_reason(&bare), Some(&DropReason::NoDigest));
1731 let asked = asked.lock().unwrap().clone();
1732 assert_eq!(
1733 asked,
1734 vec![digested],
1735 "only verifiable identities go on the wire"
1736 );
1737 }
1738
1739 #[tokio::test]
1740 async fn a_failed_verify_drops_that_providers_frames_without_touching_another() {
1741 let mut host = Host::new();
1743 host.register(Box::new(VerifyingProvider::failing("broken")));
1744 host.register(Box::new(VerifyingProvider::new(
1745 "healthy",
1746 true,
1747 &[("good", Verdict::Valid)],
1748 )));
1749
1750 let broken = held("broken", "any", Some("sha256:a"));
1751 let good = held("healthy", "good", Some("sha256:b"));
1752 let outcome = host.verify_frames(&[broken.clone(), good.clone()]).await;
1753
1754 assert_eq!(outcome.retained, vec![good], "one failure must not poison");
1755 assert!(matches!(
1756 outcome.drop_reason(&broken),
1757 Some(DropReason::VerifyFailed(_))
1758 ));
1759 assert!(outcome.requery().any(|frame| frame == &broken));
1760 }
1761
1762 #[tokio::test]
1763 async fn frames_from_an_unregistered_provider_are_dropped_not_ignored() {
1764 let host = Host::new();
1765 let orphan = held("never-registered", "f", Some("sha256:a"));
1766 let outcome = host.verify_frames(std::slice::from_ref(&orphan)).await;
1767 assert!(outcome.retained.is_empty());
1768 assert_eq!(
1769 outcome.drop_reason(&orphan),
1770 Some(&DropReason::UnknownProvider)
1771 );
1772 }
1773
1774 #[tokio::test]
1775 async fn held_frames_are_grouped_into_one_request_per_provider() {
1776 let mut host = Host::new();
1779 let docs = VerifyingProvider::new(
1780 "docs",
1781 true,
1782 &[("a", Verdict::Valid), ("b", Verdict::Valid)],
1783 );
1784 let asked = docs.asked.clone();
1785 host.register(Box::new(docs));
1786
1787 let outcome = host
1788 .verify_frames(&[
1789 held("docs", "a", Some("sha256:1")),
1790 held("docs", "b", Some("sha256:2")),
1791 ])
1792 .await;
1793 assert_eq!(outcome.retained.len(), 2);
1794 assert_eq!(asked.lock().unwrap().len(), 2);
1796 }
1797
1798 #[tokio::test]
1799 async fn the_partition_is_total_so_every_held_frame_is_accounted_for() {
1800 let mut host = Host::new();
1801 host.register(Box::new(VerifyingProvider::new(
1802 "docs",
1803 true,
1804 &[("keep", Verdict::Valid), ("drop", Verdict::Gone)],
1805 )));
1806 let input = vec![
1807 held("docs", "keep", Some("sha256:1")),
1808 held("docs", "drop", Some("sha256:2")),
1809 held("docs", "nodigest", None),
1810 held("elsewhere", "orphan", Some("sha256:3")),
1811 ];
1812 let outcome = host.verify_frames(&input).await;
1813 assert_eq!(
1814 outcome.retained.len() + outcome.dropped.len(),
1815 input.len(),
1816 "every held identity must land in exactly one bucket"
1817 );
1818 for frame in &input {
1819 assert!(
1820 outcome.retained.contains(frame) || outcome.was_dropped(frame),
1821 "{frame:?} was silently lost"
1822 );
1823 }
1824 }
1825
1826 #[tokio::test]
1827 async fn verifying_an_empty_held_set_is_a_no_op() {
1828 let host = Host::new();
1829 let outcome = host.verify_frames(&[]).await;
1830 assert_eq!(outcome, VerifyOutcome::default());
1831 }
1832}