use std::collections::BTreeSet;
use async_trait::async_trait;
use contextgraph_types::{
BYTES_PER_BUDGET_TOKEN, ContextFrame, ContextQuery, FrameId, FrameKind, budget_tokens,
};
use contextgraph_host::{ContextProvider, Host, ProviderResult, compose_for_prompt};
use crate::host_conformance::{ProbeProvider, probe_query};
use crate::report::{CheckResult, ConformanceReport};
pub const CCHECK_BUDGET_BOUND: &str = "composition-budget-bound";
pub const CCHECK_TOTAL_PARTITION: &str = "composition-total-partition";
pub const CCHECK_QUARANTINE: &str = "composition-quarantine";
pub const CCHECK_DETERMINISM: &str = "composition-determinism";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ExcludedFrame {
pub provider_id: String,
pub frame_id: String,
}
#[derive(Debug, Clone, Default)]
pub struct Composition {
pub admitted: Vec<(String, ContextFrame)>,
pub dropped: Vec<ExcludedFrame>,
}
impl Composition {
fn admitted_tokens(&self) -> u64 {
self.admitted
.iter()
.map(|(_, frame)| u64::from(frame.token_cost))
.sum()
}
fn accounted(&self) -> BTreeSet<(String, String)> {
self.admitted
.iter()
.map(|(provider, frame)| (provider.clone(), frame.id.clone()))
.chain(
self.dropped
.iter()
.map(|drop| (drop.provider_id.clone(), drop.frame_id.clone())),
)
.collect()
}
fn render_order(&self) -> Vec<(String, String)> {
self.admitted
.iter()
.map(|(provider, frame)| (provider.clone(), frame.id.clone()))
.collect()
}
}
#[async_trait]
pub trait ComposingHost: Send + Sync {
async fn compose(
&self,
providers: Vec<Box<dyn ContextProvider>>,
query: &ContextQuery,
) -> Composition;
}
pub async fn run_composition_conformance(
host: &dyn ComposingHost,
target: impl Into<String>,
) -> ConformanceReport {
let checks = vec![
check_budget_bound(host).await,
check_total_partition(host).await,
check_quarantine(host).await,
check_determinism(host).await,
];
ConformanceReport {
target: target.into(),
checks,
}
}
async fn check_budget_bound(host: &dyn ComposingHost) -> CheckResult {
let query = probe_query(); let over = host
.compose(three_providers_each_costing(400), &query)
.await;
let within_budget = over.admitted_tokens() <= u64::from(query.max_tokens);
let dropped_something = !over.dropped.is_empty();
let under = host
.compose(three_providers_each_costing(200), &query)
.await;
let kept_all = under.admitted.len() == 3 && under.dropped.is_empty();
CheckResult::from_bool(
CCHECK_BUDGET_BOUND,
within_budget && dropped_something && kept_all,
format!(
"§7 (composition): three individually-honest providers summing 1200 against a \
1000-token budget compose within budget={within_budget} \
(admitted {} tokens) and report a drop={dropped_something}; the same shape summing \
600 keeps all three frames={kept_all}",
over.admitted_tokens()
),
)
}
async fn check_total_partition(host: &dyn ComposingHost) -> CheckResult {
let query = probe_query();
let offered: BTreeSet<(String, String)> = (0..3)
.map(|i| (format!("p{i}"), format!("p{i}-f")))
.collect();
let over = host
.compose(three_providers_each_costing(400), &query)
.await;
let accounted = over.accounted();
let missing: Vec<_> = offered.difference(&accounted).collect();
let total = missing.is_empty();
let no_phantoms = accounted.difference(&offered).count() == 0;
let under = host
.compose(three_providers_each_costing(200), &query)
.await;
let no_spurious_drops = under.dropped.is_empty();
CheckResult::from_bool(
CCHECK_TOTAL_PARTITION,
total && no_phantoms && no_spurious_drops,
format!(
"issue #15 (composition): every offered frame is admitted or reported \
dropped={total} (unaccounted: {missing:?}), no frame is reported that was never \
offered={no_phantoms}; a composition that drops nothing reports \
nothing={no_spurious_drops}"
),
)
}
async fn check_quarantine(host: &dyn ComposingHost) -> CheckResult {
let query = probe_query(); let flood: Vec<ContextFrame> = (0..query.max_frames + 9)
.map(|i| honest_frame(&format!("flood-{i}"), 1))
.collect();
let providers: Vec<Box<dyn ContextProvider>> = vec![
Box::new(ProbeProvider::local("flooder", flood)),
Box::new(ProbeProvider::local(
"honest",
vec![honest_frame("honest-f", 100)],
)),
];
let composed = host.compose(providers, &query).await;
let flooder_excluded = !composed
.admitted
.iter()
.any(|(provider, _)| provider == "flooder");
let honest_admitted = composed
.admitted
.iter()
.any(|(provider, frame)| provider == "honest" && frame.id == "honest-f");
CheckResult::from_bool(
CCHECK_QUARANTINE,
flooder_excluded && honest_admitted,
format!(
"§7 B2/B4 (composition): the frames of a provider the audit rejected (a frame \
flooder, whose frames are individually cheap enough to pass any token pack) never \
reach the prompt={flooder_excluded}; an honest provider queried alongside it still \
arrives={honest_admitted}"
),
)
}
async fn check_determinism(host: &dyn ComposingHost) -> CheckResult {
let query = probe_query();
let first = host
.compose(three_providers_each_costing(100), &query)
.await;
let second = host
.compose(three_providers_each_costing(100), &query)
.await;
let stable = first.render_order() == second.render_order();
let non_empty = !first.admitted.is_empty();
CheckResult::from_bool(
CCHECK_DETERMINISM,
stable && non_empty,
format!(
"context-reuse §1 (composition): an unchanged frame set composes to the same render \
order twice={stable} (first {:?}, second {:?}); the composition is non-empty, so \
stability is not vacuous={non_empty}",
first.render_order(),
second.render_order()
),
)
}
pub struct ReferenceComposingHost;
#[async_trait]
impl ComposingHost for ReferenceComposingHost {
async fn compose(
&self,
providers: Vec<Box<dyn ContextProvider>>,
query: &ContextQuery,
) -> Composition {
let mut host = Host::new();
for provider in providers {
host.register(provider);
}
let fanout = host.query_all(query).await;
let offered: Vec<(String, ContextFrame)> = fanout
.outcomes
.iter()
.filter_map(|outcome| match &outcome.result {
ProviderResult::Frames(result) => Some(
result
.frames
.iter()
.map(|frame| (outcome.provider_id.clone(), frame.clone())),
),
_ => None,
})
.flatten()
.collect();
let composed = compose_for_prompt(
offered
.iter()
.map(|(provider, frame)| (provider.as_str(), frame)),
query.max_tokens,
);
let included: Vec<&FrameId> = composed.audit.included().collect();
let admitted = included
.iter()
.filter_map(|id| {
offered
.iter()
.find(|(provider, frame)| {
provider == &id.provider_id && frame.id == id.frame_id
})
.cloned()
})
.collect();
let dropped = composed
.audit
.excluded()
.map(|entry| ExcludedFrame {
provider_id: entry.frame.provider_id.clone(),
frame_id: entry.frame.frame_id.clone(),
})
.collect();
Composition { admitted, dropped }
}
}
fn honest_frame(id: &str, token_cost: u32) -> ContextFrame {
let content = "x".repeat(token_cost as usize * BYTES_PER_BUDGET_TOKEN);
debug_assert_eq!(
budget_tokens(&content),
token_cost,
"the fixture must satisfy B3 or the suite measures the wrong thing"
);
let mut frame = ContextFrame::full(id, FrameKind::Doc, id, &content, 0.5, token_cost);
frame.citation_label = Some(id.into());
frame
}
fn three_providers_each_costing(token_cost: u32) -> Vec<Box<dyn ContextProvider>> {
(0..3)
.map(|i| {
let id = format!("p{i}");
let frame = honest_frame(&format!("{id}-f"), token_cost);
Box::new(ProbeProvider::local(&id, vec![frame])) as Box<dyn ContextProvider>
})
.collect()
}
#[cfg(test)]
mod tests;