use std::collections::BTreeSet;
use std::fmt::Write;
use std::time::Instant;
use serde::{Deserialize, Serialize};
use crate::accounting::{
ByteMeasurement, MeasurementProvenance, ProviderUsage, ProviderUsageComponents, ProviderUsageRule, TokenMeasurement,
};
use crate::accounting::{TOKEN_ESTIMATOR_VERSION, estimate_serialized_tokens};
pub const REPLAY_FIXTURE_SCHEMA_VERSION: &str = "context-replay-v1";
pub const REPLAY_REPORT_SCHEMA_VERSION: &str = "context-replay-report-v1";
pub trait ReplayPolicy {
fn name(&self) -> &str;
fn include(&self, item: &ReplayItem) -> bool;
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct BaselinePolicy;
impl ReplayPolicy for BaselinePolicy {
fn name(&self) -> &str {
"baseline"
}
fn include(&self, _item: &ReplayItem) -> bool {
true
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CandidatePolicy {
name: String,
omitted_ids: BTreeSet<String>,
}
impl CandidatePolicy {
pub fn new(name: impl Into<String>) -> Self {
Self { name: name.into(), omitted_ids: BTreeSet::new() }
}
pub fn omit(mut self, id: impl Into<String>) -> Self {
self.omitted_ids.insert(id.into());
self
}
pub fn omit_all<I, S>(mut self, ids: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.omitted_ids.extend(ids.into_iter().map(Into::into));
self
}
pub fn omitted_ids(&self) -> &BTreeSet<String> {
&self.omitted_ids
}
}
impl ReplayPolicy for CandidatePolicy {
fn name(&self) -> &str {
&self.name
}
fn include(&self, item: &ReplayItem) -> bool {
!self.omitted_ids.contains(&item.id)
}
}
#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
pub enum ReplayError {
#[error("invalid replay fixture: {0}")]
InvalidFixture(String),
#[error("replay invariant failed for {policy}: {message}")]
InvariantViolation {
policy: String,
message: String,
},
#[error("replay serialization failed: {0}")]
Serialization(String),
}
#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ReplayItemKind {
UserTurn,
Assistant,
ToolOutput,
PassingTest,
FailingTest,
Progress,
Error,
Command,
Write,
ProtectedEvidence,
}
impl ReplayItemKind {
pub const fn label(self) -> &'static str {
match self {
Self::UserTurn => "user_turn",
Self::Assistant => "assistant",
Self::ToolOutput => "tool_output",
Self::PassingTest => "passing_test",
Self::FailingTest => "failing_test",
Self::Progress => "progress",
Self::Error => "error",
Self::Command => "command",
Self::Write => "write",
Self::ProtectedEvidence => "protected_evidence",
}
}
}
#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ReplayScenario {
RepeatedReads,
OverlappingReads,
PassingTests,
FailingTests,
NoisyProgress,
MiddlePositionError,
StateChangingCommands,
FailedLargeWrite,
ProtectedEvidence,
CacheComponents,
Recovery,
}
impl ReplayScenario {
pub const fn label(self) -> &'static str {
match self {
Self::RepeatedReads => "repeated_reads",
Self::OverlappingReads => "overlapping_reads",
Self::PassingTests => "passing_tests",
Self::FailingTests => "failing_tests",
Self::NoisyProgress => "noisy_progress",
Self::MiddlePositionError => "middle_position_error",
Self::StateChangingCommands => "state_changing_commands",
Self::FailedLargeWrite => "failed_large_write",
Self::ProtectedEvidence => "protected_evidence",
Self::CacheComponents => "cache_components",
Self::Recovery => "recovery",
}
}
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum ReplayTiming {
#[default]
Deterministic,
Measured,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct ReplayItem {
pub id: String,
pub kind: ReplayItemKind,
pub label: String,
pub content: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub original_bytes: Option<u64>,
#[serde(default)]
pub fact_ids: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub recovery_handle: Option<String>,
#[serde(default)]
pub protected: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub state_fingerprint: Option<String>,
}
impl ReplayItem {
pub fn content_bytes(&self) -> u64 {
self.content.len() as u64
}
pub fn source_bytes(&self) -> u64 {
self.original_bytes.unwrap_or_else(|| self.content_bytes())
}
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct RequiredFact {
pub id: String,
pub description: String,
#[serde(default)]
pub protected: bool,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct RecoveryCase {
pub id: String,
pub artifact_handle: String,
pub expected_available: bool,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct RecordedProviderUsage {
pub provider: String,
pub rule: ProviderUsageRule,
pub components: ProviderUsageComponents,
}
impl RecordedProviderUsage {
pub fn normalize(&self) -> ProviderUsage {
self.components.normalize(&self.provider, self.rule)
}
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct ReplayFixture {
pub schema_version: String,
pub id: String,
pub description: String,
pub scenarios: Vec<ReplayScenario>,
pub items: Vec<ReplayItem>,
pub required_facts: Vec<RequiredFact>,
#[serde(default)]
pub recovery: Vec<RecoveryCase>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_usage: Option<RecordedProviderUsage>,
}
impl ReplayFixture {
pub fn validate(&self) -> Result<(), ReplayError> {
if self.schema_version != REPLAY_FIXTURE_SCHEMA_VERSION {
return Err(ReplayError::InvalidFixture(format!(
"schema_version must be {REPLAY_FIXTURE_SCHEMA_VERSION}"
)));
}
if self.id.trim().is_empty() {
return Err(ReplayError::InvalidFixture("id must not be empty".to_string()));
}
let mut item_ids = BTreeSet::new();
for item in &self.items {
if item.id.trim().is_empty() || !item_ids.insert(&item.id) {
return Err(ReplayError::InvalidFixture(format!(
"item id is empty or duplicated: {}",
item.id
)));
}
}
let fact_ids: BTreeSet<&str> = self.required_facts.iter().map(|fact| fact.id.as_str()).collect();
if fact_ids.len() != self.required_facts.len() {
return Err(ReplayError::InvalidFixture(
"required fact ids must be unique".to_string(),
));
}
for item in &self.items {
for fact_id in &item.fact_ids {
if !fact_ids.contains(fact_id.as_str()) {
return Err(ReplayError::InvalidFixture(format!(
"item {} references unknown fact {fact_id}",
item.id
)));
}
}
}
let mut recovery_ids = BTreeSet::new();
let mut handles = BTreeSet::new();
for case in &self.recovery {
if case.id.trim().is_empty() || !recovery_ids.insert(&case.id) {
return Err(ReplayError::InvalidFixture(format!(
"recovery id is empty or duplicated: {}",
case.id
)));
}
if case.artifact_handle.trim().is_empty() || !handles.insert(&case.artifact_handle) {
return Err(ReplayError::InvalidFixture(format!(
"recovery handle is empty or duplicated: {}",
case.artifact_handle
)));
}
}
Ok(())
}
pub fn to_json(&self) -> Result<String, ReplayError> {
serde_json::to_string_pretty(self).map_err(|error| ReplayError::Serialization(error.to_string()))
}
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct ReplayReceipt {
pub item_id: String,
pub method: String,
pub version: String,
pub before_bytes: u64,
pub after_bytes: u64,
pub lossy: bool,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct RequiredFactResult {
pub id: String,
pub preserved: bool,
pub item_ids: Vec<String>,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct RecoveryOutcome {
pub id: String,
pub artifact_handle: String,
pub expected_available: bool,
pub available: bool,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ReplayProjection {
pub policy: String,
pub item_ids: Vec<String>,
pub rendered: String,
pub exact_bytes: ByteMeasurement,
pub estimated_tokens: TokenMeasurement,
pub receipts: Vec<ReplayReceipt>,
pub recovery_handles: Vec<String>,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct ProjectionReport {
pub policy: String,
pub exact_bytes: ByteMeasurement,
pub estimated_tokens: TokenMeasurement,
pub projection_digest: String,
pub item_count: usize,
pub receipts: Vec<ReplayReceipt>,
pub required_facts: Vec<RequiredFactResult>,
pub recovery: Vec<RecoveryOutcome>,
pub elapsed_micros: u64,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct ReplayComparison {
pub exact_bytes_delta: i64,
pub estimated_tokens_delta: i64,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct ReplayReport {
pub schema_version: String,
pub fixture_id: String,
pub baseline: ProjectionReport,
pub candidate: ProjectionReport,
pub comparison: ReplayComparison,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_usage: Option<ProviderUsage>,
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct ReplayEvaluator {
timing: ReplayTiming,
}
impl ReplayEvaluator {
pub const fn new() -> Self {
Self { timing: ReplayTiming::Deterministic }
}
pub const fn measured() -> Self {
Self { timing: ReplayTiming::Measured }
}
pub fn evaluate<B: ReplayPolicy, C: ReplayPolicy>(
&self, fixture: &ReplayFixture, baseline: &B, candidate: &C,
) -> Result<ReplayReport, ReplayError> {
fixture.validate()?;
let baseline = self.evaluate_policy(fixture, baseline)?;
let candidate = self.evaluate_policy(fixture, candidate)?;
let baseline_tokens = baseline.estimated_tokens.value.unwrap_or_default();
let candidate_tokens = candidate.estimated_tokens.value.unwrap_or_default();
Ok(ReplayReport {
schema_version: REPLAY_REPORT_SCHEMA_VERSION.to_string(),
fixture_id: fixture.id.clone(),
comparison: ReplayComparison {
exact_bytes_delta: signed_delta(candidate.exact_bytes.value, baseline.exact_bytes.value),
estimated_tokens_delta: signed_delta(candidate_tokens, baseline_tokens),
},
provider_usage: fixture.provider_usage.as_ref().map(RecordedProviderUsage::normalize),
baseline,
candidate,
})
}
fn evaluate_policy<P: ReplayPolicy>(
&self, fixture: &ReplayFixture, policy: &P,
) -> Result<ProjectionReport, ReplayError> {
let started = (self.timing == ReplayTiming::Measured).then(Instant::now);
let projection = project_fixture(fixture, policy);
let required_facts = required_fact_results(fixture, &projection.item_ids);
let recovery = recovery_outcomes(fixture, &projection.recovery_handles);
if let Some(fact) = required_facts.iter().find(|fact| !fact.preserved) {
return Err(ReplayError::InvariantViolation {
policy: policy.name().to_string(),
message: format!("required fact {} was not preserved", fact.id),
});
}
if let Some(outcome) = recovery
.iter()
.find(|outcome| outcome.available != outcome.expected_available)
{
return Err(ReplayError::InvariantViolation {
policy: policy.name().to_string(),
message: format!("recovery case {} availability was unexpected", outcome.id),
});
}
Ok(ProjectionReport {
policy: projection.policy,
exact_bytes: projection.exact_bytes,
estimated_tokens: projection.estimated_tokens,
projection_digest: digest(&projection.rendered),
item_count: projection.item_ids.len(),
receipts: projection.receipts,
required_facts,
recovery,
elapsed_micros: started.map_or(0, |started| started.elapsed().as_micros() as u64),
})
}
}
impl ReplayReport {
pub fn to_json(&self) -> Result<String, ReplayError> {
serde_json::to_string_pretty(self).map_err(|error| ReplayError::Serialization(error.to_string()))
}
pub fn to_markdown(&self) -> String {
let mut out = String::new();
let _ = writeln!(&mut out, "# Context replay: {}", self.fixture_id);
let _ = writeln!(&mut out, "\nSchema: `{}`", self.schema_version);
render_projection_markdown(&mut out, &self.baseline);
render_projection_markdown(&mut out, &self.candidate);
let _ = writeln!(
&mut out,
"\n## Comparison\n\n| Measure | Delta |\n| --- | ---: |\n| Exact bytes | {} |\n| Estimated tokens | {} |",
self.comparison.exact_bytes_delta, self.comparison.estimated_tokens_delta
);
render_fact_markdown(&mut out, &self.baseline, &self.candidate);
render_recovery_markdown(&mut out, &self.baseline, &self.candidate);
render_receipt_markdown(&mut out, &self.baseline);
render_receipt_markdown(&mut out, &self.candidate);
if let Some(provider_usage) = &self.provider_usage {
let _ = writeln!(
&mut out,
"\n## Provider usage\n\n- Provider: `{}`\n- Rule: `{}`\n- Inclusive input tokens: `{}`",
provider_usage.provider,
provider_usage.rule.label(),
display_token(provider_usage.inclusive_input_tokens.value)
);
}
out
}
}
pub fn load_fixture(json: &str) -> Result<ReplayFixture, ReplayError> {
let fixture: ReplayFixture =
serde_json::from_str(json).map_err(|error| ReplayError::InvalidFixture(error.to_string()))?;
fixture.validate()?;
Ok(fixture)
}
pub fn select_items<'a, P: ReplayPolicy>(fixture: &'a ReplayFixture, policy: &P) -> Vec<&'a ReplayItem> {
fixture.items.iter().filter(|item| policy.include(item)).collect()
}
pub fn project_fixture<P: ReplayPolicy>(fixture: &ReplayFixture, policy: &P) -> ReplayProjection {
let selected = select_items(fixture, policy);
let selected_ids = selected.iter().map(|item| item.id.clone()).collect::<Vec<_>>();
let mut rendered = String::new();
let mut recovery_handles = Vec::new();
let mut selected_contributions = BTreeSet::new();
for item in &selected {
let contribution = render_item(item);
rendered.push_str(&contribution);
selected_contributions.insert(item.id.as_str());
if let Some(handle) = &item.recovery_handle {
recovery_handles.push(handle.clone());
}
}
let receipts = fixture
.items
.iter()
.map(|item| ReplayReceipt {
item_id: item.id.clone(),
method: if selected_contributions.contains(item.id.as_str()) {
policy.name().to_string()
} else {
"policy_omitted".to_string()
},
version: REPLAY_REPORT_SCHEMA_VERSION.to_string(),
before_bytes: item.source_bytes(),
after_bytes: if selected_contributions.contains(item.id.as_str()) {
render_item(item).len() as u64
} else {
0
},
lossy: !selected_contributions.contains(item.id.as_str()),
})
.collect();
let exact_bytes = rendered.len() as u64;
ReplayProjection {
policy: policy.name().to_string(),
item_ids: selected_ids,
rendered,
exact_bytes: ByteMeasurement {
value: exact_bytes,
provenance: MeasurementProvenance::ExactSerialized { boundary: "replay_projection".to_string() },
},
estimated_tokens: TokenMeasurement {
value: Some(estimate_serialized_tokens(exact_bytes)),
provenance: MeasurementProvenance::Estimated {
estimator: "utf8_bytes_divisor_3_plus_item_overhead".to_string(),
version: TOKEN_ESTIMATOR_VERSION.to_string(),
},
},
receipts,
recovery_handles,
}
}
pub fn evaluate_fixture<C: ReplayPolicy>(fixture: &ReplayFixture, candidate: &C) -> Result<ReplayReport, ReplayError> {
ReplayEvaluator::new().evaluate(fixture, &BaselinePolicy, candidate)
}
fn required_fact_results(fixture: &ReplayFixture, selected_ids: &[String]) -> Vec<RequiredFactResult> {
fixture
.required_facts
.iter()
.map(|fact| {
let item_ids = fixture
.items
.iter()
.filter(|item| selected_ids.contains(&item.id) && item.fact_ids.iter().any(|id| id == &fact.id))
.map(|item| item.id.clone())
.collect::<Vec<_>>();
RequiredFactResult { id: fact.id.clone(), preserved: !item_ids.is_empty(), item_ids }
})
.collect()
}
fn recovery_outcomes(fixture: &ReplayFixture, selected_handles: &[String]) -> Vec<RecoveryOutcome> {
fixture
.recovery
.iter()
.map(|case| RecoveryOutcome {
id: case.id.clone(),
artifact_handle: case.artifact_handle.clone(),
expected_available: case.expected_available,
available: selected_handles.iter().any(|handle| handle == &case.artifact_handle),
})
.collect()
}
fn render_item(item: &ReplayItem) -> String {
format!(
"<replay_item id=\"{}\" kind=\"{}\" label=\"{}\">\n{}\n</replay_item>\n",
escape_xml(&item.id),
item.kind.label(),
escape_xml(&item.label),
item.content
)
}
fn escape_xml(value: &str) -> String {
value
.replace('&', "&")
.replace('<', "<")
.replace('>', ">")
.replace('"', """)
}
fn signed_delta(current: u64, baseline: u64) -> i64 {
(current as i128)
.saturating_sub(baseline as i128)
.clamp(i64::MIN as i128, i64::MAX as i128) as i64
}
fn digest(value: &str) -> String {
let mut hash = 0xcbf29ce484222325_u64;
for byte in value.as_bytes() {
hash ^= u64::from(*byte);
hash = hash.wrapping_mul(0x100000001b3);
}
format!("fnv1a64:{hash:016x}")
}
fn display_token(value: Option<u64>) -> String {
value.map_or_else(|| "unknown".to_string(), |value| value.to_string())
}
fn render_projection_markdown(out: &mut String, projection: &ProjectionReport) {
let _ = writeln!(
out,
"\n## {}\n\n| Measure | Value |\n| --- | ---: |\n| Exact bytes | {} |\n| Estimated tokens | {} |\n| Items | {} |\n| Projection digest | `{}` |\n| Elapsed micros | {} |",
projection.policy,
projection.exact_bytes.value,
display_token(projection.estimated_tokens.value),
projection.item_count,
projection.projection_digest,
projection.elapsed_micros
);
}
fn render_fact_markdown(out: &mut String, baseline: &ProjectionReport, candidate: &ProjectionReport) {
let _ = writeln!(
out,
"\n## Required facts\n\n| Fact | Baseline | Candidate |\n| --- | --- | --- |\n{}
",
baseline
.required_facts
.iter()
.zip(&candidate.required_facts)
.map(|(baseline, candidate)| format!(
"| `{}` | {} | {} |",
baseline.id, baseline.preserved, candidate.preserved
))
.collect::<Vec<_>>()
.join("\n")
);
}
fn render_recovery_markdown(out: &mut String, baseline: &ProjectionReport, candidate: &ProjectionReport) {
let _ = writeln!(
out,
"\n## Recovery\n\n| Case | Expected | Baseline | Candidate |\n| --- | --- | --- | --- |\n{}
",
baseline
.recovery
.iter()
.zip(&candidate.recovery)
.map(|(baseline, candidate)| format!(
"| `{}` | {} | {} | {} |",
baseline.id, baseline.expected_available, baseline.available, candidate.available
))
.collect::<Vec<_>>()
.join("\n")
);
}
fn render_receipt_markdown(out: &mut String, projection: &ProjectionReport) {
let _ = writeln!(
out,
"\n## {} receipts\n\n| Item | Method | Before bytes | After bytes | Lossy |\n| --- | --- | ---: | ---: | --- |\n{}\n",
projection.policy,
projection
.receipts
.iter()
.map(|receipt| format!(
"| `{}` | `{}` | {} | {} | {} |",
receipt.item_id, receipt.method, receipt.before_bytes, receipt.after_bytes, receipt.lossy
))
.collect::<Vec<_>>()
.join("\n")
);
}
#[cfg(test)]
mod tests {
use super::*;
const FIXTURE: &str = include_str!("../fixtures/context-replay/context.json");
#[test]
fn fixture_round_trip_and_report_are_deterministic() {
let fixture = load_fixture(FIXTURE).expect("fixture loads");
let candidate = CandidatePolicy::new("candidate");
let first = evaluate_fixture(&fixture, &candidate).expect("evaluation succeeds");
let second = evaluate_fixture(&fixture, &candidate).expect("evaluation succeeds");
assert_eq!(first, second);
assert_eq!(first.to_json().expect("json"), second.to_json().expect("json"));
assert_eq!(first.to_markdown(), second.to_markdown());
assert!(first.baseline.exact_bytes.value > 0);
assert!(!first.baseline.receipts.is_empty());
assert!(first.provider_usage.is_some());
}
#[test]
fn required_fact_omission_fails_independently_of_timing() {
let fixture = load_fixture(FIXTURE).expect("fixture loads");
let error = evaluate_fixture(&fixture, &CandidatePolicy::new("candidate").omit("error-middle"))
.expect_err("required fact omission must fail");
assert!(matches!(error, ReplayError::InvariantViolation { .. }));
}
#[test]
fn recovery_is_checked_by_opaque_handle() {
let mut fixture = load_fixture(FIXTURE).expect("fixture loads");
fixture.required_facts.retain(|fact| fact.id != "protected-write");
fixture
.items
.iter_mut()
.find(|item| item.id == "protected-output")
.expect("protected item")
.fact_ids
.clear();
let error = evaluate_fixture(&fixture, &CandidatePolicy::new("candidate").omit("protected-output"))
.expect_err("recovery omission must fail");
assert!(error.to_string().contains("recovery"));
}
#[test]
fn unknown_provider_usage_is_omitted_from_json() {
let mut fixture = load_fixture(FIXTURE).expect("fixture loads");
fixture.provider_usage = None;
let report = evaluate_fixture(&fixture, &CandidatePolicy::new("candidate")).expect("evaluation succeeds");
assert!(report.provider_usage.is_none());
assert!(!report.to_json().expect("json").contains("provider_usage"));
}
#[test]
fn malformed_fixture_is_rejected() {
let mut fixture = load_fixture(FIXTURE).expect("fixture loads");
fixture.items[0].fact_ids.push("missing".to_string());
assert!(matches!(fixture.validate(), Err(ReplayError::InvalidFixture(_))));
}
}