pub mod claude;
pub mod codex;
pub mod gemini;
pub mod kodelet;
pub mod opencode;
use crate::model::{Activity, Attribution, ContextOrigin, ContextSource, CostBreakdown, Harness, ProcNode, SpanKind, TokenUsage, ToolSpan};
use crate::process::RawProc;
use std::collections::{BTreeMap, HashSet, VecDeque};
use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime};
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct ParseHealth {
pub billable_messages: u64,
pub usage_records: u64,
pub empty_usage_records: u64,
}
impl ParseHealth {
const MIN_EVIDENCE: u64 = 3;
pub fn fields_unrecognised(&self) -> bool {
self.billable_messages >= Self::MIN_EVIDENCE && self.usage_records == self.empty_usage_records
}
}
#[derive(Debug, Clone, Default)]
pub struct SessionSummary {
pub harness: Option<Harness>,
pub session_id: Option<String>,
pub subagent: Option<crate::model::SubagentInfo>,
pub cwd: Option<PathBuf>,
pub model: Option<String>,
pub harness_version: Option<String>,
pub usage: TokenUsage,
pub cost_usd: f64,
pub cost_breakdown: CostBreakdown,
pub price_source: Option<crate::model::PriceSource>,
pub unpriced_tokens: u64,
pub turns: u64,
pub subagent_turns: u64,
pub folds_child_usage: bool,
pub tool_calls: u64,
pub tool_calls_lower_bound: bool,
pub web_searches: u64,
pub spans: SpanLog,
pub mcp: BTreeMap<String, McpUsage>,
pub context: ContextLedger,
pub health: ParseHealth,
pub activity: Activity,
pub started_at: Option<SystemTime>,
pub last_activity: Option<SystemTime>,
pub rate_limit: Option<crate::model::RateLimit>,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct McpUsage {
pub calls: u64,
pub errors: u64,
pub last_call: Option<SystemTime>,
}
impl McpUsage {
pub fn add(&mut self, o: &McpUsage) {
self.calls += o.calls;
self.errors += o.errors;
self.last_call = self.last_call.max(o.last_call);
}
}
#[derive(Debug, Clone, Default)]
pub struct ContextLedger {
shares: BTreeMap<ContextKey, ContextShare>,
live: BTreeMap<ContextKey, u64>,
pending: Vec<(String, ContextWeights)>,
prev: Option<(u64, u64)>,
}
type ContextKey = (ContextOrigin, String);
pub(super) type ContextWeights = Vec<(ContextOrigin, String, u64)>;
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct ContextShare {
pub calls: u64,
pub tokens: u64,
pub cost_usd: f64,
}
impl ContextLedger {
pub const OTHER: &str = "other";
fn other() -> ContextKey {
(ContextOrigin::Other, Self::OTHER.to_string())
}
pub fn result(&mut self, id: &str, origin: ContextOrigin, name: &str) {
self.result_weighted(id, vec![(origin, name.to_string(), 1)]);
}
pub(super) fn result_weighted(&mut self, id: &str, sources: ContextWeights) {
if !sources.is_empty() {
self.pending.push((id.to_string(), sources));
}
}
pub fn retag(&mut self, id: &str, origin: ContextOrigin, name: &str) {
if let Some((_, sources)) = self.pending.iter_mut().find(|(i, _)| i == id) {
for (source_origin, source_name, _) in sources {
*source_origin = origin;
*source_name = name.to_string();
}
}
}
pub fn response(&mut self, usage: &TokenUsage, cost: &CostBreakdown) {
let prompt = usage.prompt();
if prompt == 0 {
return;
}
let pending = std::mem::take(&mut self.pending);
match self.prev {
Some((p, _)) if prompt < p / 2 => {
self.live.clear();
self.file(prompt, 0, pending);
}
Some((p, o)) if prompt >= p => {
let growth = prompt - p;
let reply = o.min(growth);
self.file(growth - reply, reply, pending);
}
Some(_) => {
let e = self.live.entry(Self::other()).or_default();
*e = e.saturating_sub(self.prev.map(|(p, _)| p - prompt).unwrap_or(0));
self.file(0, 0, pending);
}
None => self.file(prompt, 0, pending),
}
let prompt_cost = cost.input + cost.cache_read + cost.cache_write_5m + cost.cache_write_1h;
if prompt_cost > 0.0 {
let rate = prompt_cost / prompt as f64;
for (k, t) in &self.live {
self.shares.entry(k.clone()).or_default().cost_usd += *t as f64 * rate;
}
}
self.prev = Some((prompt, usage.output));
}
fn file(&mut self, results: u64, other: u64, pending: Vec<(String, ContextWeights)>) {
if pending.is_empty() {
self.add(Self::other(), results + other, 0);
return;
}
self.add(Self::other(), other, 0);
let n = pending.len() as u64;
let (each, mut rem) = (results / n, results % n);
for (_, sources) in pending {
let t = each + u64::from(rem > 0);
rem = rem.saturating_sub(1);
let total: u128 = sources.iter().map(|(_, _, weight)| u128::from(*weight)).sum();
let denominator = if total == 0 { sources.len() as u128 } else { total };
let mut remaining = t;
let mut parts: Vec<_> = sources
.into_iter()
.map(|(origin, name, weight)| {
let weight = if total == 0 { 1 } else { u128::from(weight) };
let numerator = u128::from(t) * weight;
let tokens = (numerator / denominator) as u64;
remaining -= tokens;
((origin, name), tokens, numerator % denominator)
})
.collect();
parts.sort_by_key(|part| std::cmp::Reverse(part.2));
for (key, tokens, _) in parts {
let tokens = tokens + u64::from(remaining > 0);
remaining = remaining.saturating_sub(1);
self.add(key, tokens, 1);
}
}
}
fn add(&mut self, key: ContextKey, tokens: u64, calls: u64) {
if tokens == 0 && calls == 0 {
return;
}
let sh = self.shares.entry(key.clone()).or_default();
sh.tokens += tokens;
sh.calls += calls;
*self.live.entry(key).or_default() += tokens;
}
pub fn compacted(&mut self) {
self.live.clear();
self.prev = None;
}
pub fn merge(&mut self, other: &ContextLedger) {
for (k, sh) in &other.shares {
let e = self.shares.entry(k.clone()).or_default();
e.calls += sh.calls;
e.tokens += sh.tokens;
e.cost_usd += sh.cost_usd;
}
}
pub fn is_empty(&self) -> bool {
self.shares.is_empty()
}
pub fn sources(&self) -> Vec<ContextSource> {
let mut v: Vec<ContextSource> = self
.shares
.iter()
.map(|((origin, name), sh)| ContextSource {
name: name.clone(),
origin: *origin,
calls: sh.calls,
tokens: sh.tokens,
cost_usd: sh.cost_usd,
})
.collect();
v.sort_by(|a, b| b.tokens.cmp(&a.tokens).then_with(|| a.name.cmp(&b.name)));
v
}
}
pub fn mcp_server_of(tool_name: &str) -> Option<&str> {
let rest = tool_name.strip_prefix("mcp__")?;
let server = match rest.rfind("__") {
Some(i) => &rest[..i],
None => rest,
};
if server.is_empty() { None } else { Some(server) }
}
pub const MAX_SPANS: usize = 256;
#[derive(Debug, Clone)]
pub struct SpanLog {
spans: VecDeque<ToolSpan>,
cap: usize,
}
impl Default for SpanLog {
fn default() -> Self {
SpanLog { spans: VecDeque::new(), cap: MAX_SPANS }
}
}
impl SpanLog {
pub fn unbounded() -> Self {
SpanLog { spans: VecDeque::new(), cap: usize::MAX }
}
pub fn open(&mut self, id: String, name: String, at: SystemTime, sidechain: bool) {
self.open_kind(id, name, at, sidechain, SpanKind::Tool);
}
pub fn open_kind(&mut self, id: String, name: String, at: SystemTime, sidechain: bool, kind: SpanKind) {
if id.is_empty() || self.spans.iter().any(|s| s.is_open() && s.id == id) {
return;
}
if self.spans.len() >= self.cap {
self.spans.pop_front();
}
self.spans.push_back(ToolSpan { id, name, started_at: at, duration_ms: None, sidechain, error: false, kind });
}
pub fn end_at(&mut self, id: &str, at: SystemTime) {
let Some(s) = self.spans.iter_mut().rev().find(|s| s.id == id) else { return };
s.duration_ms = Some(at.duration_since(s.started_at).map(|d| d.as_millis() as u64).unwrap_or(0));
}
pub fn open_of_kind(&self, kind: SpanKind) -> Option<&ToolSpan> {
self.spans.iter().rev().find(|s| s.is_open() && s.kind == kind)
}
pub fn discard_open(&mut self, id: &str) {
if let Some(i) = self.spans.iter().rposition(|s| s.is_open() && s.id == id) {
self.spans.remove(i);
}
}
pub fn close(&mut self, id: &str, at: SystemTime, error: bool) {
let Some(s) = self.spans.iter_mut().rev().find(|s| s.is_open() && s.id == id) else { return };
s.duration_ms = Some(at.duration_since(s.started_at).map(|d| d.as_millis() as u64).unwrap_or(0));
s.error = error;
}
pub fn len(&self) -> usize {
self.spans.len()
}
pub fn is_empty(&self) -> bool {
self.spans.is_empty()
}
pub fn iter(&self) -> impl DoubleEndedIterator<Item = &ToolSpan> + ExactSizeIterator {
self.spans.iter()
}
pub fn to_vec(&self) -> Vec<ToolSpan> {
self.spans.iter().cloned().collect()
}
pub fn merged<'a>(logs: impl IntoIterator<Item = &'a SpanLog>, cap: usize) -> SpanLog {
let mut spans: Vec<ToolSpan> = logs.into_iter().flat_map(|l| l.spans.iter().cloned()).collect();
spans.sort_by_key(|s| s.started_at);
if spans.len() > cap {
spans.drain(..spans.len() - cap);
}
SpanLog { spans: spans.into(), cap }
}
pub fn cap(&self) -> usize {
self.cap
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum SpanRetention {
#[default]
Recent,
All,
}
impl SpanRetention {
pub(crate) fn log(self) -> SpanLog {
match self {
SpanRetention::Recent => SpanLog::default(),
SpanRetention::All => SpanLog::unbounded(),
}
}
}
pub trait SessionTracker {
fn refresh(&mut self) -> anyhow::Result<bool>;
fn summary(&self) -> &SessionSummary;
fn path(&self) -> &Path;
fn refresh_all(&mut self) -> anyhow::Result<()> {
while self.refresh()? {}
Ok(())
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct RegistryHints {
pub name: Option<String>,
pub session_id: Option<String>,
pub cwd: Option<PathBuf>,
pub version: Option<String>,
pub status: Option<String>,
}
pub struct AttributeContext<'a> {
pub cwd: Option<&'a Path>,
pub proc_start: SystemTime,
pub now: SystemTime,
pub attached: &'a HashSet<PathBuf>,
pub activity_timeout: Duration,
}
pub trait HarnessAdapter {
fn harness(&self) -> Harness;
fn rescan(&mut self, since: SystemTime);
fn prepare(&mut self, _roots: &[&ProcNode], _by_pid: &std::collections::HashMap<u32, &RawProc>) {}
fn hints(&self, _pid: u32) -> Option<RegistryHints> {
None
}
fn attribute(&self, root: &ProcNode, raw: Option<&RawProc>, ctx: &AttributeContext) -> (Vec<PathBuf>, Attribution);
fn unowned(&self, attached: &HashSet<PathBuf>) -> Vec<PathBuf>;
fn open(&self, path: &Path, spans: SpanRetention) -> Box<dyn SessionTracker>;
fn detect(&self, path: &Path) -> bool;
fn transcripts(&self) -> Vec<(String, PathBuf)>;
}
pub fn adapters() -> Vec<Box<dyn HarnessAdapter>> {
vec![
Box::new(codex::CodexAdapter::default()),
Box::new(gemini::GeminiAdapter::default()),
Box::new(opencode::OpenCodeAdapter::default()),
Box::new(kodelet::KodeletAdapter::default()),
Box::new(claude::ClaudeAdapter::default()),
]
}
pub fn adapter_for(harness: Harness) -> Option<Box<dyn HarnessAdapter>> {
adapters().into_iter().find(|a| a.harness() == harness)
}
pub fn detect(path: &Path) -> Option<Harness> {
adapters().iter().find(|a| a.detect(path)).map(|a| a.harness())
}
pub fn open_transcript(path: &Path, harness: Harness, spans: SpanRetention) -> Option<Box<dyn SessionTracker>> {
adapter_for(harness).map(|a| a.open(path, spans))
}
pub(crate) fn head_lines(path: &Path) -> Vec<serde_json::Value> {
use std::io::{BufRead, BufReader};
let Ok(f) = std::fs::File::open(path) else { return Vec::new() };
BufReader::new(f).lines().map_while(Result::ok).take(5).filter_map(|l| serde_json::from_str(&l).ok()).collect()
}
pub const REFRESH_BUDGET_BYTES: usize = 8 * 1024 * 1024;
pub fn parse_rfc3339_utc(s: &str) -> Option<SystemTime> {
let s = s.trim();
let (s, offset_secs) = match s.strip_suffix(['Z', 'z']) {
Some(rest) => (rest, 0i64),
None => {
let t_pos = s.find('T')?;
let sign_pos = s[t_pos..].rfind(['+', '-'])? + t_pos;
let (rest, zone) = s.split_at(sign_pos);
let sign = if zone.starts_with('-') { -1 } else { 1 };
let digits: String = zone[1..].chars().filter(|c| c.is_ascii_digit()).collect();
if digits.len() != 4 {
return None;
}
let oh = digits[..2].parse::<i64>().ok()?;
let om = digits[2..].parse::<i64>().ok()?;
(rest, sign * (oh * 3600 + om * 60))
}
};
let (date, time) = s.split_once('T')?;
let mut d = date.split('-');
let (y, mo, da) = (d.next()?.parse::<i64>().ok()?, d.next()?.parse::<u32>().ok()?, d.next()?.parse::<u32>().ok()?);
let mut t = time.split(':');
let (h, mi) = (t.next()?.parse::<u64>().ok()?, t.next()?.parse::<u64>().ok()?);
let sec_str = t.next()?;
let (sec, frac) = match sec_str.split_once('.') {
Some((s, f)) => (s.parse::<u64>().ok()?, f),
None => (sec_str.parse::<u64>().ok()?, ""),
};
let nanos: u32 = if frac.is_empty() {
0
} else {
let mut f = frac.to_string();
f.truncate(9);
while f.len() < 9 {
f.push('0');
}
f.parse().ok()?
};
let days = days_from_civil(y, mo, da);
let secs = days * 86_400 + (h * 3600 + mi * 60 + sec) as i64 - offset_secs;
if secs < 0 {
return None;
}
Some(SystemTime::UNIX_EPOCH + std::time::Duration::new(secs as u64, nanos))
}
fn days_from_civil(y: i64, m: u32, d: u32) -> i64 {
let y = if m <= 2 { y - 1 } else { y };
let era = if y >= 0 { y } else { y - 399 } / 400;
let yoe = y - era * 400;
let mp = (m as i64 + 9) % 12;
let doy = (153 * mp + 2) / 5 + d as i64 - 1;
let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy;
era * 146_097 + doe - 719_468
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::Duration;
fn at(secs: u64) -> SystemTime {
SystemTime::UNIX_EPOCH + Duration::from_secs(secs)
}
#[test]
fn pairs_spans_by_id_out_of_order() {
let mut log = SpanLog::default();
log.open("a".into(), "Bash".into(), at(10), false);
log.open("b".into(), "Read".into(), at(11), true);
log.open("a".into(), "Bash".into(), at(10), false);
log.close("b", at(12), false);
log.close("a", at(14), true);
log.close("zzz", at(15), false);
let v = log.to_vec();
assert_eq!(v.len(), 2);
assert_eq!(v[0].name, "Bash");
assert_eq!(v[0].duration_ms, Some(4_000));
assert!(v[0].error);
assert_eq!(v[1].duration_ms, Some(1_000));
assert!(v[1].sidechain);
assert!(!v[1].error);
}
#[test]
fn keeps_the_newest_spans_and_reports_open_ones() {
let mut log = SpanLog::default();
for i in 0..(MAX_SPANS + 10) {
log.open(format!("id{i}"), "T".into(), at(i as u64), false);
log.close(&format!("id{i}"), at(i as u64), false);
}
assert_eq!(log.len(), MAX_SPANS);
assert_eq!(log.iter().next().unwrap().id, "id10");
log.open("live".into(), "Bash".into(), at(500), false);
let last = log.to_vec().pop().unwrap();
assert!(last.is_open());
assert_eq!(last.elapsed_ms(at(503)), 3_000);
}
#[test]
fn end_at_moves_the_end_of_any_kind_of_span() {
let mut log = SpanLog::default();
log.open_kind("inference:1".into(), "inference".into(), at(10), false, SpanKind::Inference);
assert!(log.open_of_kind(SpanKind::Inference).is_some());
assert!(log.open_of_kind(SpanKind::Turn).is_none());
log.end_at("inference:1", at(11));
log.end_at("inference:1", at(13));
log.end_at("nope", at(99));
let v = log.to_vec();
assert_eq!(v[0].duration_ms, Some(3_000));
assert_eq!(v[0].kind, SpanKind::Inference);
assert!(log.open_of_kind(SpanKind::Inference).is_none());
log.discard_open("inference:1");
assert_eq!(log.len(), 1);
log.open_kind("inference:2".into(), "inference".into(), at(20), false, SpanKind::Inference);
log.discard_open("inference:2");
assert_eq!(log.len(), 1);
}
#[test]
fn unbounded_log_keeps_everything() {
let mut log = SpanLog::unbounded();
for i in 0..(MAX_SPANS * 3) {
log.open(format!("id{i}"), "T".into(), at(i as u64), false);
log.close(&format!("id{i}"), at(i as u64 + 1), false);
}
assert_eq!(log.len(), MAX_SPANS * 3);
assert_eq!(log.iter().next().unwrap().id, "id0");
assert_eq!(SpanRetention::default(), SpanRetention::Recent);
}
#[test]
fn detects_the_harness_from_the_first_lines() {
let dir = std::env::temp_dir().join(format!("agent-top-detect-{}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
let codex = dir.join("rollout.jsonl");
std::fs::write(&codex, "{\"type\":\"session_meta\",\"payload\":{\"id\":\"x\"}}\n").unwrap();
let claude = dir.join("s.jsonl");
std::fs::write(&claude, "{\"type\":\"summary\",\"leafUuid\":\"u\"}\n{\"type\":\"user\",\"sessionId\":\"abc\"}\n").unwrap();
let other = dir.join("other.jsonl");
std::fs::write(&other, "{\"hello\":1}\nnot json\n").unwrap();
assert_eq!(detect(&codex), Some(Harness::Codex));
assert_eq!(detect(&claude), Some(Harness::Claude));
assert_eq!(detect(&other), None);
assert_eq!(detect(&dir.join("missing.jsonl")), None);
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn names_the_server_behind_an_mcp_tool() {
assert_eq!(mcp_server_of("mcp__filesystem__read_file"), Some("filesystem"));
assert_eq!(mcp_server_of("mcp__chrome-devtools__take_screenshot"), Some("chrome-devtools"));
assert_eq!(mcp_server_of("mcp__claude_ai_Gmail__authenticate"), Some("claude_ai_Gmail"));
assert_eq!(mcp_server_of("mcp__odd"), Some("odd"));
assert_eq!(mcp_server_of("mcp____x"), None);
assert_eq!(mcp_server_of("Bash"), None);
}
#[test]
fn accuses_the_parser_only_with_enough_evidence() {
let h = ParseHealth { billable_messages: 40, usage_records: 40, empty_usage_records: 0 };
assert!(!h.fields_unrecognised());
let h = ParseHealth { billable_messages: 40, usage_records: 40, empty_usage_records: 39 };
assert!(!h.fields_unrecognised());
let h = ParseHealth { billable_messages: 40, usage_records: 40, empty_usage_records: 40 };
assert!(h.fields_unrecognised());
let h = ParseHealth { billable_messages: 40, usage_records: 0, empty_usage_records: 0 };
assert!(h.fields_unrecognised());
let h = ParseHealth { billable_messages: 2, usage_records: 0, empty_usage_records: 0 };
assert!(!h.fields_unrecognised());
assert!(!ParseHealth::default().fields_unrecognised());
}
fn usage(prompt: u64, output: u64) -> TokenUsage {
TokenUsage { cache_read: prompt, output, ..Default::default() }
}
fn cost(prompt: u64) -> CostBreakdown {
CostBreakdown { cache_read: prompt as f64 / 1e6, ..Default::default() }
}
fn share<'a>(v: &'a [ContextSource], name: &str) -> &'a ContextSource {
v.iter().find(|s| s.name == name).unwrap_or_else(|| panic!("no source {name}"))
}
#[test]
fn context_ledger_files_prompt_growth_under_the_results_that_caused_it() {
let mut l = ContextLedger::default();
l.response(&usage(1_000, 100), &cost(1_000));
l.result("a", ContextOrigin::Tool, "Read");
l.result("b", ContextOrigin::Mcp, "fs");
l.response(&usage(3_300, 50), &cost(3_300));
let v = l.sources();
assert_eq!(share(&v, "Read").tokens, 1_100);
assert_eq!(share(&v, "fs").tokens, 1_100);
assert_eq!(share(&v, "fs").origin, ContextOrigin::Mcp);
assert_eq!(share(&v, "fs").calls, 1);
let other = share(&v, ContextLedger::OTHER);
assert_eq!((other.tokens, other.calls), (1_100, 0));
assert!((other.cost_usd - 2_100e-6).abs() < 1e-12, "{}", other.cost_usd);
assert!((share(&v, "Read").cost_usd - 1_100e-6).abs() < 1e-12);
let total: f64 = v.iter().map(|s| s.cost_usd).sum();
assert!((total - 4_300e-6).abs() < 1e-12, "{total}");
assert_eq!(v[0].tokens, 1_100, "largest first");
}
#[test]
fn context_ledger_weights_preserve_equal_shares_between_result_records() {
let mut l = ContextLedger::default();
l.response(&usage(1_000, 100), &cost(1_000));
l.result_weighted("wrapper", vec![(ContextOrigin::Tool, "exec_command".into(), 1), (ContextOrigin::Tool, "apply_patch".into(), 3)]);
l.result("direct", ContextOrigin::Mcp, "fs");
l.response(&usage(3_100, 0), &cost(3_100));
let v = l.sources();
assert_eq!(share(&v, "exec_command").tokens, 250);
assert_eq!(share(&v, "apply_patch").tokens, 750);
assert_eq!(share(&v, "fs").tokens, 1_000);
assert_eq!(v.iter().map(|s| s.calls).sum::<u64>(), 3);
assert_eq!(v.iter().map(|s| s.tokens).sum::<u64>(), 3_100);
assert!((v.iter().map(|s| s.cost_usd).sum::<f64>() - 4_100e-6).abs() < 1e-12);
}
#[test]
fn context_ledger_weighted_rounding_handles_zero_and_equal_weights() {
for (weights, tokens, expected) in [
([0, 0, 0], 2, [1, 1, 0]),
([1, 1, 1], 2, [1, 1, 0]), ([0, 1, 3], 1, [0, 0, 1]),
([0, 1, 3], 2, [0, 1, 1]),
([0, 1, 3], 0, [0, 0, 0]),
] {
let mut l = ContextLedger::default();
l.response(&usage(100, 0), &cost(100));
l.result_weighted("wrapper", weights.into_iter().enumerate().map(|(i, w)| (ContextOrigin::Tool, i.to_string(), w)).collect());
l.response(&usage(100 + tokens, 0), &cost(100 + tokens));
let v = l.sources();
for (i, tokens) in expected.into_iter().enumerate() {
let source = share(&v, &i.to_string());
assert_eq!((source.calls, source.tokens), (1, tokens), "weights: {weights:?}");
}
assert_eq!(v.iter().map(|s| s.tokens).sum::<u64>(), 100 + tokens);
assert!((v.iter().map(|s| s.cost_usd).sum::<f64>() - (200 + tokens) as f64 / 1e6).abs() < 1e-12);
}
}
#[test]
fn context_ledger_weighted_empty_groups_do_not_lose_tokens() {
let mut l = ContextLedger::default();
l.result_weighted("empty", Vec::new());
l.response(&usage(100, 0), &cost(100));
assert_eq!(share(&l.sources(), ContextLedger::OTHER).tokens, 100);
l.result_weighted("empty", Vec::new());
l.result("direct", ContextOrigin::Tool, "Read");
l.response(&usage(200, 0), &cost(200));
let v = l.sources();
assert_eq!((share(&v, "Read").calls, share(&v, "Read").tokens), (1, 100));
assert_eq!(v.iter().map(|s| s.tokens).sum::<u64>(), 200);
}
#[test]
fn context_ledger_weights_aggregate_repeated_names_and_survive_retagging() {
let mut l = ContextLedger::default();
l.response(&usage(100, 0), &cost(100));
let sources = vec![(ContextOrigin::Tool, "Read".into(), 1), (ContextOrigin::Tool, "Read".into(), 3)];
l.result_weighted("first", sources.clone());
l.response(&usage(140, 0), &cost(140));
assert_eq!((share(&l.sources(), "Read").calls, share(&l.sources(), "Read").tokens), (2, 40));
l.result_weighted("second", sources);
l.retag("second", ContextOrigin::Mcp, "fs");
assert_eq!(l.pending[0].1, vec![(ContextOrigin::Mcp, "fs".into(), 1), (ContextOrigin::Mcp, "fs".into(), 3)]);
l.response(&usage(180, 0), &cost(180));
let v = l.sources();
assert_eq!((share(&v, "fs").calls, share(&v, "fs").tokens), (2, 40));
assert_eq!(share(&v, "fs").origin, ContextOrigin::Mcp);
assert_eq!((share(&v, "Read").calls, share(&v, "Read").tokens), (2, 40));
assert_eq!(v.iter().map(|s| s.tokens).sum::<u64>(), 180);
}
#[test]
fn context_ledger_weights_use_wide_products_and_sums() {
let mut l = ContextLedger::default();
l.result_weighted("wrapper", vec![(ContextOrigin::Tool, "a".into(), u64::MAX), (ContextOrigin::Tool, "b".into(), u64::MAX - 1)]);
l.response(&usage(u64::MAX, 0), &cost(u64::MAX));
let v = l.sources();
assert_eq!(share(&v, "a").tokens, u64::MAX / 2 + 1);
assert_eq!(share(&v, "b").tokens, u64::MAX / 2);
assert_eq!(v.iter().map(|s| s.tokens).sum::<u64>(), u64::MAX);
assert_eq!(v.iter().map(|s| s.calls).sum::<u64>(), 2);
}
#[test]
fn context_ledger_weighted_costs_conserve_totals_across_rereads_and_compaction() {
let mut l = ContextLedger::default();
l.response(&usage(1_000, 100), &cost(1_000));
l.result_weighted("wrapper", vec![(ContextOrigin::Tool, "exec_command".into(), 1), (ContextOrigin::Tool, "apply_patch".into(), 3)]);
l.response(&usage(1_500, 50), &cost(1_500));
let v = l.sources();
assert_eq!(share(&v, "exec_command").tokens, 100);
assert_eq!(share(&v, "apply_patch").tokens, 300);
l.response(&usage(1_550, 10), &cost(1_550));
let v = l.sources();
assert!((share(&v, "exec_command").cost_usd - 200e-6).abs() < 1e-12);
assert!((share(&v, "apply_patch").cost_usd - 600e-6).abs() < 1e-12);
assert!((v.iter().map(|s| s.cost_usd).sum::<f64>() - 4_050e-6).abs() < 1e-12);
l.response(&usage(500, 10), &cost(500)); l.response(&usage(510, 0), &cost(510));
l.compacted();
l.response(&usage(700, 0), &cost(700));
let v = l.sources();
assert!((share(&v, "exec_command").cost_usd - 200e-6).abs() < 1e-12);
assert!((share(&v, "apply_patch").cost_usd - 600e-6).abs() < 1e-12);
assert_eq!(v.iter().map(|s| s.calls).sum::<u64>(), 2);
assert_eq!(v.iter().map(|s| s.tokens).sum::<u64>(), 2_760);
assert!((v.iter().map(|s| s.cost_usd).sum::<f64>() - 5_760e-6).abs() < 1e-12);
}
#[test]
fn context_ledger_takes_a_shrink_off_other_and_a_halving_as_compaction() {
let mut l = ContextLedger::default();
l.response(&usage(10_000, 2_000), &cost(10_000));
l.result("a", ContextOrigin::Tool, "Bash");
l.response(&usage(12_500, 3_000), &cost(12_500)); l.response(&usage(11_500, 10), &cost(11_500));
let v = l.sources();
assert_eq!(share(&v, "Bash").tokens, 500);
assert_eq!(share(&v, ContextLedger::OTHER).tokens, 12_000);
assert!((share(&v, "Bash").cost_usd - 1_000e-6).abs() < 1e-12);
l.response(&usage(3_000, 10), &cost(3_000));
let v = l.sources();
assert!((share(&v, "Bash").cost_usd - 1_000e-6).abs() < 1e-12, "not charged after compaction");
assert_eq!(share(&v, ContextLedger::OTHER).tokens, 15_000);
l.compacted();
l.response(&usage(4_000, 10), &cost(4_000));
let v = l.sources();
assert_eq!(share(&v, ContextLedger::OTHER).tokens, 19_000);
assert!((share(&v, "Bash").cost_usd - 1_000e-6).abs() < 1e-12);
}
#[test]
fn context_ledger_retags_pending_results_and_merges() {
let mut l = ContextLedger::default();
l.response(&usage(100, 0), &cost(100));
l.result("c1", ContextOrigin::Tool, "fetch");
l.retag("c1", ContextOrigin::Mcp, "apps");
l.retag("zzz", ContextOrigin::Mcp, "nope");
l.response(&TokenUsage::default(), &CostBreakdown::default());
l.response(&usage(300, 0), &cost(300));
let v = l.sources();
assert_eq!(v.len(), 2);
assert_eq!((share(&v, "apps").origin, share(&v, "apps").tokens), (ContextOrigin::Mcp, 200));
assert!(v.iter().all(|s| s.name != "fetch"));
let mut sub = ContextLedger::default();
sub.response(&usage(50, 0), &cost(50));
sub.result("x", ContextOrigin::Mcp, "apps");
sub.response(&usage(70, 0), &cost(70));
l.merge(&sub);
let v = l.sources();
assert_eq!(share(&v, "apps").tokens, 220);
assert_eq!(share(&v, "apps").calls, 2);
assert_eq!(share(&v, ContextLedger::OTHER).tokens, 150);
assert!(!l.is_empty() && ContextLedger::default().is_empty());
}
#[test]
fn parses_timestamps() {
let t = parse_rfc3339_utc("1970-01-02T00:00:00.000Z").unwrap();
assert_eq!(t, SystemTime::UNIX_EPOCH + Duration::from_secs(86_400));
let t = parse_rfc3339_utc("2026-09-03T07:15:34.5Z").unwrap();
let secs = t.duration_since(SystemTime::UNIX_EPOCH).unwrap();
assert_eq!(secs.as_secs(), 1_788_419_734);
assert_eq!(secs.subsec_millis(), 500);
assert!(parse_rfc3339_utc("nope").is_none());
}
}
#[cfg(test)]
mod rfc3339_tests {
use super::parse_rfc3339_utc;
use std::time::{Duration, UNIX_EPOCH};
fn secs(s: &str) -> u64 {
parse_rfc3339_utc(s).unwrap().duration_since(UNIX_EPOCH).unwrap().as_secs()
}
#[test]
fn offsets_are_folded_into_utc() {
let z = secs("2026-09-03T07:15:34Z");
assert_eq!(secs("2026-09-03T08:15:34+01:00"), z);
assert_eq!(secs("2026-09-03T00:15:34-07:00"), z);
assert_eq!(secs("2026-09-03T08:15:34+0100"), z, "no colon");
assert_eq!(secs("2026-09-03T07:15:34+00:00"), z);
assert_eq!(secs("2026-09-03T12:45:34+05:30"), z, "half-hour zone");
let ms = parse_rfc3339_utc("2026-09-03T08:15:34.250+01:00").unwrap();
assert_eq!(ms, UNIX_EPOCH + Duration::new(z, 250_000_000));
assert_eq!(secs("2026-09-03T07:15:34.5Z"), z);
assert!(parse_rfc3339_utc("2026-09-03T07:15:34").is_none(), "no zone at all");
assert!(parse_rfc3339_utc("2026-09-03T07:15:34+1").is_none());
assert!(parse_rfc3339_utc("garbage").is_none());
}
}