From 1b740e4cc207daaeb98006aae4031dbc8a686fa9 Mon Sep 17 00:00:00 2001 From: makoMakoGo <48956204+makoMakoGo@users.noreply.github.com> Date: Sat, 25 Jul 2026 15:17:19 +0800 Subject: [PATCH 1/2] fix(core): make adapters authoritative for input attribution Store parser and cache messages without client identity, then inject ClientId only at adapter emission. Derive Overview Data Size from the confirmed per-client inventory and validate the invariant in TUI cache boundaries. --- crates/tokscale-cli/src/tui/cache.rs | 108 +++++-- crates/tokscale-cli/src/tui/mod.rs | 8 +- .../tokscale-core/src/adapters/antigravity.rs | 18 +- crates/tokscale-core/src/adapters/cache.rs | 43 ++- crates/tokscale-core/src/adapters/claude.rs | 30 +- .../tokscale-core/src/adapters/codebuddy.rs | 21 +- crates/tokscale-core/src/adapters/codex.rs | 68 ++++- crates/tokscale-core/src/adapters/file.rs | 39 ++- crates/tokscale-core/src/adapters/goose.rs | 10 +- crates/tokscale-core/src/adapters/hermes.rs | 29 +- crates/tokscale-core/src/adapters/junie.rs | 21 +- crates/tokscale-core/src/adapters/kilo.rs | 10 +- crates/tokscale-core/src/adapters/kiro.rs | 7 +- crates/tokscale-core/src/adapters/mod.rs | 80 +++-- crates/tokscale-core/src/adapters/omp.rs | 33 +- crates/tokscale-core/src/adapters/opencode.rs | 18 +- crates/tokscale-core/src/adapters/pi.rs | 19 +- crates/tokscale-core/src/adapters/warp.rs | 10 +- crates/tokscale-core/src/adapters/zed.rs | 9 +- crates/tokscale-core/src/input_health.rs | 11 +- crates/tokscale-core/src/lib.rs | 149 ++++++--- crates/tokscale-core/src/lib_tests.rs | 94 ++---- crates/tokscale-core/src/message_cache.rs | 69 ++--- crates/tokscale-core/src/sessions/amp.rs | 13 +- .../src/sessions/antigravity_cli.rs | 12 +- .../tokscale-core/src/sessions/claudecode.rs | 41 ++- crates/tokscale-core/src/sessions/cline.rs | 7 +- .../tokscale-core/src/sessions/codebuddy.rs | 17 +- crates/tokscale-core/src/sessions/codebuff.rs | 8 +- crates/tokscale-core/src/sessions/codex.rs | 19 +- .../tokscale-core/src/sessions/commandcode.rs | 10 +- crates/tokscale-core/src/sessions/copilot.rs | 10 +- crates/tokscale-core/src/sessions/droid.rs | 11 +- crates/tokscale-core/src/sessions/gemini.rs | 16 +- crates/tokscale-core/src/sessions/goose.rs | 5 +- crates/tokscale-core/src/sessions/grok.rs | 13 +- crates/tokscale-core/src/sessions/hermes.rs | 5 +- crates/tokscale-core/src/sessions/junie.rs | 11 +- crates/tokscale-core/src/sessions/kilo.rs | 10 +- crates/tokscale-core/src/sessions/kimi.rs | 10 +- crates/tokscale-core/src/sessions/kiro.rs | 20 +- crates/tokscale-core/src/sessions/mod.rs | 288 +++++++++++++----- crates/tokscale-core/src/sessions/mux.rs | 10 +- crates/tokscale-core/src/sessions/openclaw.rs | 10 +- crates/tokscale-core/src/sessions/opencode.rs | 20 +- crates/tokscale-core/src/sessions/pi.rs | 122 ++++---- crates/tokscale-core/src/sessions/qwen.rs | 12 +- crates/tokscale-core/src/sessions/roocode.rs | 8 +- crates/tokscale-core/src/sessions/warp.rs | 8 +- crates/tokscale-core/src/sessions/zcode.rs | 9 +- crates/tokscale-core/src/sessions/zed.rs | 10 +- crates/tokscale-core/tests/codebuff.rs | 1 - crates/tokscale-core/tests/hermes.rs | 1 - docs/adr/0007-client-identity-catalog.md | 26 +- docs/adr/0008-single-copy-memory-pipeline.md | 38 ++- 55 files changed, 1033 insertions(+), 672 deletions(-) diff --git a/crates/tokscale-cli/src/tui/cache.rs b/crates/tokscale-cli/src/tui/cache.rs index beb238c90..867eedd9a 100644 --- a/crates/tokscale-cli/src/tui/cache.rs +++ b/crates/tokscale-cli/src/tui/cache.rs @@ -29,7 +29,7 @@ use super::data::{ /// Cache staleness threshold: 5 minutes (matches TS implementation) const CACHE_STALE_THRESHOLD_MS: u64 = 5 * 60 * 1000; -const CACHE_SCHEMA_VERSION: u32 = 50; +const CACHE_SCHEMA_VERSION: u32 = 51; fn sha256_hex(bytes: &[u8]) -> String { const HEX: &[u8; 16] = b"0123456789abcdef"; @@ -126,6 +126,15 @@ mod bundle_tests { InputInventorySignature::from_bytes([0x39; 32]) } + fn health_for_client_space( + client_space: &BTreeMap, + ) -> tokscale_core::input_health::HealthReport { + tokscale_core::input_health::HealthReport { + input_data_bytes: client_space.values().copied().sum(), + ..Default::default() + } + } + fn nonempty_accumulator(home: &std::path::Path) -> TuiAcc { for (project, workspace, model, input_tokens, hour) in [ ("project-a", "/work/alpha", "claude-sonnet-4.6", 101, 10), @@ -283,6 +292,49 @@ mod bundle_tests { assert_eq!(loaded.data.health, health); } + #[test] + #[serial] + fn bundle_rejects_data_size_outside_client_space_invariant() { + let (_temp, _guard, clients, scope, sessions, client_space) = fixture(); + let invalid_health = tokscale_core::input_health::HealthReport::default(); + + let error = save_tui_bundle_cache( + &TuiAcc::new(), + &sessions, + &client_space, + &invalid_health, + &clients, + &scope, + signature(), + ) + .unwrap_err(); + assert!(error + .to_string() + .contains("TUI Data Size 0 does not match client-space total 4096")); + + save_tui_bundle_cache( + &TuiAcc::new(), + &sessions, + &client_space, + &health_for_client_space(&client_space), + &clients, + &scope, + signature(), + ) + .unwrap(); + let path = cache_file().unwrap(); + let mut value: serde_json::Value = + serde_json::from_reader(File::open(&path).unwrap()).unwrap(); + value["health"]["inputDataBytes"] = serde_json::Value::from(1_u64); + tokscale_core::fs_atomic::write_atomic(&path, &serde_json::to_vec(&value).unwrap()) + .unwrap(); + + assert!(matches!( + load_cache(&clients, &GroupBy::Model, &scope), + CacheResult::Miss + )); + } + #[test] #[serial] fn bundle_keeps_same_named_agents_separate_across_clients() { @@ -298,7 +350,7 @@ mod bundle_tests { &accumulator, &[], &client_space, - &Default::default(), + &health_for_client_space(&client_space), &clients, &scope, signature(), @@ -340,7 +392,7 @@ mod bundle_tests { &TuiAcc::new(), &sessions, &client_space, - &Default::default(), + &health_for_client_space(&client_space), &clients, &scope, signature(), @@ -584,7 +636,7 @@ mod bundle_tests { &accumulator, &[], &client_space, - &Default::default(), + &health_for_client_space(&client_space), &clients, &scope, signature(), @@ -645,7 +697,7 @@ mod bundle_tests { &accumulator, &[], &client_space, - &Default::default(), + &health_for_client_space(&client_space), &clients, &scope, signature(), @@ -684,7 +736,7 @@ mod bundle_tests { &accumulator, &[], &client_space, - &Default::default(), + &health_for_client_space(&client_space), &clients, &scope, signature(), @@ -776,7 +828,7 @@ mod bundle_tests { &accumulator, &[], &client_space, - &Default::default(), + &health_for_client_space(&client_space), &clients, &scope, signature(), @@ -813,7 +865,7 @@ mod bundle_tests { &accumulator, &[], &client_space, - &Default::default(), + &health_for_client_space(&client_space), &clients, &scope, signature(), @@ -851,7 +903,7 @@ mod bundle_tests { &accumulator, &[], &client_space, - &Default::default(), + &health_for_client_space(&client_space), &clients, &scope, signature(), @@ -878,7 +930,7 @@ mod bundle_tests { &TuiAcc::new(), &sessions, &client_space, - &Default::default(), + &health_for_client_space(&client_space), &clients, &scope, signature(), @@ -909,7 +961,7 @@ mod bundle_tests { &TuiAcc::new(), &sessions, &client_space, - &Default::default(), + &health_for_client_space(&client_space), &clients, &scope, signature(), @@ -936,7 +988,7 @@ mod bundle_tests { &TuiAcc::new(), &sessions, &client_space, - &Default::default(), + &health_for_client_space(&client_space), &clients, &scope, signature(), @@ -964,7 +1016,7 @@ mod bundle_tests { &TuiAcc::new(), &sessions, &client_space, - &Default::default(), + &health_for_client_space(&client_space), &clients, &scope, signature(), @@ -989,7 +1041,7 @@ mod bundle_tests { &TuiAcc::new(), &sessions, &client_space, - &Default::default(), + &health_for_client_space(&client_space), &clients, &scope, signature(), @@ -1018,7 +1070,7 @@ mod bundle_tests { &TuiAcc::new(), &sessions, &client_space, - &Default::default(), + &health_for_client_space(&client_space), &clients, &scope, signature(), @@ -1039,7 +1091,7 @@ mod bundle_tests { &TuiAcc::new(), &sessions, &client_space, - &Default::default(), + &health_for_client_space(&client_space), &clients, &scope, signature(), @@ -1065,7 +1117,7 @@ mod bundle_tests { &TuiAcc::new(), &sessions, &client_space, - &Default::default(), + &health_for_client_space(&client_space), &clients, &scope, signature(), @@ -1092,7 +1144,7 @@ mod bundle_tests { &TuiAcc::new(), &sessions, &client_space, - &Default::default(), + &health_for_client_space(&client_space), &clients, &scope, signature(), @@ -2198,6 +2250,24 @@ fn cache_client_space_matches_exact( .all(|client| client_space.contains_key(client.as_str())) } +fn validate_data_size_matches_client_space( + client_space: &BTreeMap, + health: &tokscale_core::input_health::HealthReport, +) -> anyhow::Result<()> { + let total = client_space + .values() + .copied() + .try_fold(0_u64, u64::checked_add) + .ok_or_else(|| anyhow::anyhow!("TUI client-space total exceeds u64::MAX"))?; + if health.input_data_bytes != total { + anyhow::bail!( + "TUI Data Size {} does not match client-space total {total}", + health.input_data_bytes + ); + } + Ok(()) +} + fn cache_session_clients_are_enabled( client_universe: &HashSet, sessions: &[TuiSessionEntry], @@ -3067,6 +3137,7 @@ fn load_bundle_from_file( if !cache_client_space_matches_exact(client_universe, &parsed.client_space) { anyhow::bail!("cached TUI client-space keys do not match the client universe"); } + validate_data_size_matches_client_space(&parsed.client_space, &parsed.health)?; if !cache_session_clients_are_enabled(client_universe, &parsed.sessions) { anyhow::bail!("cached TUI Sessions contain a client outside the client universe"); } @@ -3151,6 +3222,7 @@ pub fn save_tui_bundle_cache( if !cache_client_space_matches_exact(client_universe, client_space) { anyhow::bail!("TUI client-space keys do not match the client universe"); } + validate_data_size_matches_client_space(client_space, health)?; if !cache_session_clients_are_enabled(client_universe, sessions) { anyhow::bail!("TUI Sessions contain a client outside the client universe"); } diff --git a/crates/tokscale-cli/src/tui/mod.rs b/crates/tokscale-cli/src/tui/mod.rs index 391b4ec2d..c50d5d68e 100644 --- a/crates/tokscale-cli/src/tui/mod.rs +++ b/crates/tokscale-cli/src/tui/mod.rs @@ -1118,7 +1118,10 @@ mod tests { &tokscale_core::TuiAcc::new(), &[session("amp", "cached-session")], &std::collections::BTreeMap::from([("amp".to_string(), 512)]), - &Default::default(), + &tokscale_core::input_health::HealthReport { + input_data_bytes: 512, + ..Default::default() + }, &clients, &scope, signature, @@ -1130,13 +1133,14 @@ mod tests { } fn loaded_snapshot( - data: UsageData, + mut data: UsageData, sessions: Vec, client_space: std::collections::BTreeMap, accumulator: tokscale_core::TuiAcc, group_by: tokscale_core::GroupBy, signature: tokscale_core::InputInventorySignature, ) -> BackgroundLoad { + data.health.input_data_bytes = client_space.values().copied().sum(); BackgroundLoad::Loaded { data: Box::new(data), sessions, diff --git a/crates/tokscale-core/src/adapters/antigravity.rs b/crates/tokscale-core/src/adapters/antigravity.rs index 58c5763e6..82e2fdc99 100644 --- a/crates/tokscale-core/src/adapters/antigravity.rs +++ b/crates/tokscale-core/src/adapters/antigravity.rs @@ -131,11 +131,12 @@ fn fold_antigravity_units( ctx.input_cache.remove(&path, unit.parser_version); } let cache_write_outcome = cache_write_outcome?; - sink.extend_messages( + adapter_cache::emit_messages( + unit.client, messages .into_iter() - .filter(|message| crate::should_keep_deduped_message(seen, message)) - .collect(), + .filter(|message| crate::should_keep_deduped_message(seen, message)), + sink, ); if cache_write_outcome == adapter_cache::CacheWriteOutcome::NotPlanned && invalidate_cache { @@ -188,7 +189,8 @@ mod tests { use super::*; use crate::adapters::UnitMessagePayload; use crate::scanner::ScannerSettings; - use crate::{message_cache, TokenBreakdown, UnifiedMessage}; + use crate::sessions::ParsedMessage; + use crate::{message_cache, TokenBreakdown}; use std::collections::BTreeMap; use std::path::Path; @@ -278,7 +280,7 @@ mod tests { assert!(units.is_empty()); } - fn parsed_unit(path: &Path, meta: InputUnitMeta, message: UnifiedMessage) -> ParsedUnit { + fn parsed_unit(path: &Path, meta: InputUnitMeta, message: ParsedMessage) -> ParsedUnit { ParsedUnit::healthy( InputUnit::plain_file(ClientId::Antigravity, path.to_path_buf()).with_meta(meta), UnitMessagePayload::Fresh(vec![message]), @@ -287,9 +289,8 @@ mod tests { ) } - fn antigravity_message(session_id: &str, dedup_key: Option) -> UnifiedMessage { - UnifiedMessage::new_with_dedup( - "antigravity", + fn antigravity_message(session_id: &str, dedup_key: Option) -> ParsedMessage { + ParsedMessage::new_with_dedup( "gemini-3.1-pro", "google", session_id, @@ -332,6 +333,7 @@ mod tests { .unwrap(); assert_eq!(messages.len(), 1); + assert_eq!(messages[0].client.as_ref(), ClientId::Antigravity.as_str()); assert_eq!(messages[0].session_id.as_ref(), "first-session"); } } diff --git a/crates/tokscale-core/src/adapters/cache.rs b/crates/tokscale-core/src/adapters/cache.rs index cad6e643d..de37be907 100644 --- a/crates/tokscale-core/src/adapters/cache.rs +++ b/crates/tokscale-core/src/adapters/cache.rs @@ -4,8 +4,9 @@ use crate::adapters::{ CacheHitPlan, FingerprintPolicy, FoldContext, InputPipelineError, InputPlanningError, InputUnit, MessageSink, ParseContext, ParsedUnit, UnitMessagePayload, UnitScanHealth, }; +use crate::clients::ClientId; use crate::input_health::{InputFailure, InputHealth, InputStatus, ScannedInput}; -use crate::{message_cache, UnifiedMessage}; +use crate::{message_cache, sessions::ParsedMessage}; pub(crate) fn plan_cache_hit( mut unit: InputUnit, @@ -337,6 +338,16 @@ pub(crate) fn fold_units( fold_units_with_filter(parsed, ctx, sink, |_, messages| messages) } +pub(crate) fn emit_messages( + client: ClientId, + messages: impl IntoIterator, + sink: &mut dyn MessageSink, +) { + for message in messages { + sink.push_message(message.attribute(client)); + } +} + pub(crate) fn fold_units_with_filter( parsed: Vec, ctx: &mut FoldContext<'_>, @@ -344,7 +355,7 @@ pub(crate) fn fold_units_with_filter( mut filter: F, ) -> Result<(), InputPipelineError> where - F: FnMut(&InputUnit, Vec) -> Vec, + F: FnMut(&InputUnit, Vec) -> Vec, { for parsed_unit in parsed { let ResolvedUnit { @@ -370,7 +381,7 @@ where } let cache_write_outcome = cache_write_outcome?; let messages = filter(&unit, messages); - sink.extend_messages(messages); + emit_messages(unit.client, messages, sink); if cache_write_outcome == CacheWriteOutcome::NotPlanned && invalidate_cache { ctx.input_cache.remove(&path, parser_version); @@ -381,7 +392,7 @@ where pub(crate) struct ResolvedUnit { pub(crate) unit: InputUnit, - pub(crate) messages: Vec, + pub(crate) messages: Vec, pub(crate) cache_write: Option>, pub(crate) invalidate_cache: bool, pub(crate) status: InputStatus, @@ -470,7 +481,7 @@ pub(crate) enum CacheWriteOutcome { pub(crate) fn write_cache( cache_write: Option>, ctx: &mut FoldContext<'_>, - messages: &[UnifiedMessage], + messages: &[ParsedMessage], ) -> Result { if let Some(plan) = cache_write { ctx.input_cache.write_messages(*plan, messages)?; @@ -482,7 +493,7 @@ pub(crate) fn write_cache( pub(crate) fn resolve_messages( payload: UnitMessagePayload, ctx: &mut FoldContext<'_>, -) -> Result, message_cache::CacheReadFailure> { +) -> Result, message_cache::CacheReadFailure> { match payload { UnitMessagePayload::Fresh(messages) => Ok(messages), UnitMessagePayload::CacheHit(plan) => { @@ -503,7 +514,7 @@ mod tests { use super::*; use crate::adapters::InputUnit; use crate::clients::ClientId; - use crate::TokenBreakdown; + use crate::{TokenBreakdown, UnifiedMessage}; const PI_INPUT: &str = r#"{"type":"session","id":"input-session","timestamp":"2026-01-01T00:00:00.000Z","cwd":"/tmp"} {"type":"message","id":"msg_001","parentId":null,"timestamp":"2026-01-01T00:00:01.000Z","message":{"role":"assistant","model":"gpt-5.5","provider":"openai","usage":{"input":17,"output":3,"cacheRead":0,"cacheWrite":0,"totalTokens":20}}}"#; @@ -511,9 +522,8 @@ mod tests { const PI_REPLACEMENT_INPUT: &str = r#"{"type":"session","id":"replacement-input-session","timestamp":"2026-01-01T00:00:00.000Z","cwd":"/tmp"} {"type":"message","id":"msg_002","parentId":null,"timestamp":"2026-01-01T00:00:01.000Z","message":{"role":"assistant","model":"gpt-5.5","provider":"openai","usage":{"input":29,"output":5,"cacheRead":0,"cacheWrite":0,"totalTokens":34}}}"#; - fn cached_message() -> UnifiedMessage { - UnifiedMessage::new( - "test", + fn cached_message() -> ParsedMessage { + ParsedMessage::new( "gpt-5", "openai", "session", @@ -523,9 +533,8 @@ mod tests { ) } - fn scanned_message() -> UnifiedMessage { - UnifiedMessage::new( - "test", + fn scanned_message() -> ParsedMessage { + ParsedMessage::new( "gpt-5", "openai", "session", @@ -573,8 +582,7 @@ mod tests { &unit.path, unit.parser_version, fingerprint.clone(), - vec![UnifiedMessage::new( - "pi", + vec![ParsedMessage::new( "gpt-5.5", "openai", session_id, @@ -1097,6 +1105,7 @@ mod tests { let input_path = input_dir.path().join("session.jsonl"); std::fs::write(&input_path, PI_INPUT).unwrap(); let unit = pi_unit(&input_path); + let expected_client = unit.client; let fingerprint = seed_disk_cache(cache_dir.path(), &unit, "stale-cache-session"); message_cache::truncate_shard_after_header_for_test( cache_dir.path(), @@ -1130,6 +1139,7 @@ mod tests { ); let repaired = fold_planned_unit(parsed, &mut cache); assert_eq!(repaired.len(), 1); + assert_eq!(repaired[0].client.as_ref(), expected_client.as_str()); assert_eq!(repaired[0].session_id.as_ref(), "input-session"); assert_eq!(repaired[0].tokens.input, 17); @@ -1140,6 +1150,7 @@ mod tests { "successful recovery must atomically replace the failed shard", ); let warm_messages = fold_planned_unit(warm, &mut warm_cache); + assert_eq!(warm_messages[0].client.as_ref(), expected_client.as_str()); assert_eq!(warm_messages[0].session_id.as_ref(), "input-session"); assert_eq!( message_cache::get_input_read_stats(&input_path), @@ -1376,6 +1387,7 @@ mod tests { let input_path = input_dir.path().join("session.jsonl"); std::fs::write(&input_path, PI_INPUT).unwrap(); let unit = pi_unit(&input_path); + let expected_client = unit.client; seed_disk_cache(cache_dir.path(), &unit, "stale-cache-session"); let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); @@ -1418,6 +1430,7 @@ mod tests { ); let mut cold_cache = cold_cache; let cold_messages = fold_planned_unit(cold_parsed, &mut cold_cache); + assert_eq!(cold_messages[0].client.as_ref(), expected_client.as_str()); assert_eq!(cold_messages[0].session_id.as_ref(), "input-session"); } diff --git a/crates/tokscale-core/src/adapters/claude.rs b/crates/tokscale-core/src/adapters/claude.rs index a588ec739..0815b3c7c 100644 --- a/crates/tokscale-core/src/adapters/claude.rs +++ b/crates/tokscale-core/src/adapters/claude.rs @@ -274,11 +274,12 @@ fn fold_claude_units( ctx.input_cache.remove(&path, unit.parser_version); } let cache_write_outcome = cache_write_outcome?; - sink.extend_messages( + adapter_cache::emit_messages( + unit.client, messages .into_iter() - .filter(|message| crate::should_keep_deduped_message(seen_keys, message)) - .collect(), + .filter(|message| crate::should_keep_deduped_message(seen_keys, message)), + sink, ); if cache_write_outcome == adapter_cache::CacheWriteOutcome::NotPlanned && invalidate_cache { @@ -334,6 +335,9 @@ mod tests { CLAUDE_ADAPTER .fold(parsed, &mut FoldContext::new(cache, None), &mut messages) .unwrap(); + assert!(messages + .iter() + .all(|message| message.client.as_ref() == ClientId::Claude.as_str())); (messages, health) } @@ -350,9 +354,20 @@ mod tests { &mut messages, ) .unwrap(); + assert!(messages + .iter() + .all(|message| message.client.as_ref() == ClientId::Claude.as_str())); (messages, health) } + fn finalized(mut messages: Vec) -> Vec { + crate::finalize_token_priced_messages(&mut messages, None); + messages + .into_iter() + .map(|message| message.attribute(ClientId::Claude)) + .collect() + } + fn explore_parent(agent_id: &str) -> String { format!( r#"{{"type":"assistant","message":{{"content":[{{"type":"tool_use","id":"toolu_explore","name":"Agent","input":{{"subagent_type":"explore"}}}}]}}}} @@ -469,10 +484,14 @@ mod tests { .fold(parsed, &mut FoldContext::new(&mut cache, None), &mut actual) .unwrap(); - let expected = + let expected = finalized( sessions::claudecode::parse_claude_file_with_home(&session_path, Some(home.path())) .unwrap() - .messages; + .messages, + ); + assert!(actual + .iter() + .all(|message| message.client.as_ref() == ClientId::Claude.as_str())); assert_eq!(actual, expected); assert_eq!(actual.len(), 1); } @@ -523,6 +542,7 @@ mod tests { ) .unwrap(); assert_eq!(messages.len(), 1); + assert_eq!(messages[0].client.as_ref(), ClientId::Claude.as_str()); assert_eq!(messages[0].tokens.input, 10); assert!(cache .get_meta(&session_path, parser_version) diff --git a/crates/tokscale-core/src/adapters/codebuddy.rs b/crates/tokscale-core/src/adapters/codebuddy.rs index f6d3b3f37..174825e16 100644 --- a/crates/tokscale-core/src/adapters/codebuddy.rs +++ b/crates/tokscale-core/src/adapters/codebuddy.rs @@ -12,8 +12,7 @@ use crate::adapters::{ }; use crate::clients::ClientId; use crate::message_cache::{ParserId, ParserVersion}; -use crate::sessions; -use crate::UnifiedMessage; +use crate::sessions::{self, ParsedMessage}; const MIRROR_DEDUP_WINDOW_MS: i64 = 1000; const CODEBUDDY_JSONL_RECORD_REJECTION_REVISION: u32 = @@ -240,14 +239,14 @@ struct CodeBuddyDeduper { } impl CodeBuddyDeduper { - fn filter(&mut self, unit: &InputUnit, messages: Vec) -> Vec { + fn filter(&mut self, unit: &InputUnit, messages: Vec) -> Vec { messages .into_iter() .filter(|message| self.keep(unit, message)) .collect() } - fn keep(&mut self, unit: &InputUnit, message: &UnifiedMessage) -> bool { + fn keep(&mut self, unit: &InputUnit, message: &ParsedMessage) -> bool { if let Some(key) = message.dedup_key { return self.seen_keys.insert(key); } @@ -284,7 +283,7 @@ struct MirrorSignature { } impl MirrorSignature { - fn from_message(message: &UnifiedMessage) -> Self { + fn from_message(message: &ParsedMessage) -> Self { Self { session_id: message.session_id.clone(), model_id: message.model_id.clone(), @@ -328,8 +327,15 @@ mod tests { std::fs::write(path, content).unwrap(); } - fn finalized(mut messages: Vec) -> Vec { + fn finalized(mut messages: Vec) -> Vec { crate::finalize_token_priced_messages(&mut messages, None); + let messages: Vec<_> = messages + .into_iter() + .map(|message| message.attribute(ClientId::CodeBuddy)) + .collect(); + assert!(messages + .iter() + .all(|message| message.client.as_ref() == ClientId::CodeBuddy.as_str())); messages } @@ -340,6 +346,9 @@ mod tests { CODEBUDDY_ADAPTER .fold(parsed, &mut FoldContext::new(&mut cache, None), &mut sink) .unwrap(); + assert!(sink + .iter() + .all(|message| message.client.as_ref() == ClientId::CodeBuddy.as_str())); sink } diff --git a/crates/tokscale-core/src/adapters/codex.rs b/crates/tokscale-core/src/adapters/codex.rs index f622748ec..a4ddb5e44 100644 --- a/crates/tokscale-core/src/adapters/codex.rs +++ b/crates/tokscale-core/src/adapters/codex.rs @@ -12,7 +12,8 @@ use crate::adapters::{ }; use crate::clients::ClientId; use crate::input_health::InputStatus; -use crate::{message_cache, pricing, sessions, UnifiedMessage}; +use crate::sessions::ParsedMessage; +use crate::{message_cache, pricing, sessions}; pub(crate) struct CodexAdapter; @@ -21,7 +22,7 @@ pub(crate) struct CodexAppendInput { path: PathBuf, read_plan: message_cache::CacheReadPlan, parser_version: message_cache::ParserVersion, - tail_messages: Vec, + tail_messages: Vec, cache_write: Option>, } @@ -197,11 +198,12 @@ fn fold_codex_units( if finalization { finalize_codex_messages(&mut messages, ctx.pricing); } - sink.extend_messages( + adapter_cache::emit_messages( + client, messages .into_iter() - .filter(|message| crate::should_keep_deduped_message(seen, message)) - .collect(), + .filter(|message| crate::should_keep_deduped_message(seen, message)), + sink, ); } Ok(()) @@ -229,7 +231,7 @@ fn write_codex_cache_and_apply_recovery( path: &Path, parser_version: message_cache::ParserVersion, cache_write: Option>, - messages: &[UnifiedMessage], + messages: &[ParsedMessage], invalidate_cache: bool, recovery_requires_removal: bool, ctx: &mut FoldContext<'_>, @@ -243,7 +245,7 @@ fn write_codex_cache_and_apply_recovery( } struct CodexResolvedMessages { - messages: Vec, + messages: Vec, cache_write: Option>, finalization: bool, recovery_requires_removal: bool, @@ -325,7 +327,7 @@ fn parse_full_log_input( } fn finalize_codex_messages( - messages: &mut Vec, + messages: &mut Vec, pricing: Option<&pricing::PricingService>, ) { crate::finalize_token_priced_messages(messages, pricing); @@ -703,6 +705,7 @@ mod tests { use super::*; use crate::message_cache; use crate::pricing::{ModelPricing, PricingService}; + use crate::UnifiedMessage; const FIRST_CODEX_ENTRY: &str = concat!( r#"{"timestamp":"2026-04-27T09:59:59Z","type":"turn_context","payload":{"model":"gpt-5.4"}}"#, @@ -828,6 +831,7 @@ mod tests { &mut sink, ) .expect("valid Codex fixture must fold"); + assert_codex_attribution(&sink); sink } @@ -877,6 +881,7 @@ mod tests { CODEX_ADAPTER .fold_batches(&mut batches, &mut FoldContext::new(cache, None), &mut sink) .unwrap(); + assert_codex_attribution(&sink); sink } @@ -888,10 +893,11 @@ mod tests { CODEX_ADAPTER .fold(parsed, &mut FoldContext::new(cache, None), &mut sink) .expect("valid Codex fixture must fold"); + assert_codex_attribution(&sink); sink } - fn parser_messages(path: &Path) -> Vec { + fn parser_messages(path: &Path) -> Vec { let mut messages = sessions::codex::parse_codex_file(path).unwrap(); for message in &mut messages { message.refresh_derived_fields(); @@ -899,6 +905,25 @@ mod tests { messages } + fn assert_codex_attribution(messages: &[UnifiedMessage]) { + assert!( + messages + .iter() + .all(|message| message.client.as_ref() == ClientId::Codex.as_str()), + "Codex fold output must be attributed from its InputUnit client" + ); + } + + fn assert_output_matches_parser(actual: &[UnifiedMessage], expected: &[ParsedMessage]) { + assert_codex_attribution(actual); + let attributed_expected: Vec<_> = expected + .iter() + .cloned() + .map(|message| message.attribute(ClientId::Codex)) + .collect(); + assert_eq!(actual, attributed_expected); + } + fn assert_cached_raw_messages_match_parser(cache_home: &Path, path: &Path) { let expected = sessions::codex::parse_codex_file_incremental( path, @@ -1026,7 +1051,7 @@ mod tests { ); let expected = parser_messages(&path); - assert_eq!(actual, expected); + assert_output_matches_parser(&actual, &expected); assert!(cache .get_meta( &path, @@ -1099,6 +1124,7 @@ mod tests { let mut messages = Vec::new(); let mut ctx = FoldContext::new(&mut cache, None); CODEX_ADAPTER.fold(parsed, &mut ctx, &mut messages).unwrap(); + assert_codex_attribution(&messages); assert_eq!(messages.len(), 2); assert_eq!(messages[0].model_id.as_ref(), "gpt-5.4"); @@ -1142,6 +1168,7 @@ mod tests { let mut messages = Vec::new(); let mut ctx = FoldContext::new(&mut cache, None); CODEX_ADAPTER.fold(parsed, &mut ctx, &mut messages).unwrap(); + assert_codex_attribution(&messages); assert_eq!(messages.len(), 1); assert_eq!(messages[0].tokens.input, 8); @@ -1187,6 +1214,7 @@ mod tests { let mut messages = Vec::new(); let mut ctx = FoldContext::new(&mut cache, None); CODEX_ADAPTER.fold(parsed, &mut ctx, &mut messages).unwrap(); + assert_codex_attribution(&messages); assert_eq!(messages.len(), 1); assert_eq!(messages[0].model_id.as_ref(), "gpt-5.4"); @@ -1313,6 +1341,7 @@ mod tests { let mut sink = Vec::new(); let mut ctx = FoldContext::new(&mut warm_cache, None); CODEX_ADAPTER.fold(vec![hit], &mut ctx, &mut sink).unwrap(); + assert_codex_attribution(&sink); assert!(!sink.is_empty()); assert_eq!(ctx.health.rejected_records(), 1); @@ -1360,6 +1389,7 @@ mod tests { CODEX_ADAPTER .fold(parsed, &mut append_ctx, &mut appended_messages) .unwrap(); + assert_codex_attribution(&appended_messages); assert_eq!(append_ctx.health.rejected_records(), 1); append_ctx.input_cache.save_if_dirty().unwrap(); @@ -1373,6 +1403,7 @@ mod tests { CODEX_ADAPTER .fold(vec![hit], &mut warm_ctx, &mut warm_messages) .unwrap(); + assert_codex_attribution(&warm_messages); assert_eq!(warm_messages, appended_messages); assert_eq!(warm_ctx.health.rejected_records(), 1); @@ -1405,6 +1436,7 @@ mod tests { let mut messages = Vec::new(); let mut ctx = FoldContext::new(&mut append_cache, None); CODEX_ADAPTER.fold(parsed, &mut ctx, &mut messages).unwrap(); + assert_codex_attribution(&messages); assert_eq!(messages.len(), 2); assert_eq!(ctx.health.rejected_records(), 1); @@ -1450,6 +1482,7 @@ mod tests { let mut messages = Vec::new(); let mut ctx = FoldContext::new(&mut append_cache, None); CODEX_ADAPTER.fold(parsed, &mut ctx, &mut messages).unwrap(); + assert_codex_attribution(&messages); assert_eq!(messages.len(), 1); assert_eq!(ctx.health.partial_inputs(), 1); @@ -1495,6 +1528,7 @@ mod tests { let mut messages = Vec::new(); let mut ctx = FoldContext::new(&mut append_cache, None); CODEX_ADAPTER.fold(parsed, &mut ctx, &mut messages).unwrap(); + assert_codex_attribution(&messages); assert_eq!(messages.len(), 1); assert_eq!(messages[0].model_id.as_ref(), "gpt-5.4"); @@ -1562,6 +1596,7 @@ mod tests { CODEX_ADAPTER .fold(vec![planned], &mut ctx, &mut sink) .expect("an interrupted Codex recovery scan must stay inside its input domain"); + assert_codex_attribution(&sink); assert!(sink.is_empty()); assert_eq!(ctx.health.partial_inputs(), 1); assert_eq!(ctx.health.rejected_records(), 1); @@ -1607,6 +1642,7 @@ mod tests { CODEX_ADAPTER .fold(vec![planned], &mut ctx, &mut sink) .expect("an interrupted Codex recovery scan must stay inside its input domain"); + assert_codex_attribution(&sink); assert!(sink.is_empty()); assert_eq!(ctx.health.partial_inputs(), 1); assert_eq!(ctx.health.rejected_records(), 1); @@ -1805,7 +1841,7 @@ mod tests { "Codex append must verify the prefix and hash the tail in one pass" ); let expected = parser_messages(&path); - assert_eq!(actual, expected); + assert_output_matches_parser(&actual, &expected); assert_cached_raw_messages_match_parser(cache_home.path(), &path); } @@ -1922,16 +1958,16 @@ mod tests { )); let messages_b = fold_parsed(parsed_b, &mut cache_b); - assert_eq!(messages_b, expected); + assert_output_matches_parser(&messages_b, &expected); cache_b.save_if_dirty().unwrap(); let messages_a = fold_parsed(parsed_a, &mut cache_a); - assert_eq!(messages_a, expected); + assert_output_matches_parser(&messages_a, &expected); cache_a.save_if_dirty().unwrap(); let mut warm_cache = message_cache::InputMessageCache::load().unwrap(); let warm_messages = parse_and_fold(vec![codex_unit(&path)], &mut warm_cache); - assert_eq!(warm_messages, expected); + assert_output_matches_parser(&warm_messages, &expected); } #[test] @@ -1970,12 +2006,12 @@ mod tests { remover.save_if_dirty().unwrap(); let messages = fold_parsed(parsed, &mut cache); - assert_eq!(messages, expected); + assert_output_matches_parser(&messages, &expected); cache.save_if_dirty().unwrap(); assert_cached_raw_messages_match_parser(&cache_home.path().join("cache"), &path); let mut warm_cache = message_cache::InputMessageCache::load().unwrap(); let warm_messages = parse_and_fold(vec![codex_unit(&path)], &mut warm_cache); - assert_eq!(warm_messages, expected); + assert_output_matches_parser(&warm_messages, &expected); } } diff --git a/crates/tokscale-core/src/adapters/file.rs b/crates/tokscale-core/src/adapters/file.rs index b433705c9..546bfff56 100644 --- a/crates/tokscale-core/src/adapters/file.rs +++ b/crates/tokscale-core/src/adapters/file.rs @@ -12,9 +12,9 @@ use crate::adapters::{ use crate::clients::ClientId; use crate::input_health::ScannedInput; use crate::message_cache::{ParserId, ParserVersion, RelatedInputFailurePolicy}; +use crate::sessions; use crate::sessions::error::SessionParseResult; -use crate::sessions::WorkspaceMetadata; -use crate::{sessions, UnifiedMessage}; +use crate::sessions::{ParsedMessage, WorkspaceMetadata}; const GROK_TOTAL_ONLY_IMPUTATION_REVISION: u32 = MODEL_ID_CANONICALIZATION_REVISION + 1; const MUX_STABLE_DEDUP_REVISION: u32 = MODEL_ID_CANONICALIZATION_REVISION + 1; @@ -42,7 +42,7 @@ pub(crate) struct CachedFileAdapter { fingerprint_policy: FingerprintPolicy, dependency_failure_policy: RelatedInputFailurePolicy, dependency_path: Option Option>, - workspace_enrichment: Option, + workspace_enrichment: Option, parse: fn(&Path) -> SessionParseResult, } @@ -121,7 +121,7 @@ impl CachedFileAdapter { } } - const fn with_workspace_enrichment(mut self, enrich: fn(&Path, &mut [UnifiedMessage])) -> Self { + const fn with_workspace_enrichment(mut self, enrich: fn(&Path, &mut [ParsedMessage])) -> Self { self.workspace_enrichment = Some(enrich); self } @@ -195,7 +195,7 @@ impl LocalInputAdapter for CachedFileAdapter { } } -fn apply_workspace(messages: &mut [UnifiedMessage], workspace: Option) { +fn apply_workspace(messages: &mut [ParsedMessage], workspace: Option) { let Some(workspace) = workspace else { return; }; @@ -204,16 +204,16 @@ fn apply_workspace(messages: &mut [UnifiedMessage], workspace: Option) -> Vec { + fn finalized(client: ClientId, mut messages: Vec) -> Vec { crate::finalize_token_priced_messages(&mut messages, None); + let messages: Vec<_> = messages + .into_iter() + .map(|message| message.attribute(client)) + .collect(); + assert!(messages + .iter() + .all(|message| message.client.as_ref() == client.as_str())); messages } @@ -404,6 +412,9 @@ not-json adapter .fold(parsed, &mut FoldContext::new(cache, None), &mut sink) .unwrap(); + assert!(sink + .iter() + .all(|message| message.client.as_ref() == adapter.client().as_str())); sink } @@ -445,7 +456,10 @@ not-json let mut cache = message_cache::InputMessageCache::default(); let actual = fold_with_adapter(&_ADAPTER, units, &mut cache); - let expected = finalized(sessions::amp::parse_amp_file(&path).unwrap().messages); + let expected = finalized( + ClientId::Amp, + sessions::amp::parse_amp_file(&path).unwrap().messages, + ); assert_eq!(actual, expected); } @@ -1244,7 +1258,10 @@ model = "claude-sonnet-4" let mut cache = message_cache::InputMessageCache::default(); let actual = fold_with_adapter(&ZCODE_ADAPTER, units, &mut cache); - let expected = finalized(sessions::zcode::parse_zcode_file(&path).unwrap().messages); + let expected = finalized( + ClientId::Zcode, + sessions::zcode::parse_zcode_file(&path).unwrap().messages, + ); assert_eq!(actual, expected); } diff --git a/crates/tokscale-core/src/adapters/goose.rs b/crates/tokscale-core/src/adapters/goose.rs index 3d0d2c18c..5256a5001 100644 --- a/crates/tokscale-core/src/adapters/goose.rs +++ b/crates/tokscale-core/src/adapters/goose.rs @@ -6,7 +6,7 @@ use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ AdapterScanContext, FoldContext, InputDiscoveryError, InputUnit, LocalInputAdapter, - MessageSink, ParseContext, ParsedUnit, UnitMessagePayload, + MessageSink, ParseContext, ParsedUnit, }; use crate::clients::ClientId; use crate::message_cache::{ParserId, ParserVersion}; @@ -55,13 +55,7 @@ impl LocalInputAdapter for GooseAdapter { ctx: &mut FoldContext<'_>, sink: &mut dyn MessageSink, ) -> Result<(), crate::adapters::InputPipelineError> { - for unit in parsed { - ctx.health.record(unit.input_health()); - if let UnitMessagePayload::Fresh(messages) = unit.messages { - sink.extend_messages(messages); - } - } - Ok(()) + adapter_cache::fold_units(parsed, ctx, sink) } } diff --git a/crates/tokscale-core/src/adapters/hermes.rs b/crates/tokscale-core/src/adapters/hermes.rs index 1d076afa1..e51aea913 100644 --- a/crates/tokscale-core/src/adapters/hermes.rs +++ b/crates/tokscale-core/src/adapters/hermes.rs @@ -6,7 +6,7 @@ use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ AdapterScanContext, FingerprintPolicy, FoldContext, InputDiscoveryError, InputUnit, - LocalInputAdapter, MessageSink, ParseContext, ParsedBatchInput, ParsedUnit, UnitMessagePayload, + LocalInputAdapter, MessageSink, ParseContext, ParsedBatchInput, ParsedUnit, }; use crate::clients::ClientId; use crate::message_cache::{ParserId, ParserVersion}; @@ -73,8 +73,7 @@ impl LocalInputAdapter for HermesAdapter { sink: &mut dyn MessageSink, ) -> Result<(), crate::adapters::InputPipelineError> { let mut seen = HashSet::new(); - fold_hermes_units(parsed, ctx, sink, &mut seen); - Ok(()) + fold_hermes_units(parsed, ctx, sink, &mut seen) } fn fold_batches( @@ -85,7 +84,7 @@ impl LocalInputAdapter for HermesAdapter { ) -> Result<(), crate::adapters::InputPipelineError> { let mut seen = HashSet::new(); while let Some(parsed) = batches.next(ctx)? { - fold_hermes_units(parsed, ctx, sink, &mut seen); + fold_hermes_units(parsed, ctx, sink, &mut seen)?; } Ok(()) } @@ -96,18 +95,13 @@ fn fold_hermes_units( ctx: &mut FoldContext<'_>, sink: &mut dyn MessageSink, seen: &mut HashSet, -) { - for unit in parsed { - ctx.health.record(unit.input_health()); - if let UnitMessagePayload::Fresh(messages) = unit.messages { - sink.extend_messages( - messages - .into_iter() - .filter(|message| crate::should_keep_deduped_message(seen, message)) - .collect(), - ); - } - } +) -> Result<(), crate::adapters::InputPipelineError> { + adapter_cache::fold_units_with_filter(parsed, ctx, sink, |_, messages| { + messages + .into_iter() + .filter(|message| crate::should_keep_deduped_message(seen, message)) + .collect() + }) } pub(crate) static HERMES_ADAPTER: HermesAdapter = HermesAdapter; @@ -129,8 +123,7 @@ mod tests { &path, unit.parser_version, unit.input_policy().fingerprint().unwrap(), - vec![crate::UnifiedMessage::new( - "hermes", + vec![crate::sessions::ParsedMessage::new( "model", "provider", "cached-session", diff --git a/crates/tokscale-core/src/adapters/junie.rs b/crates/tokscale-core/src/adapters/junie.rs index 368b4c336..e186c97a8 100644 --- a/crates/tokscale-core/src/adapters/junie.rs +++ b/crates/tokscale-core/src/adapters/junie.rs @@ -113,15 +113,26 @@ mod tests { units: Vec, cache: &mut message_cache::InputMessageCache, pricing: Option<&PricingService>, - ) -> Vec { + ) -> Vec { let parsed = JUNIE_ADAPTER.parse_checked(units, &ParseContext { pricing }); let mut messages = Vec::new(); JUNIE_ADAPTER .fold(parsed, &mut FoldContext::new(cache, pricing), &mut messages) .unwrap(); + assert!(messages + .iter() + .all(|message| message.client.as_ref() == ClientId::Junie.as_str())); messages } + fn finalized(mut messages: Vec) -> Vec { + crate::finalize_token_priced_messages(&mut messages, None); + messages + .into_iter() + .map(|message| message.attribute(ClientId::Junie)) + .collect() + } + fn pricing_service() -> PricingService { let mut litellm_data = HashMap::new(); litellm_data.insert( @@ -168,8 +179,11 @@ mod tests { &mut cache, None, ); - let expected = sessions::junie::parse_junie_file(&path).unwrap().messages; + let expected = finalized(sessions::junie::parse_junie_file(&path).unwrap().messages); + assert!(actual + .iter() + .all(|message| message.client.as_ref() == ClientId::Junie.as_str())); assert_eq!(actual, expected); } @@ -199,6 +213,9 @@ mod tests { .fold(parsed, &mut FoldContext::new(&mut cache, None), &mut cached) .unwrap(); + assert!(cached + .iter() + .all(|message| message.client.as_ref() == ClientId::Junie.as_str())); assert_eq!(cached, fresh); restore_env_var("TOKSCALE_CONFIG_DIR", previous_config_dir); } diff --git a/crates/tokscale-core/src/adapters/kilo.rs b/crates/tokscale-core/src/adapters/kilo.rs index e54998494..374e54d57 100644 --- a/crates/tokscale-core/src/adapters/kilo.rs +++ b/crates/tokscale-core/src/adapters/kilo.rs @@ -4,7 +4,7 @@ use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ AdapterScanContext, FoldContext, InputDiscoveryError, InputUnit, LocalInputAdapter, - MessageSink, ParseContext, ParsedUnit, UnitMessagePayload, + MessageSink, ParseContext, ParsedUnit, }; use crate::clients::ClientId; use crate::message_cache::{ParserId, ParserVersion}; @@ -68,13 +68,7 @@ impl LocalInputAdapter for KiloAdapter { ctx: &mut FoldContext<'_>, sink: &mut dyn MessageSink, ) -> Result<(), crate::adapters::InputPipelineError> { - for unit in parsed { - ctx.health.record(unit.input_health()); - if let UnitMessagePayload::Fresh(messages) = unit.messages { - sink.extend_messages(messages); - } - } - Ok(()) + adapter_cache::fold_units(parsed, ctx, sink) } } diff --git a/crates/tokscale-core/src/adapters/kiro.rs b/crates/tokscale-core/src/adapters/kiro.rs index 3dcf39f97..984229193 100644 --- a/crates/tokscale-core/src/adapters/kiro.rs +++ b/crates/tokscale-core/src/adapters/kiro.rs @@ -490,6 +490,9 @@ mod tests { KIRO_ADAPTER.fold(parsed, &mut ctx, &mut messages).unwrap(); assert_eq!(messages.len(), 2); + assert!(messages + .iter() + .all(|message| message.client.as_ref() == ClientId::Kiro.as_str())); assert_eq!(ctx.health.rejected_records(), 1); assert_eq!(ctx.health.failed_inputs(), 0); assert_eq!(ctx.health.partial_inputs(), 0); @@ -557,6 +560,7 @@ mod tests { KIRO_ADAPTER.fold(parsed, &mut ctx, &mut messages).unwrap(); assert_eq!(messages.len(), 1); + assert_eq!(messages[0].client.as_ref(), ClientId::Kiro.as_str()); assert_eq!(messages[0].tokens.input, 13); assert_eq!(messages[0].tokens.output, 5); assert_eq!(ctx.health.partial_inputs(), 1); @@ -576,8 +580,7 @@ mod tests { &path, unit.parser_version, unit.input_policy().fingerprint().unwrap(), - vec![crate::UnifiedMessage::new( - "kiro", + vec![crate::sessions::ParsedMessage::new( "cached-model", "cached-provider", "cached-session", diff --git a/crates/tokscale-core/src/adapters/mod.rs b/crates/tokscale-core/src/adapters/mod.rs index 526caf0e6..1d54be280 100644 --- a/crates/tokscale-core/src/adapters/mod.rs +++ b/crates/tokscale-core/src/adapters/mod.rs @@ -30,7 +30,7 @@ use rayon::prelude::*; use crate::clients::ClientId; use crate::input_health::{DataHealth, InputFailure, InputHealth, InputStatus, RejectionSummary}; use crate::message_cache::{InputFileIdentity, ParserId, ParserRevision, ParserVersion}; -use crate::{message_cache, pricing, scanner, UnifiedMessage}; +use crate::{message_cache, pricing, scanner, sessions::ParsedMessage, UnifiedMessage}; pub(crate) use error::{ InputDiscoveryError, InputParseError, InputPipelineError, InputPlanningError, @@ -117,12 +117,6 @@ impl<'a> FoldContext<'a> { pub(crate) trait MessageSink { fn push_message(&mut self, message: UnifiedMessage); - - fn extend_messages(&mut self, messages: Vec) { - for message in messages { - self.push_message(message); - } - } } impl MessageSink for Vec { @@ -133,7 +127,7 @@ impl MessageSink for Vec { #[derive(Debug, Clone)] pub(crate) struct InputUnit { - pub client: ClientId, + client: ClientId, pub path: PathBuf, pub fingerprint_policy: FingerprintPolicy, pub meta: InputUnitMeta, @@ -145,6 +139,11 @@ pub(crate) struct InputUnit { } impl InputUnit { + /// Adapter-issued source provenance for this input. + pub(crate) const fn client(&self) -> ClientId { + self.client + } + pub(crate) fn plain_file(client: ClientId, path: PathBuf) -> Self { Self { client, @@ -582,8 +581,8 @@ pub(crate) enum FingerprintPolicy { #[derive(Debug)] pub(crate) enum UnitMessagePayload { - Fresh(Vec), - CodexFresh(Vec), + Fresh(Vec), + CodexFresh(Vec), CacheHit(message_cache::CacheReadPlan), CodexCacheHit(message_cache::CacheReadPlan), CodexAppend(Box), @@ -728,6 +727,22 @@ pub(crate) struct ConfirmedAdapterInputs { pub present_files: Vec<(InputFileIdentity, u64)>, } +pub(crate) fn validate_discovered_unit_sources( + adapter: &dyn LocalInputAdapter, + units: &[InputUnit], +) -> Result<(), InputPipelineError> { + let expected = adapter.client(); + if let Some(unit) = units.iter().find(|unit| unit.client != expected) { + return Err(InputPipelineError::contract(format!( + "local input adapter `{}` discovered `{}` input `{}`", + expected.as_str(), + unit.client.as_str(), + unit.path.display() + ))); + } + Ok(()) +} + pub(crate) struct ParsedBatchInput<'a> { adapter: &'a dyn LocalInputAdapter, units: Option>, @@ -933,6 +948,10 @@ pub(crate) fn run_prepared_local_input_adapters( sink: &mut dyn MessageSink, health: &mut DataHealth, ) -> Result, InputPipelineError> { + for group in &prepared { + validate_discovered_unit_sources(group.adapter, &group.units)?; + } + let mut confirmed = Vec::with_capacity(prepared.len()); for PreparedAdapterInputs { adapter, units } in prepared { let mut batches = ParsedBatchInput::new(adapter, units); @@ -1001,6 +1020,26 @@ mod tests { ); } + #[test] + fn adapter_rejects_input_unit_with_foreign_source_identity() { + let prepared = PreparedAdapterInputs { + adapter: adapter_for(ClientId::Amp).unwrap(), + units: vec![InputUnit::plain_file( + ClientId::Codex, + PathBuf::from("foreign-input.jsonl"), + )], + }; + + let error = match validate_discovered_unit_sources(prepared.adapter, &prepared.units) { + Ok(_) => panic!("foreign source identity must be rejected"), + Err(error) => error, + }; + + assert!(error + .to_string() + .contains("local input adapter `amp` discovered `codex` input `foreign-input.jsonl`")); + } + impl LocalInputAdapter for RecordingAdapter { fn client(&self) -> ClientId { ClientId::Amp @@ -1019,8 +1058,7 @@ mod tests { .into_iter() .enumerate() .map(|(index, unit)| { - let message = UnifiedMessage::new( - "amp", + let message = ParsedMessage::new( "model", "provider", unit.path.to_string_lossy(), @@ -1068,8 +1106,7 @@ mod tests { for unit in units { let session_id: Arc = Arc::from(unit.path.to_string_lossy().into_owned()); *previous = Some(Arc::downgrade(&session_id)); - let mut message = UnifiedMessage::new( - "amp", + let mut message = ParsedMessage::new( "model", "provider", "placeholder", @@ -1115,8 +1152,7 @@ mod tests { units .into_iter() .map(|unit| { - let message = UnifiedMessage::new( - "amp", + let message = ParsedMessage::new( "model", "provider", unit.path.file_name().unwrap().to_string_lossy(), @@ -1266,8 +1302,7 @@ mod tests { &unit.path, unit.parser_version, unit.input_policy().fingerprint().unwrap(), - vec![UnifiedMessage::new( - "amp", + vec![ParsedMessage::new( "model", "provider", index.to_string(), @@ -1329,8 +1364,7 @@ mod tests { &unit.path, unit.parser_version, unit.input_policy().fingerprint().unwrap(), - vec![UnifiedMessage::new( - "amp", + vec![ParsedMessage::new( "model", "provider", index.to_string(), @@ -1384,8 +1418,7 @@ mod tests { &path, unit.parser_version, unit.input_policy().fingerprint().unwrap(), - vec![UnifiedMessage::new( - "amp", + vec![ParsedMessage::new( "model", "provider", "cached-session", @@ -1426,8 +1459,7 @@ mod tests { &path, unit.parser_version, unit.input_policy().fingerprint().unwrap(), - vec![UnifiedMessage::new( - "amp", + vec![ParsedMessage::new( "model", "provider", "cached-session", diff --git a/crates/tokscale-core/src/adapters/omp.rs b/crates/tokscale-core/src/adapters/omp.rs index c65ba302d..3a01a9a8d 100644 --- a/crates/tokscale-core/src/adapters/omp.rs +++ b/crates/tokscale-core/src/adapters/omp.rs @@ -461,7 +461,7 @@ fn fold_omp_cache_hits( status, rejections, }); - sink.extend_messages(messages); + adapter_cache::emit_messages(unit.client, messages, sink); } Err(failure) => { if !failure.can_reparse_input() { @@ -584,10 +584,12 @@ mod tests { std::fs::write(path, content).unwrap(); } - fn refresh(messages: &mut [crate::UnifiedMessage]) { - for message in messages { - message.refresh_derived_fields(); - } + fn finalized(mut messages: Vec) -> Vec { + crate::finalize_token_priced_messages(&mut messages, None); + messages + .into_iter() + .map(|message| message.attribute(ClientId::Omp)) + .collect() } fn fold_with_omp_adapter( @@ -599,6 +601,9 @@ mod tests { OMP_ADAPTER .fold(parsed, &mut FoldContext::new(cache, None), &mut sink) .unwrap(); + assert!(sink + .iter() + .all(|message| message.client.as_ref() == ClientId::Omp.as_str())); sink } @@ -612,6 +617,9 @@ mod tests { OMP_ADAPTER .fold_batches(&mut batches, &mut ctx, &mut sink) .unwrap(); + assert!(sink + .iter() + .all(|message| message.client.as_ref() == ClientId::Omp.as_str())); (sink, std::mem::take(&mut ctx.health)) } @@ -625,8 +633,7 @@ mod tests { &unit.path, unit.parser_version, unit.input_policy().fingerprint().unwrap(), - vec![crate::UnifiedMessage::new( - "omp", + vec![crate::sessions::ParsedMessage::new( "gpt-5.5", "openai", session_id, @@ -731,11 +738,11 @@ mod tests { let miss_paths = vec![child_path.clone()]; let parent_index = sessions::pi::build_omp_parent_task_agent_index(&miss_paths); - let mut expected = + let expected = finalized( sessions::pi::parse_omp_file_with_parent_task_agent_index(&child_path, &parent_index) .unwrap() - .messages; - refresh(&mut expected); + .messages, + ); assert_eq!(actual, expected); assert_eq!(actual[0].agent.as_deref(), Some("OMP Reviewer")); @@ -936,8 +943,7 @@ mod tests { &cached_path, parser_version, cached_unit.input_policy().fingerprint().unwrap(), - vec![crate::UnifiedMessage::new( - "omp", + vec![crate::sessions::ParsedMessage::new( "gpt-5.5", "openai", "cached-session", @@ -1510,8 +1516,7 @@ mod tests { &path, parser_version, unit.input_policy().fingerprint().unwrap(), - vec![crate::UnifiedMessage::new( - "omp", + vec![crate::sessions::ParsedMessage::new( "gpt-5.5", "openai", "cached-session", diff --git a/crates/tokscale-core/src/adapters/opencode.rs b/crates/tokscale-core/src/adapters/opencode.rs index 1efc18478..7b73b1e6e 100644 --- a/crates/tokscale-core/src/adapters/opencode.rs +++ b/crates/tokscale-core/src/adapters/opencode.rs @@ -136,11 +136,12 @@ fn fold_opencode_unit( ctx.input_cache.remove(&path, unit.parser_version); } let cache_write_outcome = cache_write_outcome?; - sink.extend_messages( + adapter_cache::emit_messages( + unit.client, messages .into_iter() - .filter(|message| message.dedup_key.is_none_or(|key| seen.insert(key))) - .collect(), + .filter(|message| message.dedup_key.is_none_or(|key| seen.insert(key))), + sink, ); if cache_write_outcome == adapter_cache::CacheWriteOutcome::NotPlanned && invalidate_cache { @@ -156,7 +157,8 @@ mod tests { use super::*; use crate::adapters::{FoldContext, UnitMessagePayload}; use crate::message_cache; - use crate::{TokenBreakdown, UnifiedMessage}; + use crate::sessions::ParsedMessage; + use crate::TokenBreakdown; use rusqlite::Connection; use std::path::Path; @@ -229,8 +231,7 @@ mod tests { .into_iter() .enumerate() .map(|(index, name)| { - let message = UnifiedMessage::new_with_dedup( - "opencode", + let message = ParsedMessage::new_with_dedup( "gpt-5.5", "openai", format!("session-{index}"), @@ -259,6 +260,7 @@ mod tests { .fold(parsed, &mut FoldContext::new(&mut cache, None), &mut sink) .unwrap(); assert_eq!(sink.len(), 1); + assert_eq!(sink[0].client.as_ref(), ClientId::OpenCode.as_str()); assert_eq!(sink[0].session_id.as_ref(), "session-0"); } @@ -297,6 +299,7 @@ mod tests { sink }); assert_eq!(messages.len(), 1); + assert_eq!(messages[0].client.as_ref(), ClientId::OpenCode.as_str()); } #[test] @@ -321,8 +324,7 @@ mod tests { crate::adapters::MODEL_ID_CANONICALIZATION_REVISION, ), fingerprint, - vec![UnifiedMessage::new( - "opencode", + vec![ParsedMessage::new( "stale-model", "stale-provider", "stale-session", diff --git a/crates/tokscale-core/src/adapters/pi.rs b/crates/tokscale-core/src/adapters/pi.rs index 6117cdb7d..a0e8a753d 100644 --- a/crates/tokscale-core/src/adapters/pi.rs +++ b/crates/tokscale-core/src/adapters/pi.rs @@ -109,10 +109,12 @@ mod tests { } } - fn refresh(messages: &mut [crate::UnifiedMessage]) { - for message in messages { - message.refresh_derived_fields(); - } + fn finalized(mut messages: Vec) -> Vec { + crate::finalize_token_priced_messages(&mut messages, None); + messages + .into_iter() + .map(|message| message.attribute(ClientId::Pi)) + .collect() } fn fold_with_adapter( @@ -125,6 +127,9 @@ mod tests { adapter .fold(parsed, &mut FoldContext::new(cache, None), &mut sink) .unwrap(); + assert!(sink + .iter() + .all(|message| message.client.as_ref() == adapter.client().as_str())); sink } @@ -169,8 +174,7 @@ mod tests { let mut cache = message_cache::InputMessageCache::default(); let actual = fold_with_adapter(&PI_ADAPTER, units, &mut cache); - let mut expected = sessions::pi::parse_pi_file(&path).unwrap().messages; - refresh(&mut expected); + let expected = finalized(sessions::pi::parse_pi_file(&path).unwrap().messages); assert_eq!(actual, expected); } @@ -288,6 +292,9 @@ mod tests { .fold(parsed, &mut FoldContext::new(&mut cache, None), &mut second) .unwrap(); + assert!(second + .iter() + .all(|message| message.client.as_ref() == ClientId::Pi.as_str())); assert_eq!(second, first); restore_env_var("TOKSCALE_CONFIG_DIR", previous_config_dir); } diff --git a/crates/tokscale-core/src/adapters/warp.rs b/crates/tokscale-core/src/adapters/warp.rs index d28d7337b..34985f20f 100644 --- a/crates/tokscale-core/src/adapters/warp.rs +++ b/crates/tokscale-core/src/adapters/warp.rs @@ -4,7 +4,7 @@ use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ AdapterScanContext, FingerprintPolicy, FoldContext, InputDiscoveryError, InputUnit, - LocalInputAdapter, MessageSink, ParseContext, ParsedUnit, UnitMessagePayload, + LocalInputAdapter, MessageSink, ParseContext, ParsedUnit, }; use crate::clients::ClientId; use crate::local_clients; @@ -70,13 +70,7 @@ impl LocalInputAdapter for WarpAdapter { ctx: &mut FoldContext<'_>, sink: &mut dyn MessageSink, ) -> Result<(), crate::adapters::InputPipelineError> { - for unit in parsed { - ctx.health.record(unit.input_health()); - if let UnitMessagePayload::Fresh(messages) = unit.messages { - sink.extend_messages(messages); - } - } - Ok(()) + adapter_cache::fold_units(parsed, ctx, sink) } } diff --git a/crates/tokscale-core/src/adapters/zed.rs b/crates/tokscale-core/src/adapters/zed.rs index 01c20a424..592405b3f 100644 --- a/crates/tokscale-core/src/adapters/zed.rs +++ b/crates/tokscale-core/src/adapters/zed.rs @@ -175,9 +175,12 @@ mod tests { .unwrap(); } - fn finalized(mut messages: Vec) -> Vec { + fn finalized(mut messages: Vec) -> Vec { crate::finalize_token_priced_messages(&mut messages, None); messages + .into_iter() + .map(|message| message.attribute(ClientId::Zed)) + .collect() } #[test] @@ -239,6 +242,7 @@ mod tests { ZED_ADAPTER.fold(parsed, &mut fold_ctx, &mut sink).unwrap(); assert_eq!(sink.len(), 1); + assert_eq!(sink[0].client.as_ref(), ClientId::Zed.as_str()); assert_eq!(sink[0].session_id.as_ref(), "zed-thread-good"); assert_eq!(fold_ctx.health.rejected_records(), 1); assert_eq!(fold_ctx.health.failed_inputs(), 0); @@ -318,6 +322,9 @@ mod tests { .unwrap(); let expected = finalized(sessions::zed::parse_zed_sqlite(&db_path).unwrap().messages); + assert!(actual + .iter() + .all(|message| message.client.as_ref() == ClientId::Zed.as_str())); assert_eq!(actual, expected); } } diff --git a/crates/tokscale-core/src/input_health.rs b/crates/tokscale-core/src/input_health.rs index e661a83a6..22629c1ac 100644 --- a/crates/tokscale-core/src/input_health.rs +++ b/crates/tokscale-core/src/input_health.rs @@ -13,7 +13,7 @@ use serde::{Deserialize, Serialize}; use crate::clients::ClientId; use crate::sessions::error::SessionParseError; -use crate::UnifiedMessage; +use crate::sessions::ParsedMessage; /// Why a single record inside an otherwise readable input was rejected. /// @@ -342,8 +342,9 @@ pub struct HealthReport { pub rejected_records: u64, pub partial_inputs: usize, pub failed_inputs: usize, - /// Deduplicated on-disk size of scan inputs at the latest inventory - /// snapshot. Tokscale's own cache files are not included. + /// Sum of per-client input bytes at the latest inventory snapshot, + /// deduplicated within each client. Tokscale's own cache files are not + /// included. pub input_data_bytes: u64, #[serde(default)] pub issues: Vec, @@ -421,13 +422,13 @@ pub struct HealthIssueReport { /// so far. #[derive(Debug, Default)] pub struct ScannedInput { - pub messages: Vec, + pub messages: Vec, pub rejections: RejectionSummary, pub interrupted: Option, } impl ScannedInput { - pub fn complete(messages: Vec) -> Self { + pub fn complete(messages: Vec) -> Self { Self { messages, rejections: RejectionSummary::default(), diff --git a/crates/tokscale-core/src/lib.rs b/crates/tokscale-core/src/lib.rs index 14b552a21..f7629ed33 100644 --- a/crates/tokscale-core/src/lib.rs +++ b/crates/tokscale-core/src/lib.rs @@ -43,6 +43,8 @@ use std::sync::Arc; use sha2::{Digest, Sha256}; +use sessions::ParsedMessage; + /// Canonicalize a raw model string for callers that do not already hold a /// finalized `UnifiedMessage`. /// @@ -293,7 +295,7 @@ impl PreparedLocalInputs { let mut unavailable = Vec::new(); for group in &mut self.groups { group.units.retain_mut(|unit| { - let client = unit.client; + let client = unit.client(); let path = unit.path.clone(); match unit.refresh_prepared_snapshot_for_inventory_probe() { Ok(()) => true, @@ -326,9 +328,10 @@ impl PreparedLocalInputs { } fn input_data_bytes<'a>(units: impl IntoIterator) -> u64 { - let mut seen = HashSet::new(); + let mut seen_by_client = HashMap::>::new(); let mut total = 0_u64; for unit in units { + let seen = seen_by_client.entry(unit.client()).or_default(); let snapshot = unit .prepared_input_snapshot() .expect("prepared input unit must carry an inventory snapshot"); @@ -343,38 +346,35 @@ fn input_data_bytes<'a>(units: impl IntoIterator total } -fn confirmed_input_data_bytes( +fn confirmed_client_space( clients: &[String], groups: &[adapters::ConfirmedAdapterInputs], -) -> (BTreeMap, u64) { +) -> BTreeMap { let mut totals = clients .iter() .cloned() .map(|client| (client, 0_u64)) .collect::>(); - let mut seen_by_client = HashMap::>::new(); - let mut seen_globally = HashSet::new(); - let mut global_total = 0_u64; for group in groups { let client = group.client.as_str().to_string(); - let seen = seen_by_client.entry(client.clone()).or_default(); let total = totals.entry(client).or_default(); - for &(identity, size) in &group.present_files { - if seen.insert(identity) { - *total = total - .checked_add(size) - .expect("per-client input data size must fit in u64"); - } - if seen_globally.insert(identity) { - global_total = global_total - .checked_add(size) - .expect("input data size must fit in u64"); - } + for &(_, size) in &group.present_files { + *total = total + .checked_add(size) + .expect("per-client input data size must fit in u64"); } } - (totals, global_total) + totals +} + +fn client_space_total(client_space: &BTreeMap) -> u64 { + client_space.values().copied().fold(0_u64, |total, size| { + total + .checked_add(size) + .expect("input data size must fit in u64") + }) } #[derive(Debug, Clone, Default)] @@ -526,8 +526,8 @@ fn fold_prepared_local_inputs_with_pricing( .map(|confirmed| { let input_inventory_signature = confirmed_input_inventory_signature(&clients, &confirmed); - let (client_space, input_data_bytes) = confirmed_input_data_bytes(&clients, &confirmed); - health.set_input_data_bytes(input_data_bytes); + let client_space = confirmed_client_space(&clients, &confirmed); + health.set_input_data_bytes(client_space_total(&client_space)); FoldOutcome { input_inventory_signature, client_space, @@ -611,10 +611,12 @@ pub fn prepare_local_inputs( }); } }; + adapters::validate_discovered_unit_sources(adapter, &units) + .map_err(LocalReportError::operational)?; let units = units .into_iter() .filter_map(|unit| { - let client = unit.client; + let client = unit.client(); let path = unit.path.clone(); match unit.prepare_snapshot() { Ok(unit) => Some(unit), @@ -919,40 +921,99 @@ pub async fn get_usage_data( Ok(data) } -fn apply_token_pricing(message: &mut UnifiedMessage, pricing: Option<&pricing::PricingService>) { - message.cost = 0.0; +trait FinalizableMessage { + fn model_id(&self) -> &Arc; + fn set_model_id(&mut self, model_id: Arc); + fn provider_id(&self) -> &Arc; + fn set_provider_id(&mut self, provider_id: Arc); + fn tokens(&self) -> &TokenBreakdown; + fn tokens_mut(&mut self) -> &mut TokenBreakdown; + fn set_cost(&mut self, cost: f64); +} + +macro_rules! impl_finalizable_message { + ($message:ty) => { + impl FinalizableMessage for $message { + fn model_id(&self) -> &Arc { + &self.model_id + } + + fn set_model_id(&mut self, model_id: Arc) { + self.model_id = model_id; + } + + fn provider_id(&self) -> &Arc { + &self.provider_id + } + + fn set_provider_id(&mut self, provider_id: Arc) { + self.provider_id = provider_id; + } + + fn tokens(&self) -> &TokenBreakdown { + &self.tokens + } + + fn tokens_mut(&mut self) -> &mut TokenBreakdown { + &mut self.tokens + } + + fn set_cost(&mut self, cost: f64) { + self.cost = cost; + } + } + }; +} + +impl_finalizable_message!(ParsedMessage); +impl_finalizable_message!(UnifiedMessage); + +fn apply_token_pricing( + message: &mut M, + pricing: Option<&pricing::PricingService>, +) { + message.set_cost(0.0); let Some(pricing) = pricing else { return; }; let calculated_cost = pricing.calculate_cost_with_provider( - &message.model_id, - Some(message.provider_id.as_ref()), - &message.tokens, + message.model_id(), + Some(message.provider_id().as_ref()), + message.tokens(), ); if calculated_cost > 0.0 { - message.cost = calculated_cost; + message.set_cost(calculated_cost); + } +} + +fn refresh_derived_message_fields(message: &mut M) { + if let Some(provider) = provider_identity::provider_override_from_model_and_provider( + message.model_id(), + message.provider_id(), + ) { + message.set_provider_id(sessions::intern::intern(provider)); } } -fn canonicalize_message_provider(message: &mut UnifiedMessage) { +fn canonicalize_message_provider(message: &mut M) { let provider = - provider_identity::finalized_provider_id(&message.provider_id, &message.model_id); - message.provider_id = sessions::intern::intern(&provider); + provider_identity::finalized_provider_id(message.provider_id(), message.model_id()); + message.set_provider_id(sessions::intern::intern(&provider)); } -fn canonicalize_message_model( - message: &mut UnifiedMessage, +fn canonicalize_message_model( + message: &mut M, model_cache: &mut HashMap, Arc>, ) { - if let Some(canonical) = model_cache.get(&message.model_id) { - message.model_id = Arc::clone(canonical); + if let Some(canonical) = model_cache.get(message.model_id()) { + message.set_model_id(Arc::clone(canonical)); return; } - let raw = Arc::clone(&message.model_id); + let raw = Arc::clone(message.model_id()); let canonical = model_aliases::canonicalize_model_id(raw.as_ref()); let canonical = if canonical == raw.as_ref() { Arc::clone(&raw) @@ -961,21 +1022,21 @@ fn canonicalize_message_model( }; model_cache.insert(raw, Arc::clone(&canonical)); - message.model_id = canonical; + message.set_model_id(canonical); } -pub(crate) fn finalize_token_priced_messages( - messages: &mut Vec, +fn finalize_token_priced_messages( + messages: &mut Vec, pricing: Option<&pricing::PricingService>, ) { let mut model_cache = HashMap::new(); messages.retain_mut(|message| { - normalize_token_breakdown(&mut message.tokens); + normalize_token_breakdown(message.tokens_mut()); canonicalize_message_model(message, &mut model_cache); - message.refresh_derived_fields(); + refresh_derived_message_fields(message); canonicalize_message_provider(message); - if !has_positive_tokens(&message.tokens) { + if !has_positive_tokens(message.tokens()) { return false; } apply_token_pricing(message, pricing); @@ -1204,7 +1265,7 @@ pub async fn load_prepared_tui_bundle_with_diagnostics( }) } -fn should_keep_deduped_message(seen_keys: &mut HashSet, message: &UnifiedMessage) -> bool { +fn should_keep_deduped_message(seen_keys: &mut HashSet, message: &ParsedMessage) -> bool { message.dedup_key.is_none_or(|key| seen_keys.insert(key)) } diff --git a/crates/tokscale-core/src/lib_tests.rs b/crates/tokscale-core/src/lib_tests.rs index 1ef264a31..ebaf40b2c 100644 --- a/crates/tokscale-core/src/lib_tests.rs +++ b/crates/tokscale-core/src/lib_tests.rs @@ -334,6 +334,33 @@ fn test_streaming_tui_usage_matches_reference_aggregation() { assert_eq!(format!("{streaming:?}"), format!("{reference:?}")); } +#[test] +#[serial_test::serial] +fn prepared_tui_bundle_data_size_matches_real_client_inventory_sum() { + let home = tempfile::TempDir::new().unwrap(); + let _home_guard = HomeEnvGuard::set(home.path()); + let _pricing_guard = TestEnvGuard::set("TOKSCALE_PRICING_CACHE_ONLY", "1"); + write_streaming_fold_fixture(home.path()); + + let prepared = + super::prepare_local_inputs(inventory_options(home.path(), &["opencode", "codex"])) + .unwrap(); + let result = tokio::runtime::Runtime::new() + .unwrap() + .block_on(super::load_prepared_tui_bundle_with_diagnostics(prepared)) + .unwrap(); + + assert!(result.client_space["opencode"] > 0); + assert!(result.client_space["codex"] > 0); + let client_space_total = result + .client_space + .values() + .copied() + .try_fold(0_u64, u64::checked_add) + .unwrap(); + assert_eq!(result.health.input_data_bytes(), client_space_total); +} + #[test] #[serial_test::serial] fn test_streaming_tui_usage_applies_date_range() { @@ -1661,8 +1688,7 @@ fn test_input_cache_refreshes_stale_provider_on_cache_hit() { // must re-run refresh_derived_fields (dates are derived from // timestamps since schema v24, so provider identity is the // remaining derived field). - let stale_message = UnifiedMessage::new( - "opencode", + let stale_message = crate::sessions::ParsedMessage::new( "gpt-5.5", "anthropic", "session-1", @@ -1735,29 +1761,6 @@ fn prepared_test_group( } } -fn confirmed_test_group( - client: ClientId, - units: Vec, -) -> crate::adapters::ConfirmedAdapterInputs { - let prepared = prepared_test_group(client, units); - let mut present_files = Vec::new(); - let unit_digests = prepared - .units - .iter() - .map(|unit| { - unit.prepared_input_snapshot() - .expect("test unit must carry a prepared snapshot") - .visit_present_files(|identity, size| present_files.push((identity, size))); - unit.inventory_signature_digest() - }) - .collect(); - crate::adapters::ConfirmedAdapterInputs { - client, - unit_digests, - present_files, - } -} - #[test] fn input_data_size_counts_related_inputs_once_by_file_identity() { let dir = tempfile::TempDir::new().unwrap(); @@ -1777,47 +1780,6 @@ fn input_data_size_counts_related_inputs_once_by_file_identity() { assert_eq!(super::input_data_bytes([&with_dependency, &duplicate]), 13); } -#[test] -fn input_data_size_by_client_deduplicates_within_each_client() { - let dir = tempfile::TempDir::new().unwrap(); - let shared = dir.path().join("shared.jsonl"); - let amp_only = dir.path().join("amp.jsonl"); - std::fs::write(&shared, b"12345678").unwrap(); - std::fs::write(&_only, b"12345").unwrap(); - - let amp = confirmed_test_group( - ClientId::Amp, - vec![ - crate::adapters::InputUnit::plain_file(ClientId::Amp, shared.clone()), - crate::adapters::InputUnit::plain_file(ClientId::Amp, shared.clone()), - crate::adapters::InputUnit::plain_file(ClientId::Amp, amp_only), - ], - ); - let codebuddy = confirmed_test_group( - ClientId::CodeBuddy, - vec![crate::adapters::InputUnit::plain_file( - ClientId::CodeBuddy, - shared, - )], - ); - let requested = vec![ - "amp".to_string(), - "codebuddy".to_string(), - "codex".to_string(), - ]; - - let (by_client, global) = super::confirmed_input_data_bytes(&requested, &[amp, codebuddy]); - assert_eq!( - by_client, - std::collections::BTreeMap::from([ - ("amp".to_string(), 13), - ("codebuddy".to_string(), 8), - ("codex".to_string(), 0), - ]) - ); - assert_eq!(global, 13); -} - #[test] fn inventory_probe_refreshes_input_data_size_from_metadata() { let home = tempfile::TempDir::new().unwrap(); diff --git a/crates/tokscale-core/src/message_cache.rs b/crates/tokscale-core/src/message_cache.rs index 43068f460..e6c523400 100644 --- a/crates/tokscale-core/src/message_cache.rs +++ b/crates/tokscale-core/src/message_cache.rs @@ -1,5 +1,4 @@ -use crate::sessions::codex::CodexParseState; -use crate::UnifiedMessage; +use crate::sessions::{codex::CodexParseState, ParsedMessage}; use bincode::Options; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; @@ -19,7 +18,7 @@ compile_error!("input-message cache requires stable Unix or Windows file identit // Input-message cache shards split serialization layout from parser/input // semantics. Bump this only when the shard bincode layout changes; parser-only // fixes should bump the relevant InputUnit parser revision instead. -const CACHE_FORMAT_VERSION: u32 = 10; +const CACHE_FORMAT_VERSION: u32 = 11; #[cfg(test)] const UNSUPPORTED_CACHE_FORMAT_VERSION: u32 = CACHE_FORMAT_VERSION - 1; const SHARD_MAGIC: [u8; 8] = *b"TOKSHRD\0"; @@ -1303,7 +1302,7 @@ pub(crate) struct CachedInputEntry { pub path: CachedPath, pub parser_version: ParserVersion, pub fingerprint: InputFingerprint, - pub messages: Vec, + pub messages: Vec, pub codex_incremental: Option, pub rejections: crate::input_health::RejectionSummary, } @@ -1313,7 +1312,7 @@ impl CachedInputEntry { pub(crate) fn new( path: &Path, fingerprint: InputFingerprint, - messages: Vec, + messages: Vec, codex_incremental: Option, ) -> Self { Self::new_with_revision(path, 1, fingerprint, messages, codex_incremental) @@ -1324,7 +1323,7 @@ impl CachedInputEntry { path: &Path, parser_revision: ParserRevision, fingerprint: InputFingerprint, - messages: Vec, + messages: Vec, codex_incremental: Option, ) -> Self { Self::new_with_version( @@ -1341,7 +1340,7 @@ impl CachedInputEntry { path: &Path, parser_version: ParserVersion, fingerprint: InputFingerprint, - messages: Vec, + messages: Vec, codex_incremental: Option, ) -> Self { Self { @@ -1428,12 +1427,12 @@ struct CachedShardHeader { #[derive(Debug, Serialize, Deserialize)] struct CachedShardBody { - messages: Vec, + messages: Vec, } #[derive(Serialize)] struct BorrowedCachedShardBody<'a> { - messages: &'a [UnifiedMessage], + messages: &'a [ParsedMessage], } #[derive(Debug, Clone)] @@ -1560,7 +1559,7 @@ impl InputMessageCache { pub(crate) fn write_messages( &mut self, plan: CacheWritePlan, - messages: &[UnifiedMessage], + messages: &[ParsedMessage], ) -> Result<(), InputCacheError> { let key = plan.key(); ensure_cache_dir(&self.cache_dir).map_err(|source| { @@ -1584,7 +1583,7 @@ impl InputMessageCache { pub(crate) fn take_messages( &mut self, plan: &CacheReadPlan, - ) -> Result, CacheReadFailure> { + ) -> Result, CacheReadFailure> { let key = plan.key.clone(); if self.deleted_paths.contains(&key) { return Err(CacheReadFailure::new( @@ -2108,7 +2107,7 @@ fn write_shard_entry(cache_dir: &Path, entry: &CachedInputEntry) -> std::io::Res fn write_shard_borrowed( cache_dir: &Path, plan: &CacheWritePlan, - messages: &[UnifiedMessage], + messages: &[ParsedMessage], ) -> std::io::Result<()> { let final_path = shard_path_for_input_key(cache_dir, &plan.key()); let parent = final_path @@ -3181,8 +3180,7 @@ mod tests { let mut entry = CachedInputEntry::new( file.path(), fingerprint, - vec![UnifiedMessage::new( - "client", + vec![ParsedMessage::new( "gpt-5", "provider", "session-1", @@ -3236,6 +3234,13 @@ mod tests { .unwrap(); assert_eq!(messages.len(), 1); assert_eq!(messages[0].session_id.as_ref(), "session-1"); + assert!( + serde_json::to_value(&messages[0]) + .unwrap() + .get("client") + .is_none(), + "cached parsed messages must remain source-neutral" + ); restore_cache_env(prev_env); } @@ -3254,8 +3259,7 @@ mod tests { fingerprint.clone(), None, ); - let messages = vec![UnifiedMessage::new( - "client", + let messages = vec![ParsedMessage::new( "gpt-5", "provider", "session-1", @@ -3315,8 +3319,7 @@ mod tests { path, revision, InputFingerprint::from_path(path).unwrap(), - vec![UnifiedMessage::new( - "client", + vec![ParsedMessage::new( "gpt-5", "provider", format!("session-{revision}"), @@ -3524,8 +3527,7 @@ mod tests { cache.insert(CachedInputEntry::new( &path, InputFingerprint::from_path(&path).unwrap(), - vec![UnifiedMessage::new( - "opencode", + vec![ParsedMessage::new( "gpt-5", "openai", "session-1", @@ -3698,8 +3700,7 @@ mod tests { input.path(), parser_version, fingerprint.clone(), - vec![UnifiedMessage::new( - "client", + vec![ParsedMessage::new( "gpt-5", "provider", "cached-session", @@ -3730,8 +3731,7 @@ mod tests { "a failed ordinary rebuild must retain the unsupported shard" ); - let replacement = vec![UnifiedMessage::new( - "client", + let replacement = vec![ParsedMessage::new( "gpt-5", "provider", "current-session", @@ -3820,8 +3820,7 @@ mod tests { let error = cache .write_messages( CacheWritePlan::new(input.path(), parser_version, fingerprint, None), - &[UnifiedMessage::new( - "client", + &[ParsedMessage::new( "gpt-5", "provider", "session", @@ -3882,8 +3881,7 @@ mod tests { input.path(), 7, fingerprint, - vec![UnifiedMessage::new( - "client", + vec![ParsedMessage::new( "gpt-5", "provider", "session-1", @@ -3927,8 +3925,7 @@ mod tests { input.path(), ParserVersion::new(ParserId::Copilot, 1), fingerprint, - vec![UnifiedMessage::new( - "client", + vec![ParsedMessage::new( "gpt-5", "provider", "session-1", @@ -4049,8 +4046,7 @@ mod tests { input.path(), copilot_version, fingerprint.clone(), - vec![UnifiedMessage::new( - "copilot", + vec![ParsedMessage::new( "gpt-5", "openai", "copilot-session", @@ -4070,8 +4066,7 @@ mod tests { input.path(), gemini_version, fingerprint.clone(), - vec![UnifiedMessage::new( - "gemini", + vec![ParsedMessage::new( "gpt-5", "openai", "gemini-session", @@ -4138,8 +4133,7 @@ mod tests { input.path(), parser_version, initial_fingerprint.clone(), - vec![UnifiedMessage::new( - "copilot", + vec![ParsedMessage::new( "gpt-5", "openai", "initial-session", @@ -4171,8 +4165,7 @@ mod tests { input.path(), parser_version, replacement_fingerprint, - vec![UnifiedMessage::new( - "copilot", + vec![ParsedMessage::new( "gpt-5", "openai", "replacement-session", diff --git a/crates/tokscale-core/src/sessions/amp.rs b/crates/tokscale-core/src/sessions/amp.rs index a5eb6aeca..db2c648c1 100644 --- a/crates/tokscale-core/src/sessions/amp.rs +++ b/crates/tokscale-core/src/sessions/amp.rs @@ -3,7 +3,7 @@ //! Parses JSON files from ~/.local/share/amp/threads/ use super::error::{SessionParseError, SessionParseResult}; -use super::UnifiedMessage; +use super::ParsedMessage; use crate::input_health::{RecordRejectionReason, RejectionSummary, ScannedInput}; use crate::{provider_identity, TokenBreakdown}; use serde::Deserialize; @@ -112,10 +112,9 @@ impl AmpUsageRecord { self.model == other.model && self.tokens == other.tokens } - fn into_unified(self, thread_id: &str) -> UnifiedMessage { + fn into_message(self, thread_id: &str) -> ParsedMessage { let provider = provider_identity::observed_provider_id("", &self.model); - UnifiedMessage::new( - "amp", + ParsedMessage::new( &self.model, provider, thread_id, @@ -331,10 +330,10 @@ fn merge_amp_records( } } -fn build_amp_messages(records: Vec, thread_id: &str) -> Vec { +fn build_amp_messages(records: Vec, thread_id: &str) -> Vec { records .into_iter() - .map(|record| record.into_unified(thread_id)) + .map(|record| record.into_message(thread_id)) .collect() } @@ -401,7 +400,7 @@ mod tests { use super::parse_amp_file as parse_amp_file_result; use std::path::Path; - fn parse_amp_file(path: &Path) -> Vec { + fn parse_amp_file(path: &Path) -> Vec { parse_amp_file_result(path).unwrap().messages } diff --git a/crates/tokscale-core/src/sessions/antigravity_cli.rs b/crates/tokscale-core/src/sessions/antigravity_cli.rs index 5591ec58b..90a944ed1 100644 --- a/crates/tokscale-core/src/sessions/antigravity_cli.rs +++ b/crates/tokscale-core/src/sessions/antigravity_cli.rs @@ -45,7 +45,7 @@ use super::error::{SessionParseError, SessionParseResult}; use super::utils::open_readonly_sqlite; -use super::{normalize_workspace_key, workspace_label_from_key, UnifiedMessage}; +use super::{normalize_workspace_key, workspace_label_from_key, ParsedMessage}; use crate::input_health::{InputFailure, RecordRejectionReason, ScannedInput}; use crate::{provider_identity, TokenBreakdown}; use rusqlite::Connection; @@ -127,7 +127,7 @@ fn parse_gen_metadata( session_id: &str, session_timestamp: i64, seen_response_ids: &mut HashSet, -) -> SessionParseResult> { +) -> SessionParseResult> { let chat_model = message_field(blob, 1).ok_or_else(|| { SessionParseError::invalid( "decode Antigravity CLI usage protobuf", @@ -225,8 +225,7 @@ fn parse_gen_metadata( let dedup_key = response_id.as_deref().map(response_dedup_key); - Ok(Some(UnifiedMessage::new_with_dedup( - "antigravity", + Ok(Some(ParsedMessage::new_with_dedup( model_id, provider_id, session_id, @@ -576,7 +575,7 @@ mod tests { use super::*; use rusqlite::{params, Connection}; - fn parse_antigravity_cli_file(path: &Path) -> Vec { + fn parse_antigravity_cli_file(path: &Path) -> Vec { super::parse_antigravity_cli_file(path).unwrap().messages } @@ -585,7 +584,7 @@ mod tests { session_id: &str, session_timestamp: i64, seen_response_ids: &mut HashSet, - ) -> Option { + ) -> Option { super::parse_gen_metadata(blob, session_id, session_timestamp, seen_response_ids).unwrap() } @@ -702,7 +701,6 @@ mod tests { assert_eq!(messages.len(), 1); let message = &messages[0]; - assert_eq!(message.client.as_ref(), "antigravity"); assert_eq!(message.model_id.as_ref(), "gemini-3.5-flash"); assert_eq!(message.provider_id.as_ref(), "google"); assert_eq!(message.tokens.input, 1632); diff --git a/crates/tokscale-core/src/sessions/claudecode.rs b/crates/tokscale-core/src/sessions/claudecode.rs index 765958366..f315546dc 100644 --- a/crates/tokscale-core/src/sessions/claudecode.rs +++ b/crates/tokscale-core/src/sessions/claudecode.rs @@ -4,7 +4,7 @@ use super::error::{SessionParseError, SessionParseResult}; use super::utils::{extract_i64, extract_string, parse_timestamp_value}; -use super::{normalize_workspace_key, workspace_label_from_key, UnifiedMessage}; +use super::{normalize_workspace_key, workspace_label_from_key, ParsedMessage}; use crate::input_health::{InputFailure, RecordRejectionReason, RejectionSummary, ScannedInput}; use crate::{checked_token_add, model_aliases, provider_identity, TokenBreakdown}; use serde::Deserialize; @@ -727,7 +727,6 @@ fn parse_claude_file_with_cache_home_and_resolver( let mut project_candidates = ClaudeProjectCandidates::default(); let mut parent_session_id = None; let is_transcript_path = is_claude_transcripts_path(path); - let client_id = "claude".to_string(); let mut session_id = path .file_stem() .and_then(|s| s.to_str()) @@ -746,7 +745,7 @@ fn parse_claude_file_with_cache_home_and_resolver( .map_err(|source| SessionParseError::at_path(path, "open Claude session", source))?; let reader = BufReader::new(file); - let mut messages: Vec = Vec::with_capacity(64); + let mut messages: Vec = Vec::with_capacity(64); let mut rejections = RejectionSummary::default(); let mut interrupted = None; let mut provider_confidences: Vec = Vec::with_capacity(64); @@ -924,7 +923,6 @@ fn parse_claude_file_with_cache_home_and_resolver( entry: &entry, last_model: last_model.as_deref(), last_provider_hint: last_provider_hint.as_deref(), - client_id: &client_id, default_provider_hint: None, session_id: &session_id, suppress_unattributed: suppress_unattributed_tool_results, @@ -1193,8 +1191,7 @@ fn parse_claude_file_with_cache_home_and_resolver( processed_hashes.insert(*hash, messages.len()); }); - let mut unified = UnifiedMessage::new_with_dedup( - client_id.clone(), + let mut message = ParsedMessage::new_with_dedup( model, provider_choice.id, session_id.clone(), @@ -1203,23 +1200,23 @@ fn parse_claude_file_with_cache_home_and_resolver( 0.0, dedup_key, ); - unified.is_main_session = is_main_session; - unified.agent = sidechain_agent + message.is_main_session = is_main_session; + message.agent = sidechain_agent .as_deref() .map(crate::sessions::intern::intern); - unified.set_agent_instance(sidechain_agent_instance.clone()); + message.set_agent_instance(sidechain_agent_instance.clone()); let (message_workspace_key, message_workspace_label) = workspace_options_for_entry( entry_workspace.as_ref(), &workspace_key, &workspace_label, ); - unified.set_workspace(message_workspace_key, message_workspace_label); + message.set_workspace(message_workspace_key, message_workspace_label); // Mark the first assistant response after a user message as a turn start if pending_turn_start { - unified.is_turn_start = true; + message.is_turn_start = true; pending_turn_start = false; } - messages.push(unified); + messages.push(message); provider_confidences.push(provider_confidence); } } @@ -1304,7 +1301,7 @@ fn workspace_parts_from_key(raw: &str) -> Option { Some(ClaudeWorkspaceParts { key, label }) } -fn set_message_workspace(message: &mut UnifiedMessage, workspace: &ClaudeWorkspaceParts) { +fn set_message_workspace(message: &mut ParsedMessage, workspace: &ClaudeWorkspaceParts) { let (workspace_key, workspace_label) = workspace.to_options(); message.set_workspace(workspace_key, workspace_label); } @@ -1327,7 +1324,7 @@ fn apply_resolved_project_workspace( candidates: &ClaudeProjectCandidates, input_path: &Path, parent_session_id: Option<&str>, - messages: &mut [UnifiedMessage], + messages: &mut [ParsedMessage], ) -> ClaudeProjectDependency { let Some(project_key) = project_key else { return ClaudeProjectDependency::None; @@ -1603,7 +1600,7 @@ fn parse_claude_entry_timestamp_checked( } fn merge_claude_duplicate( - existing: &mut UnifiedMessage, + existing: &mut ParsedMessage, usage: &ClaudeUsage, parsed_timestamp: i64, ) { @@ -1624,7 +1621,7 @@ fn merge_claude_duplicate( } fn merge_claude_tool_result_duplicate( - existing: &mut UnifiedMessage, + existing: &mut ParsedMessage, input_tokens: i64, timestamp_ms: i64, ) { @@ -1645,7 +1642,6 @@ struct ClaudeToolResultContext<'a> { entry: &'a ClaudeEntry, last_model: Option<&'a str>, last_provider_hint: Option<&'a str>, - client_id: &'a str, default_provider_hint: Option<&'a str>, session_id: &'a str, suppress_unattributed: bool, @@ -1658,7 +1654,7 @@ struct ClaudeToolResultContext<'a> { fn extract_claude_tool_result_message( line: &str, context: ClaudeToolResultContext<'_>, -) -> SessionParseResult> { +) -> SessionParseResult> { let value: Value = serde_json::from_str(line).map_err(|source| { SessionParseError::at_path( context.input_path, @@ -1740,8 +1736,7 @@ fn extract_claude_tool_result_message( ) })?; - let mut message = UnifiedMessage::new_with_dedup( - context.client_id, + let mut message = ParsedMessage::new_with_dedup( model, provider_choice.id, context.session_id, @@ -1756,8 +1751,8 @@ fn extract_claude_tool_result_message( 0.0, usage.dedup_key.map(|key| { crate::sessions::dedup_hash_str(&format!( - "{}:tool_result:{}:{key}", - context.client_id, context.session_id + "claude:tool_result:{}:{key}", + context.session_id )) }), ); @@ -2226,7 +2221,7 @@ mod tests { // Most parser tests assert only the message projection. Health-specific // cases call `super::parse_claude_file` directly so rejections cannot be // discarded accidentally in the behavior under test. - fn parse_claude_file(path: &Path) -> SessionParseResult> { + fn parse_claude_file(path: &Path) -> SessionParseResult> { super::parse_claude_file(path).map(|scanned| scanned.messages) } diff --git a/crates/tokscale-core/src/sessions/cline.rs b/crates/tokscale-core/src/sessions/cline.rs index 7fbf90035..b15bde07c 100644 --- a/crates/tokscale-core/src/sessions/cline.rs +++ b/crates/tokscale-core/src/sessions/cline.rs @@ -5,13 +5,12 @@ //! part of this parser's input contract. use super::error::{SessionParseError, SessionParseResult}; -use super::{workspace_metadata_from_key, UnifiedMessage, WorkspaceMetadata}; +use super::{workspace_metadata_from_key, ParsedMessage, WorkspaceMetadata}; use crate::input_health::{InputFailure, RecordRejectionReason, ScannedInput}; use crate::{provider_identity, TokenBreakdown}; use serde_json::{Map, Value}; use std::path::{Path, PathBuf}; -const CLIENT_ID: &str = "cline"; const MESSAGES_SUFFIX: &str = ".messages.json"; #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -318,8 +317,7 @@ pub fn parse_cline_file(path: &Path) -> SessionParseResult { .and_then(|info| non_empty_string(info.get("provider"))) .unwrap_or_default(); let provider_id = provider_identity::observed_provider_id(raw_provider, model_id); - let mut message = UnifiedMessage::new_with_agent( - CLIENT_ID, + let mut message = ParsedMessage::new_with_agent( model_id, provider_id, &identity.session_id, @@ -431,7 +429,6 @@ mod tests { assert!(scanned.rejections.is_empty()); assert!(scanned.interrupted.is_none()); let first = &scanned.messages[0]; - assert_eq!(first.client.as_ref(), "cline"); assert_eq!(first.model_id.as_ref(), "poolside/laguna-m.1:free"); assert_eq!(first.provider_id.as_ref(), "cline"); assert_eq!(first.session_id.as_ref(), "session-a"); diff --git a/crates/tokscale-core/src/sessions/codebuddy.rs b/crates/tokscale-core/src/sessions/codebuddy.rs index 4983fff76..0ce96a4ac 100644 --- a/crates/tokscale-core/src/sessions/codebuddy.rs +++ b/crates/tokscale-core/src/sessions/codebuddy.rs @@ -5,7 +5,7 @@ //! extension writes final agent usage into extension logs. use super::error::{SessionParseError, SessionParseResult}; -use super::{dedup_hash_str, normalize_workspace_key, workspace_label_from_key, UnifiedMessage}; +use super::{dedup_hash_str, normalize_workspace_key, workspace_label_from_key, ParsedMessage}; use crate::input_health::{InputFailure, RecordRejectionReason, ScannedInput}; use crate::{provider_identity, TokenBreakdown}; use chrono::TimeZone; @@ -14,8 +14,6 @@ use std::collections::HashMap; use std::io::{BufRead, BufReader}; use std::path::Path; -const CLIENT_ID: &str = "codebuddy"; - #[derive(Debug, Deserialize)] struct CodeBuddyLine { id: Option, @@ -308,10 +306,9 @@ pub(crate) fn parse_codebuddy_jsonl_file(path: &Path) -> SessionParseResult SessionParseRes continue; }; let provider_id = provider_identity::observed_provider_id("", &model_id); - let mut message = UnifiedMessage::new_with_dedup( - CLIENT_ID, + let mut message = ParsedMessage::new_with_dedup( model_id, provider_id, agent_id, @@ -582,11 +578,11 @@ fn workspace_from_log_path(path: &Path) -> Option { mod tests { use super::*; - fn parse_codebuddy_jsonl_file(path: &Path) -> Vec { + fn parse_codebuddy_jsonl_file(path: &Path) -> Vec { super::parse_codebuddy_jsonl_file(path).unwrap().messages } - fn parse_codebuddy_extension_log_file(path: &Path) -> Vec { + fn parse_codebuddy_extension_log_file(path: &Path) -> Vec { super::parse_codebuddy_extension_log_file(path) .unwrap() .messages @@ -607,7 +603,6 @@ mod tests { let messages = parse_codebuddy_jsonl_file(&path); assert_eq!(messages.len(), 1); - assert_eq!(messages[0].client.as_ref(), "codebuddy"); assert_eq!(messages[0].model_id.as_ref(), "glm-5.2"); assert_eq!(messages[0].tokens.input, 24486); assert_eq!(messages[0].tokens.output, 3); diff --git a/crates/tokscale-core/src/sessions/codebuff.rs b/crates/tokscale-core/src/sessions/codebuff.rs index ca48daa87..c1214305d 100644 --- a/crates/tokscale-core/src/sessions/codebuff.rs +++ b/crates/tokscale-core/src/sessions/codebuff.rs @@ -15,13 +15,13 @@ use super::error::{SessionParseError, SessionParseResult}; use super::utils::{parse_timestamp_str, parse_timestamp_value, read_file}; -use super::UnifiedMessage; +use super::ParsedMessage; use crate::input_health::{RecordRejectionReason, ScannedInput}; use crate::{provider_identity, TokenBreakdown}; use serde_json::Value; use std::path::Path; -/// Parse a single `chat-messages.json` file into UnifiedMessages. +/// Parse a single `chat-messages.json` file into source-neutral messages. pub fn parse_codebuff_file(path: &Path) -> SessionParseResult { let mut bytes = read_file(path)?; let root: Value = simd_json::from_slice(&mut bytes) @@ -86,8 +86,7 @@ pub fn parse_codebuff_file(path: &Path) -> SessionParseResult { let dedup_key = upstream_message_id(msg) .unwrap_or_else(|| derive_dedup_key(&session_id, ts, &model, &usage, ordinal)); - scanned.messages.push(UnifiedMessage::new_with_dedup( - "codebuff", + scanned.messages.push(ParsedMessage::new_with_dedup( &model, provider, &session_id, @@ -667,7 +666,6 @@ mod tests { let messages = scanned.messages; assert_eq!(messages.len(), 1); let only = &messages[0]; - assert_eq!(only.client.as_ref(), "codebuff"); assert_eq!(only.model_id.as_ref(), "claude-sonnet-4-20250514"); assert_eq!(only.provider_id.as_ref(), "anthropic"); assert!(only.session_id.ends_with("/proj/2025-12-20T12-00-00.000Z")); diff --git a/crates/tokscale-core/src/sessions/codex.rs b/crates/tokscale-core/src/sessions/codex.rs index c4563c591..4a1b399f2 100644 --- a/crates/tokscale-core/src/sessions/codex.rs +++ b/crates/tokscale-core/src/sessions/codex.rs @@ -7,7 +7,7 @@ //! Note: This parser has stateful logic to track model and delta calculations. use super::error::{SessionParseError, SessionParseResult}; -use super::{normalize_workspace_key, workspace_label_from_key, UnifiedMessage}; +use super::{normalize_workspace_key, workspace_label_from_key, ParsedMessage}; use crate::input_health::{InputFailure, RecordRejectionReason, RejectionSummary}; use crate::{checked_token_sum, TokenBreakdown}; use serde::Deserialize; @@ -187,7 +187,7 @@ pub(crate) struct CodexParseState { #[derive(Debug, Clone)] pub(crate) struct ParsedCodexFile { - pub messages: Vec, + pub messages: Vec, pub rejections: RejectionSummary, pub interrupted: Option, pub consumed_offset: u64, @@ -214,9 +214,8 @@ struct PendingCodexMessage { } impl PendingCodexMessage { - fn into_message(self, model: &str) -> UnifiedMessage { - let mut message = UnifiedMessage::new_with_agent( - "codex", + fn into_message(self, model: &str) -> ParsedMessage { + let mut message = ParsedMessage::new_with_agent( model, self.provider, self.session_id, @@ -867,7 +866,7 @@ fn parse_codex_entry_timestamp(timestamp: Option<&str>) -> SessionParseResult, - messages: &mut Vec, + messages: &mut Vec, model: &str, ) { for pending in pending_model_messages.drain(..) { @@ -914,7 +913,7 @@ fn flush_pending_model_messages( } /// Parse a Codex JSONL file with stateful tracking -pub fn parse_codex_file(path: &Path) -> SessionParseResult> { +pub fn parse_codex_file(path: &Path) -> SessionParseResult> { let file = std::fs::File::open(path) .map_err(|source| SessionParseError::new("open Codex JSONL input", source))?; @@ -1157,7 +1156,7 @@ mod tests { file } - fn parse_codex_file(path: &Path) -> Vec { + fn parse_codex_file(path: &Path) -> Vec { super::parse_codex_file(path).expect("test fixture must be valid Codex JSONL") } diff --git a/crates/tokscale-core/src/sessions/commandcode.rs b/crates/tokscale-core/src/sessions/commandcode.rs index e88b7da41..0843b0325 100644 --- a/crates/tokscale-core/src/sessions/commandcode.rs +++ b/crates/tokscale-core/src/sessions/commandcode.rs @@ -9,7 +9,7 @@ use super::error::{SessionParseError, SessionParseResult}; use super::utils::parse_timestamp_str; use super::{ - normalize_workspace_key, workspace_label_from_key, workspace_metadata_from_key, UnifiedMessage, + normalize_workspace_key, workspace_label_from_key, workspace_metadata_from_key, ParsedMessage, WorkspaceMetadata, }; use crate::input_health::{InputFailure, RecordRejectionReason, ScannedInput}; @@ -19,8 +19,6 @@ use std::collections::BTreeSet; use std::io::{BufRead, BufReader}; use std::path::{Path, PathBuf}; -const CLIENT_ID: &str = "commandcode"; - #[derive(Debug, Deserialize)] struct CommandCodeEntry { role: Option, @@ -180,8 +178,7 @@ pub fn parse_commandcode_file(path: &Path) -> SessionParseResult { "commandcode:{resolved_session}:{assistant_index}" )); session_id = Some(resolved_session.clone()); - let mut message = UnifiedMessage::new_with_dedup( - CLIENT_ID, + let mut message = ParsedMessage::new_with_dedup( model_id.clone(), &provider_id, resolved_session, @@ -406,7 +403,7 @@ mod tests { use serde_json::json; use std::io::Write; - fn parse_commandcode_file(path: &Path) -> Vec { + fn parse_commandcode_file(path: &Path) -> Vec { super::parse_commandcode_file(path).unwrap().messages } @@ -457,7 +454,6 @@ mod tests { assert_eq!(messages.len(), 1); let message = &messages[0]; - assert_eq!(message.client.as_ref(), "commandcode"); assert_eq!(message.provider_id.as_ref(), "minimax"); assert_eq!(message.model_id.as_ref(), "MiniMax-M3"); assert_eq!(message.session_id.as_ref(), "sess-1"); diff --git a/crates/tokscale-core/src/sessions/copilot.rs b/crates/tokscale-core/src/sessions/copilot.rs index c61c3b1e1..c5e40fc94 100644 --- a/crates/tokscale-core/src/sessions/copilot.rs +++ b/crates/tokscale-core/src/sessions/copilot.rs @@ -5,7 +5,7 @@ //! aggregate agent records are only used as a fallback to avoid double counting. use super::error::{SessionParseError, SessionParseResult}; -use super::{workspace_metadata_from_key, UnifiedMessage, WorkspaceMetadata}; +use super::{workspace_metadata_from_key, ParsedMessage, WorkspaceMetadata}; use crate::input_health::{InputFailure, RecordRejectionReason, RejectionSummary, ScannedInput}; use crate::provider_identity::observed_provider_id; use crate::TokenBreakdown; @@ -334,9 +334,8 @@ enum SessionIdPriority { } impl CopilotUsageCandidate { - fn into_message(self) -> UnifiedMessage { - let mut message = UnifiedMessage::new_with_agent( - "copilot", + fn into_message(self) -> ParsedMessage { + let mut message = ParsedMessage::new_with_agent( self.model, self.provider_id, self.session_id, @@ -1043,7 +1042,7 @@ mod tests { use std::io::Write; use tempfile::NamedTempFile; - fn parse_copilot_file(path: &Path) -> Vec { + fn parse_copilot_file(path: &Path) -> Vec { super::parse_copilot_file(path).unwrap().messages } @@ -1176,7 +1175,6 @@ mod tests { assert_eq!(messages.len(), 1); let message = &messages[0]; - assert_eq!(message.client.as_ref(), "copilot"); assert_eq!(message.model_id.as_ref(), "claude-sonnet-4"); assert_eq!(message.provider_id.as_ref(), "anthropic"); assert_eq!(message.session_id.as_ref(), "conv-1"); diff --git a/crates/tokscale-core/src/sessions/droid.rs b/crates/tokscale-core/src/sessions/droid.rs index b798472a5..4ec42b090 100644 --- a/crates/tokscale-core/src/sessions/droid.rs +++ b/crates/tokscale-core/src/sessions/droid.rs @@ -3,7 +3,7 @@ //! Parses JSON files from ~/.factory/sessions/ use super::error::{SessionParseError, SessionParseResult}; -use super::{workspace_metadata_from_key, UnifiedMessage, WorkspaceMetadata}; +use super::{workspace_metadata_from_key, ParsedMessage, WorkspaceMetadata}; use crate::input_health::{RecordRejectionReason, ScannedInput}; use crate::{model_aliases, provider_identity, TokenBreakdown}; use serde::Deserialize; @@ -202,7 +202,7 @@ fn read_session_start(path: &Path) -> Option { /// Factory records direct delegation on the settings file's transcript header. /// Missing legacy transcripts retain the constructor default of top-level. -pub(crate) fn classify_droid_main_session(path: &Path, messages: &mut [UnifiedMessage]) { +pub(crate) fn classify_droid_main_session(path: &Path, messages: &mut [ParsedMessage]) { let is_main_session = read_session_start(path).is_none_or(|start| start.parent_session_id().is_none()); for message in messages { @@ -474,9 +474,8 @@ pub fn parse_droid_file(path: &Path) -> SessionParseResult { return Ok(scanned); } - let message = UnifiedMessage::new_with_agent( - "droid", model, provider, session_id, timestamp, tokens, 0.0, agent, - ); + let message = + ParsedMessage::new_with_agent(model, provider, session_id, timestamp, tokens, 0.0, agent); scanned.messages.push(message); classify_droid_main_session(path, &mut scanned.messages); Ok(scanned) @@ -487,7 +486,7 @@ mod tests { use super::*; use serde_json::json; - fn parse_droid_file(path: &Path) -> Vec { + fn parse_droid_file(path: &Path) -> Vec { super::parse_droid_file(path).unwrap().messages } diff --git a/crates/tokscale-core/src/sessions/gemini.rs b/crates/tokscale-core/src/sessions/gemini.rs index 9781ef2de..dbb206036 100644 --- a/crates/tokscale-core/src/sessions/gemini.rs +++ b/crates/tokscale-core/src/sessions/gemini.rs @@ -5,7 +5,7 @@ use super::error::{SessionParseError, SessionParseResult}; use super::utils::{extract_i64, extract_string, parse_timestamp_value}; -use super::{workspace_metadata_from_key, UnifiedMessage, WorkspaceMetadata}; +use super::{workspace_metadata_from_key, ParsedMessage, WorkspaceMetadata}; use crate::input_health::{InputFailure, RecordRejectionReason, RejectionSummary, ScannedInput}; use crate::{checked_token_add, checked_token_sum, TokenBreakdown}; use serde::Deserialize; @@ -347,7 +347,7 @@ fn build_gemini_token_message( session_id: &str, timestamp: i64, tokens: GeminiTokens, -) -> UnifiedMessage { +) -> ParsedMessage { let (input, cache_read) = normalize_gemini_session_input_and_cache( tokens.input.unwrap_or(0), tokens.cached.unwrap_or(0), @@ -359,8 +359,7 @@ fn build_gemini_token_message( let tool = tokens.tool.unwrap_or(0).max(0); - UnifiedMessage::new( - "gemini", + ParsedMessage::new( model, "google", session_id, @@ -394,7 +393,7 @@ fn parse_direct_gemini_token_message( value: &Value, model_hint: Option, session_id: Option<&str>, -) -> SessionParseResult> { +) -> SessionParseResult> { let Some(tokens_value) = value.get("tokens") else { return Ok(None); }; @@ -686,7 +685,7 @@ fn build_messages_from_usages( usages: Vec, session_id: &str, timestamp: i64, -) -> Vec { +) -> Vec { usages .into_iter() .map(|usage| { @@ -695,8 +694,7 @@ fn build_messages_from_usages( } else { (usage.input.max(0), usage.cached.max(0)) }; - UnifiedMessage::new( - "gemini", + ParsedMessage::new( usage.model, "google", session_id, @@ -947,7 +945,7 @@ fn extract_timestamp_from_value(value: &Value) -> Option { mod tests { use super::*; - fn parse_gemini_file(path: &Path) -> Vec { + fn parse_gemini_file(path: &Path) -> Vec { super::parse_gemini_file(path).unwrap().messages } use std::io::Write; diff --git a/crates/tokscale-core/src/sessions/goose.rs b/crates/tokscale-core/src/sessions/goose.rs index 63d1db3f4..11d812a33 100644 --- a/crates/tokscale-core/src/sessions/goose.rs +++ b/crates/tokscale-core/src/sessions/goose.rs @@ -8,7 +8,7 @@ use super::error::{SessionParseError, SessionParseResult}; use super::utils::open_readonly_sqlite; -use super::UnifiedMessage; +use super::ParsedMessage; use crate::input_health::{InputFailure, RecordRejectionReason, ScannedInput}; use crate::{provider_identity, TokenBreakdown}; use serde::Deserialize; @@ -198,8 +198,7 @@ pub fn parse_goose_sqlite(db_path: &Path) -> SessionParseResult { .record(RecordRejectionReason::MalformedRecord); continue; } - let mut msg = UnifiedMessage::new( - "goose", + let mut msg = ParsedMessage::new( model_id, provider, session_id, diff --git a/crates/tokscale-core/src/sessions/grok.rs b/crates/tokscale-core/src/sessions/grok.rs index 41753e4d4..fbac0dd7c 100644 --- a/crates/tokscale-core/src/sessions/grok.rs +++ b/crates/tokscale-core/src/sessions/grok.rs @@ -8,14 +8,13 @@ use super::error::{SessionParseError, SessionParseResult}; use super::utils::{extract_string, parse_timestamp_value}; -use super::{normalize_workspace_key, workspace_label_from_key, UnifiedMessage}; +use super::{normalize_workspace_key, workspace_label_from_key, ParsedMessage}; use crate::input_health::{InputFailure, RecordRejectionReason, RejectionSummary, ScannedInput}; use crate::{model_aliases, token_imputation}; use serde_json::Value; use std::io::{BufRead, BufReader}; use std::path::{Path, PathBuf}; -const CLIENT_ID: &str = "grok"; const PROVIDER_ID: &str = "xai"; #[derive(Debug, Clone)] @@ -367,7 +366,7 @@ fn build_messages( metadata: GrokMetadata, session_id: String, pending_messages: Vec, -) -> SessionParseResult> { +) -> SessionParseResult> { let totals: Vec = pending_messages .iter() .map(|message| message.token_delta) @@ -386,7 +385,7 @@ fn build_messages( let messages = pending_messages .into_iter() .zip(token_rows) - .map(|(pending, tokens)| -> SessionParseResult { + .map(|(pending, tokens)| -> SessionParseResult { let token_total = tokens.checked_total().ok_or_else(|| { SessionParseError::invalid( "validate imputed token breakdown", @@ -399,8 +398,7 @@ fn build_messages( "Grok imputed token total does not match its cumulative delta", )); } - let mut message = UnifiedMessage::new_with_dedup( - CLIENT_ID, + let mut message = ParsedMessage::new_with_dedup( pending.model_id, PROVIDER_ID, session_id.clone(), @@ -667,7 +665,7 @@ fn hex_value(byte: u8) -> Option { mod tests { use super::*; - fn parse_grok_updates_file(path: &Path) -> Vec { + fn parse_grok_updates_file(path: &Path) -> Vec { super::parse_grok_updates_file(path).unwrap().messages } @@ -776,7 +774,6 @@ not-json let expected_tokens = crate::token_imputation::impute_total_only_token_breakdowns(&[200, 150]); assert_eq!(messages.len(), 2); - assert_eq!(messages[0].client.as_ref(), "grok"); assert_eq!(messages[0].model_id.as_ref(), "composer-2.5-fast"); assert_eq!(messages[0].provider_id.as_ref(), "xai"); assert_eq!(messages[0].session_id.as_ref(), "session-1"); diff --git a/crates/tokscale-core/src/sessions/hermes.rs b/crates/tokscale-core/src/sessions/hermes.rs index d1a31d206..7cfa6682c 100644 --- a/crates/tokscale-core/src/sessions/hermes.rs +++ b/crates/tokscale-core/src/sessions/hermes.rs @@ -9,7 +9,7 @@ use super::error::{SessionParseError, SessionParseResult}; use super::utils::{open_readonly_sqlite, parse_epoch_f64_millis}; -use super::UnifiedMessage; +use super::ParsedMessage; use crate::input_health::{InputFailure, RecordRejectionReason, ScannedInput}; use crate::{provider_identity, TokenBreakdown}; use std::path::Path; @@ -147,8 +147,7 @@ pub fn parse_hermes_sqlite(db_path: &Path) -> SessionParseResult { continue; }; let provider = resolved_provider(billing_provider, &model_id); - let mut msg = UnifiedMessage::new_with_agent( - "hermes", + let mut msg = ParsedMessage::new_with_agent( model_id, provider, session_id.clone(), diff --git a/crates/tokscale-core/src/sessions/junie.rs b/crates/tokscale-core/src/sessions/junie.rs index b8b58ea9d..e4ef3748e 100644 --- a/crates/tokscale-core/src/sessions/junie.rs +++ b/crates/tokscale-core/src/sessions/junie.rs @@ -3,7 +3,7 @@ //! Junie stores local sessions under `~/.junie/sessions//events.jsonl`. use super::error::{SessionParseError, SessionParseResult}; -use super::{dedup_hash_str, normalize_agent_name, UnifiedMessage}; +use super::{dedup_hash_str, normalize_agent_name, ParsedMessage}; use crate::input_health::{InputFailure, RecordRejectionReason, ScannedInput}; use crate::{model_aliases, provider_identity, TokenBreakdown}; use chrono::{Local, LocalResult, NaiveDateTime, TimeZone}; @@ -12,7 +12,6 @@ use std::collections::HashSet; use std::io::{BufRead, BufReader}; use std::path::Path; -const CLIENT_ID: &str = "junie"; const USAGE_EVENT_KIND: &str = "LlmResponseMetadataEvent"; const USER_PROMPT_KIND: &str = "UserPromptEvent"; const SKIP_EVENT_KINDS: &[&str] = &[ @@ -130,7 +129,7 @@ pub fn parse_junie_file(path: &Path) -> SessionParseResult { let provider_id = provider_from_usage(usage, &model_id); let dedup_key = format!( - "{CLIENT_ID}:{session_id}:{timestamp}:{model_id}:{}:{}:{}:{}:{}:{usage_index}", + "junie:{session_id}:{timestamp}:{model_id}:{}:{}:{}:{}:{}:{usage_index}", tokens.input, tokens.output, tokens.cache_read, @@ -141,8 +140,7 @@ pub fn parse_junie_file(path: &Path) -> SessionParseResult { continue; } - let mut message = UnifiedMessage::new_with_agent( - CLIENT_ID, + let mut message = ParsedMessage::new_with_agent( model_id, provider_id, &session_id, @@ -401,7 +399,7 @@ mod tests { parse_junie_file(&path) } - fn parse_events(content: &str) -> Vec { + fn parse_events(content: &str) -> Vec { parse_events_result(content).unwrap().messages } @@ -422,7 +420,6 @@ mod tests { assert_eq!(messages.len(), 1); let message = &messages[0]; - assert_eq!(message.client.as_ref(), "junie"); assert_eq!(message.session_id.as_ref(), "session-250622-101010"); assert_eq!(message.model_id.as_ref(), "gpt-4.1"); assert_eq!(message.provider_id.as_ref(), "openai"); diff --git a/crates/tokscale-core/src/sessions/kilo.rs b/crates/tokscale-core/src/sessions/kilo.rs index b4b90ae4f..973ea01ea 100644 --- a/crates/tokscale-core/src/sessions/kilo.rs +++ b/crates/tokscale-core/src/sessions/kilo.rs @@ -7,7 +7,7 @@ use super::error::{SessionParseError, SessionParseResult}; use super::utils::open_readonly_sqlite; -use super::{normalize_agent_name, UnifiedMessage}; +use super::{normalize_agent_name, ParsedMessage}; use crate::input_health::{InputFailure, RecordRejectionReason, ScannedInput}; use crate::{provider_identity, TokenBreakdown}; use serde::Deserialize; @@ -237,8 +237,7 @@ pub fn parse_kilo_sqlite(db_path: &Path) -> SessionParseResult { .map(str::to_string) .unwrap_or_else(|| provider_identity::observed_provider_id("", &model_id)); - let mut unified = UnifiedMessage::new_with_agent( - "kilo", + let mut message = ParsedMessage::new_with_agent( model_id, provider, session_id, @@ -247,9 +246,9 @@ pub fn parse_kilo_sqlite(db_path: &Path) -> SessionParseResult { 0.0, agent, ); - unified.dedup_key = dedup_key; + message.dedup_key = dedup_key; - scanned.messages.push(unified); + scanned.messages.push(message); } Ok(scanned) @@ -337,7 +336,6 @@ mod tests { assert_eq!(messages.len(), 1); let msg = &messages[0]; - assert_eq!(msg.client.as_ref(), "kilo"); assert_eq!(msg.session_id.as_ref(), "sess-1"); assert_eq!(msg.model_id.as_ref(), "claude-sonnet-4"); assert_eq!(msg.provider_id.as_ref(), "anthropic"); diff --git a/crates/tokscale-core/src/sessions/kimi.rs b/crates/tokscale-core/src/sessions/kimi.rs index ac8f61d58..d7a95f771 100644 --- a/crates/tokscale-core/src/sessions/kimi.rs +++ b/crates/tokscale-core/src/sessions/kimi.rs @@ -5,7 +5,7 @@ //! `~/.kimi-code/sessions///agents//wire.jsonl`. use super::error::{SessionParseError, SessionParseResult}; -use super::{workspace_metadata_from_key, UnifiedMessage, WorkspaceMetadata}; +use super::{workspace_metadata_from_key, ParsedMessage, WorkspaceMetadata}; use crate::input_health::{InputFailure, RecordRejectionReason, ScannedInput}; use crate::{provider_identity, TokenBreakdown}; use serde::Deserialize; @@ -13,8 +13,6 @@ use std::collections::HashMap; use std::io::{BufRead, BufReader}; use std::path::{Path, PathBuf}; -const CLIENT_ID: &str = "kimi"; - #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] struct TokenUsage { @@ -189,8 +187,7 @@ pub fn parse_kimi_file(path: &Path) -> SessionParseResult { continue; }; let (provider_id, model_id) = resolve_model(raw_model, &observed_aliases, &aliases); - let mut message = UnifiedMessage::new_with_agent( - CLIENT_ID, + let mut message = ParsedMessage::new_with_agent( model_id, provider_id, session_id.clone(), @@ -398,7 +395,7 @@ fn workspace_from_initial_wire_config(path: &Path) -> Option mod tests { use super::*; - fn parse_kimi_file(path: &Path) -> Vec { + fn parse_kimi_file(path: &Path) -> Vec { super::parse_kimi_file(path).unwrap().messages } use std::io::Write; @@ -458,7 +455,6 @@ model = "gpt-5.5" let messages = parse_kimi_file(&wire); assert_eq!(messages.len(), 1); - assert_eq!(messages[0].client.as_ref(), "kimi"); assert_eq!(messages[0].provider_id.as_ref(), "openai-pro"); assert_eq!(messages[0].model_id.as_ref(), "gpt-5.5"); assert_eq!(messages[0].session_id.as_ref(), "session_123"); diff --git a/crates/tokscale-core/src/sessions/kiro.rs b/crates/tokscale-core/src/sessions/kiro.rs index 37072e7a6..50ad4b3f0 100644 --- a/crates/tokscale-core/src/sessions/kiro.rs +++ b/crates/tokscale-core/src/sessions/kiro.rs @@ -11,7 +11,7 @@ use super::error::{SessionParseError, SessionParseResult}; use super::utils::{open_readonly_sqlite, parse_epoch_f64_millis, read_file}; -use super::{normalize_workspace_key, workspace_label_from_key, UnifiedMessage}; +use super::{normalize_workspace_key, workspace_label_from_key, ParsedMessage}; use crate::input_health::{InputFailure, RecordRejectionReason, ScannedInput}; use crate::TokenBreakdown; use serde::de::{IgnoredAny, MapAccess, SeqAccess, Visitor}; @@ -21,7 +21,6 @@ use std::collections::HashMap; use std::io::{BufRead, BufReader}; use std::path::Path; -const CLIENT_ID: &str = "kiro"; const PROVIDER_ID: &str = "amazon-bedrock"; #[derive(Debug, Deserialize)] @@ -350,7 +349,7 @@ pub fn parse_kiro_file(path: &Path) -> SessionParseResult { } for result in turns.into_iter().enumerate().map( - |(index, turn)| -> SessionParseResult> { + |(index, turn)| -> SessionParseResult> { let turn = serde_json::from_str::(turn.get()).map_err(|source| { SessionParseError::at_path(path, "decode Kiro turn metadata", source) })?; @@ -509,8 +508,7 @@ pub fn parse_kiro_file(path: &Path) -> SessionParseResult { ) })?; - let mut message = UnifiedMessage::new_with_dedup( - CLIENT_ID, + let mut message = ParsedMessage::new_with_dedup( model_id, PROVIDER_ID, session_id, @@ -721,8 +719,7 @@ fn parse_kiro_global_storage_file(path: &Path) -> SessionParseResult SessionParseResult { continue; } - let mut message = UnifiedMessage::new_with_dedup( - CLIENT_ID, + let mut message = ParsedMessage::new_with_dedup( model_id, PROVIDER_ID, conversation_id, @@ -953,11 +949,11 @@ mod tests { use std::io::Write; use tempfile::TempDir; - fn parse_kiro_file(path: &Path) -> Vec { + fn parse_kiro_file(path: &Path) -> Vec { super::parse_kiro_file(path).unwrap().messages } - fn parse_kiro_sqlite(path: &Path) -> Vec { + fn parse_kiro_sqlite(path: &Path) -> Vec { super::parse_kiro_sqlite(path).unwrap().messages } @@ -987,7 +983,6 @@ mod tests { let messages = parse_kiro_file(&path); assert_eq!(messages.len(), 1); - assert_eq!(messages[0].client.as_ref(), "kiro"); assert_eq!(messages[0].provider_id.as_ref(), "amazon-bedrock"); assert_eq!(messages[0].model_id.as_ref(), "claude-sonnet-4-5"); assert_eq!(messages[0].session_id.as_ref(), "session-1"); @@ -1565,7 +1560,6 @@ not json let messages = parse_kiro_file(&path); assert_eq!(messages.len(), 1); - assert_eq!(messages[0].client.as_ref(), "kiro"); assert_eq!(messages[0].model_id.as_ref(), "claude-sonnet-4-5"); assert_eq!(messages[0].session_id.as_ref(), "ide-session-1"); assert_eq!(messages[0].timestamp, 1_770_983_426_420); diff --git a/crates/tokscale-core/src/sessions/mod.rs b/crates/tokscale-core/src/sessions/mod.rs index 9693baa83..3b9d1ba28 100644 --- a/crates/tokscale-core/src/sessions/mod.rs +++ b/crates/tokscale-core/src/sessions/mod.rs @@ -1,6 +1,7 @@ //! Session parsers for different AI coding assistant formats //! -//! Each client has its own parser that converts to a unified message format. +//! Parsers produce source-neutral messages. Local-input adapters attach the +//! owning client only when emitting the final unified message stream. pub mod amp; pub mod antigravity_cli; @@ -33,41 +34,60 @@ pub mod warp; pub mod zcode; pub mod zed; -use crate::TokenBreakdown; +use crate::{clients::ClientId, TokenBreakdown}; + +macro_rules! define_message { + ($(#[$meta:meta])* $name:ident, $($source_field:tt)*) => { + $(#[$meta])* + #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] + pub struct $name { + $($source_field)* + #[serde(deserialize_with = "intern::de_intern")] + pub model_id: std::sync::Arc, + #[serde(deserialize_with = "intern::de_intern")] + pub provider_id: std::sync::Arc, + #[serde(deserialize_with = "intern::de_intern")] + pub session_id: std::sync::Arc, + /// Whether the input record directly identifies this usage as belonging to a + /// top-level session. `session_id` remains the legacy Total grouping key. + #[serde(default = "default_true")] + pub is_main_session: bool, + #[serde(default, deserialize_with = "intern::de_intern_opt")] + pub workspace_key: Option>, + #[serde(default, deserialize_with = "intern::de_intern_opt")] + pub workspace_label: Option>, + pub timestamp: i64, + pub tokens: TokenBreakdown, + pub cost: f64, + #[serde(default = "default_message_count")] + pub message_count: i32, + #[serde(default, deserialize_with = "intern::de_intern_opt")] + pub agent: Option>, + #[serde(default, deserialize_with = "intern::de_intern_opt")] + pub agent_instance: Option>, + pub dedup_key: Option, + /// True if this message is the first assistant response after a user turn. + /// Used to count user interaction turns (as opposed to API message count). + #[serde(default)] + pub is_turn_start: bool, + } + }; +} -#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] -pub struct UnifiedMessage { +define_message!( + /// Source-neutral usage produced by a session parser. + /// + /// The local-input adapter owns source identity and attributes it only after + /// parsing, cache persistence, and adapter-specific filtering have completed. + ParsedMessage, +); + +define_message!( + /// Usage attributed to a concrete client and ready for public aggregation. + UnifiedMessage, #[serde(deserialize_with = "intern::de_intern")] pub client: std::sync::Arc, - #[serde(deserialize_with = "intern::de_intern")] - pub model_id: std::sync::Arc, - #[serde(deserialize_with = "intern::de_intern")] - pub provider_id: std::sync::Arc, - #[serde(deserialize_with = "intern::de_intern")] - pub session_id: std::sync::Arc, - /// Whether the input record directly identifies this usage as belonging to a - /// top-level session. `session_id` remains the legacy Total grouping key. - #[serde(default = "default_true")] - pub is_main_session: bool, - #[serde(default, deserialize_with = "intern::de_intern_opt")] - pub workspace_key: Option>, - #[serde(default, deserialize_with = "intern::de_intern_opt")] - pub workspace_label: Option>, - pub timestamp: i64, - pub tokens: TokenBreakdown, - pub cost: f64, - #[serde(default = "default_message_count")] - pub message_count: i32, - #[serde(default, deserialize_with = "intern::de_intern_opt")] - pub agent: Option>, - #[serde(default, deserialize_with = "intern::de_intern_opt")] - pub agent_instance: Option>, - pub dedup_key: Option, - /// True if this message is the first assistant response after a user turn. - /// Used to count user interaction turns (as opposed to API message count). - #[serde(default)] - pub is_turn_start: bool, -} +); const fn default_message_count() -> i32 { 1 @@ -240,9 +260,8 @@ fn titlecase_agent(name: &str) -> String { .join(" ") } -impl UnifiedMessage { +impl ParsedMessage { pub fn new( - client: impl AsRef, model_id: impl AsRef, provider_id: impl AsRef, session_id: impl AsRef, @@ -251,7 +270,6 @@ impl UnifiedMessage { cost: f64, ) -> Self { Self::new_full( - client, model_id, provider_id, session_id, @@ -265,7 +283,6 @@ impl UnifiedMessage { #[allow(clippy::too_many_arguments)] pub fn new_with_agent( - client: impl AsRef, model_id: impl AsRef, provider_id: impl AsRef, session_id: impl AsRef, @@ -275,7 +292,6 @@ impl UnifiedMessage { agent: Option, ) -> Self { Self::new_full( - client, model_id, provider_id, session_id, @@ -289,7 +305,6 @@ impl UnifiedMessage { #[allow(clippy::too_many_arguments)] pub fn new_with_dedup( - client: impl AsRef, model_id: impl AsRef, provider_id: impl AsRef, session_id: impl AsRef, @@ -299,7 +314,6 @@ impl UnifiedMessage { dedup_key: Option, ) -> Self { Self::new_full( - client, model_id, provider_id, session_id, @@ -313,7 +327,6 @@ impl UnifiedMessage { #[allow(clippy::too_many_arguments)] fn new_full( - client: impl AsRef, model_id: impl AsRef, provider_id: impl AsRef, session_id: impl AsRef, @@ -324,7 +337,6 @@ impl UnifiedMessage { dedup_key: Option, ) -> Self { Self { - client: intern::intern(client.as_ref()), model_id: intern::intern(model_id.as_ref()), provider_id: intern::intern(provider_id.as_ref()), session_id: intern::intern(session_id.as_ref()), @@ -342,44 +354,159 @@ impl UnifiedMessage { } } - /// Local calendar date derived from `timestamp`; `None` when the - /// timestamp is outside the representable range. - pub fn local_date(&self) -> Option { - use chrono::TimeZone; - match chrono::Local.timestamp_millis_opt(self.timestamp) { - chrono::LocalResult::Single(dt) => Some(dt.date_naive()), - _ => None, + pub(crate) fn attribute(self, client: ClientId) -> UnifiedMessage { + self.with_client(client.as_str()) + } + + fn with_client(self, client: impl AsRef) -> UnifiedMessage { + let Self { + model_id, + provider_id, + session_id, + is_main_session, + workspace_key, + workspace_label, + timestamp, + tokens, + cost, + message_count, + agent, + agent_instance, + dedup_key, + is_turn_start, + } = self; + UnifiedMessage { + client: intern::intern(client.as_ref()), + model_id, + provider_id, + session_id, + is_main_session, + workspace_key, + workspace_label, + timestamp, + tokens, + cost, + message_count, + agent, + agent_instance, + dedup_key, + is_turn_start, } } +} - /// Local `YYYY-MM-DD` string derived from `timestamp` (empty when out of - /// range). Allocates; prefer [`Self::local_date`] in hot paths. - pub fn date_string(&self) -> String { - timestamp_to_date(self.timestamp) +impl UnifiedMessage { + pub fn new( + client: impl AsRef, + model_id: impl AsRef, + provider_id: impl AsRef, + session_id: impl AsRef, + timestamp: i64, + tokens: TokenBreakdown, + cost: f64, + ) -> Self { + ParsedMessage::new(model_id, provider_id, session_id, timestamp, tokens, cost) + .with_client(client) } - pub fn set_workspace( - &mut self, - workspace_key: Option, - workspace_label: Option, - ) { - self.workspace_key = workspace_key.as_deref().map(intern::intern); - self.workspace_label = workspace_label.as_deref().map(intern::intern); + #[allow(clippy::too_many_arguments)] + pub fn new_with_agent( + client: impl AsRef, + model_id: impl AsRef, + provider_id: impl AsRef, + session_id: impl AsRef, + timestamp: i64, + tokens: TokenBreakdown, + cost: f64, + agent: Option, + ) -> Self { + ParsedMessage::new_with_agent( + model_id, + provider_id, + session_id, + timestamp, + tokens, + cost, + agent, + ) + .with_client(client) } - pub fn set_agent_instance(&mut self, agent_instance: Option) { - self.agent_instance = agent_instance.as_deref().map(intern::intern); + #[allow(clippy::too_many_arguments)] + pub fn new_with_dedup( + client: impl AsRef, + model_id: impl AsRef, + provider_id: impl AsRef, + session_id: impl AsRef, + timestamp: i64, + tokens: TokenBreakdown, + cost: f64, + dedup_key: Option, + ) -> Self { + ParsedMessage::new_with_dedup( + model_id, + provider_id, + session_id, + timestamp, + tokens, + cost, + dedup_key, + ) + .with_client(client) } +} - pub(crate) fn refresh_derived_fields(&mut self) { - if let Some(provider) = crate::provider_identity::provider_override_from_model_and_provider( - &self.model_id, - &self.provider_id, - ) { - self.provider_id = intern::intern(provider); +macro_rules! impl_message_methods { + ($name:ident) => { + impl $name { + /// Local calendar date derived from `timestamp`; `None` when the + /// timestamp is outside the representable range. + pub fn local_date(&self) -> Option { + use chrono::TimeZone; + match chrono::Local.timestamp_millis_opt(self.timestamp) { + chrono::LocalResult::Single(dt) => Some(dt.date_naive()), + _ => None, + } + } + + /// Local `YYYY-MM-DD` string derived from `timestamp` (empty when out of + /// range). Allocates; prefer [`Self::local_date`] in hot paths. + pub fn date_string(&self) -> String { + timestamp_to_date(self.timestamp) + } + + pub fn set_workspace( + &mut self, + workspace_key: Option, + workspace_label: Option, + ) { + self.workspace_key = workspace_key.as_deref().map(intern::intern); + self.workspace_label = workspace_label.as_deref().map(intern::intern); + } + + pub fn set_agent_instance(&mut self, agent_instance: Option) { + self.agent_instance = agent_instance.as_deref().map(intern::intern); + } + + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn refresh_derived_fields(&mut self) { + if let Some(provider) = + crate::provider_identity::provider_override_from_model_and_provider( + &self.model_id, + &self.provider_id, + ) + { + self.provider_id = intern::intern(provider); + } + } } - } + }; +} + +impl_message_methods!(ParsedMessage); +impl_message_methods!(UnifiedMessage); +impl ParsedMessage { pub(crate) fn set_timestamp(&mut self, timestamp: i64) { self.timestamp = timestamp; self.refresh_derived_fields(); @@ -521,6 +648,29 @@ mod tests { assert_eq!(msg.workspace_label, None); } + #[test] + fn parsed_message_has_no_source_until_adapter_attribution() { + let parsed = ParsedMessage::new( + "gpt-5.5", + "openai", + "session", + 1_733_011_200_000, + TokenBreakdown { + input: 10, + output: 2, + ..Default::default() + }, + 0.0, + ); + + let serialized = serde_json::to_value(&parsed).unwrap(); + assert!(serialized.get("client").is_none()); + + let attributed = parsed.attribute(ClientId::Codex); + assert_eq!(attributed.client.as_ref(), ClientId::Codex.as_str()); + assert_eq!(attributed.model_id.as_ref(), "gpt-5.5"); + } + #[test] fn test_refresh_derived_fields_overrides_historical_model_provider() { let tokens = TokenBreakdown::default(); diff --git a/crates/tokscale-core/src/sessions/mux.rs b/crates/tokscale-core/src/sessions/mux.rs index 8aeb37265..2659c72ee 100644 --- a/crates/tokscale-core/src/sessions/mux.rs +++ b/crates/tokscale-core/src/sessions/mux.rs @@ -3,7 +3,7 @@ //! Parses session-usage.json files from ~/.mux/sessions//session-usage.json use super::error::{SessionParseError, SessionParseResult}; -use super::UnifiedMessage; +use super::ParsedMessage; use crate::input_health::{RecordRejectionReason, ScannedInput}; use crate::{model_aliases, provider_identity, TokenBreakdown}; use serde::Deserialize; @@ -42,7 +42,7 @@ pub struct MuxLastRequest { } /// Parse a mux session-usage.json file. -/// Returns one UnifiedMessage per model entry in byModel. +/// Returns one ParsedMessage per model entry in byModel. pub fn parse_mux_file(path: &Path) -> SessionParseResult { let data = std::fs::read(path) .map_err(|error| SessionParseError::at_path(path, "read file", error))?; @@ -161,8 +161,7 @@ pub fn parse_mux_file(path: &Path) -> SessionParseResult { .unwrap_or_else(|| raw_model_id.to_string()); let provider = provider_identity::observed_provider_id(raw_provider, &model_id); - scanned.messages.push(UnifiedMessage::new_with_dedup( - "mux", + scanned.messages.push(ParsedMessage::new_with_dedup( model_id, provider, session_id.clone(), @@ -210,7 +209,7 @@ fn invalid_at_path( mod tests { use super::*; - fn parse_mux_file(path: &Path) -> Vec { + fn parse_mux_file(path: &Path) -> Vec { super::parse_mux_file(path).unwrap().messages } use std::io::Write; @@ -272,7 +271,6 @@ mod tests { .iter() .find(|m| m.model_id.as_ref() == "claude-opus-4.6") .unwrap(); - assert_eq!(claude.client.as_ref(), "mux"); assert_eq!(claude.provider_id.as_ref(), "anthropic"); assert_eq!(claude.tokens.input, 100); assert_eq!(claude.tokens.cache_read, 5000); diff --git a/crates/tokscale-core/src/sessions/openclaw.rs b/crates/tokscale-core/src/sessions/openclaw.rs index 30e30fb7d..fa1afab18 100644 --- a/crates/tokscale-core/src/sessions/openclaw.rs +++ b/crates/tokscale-core/src/sessions/openclaw.rs @@ -4,7 +4,7 @@ //! Current-format inputs are individual transcript files. use super::error::{SessionParseError, SessionParseResult}; -use super::UnifiedMessage; +use super::ParsedMessage; use crate::input_health::{InputFailure, RecordRejectionReason, ScannedInput}; use crate::{model_aliases, provider_identity, TokenBreakdown}; use serde::Deserialize; @@ -234,8 +234,8 @@ fn parse_openclaw_session( current_model = Some(model.clone()); current_provider = Some(provider.clone()); - scanned.messages.push(UnifiedMessage::new( - "openclaw", model, provider, session_id, timestamp, tokens, 0.0, + scanned.messages.push(ParsedMessage::new( + model, provider, session_id, timestamp, tokens, 0.0, )); } } @@ -300,13 +300,13 @@ mod tests { use std::io::Write; use tempfile::TempDir; - fn parse_openclaw_session(path: &Path, session_id: &str) -> Vec { + fn parse_openclaw_session(path: &Path, session_id: &str) -> Vec { super::parse_openclaw_session(path, session_id) .unwrap() .messages } - fn parse_openclaw_transcript(path: &Path) -> Vec { + fn parse_openclaw_transcript(path: &Path) -> Vec { super::parse_openclaw_transcript(path).unwrap().messages } diff --git a/crates/tokscale-core/src/sessions/opencode.rs b/crates/tokscale-core/src/sessions/opencode.rs index 050867aaf..a3b200d80 100644 --- a/crates/tokscale-core/src/sessions/opencode.rs +++ b/crates/tokscale-core/src/sessions/opencode.rs @@ -1,8 +1,7 @@ //! OpenCode current-format SQLite session parser. use super::{ - normalize_opencode_agent_name, normalize_workspace_key, workspace_label_from_key, - UnifiedMessage, + normalize_opencode_agent_name, normalize_workspace_key, workspace_label_from_key, ParsedMessage, }; use crate::input_health::{InputFailure, RecordRejectionReason, ScannedInput}; use crate::TokenBreakdown; @@ -304,13 +303,13 @@ fn workspace_from_root(root: Option<&str>) -> (Option, Option) { (workspace_key, workspace_label) } -fn set_workspace_from_root(message: &mut UnifiedMessage, root: Option<&str>) { +fn set_workspace_from_root(message: &mut ParsedMessage, root: Option<&str>) { let (workspace_key, workspace_label) = workspace_from_root(root); message.set_workspace(workspace_key, workspace_label); } fn merge_duplicate_workspace( - message: &mut UnifiedMessage, + message: &mut ParsedMessage, state: &mut OpenCodeSqliteDedupState, root: Option<&str>, ) { @@ -615,8 +614,7 @@ pub fn parse_opencode_sqlite(db_path: &Path) -> Result Result Result, } +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum PiFormat { + Pi, + Omp, +} + +impl PiFormat { + const fn usage_validation_operation(self) -> &'static str { + match self { + Self::Pi => "validate Pi assistant message", + Self::Omp => "validate OMP assistant message", + } + } + + const fn is_omp(self) -> bool { + matches!(self, Self::Omp) + } +} + /// Parse a Pi JSONL session file pub fn parse_pi_file(path: &Path) -> SessionParseResult { - parse_pi_format_file(path, "pi", None) + parse_pi_format_file(path, PiFormat::Pi, None) } /// Parse an OMP JSONL session file. pub fn parse_omp_file(path: &Path) -> SessionParseResult { let parent_index = build_omp_parent_task_agent_index(&[path.to_path_buf()]); - parse_pi_format_file(path, "omp", Some(&parent_index)) + parse_pi_format_file(path, PiFormat::Omp, Some(&parent_index)) } pub fn parse_omp_file_with_parent_task_agent_index( path: &Path, parent_task_agent_index: &OmpParentTaskAgentIndex, ) -> SessionParseResult { - parse_pi_format_file(path, "omp", Some(parent_task_agent_index)) -} - -fn usage_validation_operation(client: &str) -> &'static str { - match client { - "pi" => "validate Pi assistant message", - "omp" => "validate OMP assistant message", - _ => unreachable!("Pi-format parser only supports Pi and OMP"), - } + parse_pi_format_file(path, PiFormat::Omp, Some(parent_task_agent_index)) } -fn required_usage_value(value: Option, field: &str, client: &str) -> SessionParseResult { +fn required_usage_value( + value: Option, + field: &str, + pi_format: PiFormat, +) -> SessionParseResult { value.ok_or_else(|| { SessionParseError::invalid( - usage_validation_operation(client), + pi_format.usage_validation_operation(), format!("current usage is missing `{field}`"), ) }) @@ -277,11 +292,11 @@ fn required_usage_value(value: Option, field: &str, client: &str) -> Sessio fn validate_nonnegative_usage_value( field: &str, value: i64, - client: &str, + pi_format: PiFormat, ) -> SessionParseResult<()> { if value < 0 { return Err(SessionParseError::invalid( - usage_validation_operation(client), + pi_format.usage_validation_operation(), format!("token counts must not be negative: `{field}`"), )); } @@ -291,12 +306,12 @@ fn validate_nonnegative_usage_value( fn checked_usage_sum( values: impl IntoIterator, description: &str, - client: &str, + pi_format: PiFormat, ) -> SessionParseResult { values.into_iter().try_fold(0_i64, |total, value| { total.checked_add(value).ok_or_else(|| { SessionParseError::invalid( - usage_validation_operation(client), + pi_format.usage_validation_operation(), format!("{description} exceeds i64::MAX"), ) }) @@ -305,13 +320,13 @@ fn checked_usage_sum( fn token_breakdown_from_pi_usage( usage: &PiUsage, - client: &str, + pi_format: PiFormat, ) -> SessionParseResult { - let input = required_usage_value(usage.input, "input", client)?; - let raw_output = required_usage_value(usage.output, "output", client)?; - let cache_read = required_usage_value(usage.cache_read, "cacheRead", client)?; - let cache_write = required_usage_value(usage.cache_write, "cacheWrite", client)?; - let input_total = required_usage_value(usage.total_tokens, "totalTokens", client)?; + let input = required_usage_value(usage.input, "input", pi_format)?; + let raw_output = required_usage_value(usage.output, "output", pi_format)?; + let cache_read = required_usage_value(usage.cache_read, "cacheRead", pi_format)?; + let cache_write = required_usage_value(usage.cache_write, "cacheWrite", pi_format)?; + let input_total = required_usage_value(usage.total_tokens, "totalTokens", pi_format)?; let ( reasoning_field, @@ -319,26 +334,26 @@ fn token_breakdown_from_pi_usage( orchestration_input, orchestration_output, orchestration_cache_read, - ) = match client { - "pi" => { + ) = match pi_format { + PiFormat::Pi => { if usage.reasoning_tokens.is_some() { return Err(SessionParseError::invalid( - usage_validation_operation(client), + pi_format.usage_validation_operation(), "Pi usage must use `reasoning`, not OMP `reasoningTokens`", )); } if usage.orchestration.is_some() { return Err(SessionParseError::invalid( - usage_validation_operation(client), + pi_format.usage_validation_operation(), "Pi usage must not contain OMP `orchestration` tokens", )); } ("reasoning", usage.reasoning.unwrap_or(0), 0, 0, 0) } - "omp" => { + PiFormat::Omp => { if usage.reasoning.is_some() { return Err(SessionParseError::invalid( - usage_validation_operation(client), + pi_format.usage_validation_operation(), "OMP usage must use `reasoningTokens`, not Pi `reasoning`", )); } @@ -353,7 +368,6 @@ fn token_breakdown_from_pi_usage( .unwrap_or(0), ) } - _ => unreachable!("Pi-format parser only supports Pi and OMP"), }; for (field, value) in [ @@ -367,7 +381,7 @@ fn token_breakdown_from_pi_usage( ("orchestration.output", orchestration_output), ("orchestration.cacheRead", orchestration_cache_read), ] { - validate_nonnegative_usage_value(field, value, client)?; + validate_nonnegative_usage_value(field, value, pi_format)?; } // Xiaomi MiMo token-plan has emitted a length-stopped response with @@ -388,11 +402,11 @@ fn token_breakdown_from_pi_usage( orchestration_cache_read, ], "reported totalTokens", - client, + pi_format, )?; if input_total != expected_input_total { return Err(SessionParseError::invalid( - usage_validation_operation(client), + pi_format.usage_validation_operation(), format!( "reported totalTokens is {input_total}, expected {expected_input_total} from usage buckets" ), @@ -403,30 +417,30 @@ fn token_breakdown_from_pi_usage( input: checked_usage_sum( [input, orchestration_input], "normalized input token count", - client, + pi_format, )?, output: checked_usage_sum( [raw_output - reasoning, orchestration_output], "normalized output token count", - client, + pi_format, )?, cache_read: checked_usage_sum( [cache_read, orchestration_cache_read], "normalized cache-read token count", - client, + pi_format, )?, cache_write, reasoning, }; let normalized_total = tokens.checked_total().ok_or_else(|| { SessionParseError::invalid( - usage_validation_operation(client), + pi_format.usage_validation_operation(), "normalized token total exceeds i64::MAX", ) })?; if normalized_total != input_total { return Err(SessionParseError::invalid( - usage_validation_operation(client), + pi_format.usage_validation_operation(), format!( "normalized token total is {normalized_total}, expected reported totalTokens {input_total}" ), @@ -848,7 +862,7 @@ fn parse_pi_header_line( fn parse_pi_format_file( path: &Path, - client: &'static str, + pi_format: PiFormat, omp_parent_task_agent_index: Option<&OmpParentTaskAgentIndex>, ) -> SessionParseResult { let file = std::fs::File::open(path) @@ -864,12 +878,12 @@ fn parse_pi_format_file( .file_stem() .and_then(|stem| stem.to_str()) .map(str::to_string); - let omp_parent_scan = if client == "omp" { + let omp_parent_scan = if pi_format.is_omp() { omp_parent_task_agent_index.and_then(|index| index.parent_scan_for_child(path)) } else { None }; - let omp_subagent_label = if client == "omp" { + let omp_subagent_label = if pi_format.is_omp() { match child_stem.as_deref() { Some(stem) => { if let Some(label) = normalize_omp_advisor_label(stem) { @@ -895,7 +909,7 @@ fn parse_pi_format_file( None }; let is_main_session = - client != "omp" || (omp_parent_scan.is_none() && omp_subagent_label.is_none()); + !pi_format.is_omp() || (omp_parent_scan.is_none() && omp_subagent_label.is_none()); let mut session_id: Option = None; let mut workspace_key: Option = None; @@ -926,7 +940,7 @@ fn parse_pi_format_file( let parsed_header = match parse_pi_header_line( trimmed, &mut buffer, - client == "omp" && !saw_omp_title_slot, + pi_format.is_omp() && !saw_omp_title_slot, ) { Ok(parsed) => parsed, Err(_error) => { @@ -980,7 +994,7 @@ fn parse_pi_format_file( None => continue, }; - let tokens = match token_breakdown_from_pi_usage(&usage, client) { + let tokens = match token_breakdown_from_pi_usage(&usage, pi_format) { Ok(tokens) => tokens, Err(_) => { record_pi_rejection(&mut scanned); @@ -1021,8 +1035,7 @@ fn parse_pi_format_file( } }; - let mut unified = UnifiedMessage::new( - client, + let mut parsed = ParsedMessage::new( model, provider, session_id.clone().expect( @@ -1032,13 +1045,13 @@ fn parse_pi_format_file( tokens, 0.0, ); - unified.is_main_session = is_main_session; - unified.set_workspace(workspace_key.clone(), workspace_label.clone()); - unified.agent = omp_subagent_label + parsed.is_main_session = is_main_session; + parsed.set_workspace(workspace_key.clone(), workspace_label.clone()); + parsed.agent = omp_subagent_label .as_deref() .map(crate::sessions::intern::intern); - unified.set_agent_instance(child_stem.clone()); - scanned.messages.push(unified); + parsed.set_agent_instance(child_stem.clone()); + scanned.messages.push(parsed); } if session_id.is_none() { @@ -1105,7 +1118,6 @@ mod tests { // then assert_eq!(messages.len(), 1); - assert_eq!(messages[0].client.as_ref(), "pi"); assert_eq!(messages[0].session_id.as_ref(), "pi_ses_001"); assert_eq!(messages[0].model_id.as_ref(), "claude-sonnet-4.6"); assert_eq!(messages[0].provider_id.as_ref(), "anthropic"); @@ -1181,7 +1193,7 @@ mod tests { } #[test] - fn test_parse_omp_jsonl_uses_omp_client() { + fn test_parse_omp_jsonl_uses_omp_format() { // given let content = r#"{"type":"session","id":"omp_ses_001","timestamp":"2026-01-01T00:00:00.000Z","cwd":"/tmp"} {"type":"message","id":"msg_001","parentId":null,"timestamp":"2026-01-01T00:00:01.000Z","message":{"role":"assistant","model":"gpt-5.5","provider":"openai","usage":{"input":20,"output":10,"cacheRead":0,"cacheWrite":0,"totalTokens":30}}}"#; @@ -1192,7 +1204,6 @@ mod tests { // then assert_eq!(messages.len(), 1); - assert_eq!(messages[0].client.as_ref(), "omp"); assert_eq!(messages[0].session_id.as_ref(), "omp_ses_001"); assert_eq!(messages[0].model_id.as_ref(), "gpt-5.5"); assert_eq!(messages[0].provider_id.as_ref(), "openai"); @@ -1209,7 +1220,6 @@ mod tests { let messages = parse_omp_file(file.path()).unwrap().messages; assert_eq!(messages.len(), 1); - assert_eq!(messages[0].client.as_ref(), "omp"); assert_eq!(messages[0].session_id.as_ref(), "omp_ses_title"); assert_eq!(messages[0].tokens.output, 8); assert_eq!(messages[0].tokens.reasoning, 2); diff --git a/crates/tokscale-core/src/sessions/qwen.rs b/crates/tokscale-core/src/sessions/qwen.rs index 380927c8d..06a493dda 100644 --- a/crates/tokscale-core/src/sessions/qwen.rs +++ b/crates/tokscale-core/src/sessions/qwen.rs @@ -5,7 +5,7 @@ use super::error::{SessionParseError, SessionParseResult}; use super::utils::parse_timestamp_str; -use super::{normalize_workspace_key, workspace_label_from_key, UnifiedMessage}; +use super::{normalize_workspace_key, workspace_label_from_key, ParsedMessage}; use crate::input_health::{InputFailure, RecordRejectionReason, ScannedInput}; use crate::model_aliases; use crate::TokenBreakdown; @@ -155,8 +155,7 @@ pub fn parse_qwen_file(path: &Path) -> SessionParseResult { crate::sessions::dedup_hash_str(&format!("qwen:{line_session_id}:{message_index}")); message_index += 1; - let mut unified = UnifiedMessage::new_with_dedup( - "qwen", + let mut message = ParsedMessage::new_with_dedup( model, DEFAULT_PROVIDER, line_session_id, @@ -171,8 +170,8 @@ pub fn parse_qwen_file(path: &Path) -> SessionParseResult { 0.0, // Cost calculated later by pricing resolver Some(dedup_key), ); - unified.set_workspace(workspace_key.clone(), workspace_label.clone()); - scanned.messages.push(unified); + message.set_workspace(workspace_key.clone(), workspace_label.clone()); + scanned.messages.push(message); } Ok(scanned) @@ -202,7 +201,7 @@ mod tests { use std::path::Path; use tempfile::{NamedTempFile, TempDir}; - fn parse_qwen_file(path: &Path) -> Vec { + fn parse_qwen_file(path: &Path) -> Vec { super::parse_qwen_file(path).unwrap().messages } @@ -247,7 +246,6 @@ mod tests { let messages = parse_qwen_file(file.path()); assert_eq!(messages.len(), 1); - assert_eq!(messages[0].client.as_ref(), "qwen"); assert_eq!(messages[0].model_id.as_ref(), "qwen3.5-plus"); assert_eq!(messages[0].provider_id.as_ref(), "qwen"); // Session ID comes from filename, not JSON content (temp file has random name) diff --git a/crates/tokscale-core/src/sessions/roocode.rs b/crates/tokscale-core/src/sessions/roocode.rs index 23ec4f496..452f429ca 100644 --- a/crates/tokscale-core/src/sessions/roocode.rs +++ b/crates/tokscale-core/src/sessions/roocode.rs @@ -6,7 +6,7 @@ use super::error::{SessionParseError, SessionParseResult}; use super::utils::parse_timestamp_str; -use super::{normalize_agent_name, UnifiedMessage}; +use super::{normalize_agent_name, ParsedMessage}; use crate::input_health::{RecordRejectionReason, RejectionSummary, ScannedInput}; use crate::{provider_identity, TokenBreakdown}; use serde_json::Value; @@ -104,8 +104,7 @@ pub fn parse_roocode_file(path: &Path) -> SessionParseResult { provider.as_deref().unwrap_or_default(), &model_id, ); - messages.push(UnifiedMessage::new_with_agent( - "roocode", + messages.push(ParsedMessage::new_with_agent( model_id.clone(), provider, session_id.clone(), @@ -301,7 +300,7 @@ mod tests { use std::fs; use tempfile::TempDir; - fn parse_roocode_file(path: &Path) -> Vec { + fn parse_roocode_file(path: &Path) -> Vec { super::parse_roocode_file(path).unwrap().messages } @@ -348,7 +347,6 @@ after"#; let messages = parse_roocode_file(&path); assert_eq!(messages.len(), 1); - assert_eq!(messages[0].client.as_ref(), "roocode"); assert_eq!(messages[0].model_id.as_ref(), "claude-sonnet-4"); assert_eq!(messages[0].provider_id.as_ref(), "anthropic"); assert_eq!(messages[0].session_id.as_ref(), "task-abc"); diff --git a/crates/tokscale-core/src/sessions/warp.rs b/crates/tokscale-core/src/sessions/warp.rs index bc3ae22ba..50889ac82 100644 --- a/crates/tokscale-core/src/sessions/warp.rs +++ b/crates/tokscale-core/src/sessions/warp.rs @@ -7,7 +7,7 @@ use super::error::{SessionParseError, SessionParseResult}; use super::utils::open_readonly_sqlite; -use super::{normalize_workspace_key, workspace_label_from_key, UnifiedMessage}; +use super::{normalize_workspace_key, workspace_label_from_key, ParsedMessage}; use crate::input_health::{InputFailure, RecordRejectionReason, ScannedInput}; use crate::{model_aliases, provider_identity, token_imputation}; use chrono::TimeZone; @@ -16,8 +16,6 @@ use serde_json::Value; use std::collections::HashMap; use std::path::Path; -const CLIENT_ID: &str = "warp"; - #[derive(Debug, Clone, Default)] struct ConversationMeta { workspace_key: Option, @@ -208,8 +206,7 @@ pub fn parse_warp_sqlite(db_path: &Path) -> SessionParseResult { .into_iter() .zip(token_rows) .map(|(pending, tokens)| { - let mut message = UnifiedMessage::new_with_dedup( - CLIENT_ID, + let mut message = ParsedMessage::new_with_dedup( pending.model_id, pending.provider_id, pending.conversation_id, @@ -495,7 +492,6 @@ mod tests { crate::finalize_token_priced_messages(&mut messages, None); assert_eq!(messages.len(), 2); - assert_eq!(messages[0].client.as_ref(), "warp"); assert_eq!(messages[0].session_id.as_ref(), "conversation-1"); assert_eq!(messages[0].model_id.as_ref(), "claude-opus-4.6"); assert_eq!(messages[0].provider_id.as_ref(), "anthropic"); diff --git a/crates/tokscale-core/src/sessions/zcode.rs b/crates/tokscale-core/src/sessions/zcode.rs index 30812a984..159dd6f1c 100644 --- a/crates/tokscale-core/src/sessions/zcode.rs +++ b/crates/tokscale-core/src/sessions/zcode.rs @@ -6,14 +6,13 @@ use super::error::{SessionParseError, SessionParseResult}; use super::utils::parse_timestamp_str; -use super::{dedup_hash_str, normalize_workspace_key, workspace_label_from_key, UnifiedMessage}; +use super::{dedup_hash_str, normalize_workspace_key, workspace_label_from_key, ParsedMessage}; use crate::input_health::{InputFailure, RecordRejectionReason, ScannedInput}; use crate::{checked_token_sum, TokenBreakdown}; use serde::Deserialize; use std::io::{BufRead, BufReader}; use std::path::Path; -const CLIENT_ID: &str = "zcode"; const PROVIDER_ID: &str = "zai"; #[derive(Debug, Deserialize)] @@ -365,8 +364,7 @@ pub fn parse_zcode_file(path: &Path) -> SessionParseResult { let dedup_key = dedup_hash_str(&format!("zcode:{resolved_session}:{assistant_index}")); - let mut message = UnifiedMessage::new_with_dedup( - CLIENT_ID, + let mut message = ParsedMessage::new_with_dedup( resolved_model, PROVIDER_ID, resolved_session, @@ -445,7 +443,7 @@ fn workspace_key_from_path(path: &Path) -> Option { mod tests { use super::*; - fn parse_zcode_file(path: &Path) -> Vec { + fn parse_zcode_file(path: &Path) -> Vec { super::parse_zcode_file(path).unwrap().messages } use serde_json::json; @@ -490,7 +488,6 @@ mod tests { assert_eq!(messages.len(), 1); let msg = &messages[0]; - assert_eq!(msg.client.as_ref(), "zcode"); assert_eq!(msg.provider_id.as_ref(), "zai"); assert_eq!(msg.model_id.as_ref(), "glm-5.2"); assert_eq!(msg.session_id.as_ref(), "s1"); diff --git a/crates/tokscale-core/src/sessions/zed.rs b/crates/tokscale-core/src/sessions/zed.rs index c1af2b177..5929cd748 100644 --- a/crates/tokscale-core/src/sessions/zed.rs +++ b/crates/tokscale-core/src/sessions/zed.rs @@ -11,7 +11,7 @@ use super::error::{SessionParseError, SessionParseResult}; use super::utils::{open_readonly_sqlite, parse_timestamp_str}; -use super::{normalize_workspace_key, workspace_label_from_key, UnifiedMessage}; +use super::{normalize_workspace_key, workspace_label_from_key, ParsedMessage}; use crate::input_health::{InputFailure, RecordRejectionReason, ScannedInput}; use crate::TokenBreakdown; use serde_json::Value; @@ -34,7 +34,7 @@ struct ZedThreadRow { /// What one thread row contributed to the scan. enum ThreadOutcome { - Message(Box), + Message(Box), /// Skipped by the input contract (imported, non-hosted, zero usage). Filtered, Rejected(RecordRejectionReason, String), @@ -205,8 +205,7 @@ fn parse_thread_row(row: ZedThreadRow) -> ThreadOutcome { ); }; - let mut message = UnifiedMessage::new_with_dedup( - "zed", + let mut message = ParsedMessage::new_with_dedup( model_id, ZED_HOSTED_PROVIDER, row.id.clone(), @@ -430,7 +429,7 @@ mod tests { use std::fs; use tempfile::TempDir; - fn parse_zed_sqlite(path: &Path) -> Vec { + fn parse_zed_sqlite(path: &Path) -> Vec { let scanned = super::parse_zed_sqlite(path).unwrap(); assert!(scanned.interrupted.is_none()); assert!(scanned.rejections.is_empty()); @@ -569,7 +568,6 @@ mod tests { assert_eq!(messages.len(), 1); let message = &messages[0]; - assert_eq!(message.client.as_ref(), "zed"); assert_eq!(message.provider_id.as_ref(), ZED_HOSTED_PROVIDER); assert_eq!(message.model_id.as_ref(), "claude-sonnet-4-5"); assert_eq!(message.session_id.as_ref(), "thread-1"); diff --git a/crates/tokscale-core/tests/codebuff.rs b/crates/tokscale-core/tests/codebuff.rs index aa8b71642..ac0ab7c95 100644 --- a/crates/tokscale-core/tests/codebuff.rs +++ b/crates/tokscale-core/tests/codebuff.rs @@ -72,7 +72,6 @@ fn test_parse_codebuff_emits_one_event_per_assistant_message_with_usage() { assert_eq!(msgs.len(), 2); let first = &msgs[0]; - assert_eq!(first.client.as_ref(), "codebuff"); assert_eq!(first.model_id.as_ref(), "claude-sonnet-4-20250514"); assert_eq!(first.provider_id.as_ref(), "anthropic"); assert_eq!(first.tokens.input, 500); diff --git a/crates/tokscale-core/tests/hermes.rs b/crates/tokscale-core/tests/hermes.rs index b1c7c85c7..161d1b939 100644 --- a/crates/tokscale-core/tests/hermes.rs +++ b/crates/tokscale-core/tests/hermes.rs @@ -64,7 +64,6 @@ fn test_parse_hermes_sqlite_reads_session_rows_and_preserves_message_count() { assert_eq!(messages.len(), 1); let msg = &messages[0]; - assert_eq!(msg.client.as_ref(), "hermes"); assert_eq!(msg.agent.as_deref(), Some("Hermes Agent")); assert_eq!(msg.session_id.as_ref(), "session-1"); assert_eq!(msg.model_id.as_ref(), "claude-sonnet-4"); diff --git a/docs/adr/0007-client-identity-catalog.md b/docs/adr/0007-client-identity-catalog.md index 5fe452324..854a715cc 100644 --- a/docs/adr/0007-client-identity-catalog.md +++ b/docs/adr/0007-client-identity-catalog.md @@ -45,11 +45,21 @@ definition. These three sets must have exact parity and no duplicate entries. The scan definition and the adapter's discovery implementation are the sole authority for fixed default roots beneath the selected home, filename -selection, companion files, database sidecars, and custom-root support. +selection, companion files, database sidecars, custom-root support, and source +identity. An adapter stamps its generated `InputUnit` with its catalog +`ClientId`; that value is the only source identity carried through acquisition. Additional roots come only from `scanner.extraScanPaths`; OpenCode database -files come only from `scanner.opencodeDbPaths`. The adapter's session -schema/parser is the sole authority for accepted envelopes, database schemas, -required fields, record semantics, deduplication, and token interpretation. +files come only from `scanner.opencodeDbPaths`. + +The adapter's session schema/parser is the sole authority for accepted +envelopes, database schemas, required fields, record semantics, deduplication, +and token interpretation, but not source identity. Parsers produce +source-neutral `ParsedMessage` values. Derived message-cache shards also store +only source-neutral messages. The sequential adapter fold attributes each +message from `InputUnit.client` when it emits the final `UnifiedMessage`; parser +modules do not repeat catalog IDs as message data. Before any adapter executes, +the pipeline verifies that every discovered unit carries that adapter's +`ClientId` and rejects a mismatched batch as a contract error. `docs/clients.md` is the user-facing discovery map generated from that contract. It does not create a second path or schema authority. @@ -93,7 +103,7 @@ Adding a client requires one atomic contract change containing: 1. one catalog identity; 2. one scan definition; 3. one registered local-input adapter; -4. one current session schema/parser; +4. one current source-neutral session schema/parser; 5. focused discovery, parser, and health tests; 6. catalog/scan/adapter parity checks; and 7. a current discovery row in `docs/clients.md`. @@ -107,5 +117,7 @@ requires a parser revision change. Each accepted local integration has one public identity and one executable input contract. Reports, filters, scanner configuration, caches, and TUI views -therefore share the same client namespace, while path and format evolution stays -owned by the adapter and schema that can validate it. +therefore share the same client namespace. A parser cannot silently create a +second source namespace because its output and persisted cache representation +have no client field; path and format evolution stays owned by the adapter and +schema that can validate it. diff --git a/docs/adr/0008-single-copy-memory-pipeline.md b/docs/adr/0008-single-copy-memory-pipeline.md index af5e2ca8e..73b0815af 100644 --- a/docs/adr/0008-single-copy-memory-pipeline.md +++ b/docs/adr/0008-single-copy-memory-pipeline.md @@ -34,6 +34,13 @@ Discovery produces one consumptive `PreparedLocalInputs` inventory. It records: inputs; and - compact pre-execution metadata snapshots. +Each `InputUnit.client` is adapter-issued provenance from ADR 0007. Session +parsers produce `ParsedMessage`, which intentionally has no source identity, and +message-cache bodies persist that same source-neutral type. After cache writes, +adapter-specific enrichment, filtering, and deduplication, the sequential fold +combines each parsed message with the unit's `ClientId` and emits a +`UnifiedMessage`. No parser or cache shard can override that attribution. + Freshness probing and execution consume the same inventory and never rediscover inputs. A file added after preparation belongs to the next inventory. Related inputs include SQLite WAL files, Claude `.meta.json`, optional workspace @@ -112,10 +119,12 @@ hit bodies before that index is built, and then folds hits and misses in bounded order. Codex retains its dedicated exact-hit, stale, append, and recovery path because its incremental parse state is not the generic adapter contract. -Codex cold parses, append merges, and race reparses own one raw message vector. +Codex cold parses, append merges, and race reparses own one raw +`ParsedMessage` vector. When cacheable, the fold serializes a borrowed raw slice before applying timestamp completion, token normalization/filtering, canonical identity, -pricing, and exec-session attribution in place. +pricing, and exec-session attribution in place. Client attribution occurs only +after those operations when the adapter emits the messages. OpenCode borrows potentially large message TEXT. It stream-validates the full JSON document and required role envelope, then fully decodes assistant payloads @@ -124,7 +133,9 @@ instead of becoming empty usage. ### Message representation and aggregation -`UnifiedMessage` stores no redundant derivable value: +`ParsedMessage` is the source-neutral parser and shard representation. +`UnifiedMessage` adds the adapter-attributed client for public aggregation and +otherwise stores no redundant derivable value: - date is derived from timestamp; - `dedup_key` is a 64-bit hash rather than a formatted string; and @@ -143,10 +154,16 @@ buckets. Composite identities never use delimiter-concatenated public keys. Persisted map keys are versioned, variant-tagged, byte-length-prefixed values with a distinct unknown-workspace tag. -Session client-space accounting is the deduplicated byte size of the input -snapshots confirmed at the final cache-decision/fold boundary. Usage, Sessions, -Data Health, input space, and the inventory signature all derive from those -same confirmed snapshots. +Session client-space accounting is the byte size of the input snapshots +confirmed at the final cache-decision/fold boundary, deduplicated within each +client inventory. Overview Data Size is exactly the checked sum of that +per-client map; it does not apply a second cross-client physical-file +deduplication rule. The TUI generation schema changes with this accounting +contract so an older generation cannot preserve the previous total alongside +current per-client values. Current-schema cache writes and reads both reject a +Data Size that differs from the checked sum of `clientSpace`. Usage, Sessions, +Data Health, input space, and the inventory signature all derive from those same +confirmed snapshots. ### Shard contract and recovery @@ -167,9 +184,10 @@ bypassed, but a potentially valid replacement shard remains unless the reparse independently proves it invalid or non-cacheable. Partial and unavailable input never publishes a shard. -Cache writes serialize borrowed message slices and do not clone messages merely -to construct a cache representation. Parser-semantic changes bump the owning -parser revision; serialization-layout changes bump the shard format. +Cache writes serialize borrowed source-neutral `ParsedMessage` slices and do +not clone messages merely to construct a cache representation. +Parser-semantic changes bump the owning parser revision; +serialization-layout changes bump the shard format. The message-shard envelope is the `TOKSHRD\0` magic, a little-endian format version, a little-endian `u64` header length, the bincode header, and the From cb02f77c86136f885b966201f21811aefb07513a Mon Sep 17 00:00:00 2001 From: makoMakoGo <48956204+makoMakoGo@users.noreply.github.com> Date: Sat, 25 Jul 2026 19:12:16 +0800 Subject: [PATCH 2/2] refactor(core): make client attribution structural Keep inputs, decoder output, cache shards, and pipeline errors source-neutral; attach ClientId only through AdapterBinding. Use typed client identities through aggregation and derive all input-size views from one confirmed InputFootprint. --- crates/tokscale-cli/src/commands/cache.rs | 2 +- crates/tokscale-cli/src/commands/models.rs | 13 +- crates/tokscale-cli/src/commands/shared.rs | 3 + crates/tokscale-cli/src/tui/actions.rs | 14 +- crates/tokscale-cli/src/tui/app.rs | 24 +- crates/tokscale-cli/src/tui/cache.rs | 300 ++++------ crates/tokscale-cli/src/tui/data/mod.rs | 7 +- crates/tokscale-cli/src/tui/mod.rs | 64 ++- crates/tokscale-cli/src/tui/presentation.rs | 4 +- crates/tokscale-cli/src/tui/session_data.rs | 35 +- crates/tokscale-cli/src/tui/ui/mod.rs | 13 +- crates/tokscale-cli/src/tui/ui/models.rs | 4 +- .../src/tui/ui/overview_snapshot.rs | 4 +- crates/tokscale-cli/src/tui/ui/view_footer.rs | 4 +- crates/tokscale-cli/tests/cli_tests.rs | 8 + crates/tokscale-core/benches/aggregation.rs | 7 +- crates/tokscale-core/build.rs | 44 ++ .../tokscale-core/src/adapters/antigravity.rs | 122 ++-- crates/tokscale-core/src/adapters/cache.rs | 263 ++++----- crates/tokscale-core/src/adapters/claude.rs | 192 +++---- crates/tokscale-core/src/adapters/cline.rs | 50 +- .../tokscale-core/src/adapters/codebuddy.rs | 240 ++++---- crates/tokscale-core/src/adapters/codebuff.rs | 49 +- crates/tokscale-core/src/adapters/codex.rs | 417 +++++++------- crates/tokscale-core/src/adapters/decoder.rs | 123 ++++ crates/tokscale-core/src/adapters/discover.rs | 54 +- crates/tokscale-core/src/adapters/error.rs | 67 ++- crates/tokscale-core/src/adapters/file.rs | 484 ++++++++-------- crates/tokscale-core/src/adapters/goose.rs | 57 +- crates/tokscale-core/src/adapters/hermes.rs | 63 +- crates/tokscale-core/src/adapters/junie.rs | 67 ++- crates/tokscale-core/src/adapters/kilo.rs | 48 +- crates/tokscale-core/src/adapters/kiro.rs | 280 ++++----- crates/tokscale-core/src/adapters/mod.rs | 539 +++++++----------- crates/tokscale-core/src/adapters/omp.rs | 401 ++++++------- crates/tokscale-core/src/adapters/openclaw.rs | 47 +- crates/tokscale-core/src/adapters/opencode.rs | 176 +++--- crates/tokscale-core/src/adapters/pi.rs | 95 +-- crates/tokscale-core/src/adapters/roocode.rs | 72 ++- crates/tokscale-core/src/adapters/runtime.rs | 100 ++++ crates/tokscale-core/src/adapters/warp.rs | 42 +- crates/tokscale-core/src/adapters/zed.rs | 99 ++-- crates/tokscale-core/src/aggregate/keys.rs | 83 +-- crates/tokscale-core/src/aggregate/tui.rs | 246 ++++---- .../src/aggregate/tui_sessions.rs | 26 +- crates/tokscale-core/src/input_footprint.rs | 173 ++++++ crates/tokscale-core/src/input_health.rs | 34 +- crates/tokscale-core/src/lib.rs | 358 ++++++------ crates/tokscale-core/src/lib_tests.rs | 537 +++++++++-------- crates/tokscale-core/src/message_cache.rs | 476 ++++++++-------- crates/tokscale-core/src/sessions/mod.rs | 277 +++++---- docs/adr/0007-client-identity-catalog.md | 61 +- docs/adr/0008-single-copy-memory-pipeline.md | 82 +-- .../0010-period-views-derive-from-daily.md | 2 +- docs/cli.md | 10 +- docs/configuration.md | 6 +- 56 files changed, 3740 insertions(+), 3328 deletions(-) create mode 100644 crates/tokscale-core/src/adapters/decoder.rs create mode 100644 crates/tokscale-core/src/adapters/runtime.rs create mode 100644 crates/tokscale-core/src/input_footprint.rs diff --git a/crates/tokscale-cli/src/commands/cache.rs b/crates/tokscale-cli/src/commands/cache.rs index 91e71a45c..b2c01f18e 100644 --- a/crates/tokscale-cli/src/commands/cache.rs +++ b/crates/tokscale-cli/src/commands/cache.rs @@ -36,7 +36,7 @@ pub(crate) fn run_warm_tui_cache( let _ = save_tui_bundle_cache( &result.accumulator, &result.sessions, - &result.client_space, + &result.input_footprint, &health, &enabled_set, &report_scope, diff --git a/crates/tokscale-cli/src/commands/models.rs b/crates/tokscale-cli/src/commands/models.rs index c8393e0f8..56d56c708 100644 --- a/crates/tokscale-cli/src/commands/models.rs +++ b/crates/tokscale-cli/src/commands/models.rs @@ -60,8 +60,8 @@ pub(crate) fn run_models_report( let scanner_settings = tui::settings::load_scanner_settings_for_home(&home_dir)?; let start = Instant::now(); let rt = Runtime::new()?; - let data = rt - .block_on(tokscale_core::get_usage_data(ReportOptions { + let report = rt + .block_on(tokscale_core::get_usage_report(ReportOptions { home_dir: home_dir.clone(), clients: clients.clone(), since, @@ -71,6 +71,8 @@ pub(crate) fn run_models_report( scanner_settings, })) .map_err(anyhow::Error::new)?; + let data = report.data; + let input_footprint = report.metadata.input_footprint; if let Some(spinner) = spinner { spinner.stop(); @@ -96,7 +98,12 @@ pub(crate) fn run_models_report( if json { let health = data.health.clone(); let report_data = crate::tui::build_models_export_value(&data, &group_by); - let output = ReportEnvelope::new(report_data, health, processing_time_ms as u64); + let output = ReportEnvelope::new( + report_data, + health, + input_footprint, + processing_time_ms as u64, + ); println!("{}", serde_json::to_string_pretty(&output)?); } else { render_models_table(&data, &group_by, date_range.as_deref())?; diff --git a/crates/tokscale-cli/src/commands/shared.rs b/crates/tokscale-cli/src/commands/shared.rs index 07e854527..61dcc2f58 100644 --- a/crates/tokscale-cli/src/commands/shared.rs +++ b/crates/tokscale-cli/src/commands/shared.rs @@ -277,6 +277,7 @@ pub(crate) struct ReportEnvelope { #[derive(Debug, serde::Serialize)] #[serde(rename_all = "camelCase")] pub(crate) struct ReportMetadata { + pub(crate) input_footprint: tokscale_core::InputFootprint, pub(crate) processing_time_ms: u64, } @@ -284,12 +285,14 @@ impl ReportEnvelope { pub(crate) fn new( data: T, health: tokscale_core::input_health::HealthReport, + input_footprint: tokscale_core::InputFootprint, processing_time_ms: impl Into, ) -> Self { Self { data, health, metadata: ReportMetadata { + input_footprint, processing_time_ms: processing_time_ms.into(), }, } diff --git a/crates/tokscale-cli/src/tui/actions.rs b/crates/tokscale-cli/src/tui/actions.rs index 1b26ab3f2..2858dc5ba 100644 --- a/crates/tokscale-cli/src/tui/actions.rs +++ b/crates/tokscale-cli/src/tui/actions.rs @@ -245,7 +245,7 @@ mod tests { use chrono::{NaiveDate, NaiveDateTime}; use crossterm::event::KeyModifiers; - use tokscale_core::{GroupBy, TuiAcc, TuiSessionEntry}; + use tokscale_core::{ClientId, GroupBy, InputFootprint, TuiAcc, TuiSessionEntry}; use super::*; use crate::tui::app::{ProjectionBackend, SortDirection, TuiConfig}; @@ -279,7 +279,7 @@ mod tests { app.install_tui_snapshot( data, Vec::new(), - BTreeMap::new(), + InputFootprint::default(), ProjectionBackend::Memory(accumulator), GroupBy::Model, ); @@ -401,8 +401,10 @@ mod tests { #[test] fn zero_session_client_row_cannot_open_detail() { let mut app = make_app(Tab::Sessions, true); - app.session_snapshot = - SessionSnapshot::new(Vec::new(), BTreeMap::from([("codex".to_string(), 0)])); + app.session_snapshot = SessionSnapshot::new( + Vec::new(), + tokscale_core::InputFootprint::from_client_bytes([(ClientId::Codex, 0)]).unwrap(), + ); let set = action_set(&app, &ViewState::default()); assert!(!set.is_empty_view()); assert!(!set.contains(Action::OpenDetails)); @@ -514,7 +516,7 @@ mod tests { session_id: "session-1".to_string(), ..TuiSessionEntry::default() }], - BTreeMap::new(), + tokscale_core::InputFootprint::default(), ); let set = action_set(&app, &ViewState::default()); @@ -533,7 +535,7 @@ mod tests { session_id: "session-1".to_string(), ..TuiSessionEntry::default() }], - BTreeMap::from([("codex".to_string(), 0)]), + tokscale_core::InputFootprint::from_client_bytes([(ClientId::Codex, 0)]).unwrap(), ); app.sort_field = SortField::Tokens; app.sort_direction = SortDirection::Ascending; diff --git a/crates/tokscale-cli/src/tui/app.rs b/crates/tokscale-cli/src/tui/app.rs index 31ffd6e42..9308a788e 100644 --- a/crates/tokscale-cli/src/tui/app.rs +++ b/crates/tokscale-cli/src/tui/app.rs @@ -1021,11 +1021,11 @@ impl App { &mut self, data: UsageData, sessions: Vec, - client_space: BTreeMap, + input_footprint: tokscale_core::InputFootprint, projection_backend: ProjectionBackend, group_by: tokscale_core::GroupBy, ) { - let session_snapshot = SessionSnapshot::new(sessions, client_space); + let session_snapshot = SessionSnapshot::new(sessions, input_footprint); let overview_summary = OverviewSummary::derive( &data, self.selected_main_session_count_for(&session_snapshot), @@ -3312,21 +3312,33 @@ mod tests { fn model_detail_accumulator() -> tokscale_core::TuiAcc { let messages = [ ( - "claude", + tokscale_core::ClientId::Claude, "shared-model", "anthropic", "claude-anthropic", 11, ), ( - "claude", + tokscale_core::ClientId::Claude, "shared-model", "openrouter", "claude-openrouter", 22, ), - ("codex", "shared-model", "openai", "codex-openai", 33), - ("claude", "other-model", "anthropic", "claude-other", 7), + ( + tokscale_core::ClientId::Codex, + "shared-model", + "openai", + "codex-openai", + 33, + ), + ( + tokscale_core::ClientId::Claude, + "other-model", + "anthropic", + "claude-other", + 7, + ), ] .map(|(client, model, provider, session, input)| { tokscale_core::UnifiedMessage::new( diff --git a/crates/tokscale-cli/src/tui/cache.rs b/crates/tokscale-cli/src/tui/cache.rs index 867eedd9a..c86bb9f13 100644 --- a/crates/tokscale-cli/src/tui/cache.rs +++ b/crates/tokscale-cli/src/tui/cache.rs @@ -17,7 +17,7 @@ use serde::{Deserialize, Deserializer, Serialize, Serializer}; use serde_json::value::RawValue; use sha2::{Digest, Sha256}; use tokscale_core::usage_views::ContributionGrade; -use tokscale_core::{GroupBy, InputInventorySignature, TuiAcc, TuiSessionEntry}; +use tokscale_core::{GroupBy, InputFootprint, InputInventorySignature, TuiAcc, TuiSessionEntry}; use tokscale_core::ClientId; @@ -29,7 +29,7 @@ use super::data::{ /// Cache staleness threshold: 5 minutes (matches TS implementation) const CACHE_STALE_THRESHOLD_MS: u64 = 5 * 60 * 1000; -const CACHE_SCHEMA_VERSION: u32 = 51; +const CACHE_SCHEMA_VERSION: u32 = 52; fn sha256_hex(bytes: &[u8]) -> String { const HEX: &[u8; 16] = b"0123456789abcdef"; @@ -92,7 +92,7 @@ mod bundle_tests { HashSet, CacheReportScope, Vec, - BTreeMap, + InputFootprint, ) { let temp = TempDir::new().unwrap(); let guard = EnvVarGuard::set("TOKSCALE_CONFIG_DIR", temp.path().as_os_str()); @@ -118,21 +118,19 @@ mod bundle_tests { first_seen: 100, last_seen: 200, }]; - let client_space = BTreeMap::from([("claude".into(), 4096)]); - (temp, guard, clients, scope, sessions, client_space) + let input_footprint = + InputFootprint::from_client_bytes([(ClientId::Claude, 4096)]).unwrap(); + (temp, guard, clients, scope, sessions, input_footprint) } fn signature() -> InputInventorySignature { InputInventorySignature::from_bytes([0x39; 32]) } - fn health_for_client_space( - client_space: &BTreeMap, + fn health_for_input_footprint( + _input_footprint: &InputFootprint, ) -> tokscale_core::input_health::HealthReport { - tokscale_core::input_health::HealthReport { - input_data_bytes: client_space.values().copied().sum(), - ..Default::default() - } + tokscale_core::input_health::HealthReport::default() } fn nonempty_accumulator(home: &std::path::Path) -> TuiAcc { @@ -214,7 +212,7 @@ mod bundle_tests { }; let messages = [ UnifiedMessage::new_with_agent( - "opencode", + ClientId::OpenCode, "gpt-5.5", "openai", "open-session", @@ -224,7 +222,7 @@ mod bundle_tests { Some("Builder".into()), ), UnifiedMessage::new_with_agent( - "roocode", + ClientId::RooCode, "gpt-5.5", "openai", "roo-session", @@ -240,19 +238,18 @@ mod bundle_tests { #[test] #[serial] fn bundle_round_trips_sessions_and_metadata() { - let (_temp, _guard, clients, scope, sessions, client_space) = fixture(); + let (_temp, _guard, clients, scope, sessions, input_footprint) = fixture(); let accumulator = TuiAcc::new(); let expected_signature = signature(); let health = tokscale_core::input_health::HealthReport { clean_inputs: 1, - input_data_bytes: 4096, ..Default::default() }; let store = save_tui_bundle_cache( &accumulator, &sessions, - &client_space, + &input_footprint, &health, &clients, &scope, @@ -263,7 +260,7 @@ mod bundle_tests { .load_snapshot(&clients, &GroupBy::Model, &scope) .unwrap(); assert_eq!(saved.sessions, sessions); - assert_eq!(saved.client_space, client_space); + assert_eq!(saved.input_footprint, input_footprint); assert_eq!(saved.data.health, health); assert_eq!( saved.input_inventory_signature.process_digest(), @@ -277,7 +274,7 @@ mod bundle_tests { assert!(raw.get("sourceUniverse").is_none()); assert!(raw.get("sourceSpace").is_none()); assert!(raw["sessions"][0].get("source").is_none()); - assert_eq!(raw["health"]["inputDataBytes"], 4096); + assert!(raw["health"].get("inputDataBytes").is_none()); assert_eq!(raw["canonicalDigest"].as_str().unwrap().len(), 64); assert!(raw["common"].get("health").is_none()); assert!(raw["common"].get("models").is_none()); @@ -288,69 +285,27 @@ mod bundle_tests { panic!("expected a fresh cache bundle"); }; assert_eq!(loaded.sessions, sessions); - assert_eq!(loaded.client_space, client_space); + assert_eq!(loaded.input_footprint, input_footprint); assert_eq!(loaded.data.health, health); } - #[test] - #[serial] - fn bundle_rejects_data_size_outside_client_space_invariant() { - let (_temp, _guard, clients, scope, sessions, client_space) = fixture(); - let invalid_health = tokscale_core::input_health::HealthReport::default(); - - let error = save_tui_bundle_cache( - &TuiAcc::new(), - &sessions, - &client_space, - &invalid_health, - &clients, - &scope, - signature(), - ) - .unwrap_err(); - assert!(error - .to_string() - .contains("TUI Data Size 0 does not match client-space total 4096")); - - save_tui_bundle_cache( - &TuiAcc::new(), - &sessions, - &client_space, - &health_for_client_space(&client_space), - &clients, - &scope, - signature(), - ) - .unwrap(); - let path = cache_file().unwrap(); - let mut value: serde_json::Value = - serde_json::from_reader(File::open(&path).unwrap()).unwrap(); - value["health"]["inputDataBytes"] = serde_json::Value::from(1_u64); - tokscale_core::fs_atomic::write_atomic(&path, &serde_json::to_vec(&value).unwrap()) - .unwrap(); - - assert!(matches!( - load_cache(&clients, &GroupBy::Model, &scope), - CacheResult::Miss - )); - } - #[test] #[serial] fn bundle_keeps_same_named_agents_separate_across_clients() { - let (_temp, _guard, _fixture_clients, scope, _sessions, _client_space) = fixture(); + let (_temp, _guard, _fixture_clients, scope, _sessions, _input_footprint) = fixture(); let accumulator = same_named_cross_client_agent_accumulator(); let clients = HashSet::from([ClientId::OpenCode, ClientId::RooCode]); - let client_space = BTreeMap::from([ - ("opencode".to_string(), 1024), - ("roocode".to_string(), 2048), - ]); + let input_footprint = InputFootprint::from_client_bytes([ + (ClientId::OpenCode, 1024), + (ClientId::RooCode, 2048), + ]) + .unwrap(); save_tui_bundle_cache( &accumulator, &[], - &client_space, - &health_for_client_space(&client_space), + &input_footprint, + &health_for_input_footprint(&input_footprint), &clients, &scope, signature(), @@ -384,15 +339,15 @@ mod bundle_tests { #[test] #[serial] fn missing_or_null_projection_graph_is_an_explicit_miss() { - let (_temp, _guard, clients, scope, sessions, client_space) = fixture(); + let (_temp, _guard, clients, scope, sessions, input_footprint) = fixture(); let path = cache_file().unwrap(); for graph in [None, Some(serde_json::Value::Null)] { save_tui_bundle_cache( &TuiAcc::new(), &sessions, - &client_space, - &health_for_client_space(&client_space), + &input_footprint, + &health_for_input_footprint(&input_footprint), &clients, &scope, signature(), @@ -419,7 +374,7 @@ mod bundle_tests { #[test] #[serial] fn nonempty_bundle_round_trips_all_four_public_groupings() { - let (temp, _guard, _clients, scope, _sessions, _client_space) = fixture(); + let (temp, _guard, _clients, scope, _sessions, _input_footprint) = fixture(); let _pricing_guard = EnvVarGuard::set("TOKSCALE_PRICING_CACHE_ONLY", OsStr::new("1")); let accumulator = nonempty_accumulator(temp.path()); let clients = HashSet::from([ClientId::Claude, ClientId::OpenCode]); @@ -443,11 +398,13 @@ mod bundle_tests { ..Default::default() }, ]; - let client_space = - BTreeMap::from([("claude".to_string(), 8192), ("opencode".to_string(), 4096)]); + let input_footprint = InputFootprint::from_client_bytes([ + (ClientId::Claude, 8192), + (ClientId::OpenCode, 4096), + ]) + .unwrap(); let mut health = tokscale_core::input_health::HealthReport { clean_inputs: 2, - input_data_bytes: 12_288, ..Default::default() }; health.record_unavailable_input("unrelated-test-client"); @@ -478,7 +435,7 @@ mod bundle_tests { let mut store = save_tui_bundle_cache( &accumulator, &sessions, - &client_space, + &input_footprint, &health, &clients, &scope, @@ -534,7 +491,7 @@ mod bundle_tests { assert_projection_eq(&loaded.data, &expected); assert_eq!(loaded.data.health, health); assert_eq!(loaded.sessions, sessions); - assert_eq!(loaded.client_space, client_space); + assert_eq!(loaded.input_footprint, input_footprint); assert!(loaded.data.error.is_none()); } let claude_only = HashSet::from([ClientId::Claude]); @@ -607,12 +564,15 @@ mod bundle_tests { day["grade"] = serde_json::Value::from("peak"); } - let (temp, _guard, _fixture_clients, scope, _sessions, _client_space) = fixture(); + let (temp, _guard, _fixture_clients, scope, _sessions, _input_footprint) = fixture(); let _pricing_guard = EnvVarGuard::set("TOKSCALE_PRICING_CACHE_ONLY", OsStr::new("1")); let accumulator = nonempty_accumulator(temp.path()); let clients = HashSet::from([ClientId::Claude, ClientId::OpenCode]); - let client_space = - BTreeMap::from([("claude".to_string(), 8192), ("opencode".to_string(), 4096)]); + let input_footprint = InputFootprint::from_client_bytes([ + (ClientId::Claude, 8192), + (ClientId::OpenCode, 4096), + ]) + .unwrap(); let path = cache_file().unwrap(); for (boundary, mutate) in [ @@ -635,8 +595,8 @@ mod bundle_tests { save_tui_bundle_cache( &accumulator, &[], - &client_space, - &health_for_client_space(&client_space), + &input_footprint, + &health_for_input_footprint(&input_footprint), &clients, &scope, signature(), @@ -675,12 +635,15 @@ mod bundle_tests { .remove("modelId"); } - let (temp, _guard, _fixture_clients, scope, _sessions, _client_space) = fixture(); + let (temp, _guard, _fixture_clients, scope, _sessions, _input_footprint) = fixture(); let _pricing_guard = EnvVarGuard::set("TOKSCALE_PRICING_CACHE_ONLY", OsStr::new("1")); let accumulator = nonempty_accumulator(temp.path()); let clients = HashSet::from([ClientId::Claude, ClientId::OpenCode]); - let client_space = - BTreeMap::from([("claude".to_string(), 8192), ("opencode".to_string(), 4096)]); + let input_footprint = InputFootprint::from_client_bytes([ + (ClientId::Claude, 8192), + (ClientId::OpenCode, 4096), + ]) + .unwrap(); let path = cache_file().unwrap(); for (boundary, mutate) in [ @@ -696,8 +659,8 @@ mod bundle_tests { save_tui_bundle_cache( &accumulator, &[], - &client_space, - &health_for_client_space(&client_space), + &input_footprint, + &health_for_input_footprint(&input_footprint), &clients, &scope, signature(), @@ -722,21 +685,22 @@ mod bundle_tests { #[test] #[serial] fn bundle_rejects_model_clients_outside_universe_in_active_and_inactive_projections() { - let (_temp, _guard, _fixture_clients, scope, _sessions, _client_space) = fixture(); + let (_temp, _guard, _fixture_clients, scope, _sessions, _input_footprint) = fixture(); let accumulator = same_named_cross_client_agent_accumulator(); let clients = HashSet::from([ClientId::OpenCode, ClientId::RooCode]); - let client_space = BTreeMap::from([ - ("opencode".to_string(), 1024), - ("roocode".to_string(), 2048), - ]); + let input_footprint = InputFootprint::from_client_bytes([ + (ClientId::OpenCode, 1024), + (ClientId::RooCode, 2048), + ]) + .unwrap(); let path = cache_file().unwrap(); for (boundary, projection) in [("active", "model"), ("inactive", "workspaceModel")] { save_tui_bundle_cache( &accumulator, &[], - &client_space, - &health_for_client_space(&client_space), + &input_footprint, + &health_for_input_footprint(&input_footprint), &clients, &scope, signature(), @@ -803,12 +767,15 @@ mod bundle_tests { .remove("modelId"); } - let (temp, _guard, _fixture_clients, scope, _sessions, _client_space) = fixture(); + let (temp, _guard, _fixture_clients, scope, _sessions, _input_footprint) = fixture(); let _pricing_guard = EnvVarGuard::set("TOKSCALE_PRICING_CACHE_ONLY", OsStr::new("1")); let accumulator = nonempty_accumulator(temp.path()); let clients = HashSet::from([ClientId::Claude, ClientId::OpenCode]); - let client_space = - BTreeMap::from([("claude".to_string(), 8192), ("opencode".to_string(), 4096)]); + let input_footprint = InputFootprint::from_client_bytes([ + (ClientId::Claude, 8192), + (ClientId::OpenCode, 4096), + ]) + .unwrap(); let path = cache_file().unwrap(); for (field, mutate) in [ @@ -827,8 +794,8 @@ mod bundle_tests { save_tui_bundle_cache( &accumulator, &[], - &client_space, - &health_for_client_space(&client_space), + &input_footprint, + &health_for_input_footprint(&input_footprint), &clients, &scope, signature(), @@ -853,19 +820,20 @@ mod bundle_tests { #[test] #[serial] fn bundle_rejects_duplicate_client_agent_identity() { - let (_temp, _guard, _fixture_clients, scope, _sessions, _client_space) = fixture(); + let (_temp, _guard, _fixture_clients, scope, _sessions, _input_footprint) = fixture(); let accumulator = same_named_cross_client_agent_accumulator(); let clients = HashSet::from([ClientId::OpenCode, ClientId::RooCode]); - let client_space = BTreeMap::from([ - ("opencode".to_string(), 1024), - ("roocode".to_string(), 2048), - ]); + let input_footprint = InputFootprint::from_client_bytes([ + (ClientId::OpenCode, 1024), + (ClientId::RooCode, 2048), + ]) + .unwrap(); save_tui_bundle_cache( &accumulator, &[], - &client_space, - &health_for_client_space(&client_space), + &input_footprint, + &health_for_input_footprint(&input_footprint), &clients, &scope, signature(), @@ -891,19 +859,20 @@ mod bundle_tests { #[test] #[serial] fn bundle_rejects_agent_outside_client_universe() { - let (_temp, _guard, _fixture_clients, scope, _sessions, _client_space) = fixture(); + let (_temp, _guard, _fixture_clients, scope, _sessions, _input_footprint) = fixture(); let accumulator = same_named_cross_client_agent_accumulator(); let clients = HashSet::from([ClientId::OpenCode, ClientId::RooCode]); - let client_space = BTreeMap::from([ - ("opencode".to_string(), 1024), - ("roocode".to_string(), 2048), - ]); + let input_footprint = InputFootprint::from_client_bytes([ + (ClientId::OpenCode, 1024), + (ClientId::RooCode, 2048), + ]) + .unwrap(); save_tui_bundle_cache( &accumulator, &[], - &client_space, - &health_for_client_space(&client_space), + &input_footprint, + &health_for_input_footprint(&input_footprint), &clients, &scope, signature(), @@ -925,12 +894,12 @@ mod bundle_tests { #[test] #[serial] fn unsupported_schema_version_is_an_explicit_miss() { - let (_temp, _guard, clients, scope, sessions, client_space) = fixture(); + let (_temp, _guard, clients, scope, sessions, input_footprint) = fixture(); save_tui_bundle_cache( &TuiAcc::new(), &sessions, - &client_space, - &health_for_client_space(&client_space), + &input_footprint, + &health_for_input_footprint(&input_footprint), &clients, &scope, signature(), @@ -953,15 +922,15 @@ mod bundle_tests { #[test] #[serial] fn missing_or_incomplete_canonical_state_is_an_explicit_miss() { - let (_temp, _guard, clients, scope, sessions, client_space) = fixture(); + let (_temp, _guard, clients, scope, sessions, input_footprint) = fixture(); let path = cache_file().unwrap(); for replacement in [None, Some(serde_json::json!({ "model_map": [] }))] { save_tui_bundle_cache( &TuiAcc::new(), &sessions, - &client_space, - &health_for_client_space(&client_space), + &input_footprint, + &health_for_input_footprint(&input_footprint), &clients, &scope, signature(), @@ -987,8 +956,8 @@ mod bundle_tests { save_tui_bundle_cache( &TuiAcc::new(), &sessions, - &client_space, - &health_for_client_space(&client_space), + &input_footprint, + &health_for_input_footprint(&input_footprint), &clients, &scope, signature(), @@ -1015,8 +984,8 @@ mod bundle_tests { save_tui_bundle_cache( &TuiAcc::new(), &sessions, - &client_space, - &health_for_client_space(&client_space), + &input_footprint, + &health_for_input_footprint(&input_footprint), &clients, &scope, signature(), @@ -1036,12 +1005,12 @@ mod bundle_tests { #[test] #[serial] fn canonical_content_corruption_is_an_explicit_miss() { - let (_temp, _guard, clients, scope, sessions, client_space) = fixture(); + let (_temp, _guard, clients, scope, sessions, input_footprint) = fixture(); save_tui_bundle_cache( &TuiAcc::new(), &sessions, - &client_space, - &health_for_client_space(&client_space), + &input_footprint, + &health_for_input_footprint(&input_footprint), &clients, &scope, signature(), @@ -1063,14 +1032,14 @@ mod bundle_tests { #[test] #[serial] fn partial_or_foreign_client_membership_is_a_miss() { - let (_temp, _guard, clients, scope, sessions, client_space) = fixture(); + let (_temp, _guard, clients, scope, sessions, input_footprint) = fixture(); let path = cache_file().unwrap(); save_tui_bundle_cache( &TuiAcc::new(), &sessions, - &client_space, - &health_for_client_space(&client_space), + &input_footprint, + &health_for_input_footprint(&input_footprint), &clients, &scope, signature(), @@ -1090,8 +1059,8 @@ mod bundle_tests { save_tui_bundle_cache( &TuiAcc::new(), &sessions, - &client_space, - &health_for_client_space(&client_space), + &input_footprint, + &health_for_input_footprint(&input_footprint), &clients, &scope, signature(), @@ -1112,12 +1081,12 @@ mod bundle_tests { #[test] #[serial] fn writer_is_atomic_and_projection_store_pins_the_old_inode() { - let (_temp, _guard, clients, scope, sessions, client_space) = fixture(); + let (_temp, _guard, clients, scope, sessions, input_footprint) = fixture(); let mut store = save_tui_bundle_cache( &TuiAcc::new(), &sessions, - &client_space, - &health_for_client_space(&client_space), + &input_footprint, + &health_for_input_footprint(&input_footprint), &clients, &scope, signature(), @@ -1139,12 +1108,12 @@ mod bundle_tests { #[test] #[serial] fn fresh_and_stale_are_derived_from_the_bundle_timestamp() { - let (_temp, _guard, clients, scope, sessions, client_space) = fixture(); + let (_temp, _guard, clients, scope, sessions, input_footprint) = fixture(); save_tui_bundle_cache( &TuiAcc::new(), &sessions, - &client_space, - &health_for_client_space(&client_space), + &input_footprint, + &health_for_input_footprint(&input_footprint), &clients, &scope, signature(), @@ -2240,32 +2209,14 @@ fn cache_clients_match_exact( cached.len() == cached_clients.len() && enabled == cached } -fn cache_client_space_matches_exact( +fn cache_input_footprint_matches_exact( client_universe: &HashSet, - client_space: &BTreeMap, + input_footprint: &InputFootprint, ) -> bool { - client_space.len() == client_universe.len() + input_footprint.len() == client_universe.len() && client_universe .iter() - .all(|client| client_space.contains_key(client.as_str())) -} - -fn validate_data_size_matches_client_space( - client_space: &BTreeMap, - health: &tokscale_core::input_health::HealthReport, -) -> anyhow::Result<()> { - let total = client_space - .values() - .copied() - .try_fold(0_u64, u64::checked_add) - .ok_or_else(|| anyhow::anyhow!("TUI client-space total exceeds u64::MAX"))?; - if health.input_data_bytes != total { - anyhow::bail!( - "TUI Data Size {} does not match client-space total {total}", - health.input_data_bytes - ); - } - Ok(()) + .all(|client| input_footprint.contains_client(*client)) } fn cache_session_clients_are_enabled( @@ -2314,7 +2265,7 @@ fn cache_usage_clients_are_enabled(client_universe: &HashSet, data: &U pub struct LoadedTuiCache { pub data: UsageData, pub sessions: Vec, - pub client_space: BTreeMap, + pub input_footprint: InputFootprint, pub projection_store: ProjectionStore, pub input_inventory_signature: InputInventorySignature, } @@ -2426,7 +2377,8 @@ struct CachedTuiBundleRef<'a> { input_inventory_signature: &'a InputInventorySignature, health: &'a tokscale_core::input_health::HealthReport, sessions: &'a [TuiSessionEntry], - client_space: &'a BTreeMap, + #[serde(rename = "clientSpace")] + input_footprint: &'a InputFootprint, canonical_digest: &'a str, canonical: &'a RawValue, common: CachedCommonProjectionRef<'a>, @@ -2750,7 +2702,7 @@ struct ParsedTuiBundle { input_inventory_signature: InputInventorySignature, health: tokscale_core::input_health::HealthReport, sessions: Vec, - client_space: BTreeMap, + input_footprint: InputFootprint, common: UsageCommonData, grouped: UsageGroupedData, } @@ -2863,7 +2815,7 @@ impl<'de> Visitor<'de> for FullBundleVisitor<'_> { let mut input_inventory_signature = None; let mut health = None; let mut sessions = None; - let mut client_space = None; + let mut input_footprint = None; let mut expected_canonical_digest: Option = None; let mut actual_canonical_digest: Option = None; let mut common: Option = None; @@ -2886,7 +2838,7 @@ impl<'de> Visitor<'de> for FullBundleVisitor<'_> { )?, "health" => set_once(&mut health, map.next_value()?, "health")?, "sessions" => set_once(&mut sessions, map.next_value()?, "sessions")?, - "clientSpace" => set_once(&mut client_space, map.next_value()?, "clientSpace")?, + "clientSpace" => set_once(&mut input_footprint, map.next_value()?, "clientSpace")?, "canonicalDigest" => set_once( &mut expected_canonical_digest, map.next_value()?, @@ -2958,7 +2910,7 @@ impl<'de> Visitor<'de> for FullBundleVisitor<'_> { )?, health: required(health, "health")?, sessions: required(sessions, "sessions")?, - client_space: required(client_space, "clientSpace")?, + input_footprint: required(input_footprint, "clientSpace")?, common, grouped: grouped.selected, }) @@ -3134,10 +3086,9 @@ fn load_bundle_from_file( if !cache_clients_match_exact(client_universe, &parsed.client_universe) { anyhow::bail!("cached TUI client universe does not match the request"); } - if !cache_client_space_matches_exact(client_universe, &parsed.client_space) { - anyhow::bail!("cached TUI client-space keys do not match the client universe"); + if !cache_input_footprint_matches_exact(client_universe, &parsed.input_footprint) { + anyhow::bail!("cached TUI input-footprint keys do not match the client universe"); } - validate_data_size_matches_client_space(&parsed.client_space, &parsed.health)?; if !cache_session_clients_are_enabled(client_universe, &parsed.sessions) { anyhow::bail!("cached TUI Sessions contain a client outside the client universe"); } @@ -3156,7 +3107,7 @@ fn load_bundle_from_file( loaded: LoadedTuiCache { data, sessions: parsed.sessions, - client_space: parsed.client_space, + input_footprint: parsed.input_footprint, projection_store: ProjectionStore { file, health: parsed.health, @@ -3213,16 +3164,15 @@ pub fn load_cache( pub fn save_tui_bundle_cache( accumulator: &TuiAcc, sessions: &[TuiSessionEntry], - client_space: &BTreeMap, + input_footprint: &InputFootprint, health: &tokscale_core::input_health::HealthReport, client_universe: &HashSet, report_scope: &CacheReportScope, input_inventory_signature: InputInventorySignature, ) -> anyhow::Result { - if !cache_client_space_matches_exact(client_universe, client_space) { - anyhow::bail!("TUI client-space keys do not match the client universe"); + if !cache_input_footprint_matches_exact(client_universe, input_footprint) { + anyhow::bail!("TUI input-footprint keys do not match the client universe"); } - validate_data_size_matches_client_space(client_space, health)?; if !cache_session_clients_are_enabled(client_universe, sessions) { anyhow::bail!("TUI Sessions contain a client outside the client universe"); } @@ -3246,7 +3196,7 @@ pub fn save_tui_bundle_cache( input_inventory_signature: &input_inventory_signature, health, sessions, - client_space, + input_footprint, canonical_digest: &canonical_digest, canonical: canonical.as_ref(), common: CachedCommonProjectionRef(accumulator), diff --git a/crates/tokscale-cli/src/tui/data/mod.rs b/crates/tokscale-cli/src/tui/data/mod.rs index 669727dcd..bb5fce394 100644 --- a/crates/tokscale-cli/src/tui/data/mod.rs +++ b/crates/tokscale-cli/src/tui/data/mod.rs @@ -10,7 +10,8 @@ use chrono::NaiveDate; use tokscale_core::GroupBy; use tokscale_core::{ load_prepared_tui_bundle_with_diagnostics, prepare_local_inputs, ClientId, DataHealth, - InputInventorySignature, LocalParseOptions, PreparedLocalInputs, TuiAcc, TuiSessionEntry, + InputFootprint, InputInventorySignature, LocalParseOptions, PreparedLocalInputs, TuiAcc, + TuiSessionEntry, }; mod overview; @@ -84,7 +85,7 @@ pub struct DataLoader { pub struct TuiBundleLoadResult { pub accumulator: TuiAcc, pub sessions: Vec, - pub client_space: std::collections::BTreeMap, + pub input_footprint: InputFootprint, pub pricing_diagnostics: Vec, pub input_inventory_signature: InputInventorySignature, pub input_digest: u64, @@ -170,7 +171,7 @@ impl DataLoader { bundle.map(|result| TuiBundleLoadResult { accumulator: result.accumulator, sessions: result.sessions, - client_space: result.client_space, + input_footprint: result.input_footprint, pricing_diagnostics: result.pricing_diagnostics, input_inventory_signature: result.input_inventory_signature, input_digest: result.input_inventory_signature.process_digest(), diff --git a/crates/tokscale-cli/src/tui/mod.rs b/crates/tokscale-cli/src/tui/mod.rs index c50d5d68e..ac9d21db4 100644 --- a/crates/tokscale-cli/src/tui/mod.rs +++ b/crates/tokscale-cli/src/tui/mod.rs @@ -58,7 +58,7 @@ use crossterm::{ }, }; use ratatui::prelude::*; -use tokscale_core::ClientId; +use tokscale_core::{ClientId, InputFootprint}; fn decide_initial_data(load_result: CacheResult) -> (Option, bool, Option) { match load_result { @@ -99,7 +99,7 @@ enum BackgroundLoad { Loaded { data: Box, sessions: Vec, - client_space: std::collections::BTreeMap, + input_footprint: InputFootprint, projection_backend: Box, digest: u64, /// The grouping this `data` projection was aggregated with; the App @@ -143,7 +143,7 @@ fn load_background_data( BackgroundLoad::Loaded { data: Box::new(data), sessions: result.sessions, - client_space: result.client_space, + input_footprint: result.input_footprint, projection_backend: Box::new(ProjectionBackend::Memory(result.accumulator)), digest: result.input_digest, group_by: group_by.clone(), @@ -163,7 +163,7 @@ fn persist_background_load( let BackgroundLoad::Loaded { data, sessions, - client_space, + input_footprint, projection_backend, digest, group_by, @@ -188,7 +188,7 @@ fn persist_background_load( match save_tui_bundle_cache( &accumulator, &sessions, - &client_space, + &input_footprint, &health, client_universe, report_scope, @@ -201,7 +201,7 @@ fn persist_background_load( // snapshot that will be published to the UI thread. drop(accumulator); drop(sessions); - drop(client_space); + drop(input_footprint); data::trim_allocator(); Ok(BackgroundLoad::Persisted { @@ -223,7 +223,7 @@ fn persist_background_load( Ok(BackgroundLoad::Loaded { data: Box::new(data), sessions, - client_space, + input_footprint, projection_backend: Box::new(ProjectionBackend::Memory(accumulator)), digest, group_by, @@ -273,7 +273,7 @@ fn apply_background_result(app: &mut App, result: Result) { app.install_tui_snapshot( cached.data, cached.sessions, - cached.client_space, + cached.input_footprint, ProjectionBackend::Cache(cached.projection_store), selected_group_by, ); @@ -285,7 +285,7 @@ fn apply_background_result(app: &mut App, result: Result) { Ok(BackgroundLoad::Loaded { data, sessions, - client_space, + input_footprint, mut projection_backend, digest, group_by, @@ -315,7 +315,7 @@ fn apply_background_result(app: &mut App, result: Result) { app.install_tui_snapshot( current_data, sessions, - client_space, + input_footprint, *projection_backend, selected_group_by, ); @@ -462,7 +462,7 @@ pub fn run( app.install_tui_snapshot( cached.data, cached.sessions, - cached.client_space, + cached.input_footprint, ProjectionBackend::Cache(cached.projection_store), initial_group_by, ); @@ -1023,7 +1023,7 @@ mod tests { app.projection_backend = Some(ProjectionBackend::Memory(tokscale_core::TuiAcc::new())); app.session_snapshot = session_data::SessionSnapshot::new( Vec::new(), - std::collections::BTreeMap::from([("junie".to_string(), 0)]), + InputFootprint::from_client_bytes([(ClientId::Junie, 0)]).unwrap(), ); let mut view_state = view_state::ViewState::default(); @@ -1043,7 +1043,7 @@ mod tests { let (mut app, mut view_state) = app_with_codex_session_detail(); app.session_snapshot = session_data::SessionSnapshot::new( Vec::new(), - std::collections::BTreeMap::from([(ClientId::Codex.as_str().to_string(), 0)]), + InputFootprint::from_client_bytes([(ClientId::Codex, 0)]).unwrap(), ); view_state.reconcile_session_snapshot(&app); @@ -1100,7 +1100,7 @@ mod tests { } } - fn client_space_for(app: &App, client: &str) -> Option { + fn input_bytes_for(app: &App, client: &str) -> Option { app.session_snapshot .client_summaries() .iter() @@ -1117,11 +1117,8 @@ mod tests { let store = save_tui_bundle_cache( &tokscale_core::TuiAcc::new(), &[session("amp", "cached-session")], - &std::collections::BTreeMap::from([("amp".to_string(), 512)]), - &tokscale_core::input_health::HealthReport { - input_data_bytes: 512, - ..Default::default() - }, + &InputFootprint::from_client_bytes([(ClientId::Amp, 512)]).unwrap(), + &tokscale_core::input_health::HealthReport::default(), &clients, &scope, signature, @@ -1133,18 +1130,25 @@ mod tests { } fn loaded_snapshot( - mut data: UsageData, + data: UsageData, sessions: Vec, - client_space: std::collections::BTreeMap, + client_bytes: std::collections::BTreeMap, accumulator: tokscale_core::TuiAcc, group_by: tokscale_core::GroupBy, signature: tokscale_core::InputInventorySignature, ) -> BackgroundLoad { - data.health.input_data_bytes = client_space.values().copied().sum(); + let input_footprint = + InputFootprint::from_client_bytes(client_bytes.into_iter().map(|(client, bytes)| { + ( + ClientId::from_str(&client).expect("test client must be canonical"), + bytes, + ) + })) + .unwrap(); BackgroundLoad::Loaded { data: Box::new(data), sessions, - client_space, + input_footprint, projection_backend: Box::new(ProjectionBackend::Memory(accumulator)), digest: signature.process_digest(), group_by, @@ -1169,7 +1173,7 @@ mod tests { assert!(!needs_background_load); assert_eq!(digest, Some(signature.process_digest())); assert_eq!(cached_data.sessions[0].session_id, "cached-session"); - assert_eq!(cached_data.client_space.get("amp"), Some(&512)); + assert_eq!(cached_data.input_footprint.bytes_for(ClientId::Amp), 512); } #[test] @@ -1380,7 +1384,7 @@ mod tests { assert!(!app.background_loading); assert_eq!(app.data.total_tokens, 99); assert_eq!(app.session_snapshot.sessions()[0].session_id, "new-session"); - assert_eq!(client_space_for(&app, "codex"), Some(4096)); + assert_eq!(input_bytes_for(&app, "codex"), Some(4096)); assert!(matches!( app.projection_backend, Some(ProjectionBackend::Memory(_)) @@ -1420,7 +1424,7 @@ mod tests { app.session_snapshot.sessions()[0].session_id, "retained-session" ); - assert_eq!(client_space_for(&app, "amp"), Some(2048)); + assert_eq!(input_bytes_for(&app, "amp"), Some(2048)); assert!(app.projection_backend.is_some()); assert_eq!(app.last_input_digest, Some(signature.process_digest())); assert_eq!( @@ -1520,14 +1524,14 @@ mod tests { apply_background_result(&mut app, Ok(loaded)); let old_tokens = app.data.total_tokens; let old_sessions = app.session_snapshot.sessions().to_vec(); - let old_client_space = client_space_for(&app, "amp"); + let old_input_bytes = input_bytes_for(&app, "amp"); apply_background_result(&mut app, Err(anyhow::anyhow!("load failed"))); assert_eq!(app.data.total_tokens, old_tokens); assert_eq!(app.data.models[0].model_id, "retained-model"); assert_eq!(app.session_snapshot.sessions(), old_sessions); - assert_eq!(client_space_for(&app, "amp"), old_client_space); + assert_eq!(input_bytes_for(&app, "amp"), old_input_bytes); assert!(app.projection_backend.is_some()); assert!(matches!( app.session_projection_status, @@ -1548,7 +1552,7 @@ mod tests { let digest = signature.process_digest(); let accumulator = tokscale_core::build_tui_accumulator( &[tokscale_core::UnifiedMessage::new( - "amp", + tokscale_core::ClientId::Amp, "test-model", "test-provider", "loaded-despite-cache-error", @@ -1734,7 +1738,7 @@ mod tests { ..UsageData::default() }), sessions: vec![session("amp", "warning-session")], - client_space: std::collections::BTreeMap::from([("amp".to_string(), 99)]), + input_footprint: InputFootprint::from_client_bytes([(ClientId::Amp, 99)]).unwrap(), projection_backend: Box::new(ProjectionBackend::Memory( tokscale_core::TuiAcc::new(), )), diff --git a/crates/tokscale-cli/src/tui/presentation.rs b/crates/tokscale-cli/src/tui/presentation.rs index 4c6d8c84d..a824ba58b 100644 --- a/crates/tokscale-cli/src/tui/presentation.rs +++ b/crates/tokscale-cli/src/tui/presentation.rs @@ -142,7 +142,7 @@ mod tests { use crate::tui::app::{ProjectionBackend, TuiConfig}; use crate::tui::settings::Settings; use tokscale_core::{ - build_tui_accumulator, DateRange, GroupBy, TokenBreakdown, TuiAcc, UnifiedMessage, + build_tui_accumulator, ClientId, DateRange, GroupBy, TokenBreakdown, TuiAcc, UnifiedMessage, }; fn app(tab: Tab, installed: bool) -> App { @@ -187,7 +187,7 @@ mod tests { fn populated_accumulator() -> TuiAcc { build_tui_accumulator( &[UnifiedMessage::new( - "codex", + ClientId::Codex, "gpt-5", "openai", "session-1", diff --git a/crates/tokscale-cli/src/tui/session_data.rs b/crates/tokscale-cli/src/tui/session_data.rs index 0ba267a70..2fb4e2ec9 100644 --- a/crates/tokscale-cli/src/tui/session_data.rs +++ b/crates/tokscale-cli/src/tui/session_data.rs @@ -1,6 +1,7 @@ use std::collections::{BTreeMap, BTreeSet}; pub(crate) use tokscale_core::TuiSessionEntry as SessionEntry; +use tokscale_core::{ClientId, InputFootprint}; #[derive(Debug, Clone, Default)] pub(crate) struct ClientSummary { @@ -22,6 +23,7 @@ pub(crate) struct SessionSnapshot { sessions: Vec, client_summaries: Vec, session_indices_by_client: BTreeMap>, + input_footprint: InputFootprint, } #[derive(Debug, Clone, Default, PartialEq, Eq)] @@ -38,10 +40,7 @@ pub(crate) enum SessionProjectionStatus { } impl SessionSnapshot { - pub(crate) fn new( - mut sessions: Vec, - client_space: BTreeMap, - ) -> Self { + pub(crate) fn new(mut sessions: Vec, input_footprint: InputFootprint) -> Self { sessions.sort_by(|left, right| { right .last_seen @@ -81,9 +80,9 @@ impl SessionSnapshot { entry.3 = entry.3.max(session.last_seen); } - for client in client_space.keys() { + for (client, _) in input_footprint.iter() { summaries - .entry(client.clone()) + .entry(client.as_str().to_string()) .or_insert_with(|| (0, 0, BTreeSet::new(), 0)); } @@ -92,7 +91,9 @@ impl SessionSnapshot { .map( |(client, (session_count, main_session_count, workspaces, last_seen))| { ClientSummary { - space_bytes: client_space.get(&client).copied().unwrap_or(0), + space_bytes: ClientId::from_str(&client) + .map(|client| input_footprint.bytes_for(client)) + .unwrap_or(0), client, main_session_count, session_count, @@ -107,9 +108,16 @@ impl SessionSnapshot { sessions, client_summaries, session_indices_by_client, + input_footprint, } } + pub(crate) fn total_input_bytes(&self) -> u64 { + self.input_footprint + .total_bytes() + .expect("validated input footprint must fit in u64") + } + #[cfg(test)] pub(crate) fn sessions(&self) -> &[SessionEntry] { &self.sessions @@ -179,7 +187,7 @@ mod tests { session("opencode", "o-new", true, Some("repo-b"), None, 30), session("codex", "c-new", false, Some("repo-a"), None, 30), ], - BTreeMap::new(), + InputFootprint::default(), ); assert_eq!( @@ -211,11 +219,12 @@ mod tests { session("codex", "c-2", false, Some("repo-a"), None, 30), session("opencode", "o-1", true, Some(""), Some("repo-b"), 20), ], - BTreeMap::from([ - ("claude".to_string(), 7), - ("codex".to_string(), 42), - ("opencode".to_string(), 99), - ]), + InputFootprint::from_client_bytes([ + (ClientId::Claude, 7), + (ClientId::Codex, 42), + (ClientId::OpenCode, 99), + ]) + .unwrap(), ); let codex = snapshot diff --git a/crates/tokscale-cli/src/tui/ui/mod.rs b/crates/tokscale-cli/src/tui/ui/mod.rs index e9658d893..3bc9a2a2e 100644 --- a/crates/tokscale-cli/src/tui/ui/mod.rs +++ b/crates/tokscale-cli/src/tui/ui/mod.rs @@ -274,13 +274,22 @@ mod tests { app: &mut App, clients: &[ClientId], data: UsageData, - client_space: BTreeMap, + client_bytes: BTreeMap, ) { + let input_footprint = tokscale_core::InputFootprint::from_client_bytes( + client_bytes.into_iter().map(|(client, bytes)| { + ( + ClientId::from_str(&client).expect("test client must be canonical"), + bytes, + ) + }), + ) + .unwrap(); *app.selected_clients.borrow_mut() = clients.iter().copied().collect(); app.install_tui_snapshot( data, Vec::new(), - client_space, + input_footprint, ProjectionBackend::Memory(TuiAcc::default()), GroupBy::Model, ); diff --git a/crates/tokscale-cli/src/tui/ui/models.rs b/crates/tokscale-cli/src/tui/ui/models.rs index b6b74066f..7570c0da6 100644 --- a/crates/tokscale-cli/src/tui/ui/models.rs +++ b/crates/tokscale-cli/src/tui/ui/models.rs @@ -647,7 +647,7 @@ mod tests { fn make_model_detail_app(group_by: GroupBy) -> App { let messages = [ tokscale_core::UnifiedMessage::new( - "claude", + tokscale_core::ClientId::Claude, "shared-model", "anthropic", "anthropic-session", @@ -659,7 +659,7 @@ mod tests { 0.1, ), tokscale_core::UnifiedMessage::new( - "claude", + tokscale_core::ClientId::Claude, "shared-model", "openrouter", "openrouter-session", diff --git a/crates/tokscale-cli/src/tui/ui/overview_snapshot.rs b/crates/tokscale-cli/src/tui/ui/overview_snapshot.rs index 0af39e45b..18692e3c5 100644 --- a/crates/tokscale-cli/src/tui/ui/overview_snapshot.rs +++ b/crates/tokscale-cli/src/tui/ui/overview_snapshot.rs @@ -170,7 +170,7 @@ fn render_core(frame: &mut Frame, app: &App, area: Rect, data: &OverviewSummary) metric_line( app, "Data Size", - format_bytes(app.data.health.input_data_bytes), + format_bytes(app.session_snapshot.total_input_bytes()), app.theme.text.primary, ), ]; @@ -681,7 +681,7 @@ fn left_lines( metric_line( app, "Input Data", - format_bytes(app.data.health.input_data_bytes), + format_bytes(app.session_snapshot.total_input_bytes()), app.theme.text.primary, ), // The narrow fallback omits input health; Active Days keeps this diff --git a/crates/tokscale-cli/src/tui/ui/view_footer.rs b/crates/tokscale-cli/src/tui/ui/view_footer.rs index eb3fcacc9..c602ad4c7 100644 --- a/crates/tokscale-cli/src/tui/ui/view_footer.rs +++ b/crates/tokscale-cli/src/tui/ui/view_footer.rs @@ -154,7 +154,7 @@ mod tests { use chrono::NaiveDate; use crossterm::event::{KeyCode, KeyEvent, KeyModifiers}; use ratatui::{backend::TestBackend, Terminal}; - use tokscale_core::{GroupBy, TuiAcc, TuiSessionEntry}; + use tokscale_core::{GroupBy, InputFootprint, TuiAcc, TuiSessionEntry}; use unicode_width::UnicodeWidthStr; fn make_app(width: u16) -> App { @@ -187,7 +187,7 @@ mod tests { session_id: "session-1".to_string(), ..TuiSessionEntry::default() }], - BTreeMap::new(), + InputFootprint::default(), ProjectionBackend::Memory(TuiAcc::default()), GroupBy::Model, ); diff --git a/crates/tokscale-cli/tests/cli_tests.rs b/crates/tokscale-cli/tests/cli_tests.rs index 55b3d8126..cbded60df 100644 --- a/crates/tokscale-cli/tests/cli_tests.rs +++ b/crates/tokscale-cli/tests/cli_tests.rs @@ -1178,6 +1178,13 @@ fn test_models_json_output() { json["metadata"].get("processingTimeMs").is_some(), "Missing processingTimeMs" ); + assert!( + json["metadata"]["inputFootprint"]["opencode"] + .as_u64() + .is_some_and(|bytes| bytes > 0), + "Missing confirmed OpenCode input footprint" + ); + assert!(json["health"].get("inputDataBytes").is_none()); let models = model_rows(&json); assert!(!models.is_empty(), "Should have models from fixture data"); @@ -1237,6 +1244,7 @@ fn test_every_local_json_command_uses_the_common_envelope() { invocation.join(" ") ); assert!(document["metadata"]["processingTimeMs"].is_number()); + assert!(document["metadata"]["inputFootprint"].is_object()); } } diff --git a/crates/tokscale-core/benches/aggregation.rs b/crates/tokscale-core/benches/aggregation.rs index 71c72591f..79d70cf60 100644 --- a/crates/tokscale-core/benches/aggregation.rs +++ b/crates/tokscale-core/benches/aggregation.rs @@ -10,7 +10,12 @@ use tokscale_core::{ const MESSAGE_COUNT: usize = 100_000; -const CLIENTS: &[&str] = &["opencode", "claude", "codex", "zed"]; +const CLIENTS: &[ClientId] = &[ + ClientId::OpenCode, + ClientId::Claude, + ClientId::Codex, + ClientId::Zed, +]; const MODELS: &[&str] = &["gpt-5.5", "claude-sonnet-4.5", "qwen3-coder", "kimi-k2.5"]; const PROVIDERS: &[&str] = &["openai", "anthropic", "qwen", "kimi"]; const WORKSPACES: &[&str] = &[ diff --git a/crates/tokscale-core/build.rs b/crates/tokscale-core/build.rs index e993e39ab..fa39722c1 100644 --- a/crates/tokscale-core/build.rs +++ b/crates/tokscale-core/build.rs @@ -178,6 +178,50 @@ impl ClientId {{ }} }} +impl std::fmt::Display for ClientId {{ + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {{ + formatter.write_str(self.as_str()) + }} +}} + +impl PartialOrd for ClientId {{ + fn partial_cmp(&self, other: &Self) -> Option {{ + Some(self.cmp(other)) + }} +}} + +impl Ord for ClientId {{ + fn cmp(&self, other: &Self) -> std::cmp::Ordering {{ + self.as_str().cmp(other.as_str()) + }} +}} + +impl AsRef for ClientId {{ + fn as_ref(&self) -> &str {{ + self.as_str() + }} +}} + +impl serde::Serialize for ClientId {{ + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + {{ + serializer.serialize_str(self.as_str()) + }} +}} + +impl<'de> serde::Deserialize<'de> for ClientId {{ + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + {{ + let id = ::deserialize(deserializer)?; + Self::from_str(&id) + .ok_or_else(|| serde::de::Error::custom(format!("unknown local client `{{id}}`"))) + }} +}} + pub const CLIENT_IDENTITIES: [ClientIdentity; ClientId::COUNT] = [ {identities}]; "# diff --git a/crates/tokscale-core/src/adapters/antigravity.rs b/crates/tokscale-core/src/adapters/antigravity.rs index 82e2fdc99..2237e56df 100644 --- a/crates/tokscale-core/src/adapters/antigravity.rs +++ b/crates/tokscale-core/src/adapters/antigravity.rs @@ -6,11 +6,11 @@ use rayon::prelude::*; use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ - AdapterScanContext, FingerprintPolicy, FoldContext, InputDiscoveryError, InputUnit, - InputUnitMeta, LocalInputAdapter, MessageSink, ParseContext, ParsedBatchInput, ParsedUnit, + AdapterScanContext, BoundMessageSink, DecoderRoute, DecoderSpec, FingerprintPolicy, + FoldContext, InputDiscoveryError, InputUnit, LocalInputAdapter, ParseContext, ParsedBatchInput, + ParsedUnit, }; use crate::clients::ClientId; -use crate::message_cache::{ParserId, ParserVersion}; use crate::sessions; const ANTIGRAVITY_CLI_RECORD_REJECTION_REVISION: u32 = @@ -19,53 +19,42 @@ const ANTIGRAVITY_CLI_RECORD_REJECTION_REVISION: u32 = pub(crate) struct AntigravityAdapter; impl LocalInputAdapter for AntigravityAdapter { - fn client(&self) -> ClientId { - ClientId::Antigravity - } - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { - let def = ClientId::Antigravity + let def = client .local_def() .expect("Antigravity adapter requires a local scan policy"); let mut roots = vec![def.resolve_path(ctx.home_dir)]; - roots.extend(antigravity_extra_roots(ctx)?); + roots.extend(antigravity_extra_roots(client, ctx)?); - Ok(adapter_discover::input_units_from_paths( - ClientId::Antigravity, - adapter_discover::scan_roots(ClientId::Antigravity, roots, def.pattern)?, + adapter_discover::input_units_from_paths( + client, + adapter_discover::scan_roots(client, roots, def.pattern)?, FingerprintPolicy::SqliteWithWal, - )? - .into_iter() - .map(|unit| { - unit.with_meta(InputUnitMeta::AntigravityCliSqlite) - .with_parser_version(ParserVersion::new( - ParserId::AntigravityCliSqlite, - ANTIGRAVITY_CLI_RECORD_REJECTION_REVISION, - )) - }) - .collect()) + DecoderSpec::antigravity_cli_sqlite(ANTIGRAVITY_CLI_RECORD_REJECTION_REVISION), + ) } fn parse_checked(&self, units: Vec, ctx: &ParseContext<'_>) -> Vec { units .into_par_iter() - .map(|unit| match unit.meta { - InputUnitMeta::AntigravityCliSqlite => adapter_cache::load_or_scan_unit_with( + .map(|unit| match unit.decoder.route() { + DecoderRoute::AntigravityCliSqlite => adapter_cache::load_or_scan_unit_with( unit, ctx, sessions::antigravity_cli::parse_antigravity_cli_file, ), - InputUnitMeta::None - | InputUnitMeta::OpenCodeSqlite - | InputUnitMeta::KiroFile - | InputUnitMeta::KiroSqlite - | InputUnitMeta::KiroGlobalStorage - | InputUnitMeta::CodeBuddyJsonl - | InputUnitMeta::CodeBuddyExtensionLog { .. } - | InputUnitMeta::Codex => { + DecoderRoute::None + | DecoderRoute::OpenCodeSqlite + | DecoderRoute::KiroFile + | DecoderRoute::KiroSqlite + | DecoderRoute::KiroGlobalStorage + | DecoderRoute::CodeBuddyJsonl + | DecoderRoute::CodeBuddyExtensionLog { .. } + | DecoderRoute::Codex => { unreachable!("unexpected Antigravity input unit meta") } }) @@ -84,7 +73,7 @@ impl LocalInputAdapter for AntigravityAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), crate::adapters::InputPipelineError> { let mut seen = HashSet::new(); fold_antigravity_units(parsed, ctx, sink, &mut seen) @@ -94,7 +83,7 @@ impl LocalInputAdapter for AntigravityAdapter { &self, batches: &mut ParsedBatchInput<'_>, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), crate::adapters::InputPipelineError> { let mut seen = HashSet::new(); while let Some(parsed) = batches.next(ctx)? { @@ -107,7 +96,7 @@ impl LocalInputAdapter for AntigravityAdapter { fn fold_antigravity_units( parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, seen: &mut HashSet, ) -> Result<(), crate::adapters::InputPipelineError> { for parsed_unit in parsed { @@ -119,20 +108,14 @@ fn fold_antigravity_units( status, rejections, } = adapter_cache::resolve_unit(parsed_unit, ctx)?; - ctx.health.record(crate::input_health::InputHealth { - client: unit.client, - path: unit.path.clone(), - status, - rejections, - }); + ctx.record_health(unit.path.clone(), status, rejections); let path = unit.path.clone(); let cache_write_outcome = adapter_cache::write_cache(cache_write, ctx, &messages); if cache_write_outcome.is_err() && invalidate_cache { - ctx.input_cache.remove(&path, unit.parser_version); + ctx.input_cache.remove(&path, unit.decoder.version()); } let cache_write_outcome = cache_write_outcome?; adapter_cache::emit_messages( - unit.client, messages .into_iter() .filter(|message| crate::should_keep_deduped_message(seen, message)), @@ -140,7 +123,7 @@ fn fold_antigravity_units( ); if cache_write_outcome == adapter_cache::CacheWriteOutcome::NotPlanned && invalidate_cache { - ctx.input_cache.remove(&path, unit.parser_version); + ctx.input_cache.remove(&path, unit.decoder.version()); } } Ok(()) @@ -149,12 +132,13 @@ fn fold_antigravity_units( pub(crate) static ANTIGRAVITY_ADAPTER: AntigravityAdapter = AntigravityAdapter; fn antigravity_extra_roots( + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { let mut roots = Vec::new(); let mut seen = HashSet::new(); - for root in adapter_discover::extra_roots_for_client(ClientId::Antigravity, ctx)? { + for root in adapter_discover::extra_roots_for_client(client, ctx)? { push_unique_root(&mut roots, &mut seen, root)?; } @@ -171,7 +155,6 @@ fn push_unique_root( Err(source) if source.kind() == std::io::ErrorKind::NotFound => return Ok(()), Err(source) => { return Err(InputDiscoveryError::new( - ClientId::Antigravity, &root, "canonicalize configured scan root", source, @@ -188,6 +171,7 @@ fn push_unique_root( mod tests { use super::*; use crate::adapters::UnitMessagePayload; + use crate::message_cache::{DecoderId, DecoderVersion}; use crate::scanner::ScannerSettings; use crate::sessions::ParsedMessage; use crate::{message_cache, TokenBreakdown}; @@ -215,22 +199,24 @@ mod tests { let settings = ScannerSettings::default(); let units = ANTIGRAVITY_ADAPTER - .discover_checked(&scan_context(home.path(), &settings)) + .discover_checked(ClientId::Antigravity, &scan_context(home.path(), &settings)) .unwrap(); assert_eq!(units.len(), 1); let unit = &units[0]; - assert_eq!(unit.client, ClientId::Antigravity); assert_eq!(unit.path, cli_path); assert_eq!(unit.fingerprint_policy, FingerprintPolicy::SqliteWithWal); assert_eq!( - unit.parser_version, - ParserVersion::new( - ParserId::AntigravityCliSqlite, + unit.decoder.version(), + DecoderVersion::new( + DecoderId::AntigravityCliSqlite, ANTIGRAVITY_CLI_RECORD_REJECTION_REVISION, ) ); - assert!(matches!(unit.meta, InputUnitMeta::AntigravityCliSqlite)); + assert!(matches!( + unit.decoder.route(), + DecoderRoute::AntigravityCliSqlite + )); } #[test] @@ -249,7 +235,7 @@ mod tests { ..ScannerSettings::default() }; let units = ANTIGRAVITY_ADAPTER - .discover_checked(&scan_context(home.path(), &settings)) + .discover_checked(ClientId::Antigravity, &scan_context(home.path(), &settings)) .unwrap(); assert_eq!(units.len(), 1); @@ -258,7 +244,10 @@ mod tests { units[0].fingerprint_policy, FingerprintPolicy::SqliteWithWal ); - assert!(matches!(units[0].meta, InputUnitMeta::AntigravityCliSqlite)); + assert!(matches!( + units[0].decoder.route(), + DecoderRoute::AntigravityCliSqlite + )); assert!(!units.iter().any(|unit| unit.path == jsonl_path)); } @@ -274,15 +263,19 @@ mod tests { }; let units = ANTIGRAVITY_ADAPTER - .discover_checked(&scan_context(home.path(), &settings)) + .discover_checked(ClientId::Antigravity, &scan_context(home.path(), &settings)) .unwrap(); assert!(units.is_empty()); } - fn parsed_unit(path: &Path, meta: InputUnitMeta, message: ParsedMessage) -> ParsedUnit { + fn parsed_unit(path: &Path, route: DecoderRoute, message: ParsedMessage) -> ParsedUnit { + assert_eq!(route, DecoderRoute::AntigravityCliSqlite); ParsedUnit::healthy( - InputUnit::plain_file(ClientId::Antigravity, path.to_path_buf()).with_meta(meta), + InputUnit::plain_file( + path.to_path_buf(), + DecoderSpec::antigravity_cli_sqlite(ANTIGRAVITY_CLI_RECORD_REJECTION_REVISION), + ), UnitMessagePayload::Fresh(vec![message]), None, false, @@ -313,27 +306,26 @@ mod tests { let dedup_key = sessions::antigravity_cli::response_dedup_key("resp-shared"); let first = parsed_unit( &dir.path().join("first.db"), - InputUnitMeta::AntigravityCliSqlite, + DecoderRoute::AntigravityCliSqlite, antigravity_message("first-session", Some(dedup_key)), ); let second = parsed_unit( &dir.path().join("second.db"), - InputUnitMeta::AntigravityCliSqlite, + DecoderRoute::AntigravityCliSqlite, antigravity_message("second-session", Some(dedup_key)), ); let mut cache = message_cache::InputMessageCache::default(); let mut messages = Vec::new(); + let binding = crate::adapters::adapter_for(ClientId::Antigravity).unwrap(); + let mut fold_ctx = FoldContext::new(binding, &mut cache, None); + let mut sink = BoundMessageSink::new(binding, &mut messages); ANTIGRAVITY_ADAPTER - .fold( - vec![first, second], - &mut FoldContext::new(&mut cache, None), - &mut messages, - ) + .fold(vec![first, second], &mut fold_ctx, &mut sink) .unwrap(); assert_eq!(messages.len(), 1); - assert_eq!(messages[0].client.as_ref(), ClientId::Antigravity.as_str()); + assert_eq!(messages[0].client, ClientId::Antigravity); assert_eq!(messages[0].session_id.as_ref(), "first-session"); } } diff --git a/crates/tokscale-core/src/adapters/cache.rs b/crates/tokscale-core/src/adapters/cache.rs index de37be907..89bc7f778 100644 --- a/crates/tokscale-core/src/adapters/cache.rs +++ b/crates/tokscale-core/src/adapters/cache.rs @@ -1,11 +1,10 @@ use std::path::Path; use crate::adapters::{ - CacheHitPlan, FingerprintPolicy, FoldContext, InputPipelineError, InputPlanningError, - InputUnit, MessageSink, ParseContext, ParsedUnit, UnitMessagePayload, UnitScanHealth, + BoundMessageSink, CacheHitPlan, FingerprintPolicy, FoldContext, InputPipelineError, + InputPlanningError, InputUnit, ParseContext, ParsedUnit, UnitMessagePayload, UnitScanHealth, }; -use crate::clients::ClientId; -use crate::input_health::{InputFailure, InputHealth, InputStatus, ScannedInput}; +use crate::input_health::{InputFailure, InputStatus, ScannedInput}; use crate::{message_cache, sessions::ParsedMessage}; pub(crate) fn plan_cache_hit( @@ -16,7 +15,7 @@ pub(crate) fn plan_cache_hit( return Ok(CacheHitPlan::Miss(unit)); } unit.revalidate_snapshot_for_cache_decision()?; - let cached = match input_cache.get_meta(&unit.path, unit.parser_version) { + let cached = match input_cache.get_meta(&unit.path, unit.decoder.version()) { Ok(Some(cached)) => cached, Ok(None) => { unit.mark_cache_lookup_completed_no_hit(); @@ -48,7 +47,7 @@ pub(crate) fn plan_cache_hit( unit.release_prepared_snapshot(); let read_plan = - message_cache::CacheReadPlan::new(&unit.path, unit.parser_version, cached.fingerprint); + message_cache::CacheReadPlan::new(&unit.path, unit.decoder.version(), cached.fingerprint); let mut parsed = ParsedUnit::healthy(unit, UnitMessagePayload::CacheHit(read_plan), None, false); parsed.health.rejections = cached.rejections; @@ -250,7 +249,7 @@ where Some(Box::new( message_cache::CacheWritePlan::new( &unit.path, - unit.parser_version, + unit.decoder.version(), fingerprint, None, ) @@ -333,25 +332,22 @@ fn preserves_primary_on_related_failure( pub(crate) fn fold_units( parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), InputPipelineError> { fold_units_with_filter(parsed, ctx, sink, |_, messages| messages) } pub(crate) fn emit_messages( - client: ClientId, messages: impl IntoIterator, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) { - for message in messages { - sink.push_message(message.attribute(client)); - } + sink.emit_all(messages); } pub(crate) fn fold_units_with_filter( parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, mut filter: F, ) -> Result<(), InputPipelineError> where @@ -366,25 +362,19 @@ where status, rejections, } = resolve_unit(parsed_unit, ctx)?; - debug_assert!(unit.client.local_def().is_some()); - ctx.health.record(InputHealth { - client: unit.client, - path: unit.path.clone(), - status, - rejections, - }); + ctx.record_health(unit.path.clone(), status, rejections); let path = unit.path.clone(); - let parser_version = unit.parser_version; + let decoder_version = unit.decoder.version(); let cache_write_outcome = write_cache(cache_write, ctx, &messages); if cache_write_outcome.is_err() && invalidate_cache { - ctx.input_cache.remove(&path, parser_version); + ctx.input_cache.remove(&path, decoder_version); } let cache_write_outcome = cache_write_outcome?; let messages = filter(&unit, messages); - emit_messages(unit.client, messages, sink); + emit_messages(messages, sink); if cache_write_outcome == CacheWriteOutcome::NotPlanned && invalidate_cache { - ctx.input_cache.remove(&path, parser_version); + ctx.input_cache.remove(&path, decoder_version); } } Ok(()) @@ -432,34 +422,18 @@ pub(crate) fn resolve_unit( return Err(failure.into()); } debug_assert_eq!(failure.input_path, unit.path); - debug_assert_eq!(failure.parser_version, unit.parser_version); + debug_assert_eq!(failure.decoder_version, unit.decoder.version()); let remove_failed_shard = failure.requires_shard_removal(); if remove_failed_shard { - ctx.input_cache.remove(&unit.path, unit.parser_version); + ctx.input_cache.remove(&unit.path, unit.decoder.version()); } else { ctx.input_cache - .invalidate_read(&unit.path, unit.parser_version); + .invalidate_read(&unit.path, unit.decoder.version()); } unit.mark_cache_lookup_completed_no_hit(); recovery_requires_removal |= remove_failed_shard; - let adapter = super::adapter_for(unit.client) - .expect("cacheable input unit must have a registered local adapter"); - let mut reparsed = adapter.parse_checked( - vec![unit], - &ParseContext { - pricing: ctx.pricing, - }, - ); - if reparsed.len() != 1 { - return Err(InputPipelineError::contract(format!( - "single-input cache recovery returned {} parsed units instead of one", - reparsed.len() - ))); - } - parsed = reparsed.pop().ok_or_else(|| { - InputPipelineError::contract("single-input cache recovery result disappeared") - })?; + parsed = ctx.reparse_one(unit)?; } } } @@ -511,9 +485,12 @@ pub(crate) fn resolve_messages( #[cfg(test)] mod tests { + use std::path::PathBuf; + use super::*; - use crate::adapters::InputUnit; + use crate::adapters::{adapter_for, DecoderSpec, InputUnit}; use crate::clients::ClientId; + use crate::message_cache::DecoderId; use crate::{TokenBreakdown, UnifiedMessage}; const PI_INPUT: &str = r#"{"type":"session","id":"input-session","timestamp":"2026-01-01T00:00:00.000Z","cwd":"/tmp"} @@ -548,7 +525,15 @@ mod tests { } fn pi_unit(path: &Path) -> InputUnit { - InputUnit::plain_file(ClientId::Pi, path.to_path_buf()) + InputUnit::plain_file(path.to_path_buf(), DecoderSpec::plain(DecoderId::Pi, 1)) + } + + fn plain_unit(path: impl Into, decoder_id: DecoderId) -> InputUnit { + InputUnit::plain_file(path.into(), DecoderSpec::plain(decoder_id, 1)) + } + + fn binding(client: ClientId) -> crate::adapters::AdapterBinding<'static> { + adapter_for(client).expect("test client must have an adapter binding") } fn expect_cache_hit( @@ -580,7 +565,7 @@ mod tests { let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_dir); cache.insert(message_cache::CachedInputEntry::new_with_version( &unit.path, - unit.parser_version, + unit.decoder.version(), fingerprint.clone(), vec![ParsedMessage::new( "gpt-5.5", @@ -600,19 +585,27 @@ mod tests { } fn fold_planned_unit( + client: ClientId, parsed: ParsedUnit, cache: &mut message_cache::InputMessageCache, ) -> Vec { - fold_planned_unit_result(parsed, cache).unwrap() + fold_planned_unit_result(client, parsed, cache).unwrap() } fn fold_planned_unit_result( + client: ClientId, parsed: ParsedUnit, cache: &mut message_cache::InputMessageCache, ) -> Result, InputPipelineError> { - let mut sink = Vec::new(); - fold_units(vec![parsed], &mut FoldContext::new(cache, None), &mut sink)?; - Ok(sink) + let binding = binding(client); + let mut messages = Vec::new(); + let mut sink = BoundMessageSink::new(binding, &mut messages); + fold_units( + vec![parsed], + &mut FoldContext::new(binding, cache, None), + &mut sink, + )?; + Ok(messages) } fn assert_warm_hit_reads_no_input_bytes(unit: InputUnit) { @@ -624,7 +617,7 @@ mod tests { let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); cache.insert(message_cache::CachedInputEntry::new_with_version( &unit.path, - unit.parser_version, + unit.decoder.version(), fingerprint, vec![cached_message()], None, @@ -661,7 +654,7 @@ mod tests { let path = dir.path().join("session.json"); std::fs::write(&path, b"input contents").unwrap(); - assert_warm_hit_reads_no_input_bytes(InputUnit::plain_file(ClientId::Amp, path)); + assert_warm_hit_reads_no_input_bytes(plain_unit(path, DecoderId::Amp)); } #[test] @@ -672,7 +665,10 @@ mod tests { std::fs::write(&path, b"sqlite contents").unwrap(); std::fs::write(&wal_path, b"wal contents").unwrap(); - assert_warm_hit_reads_no_input_bytes(InputUnit::sqlite_with_wal(ClientId::Zed, path)); + assert_warm_hit_reads_no_input_bytes(InputUnit::sqlite_with_wal( + path, + DecoderSpec::plain(DecoderId::Zed, 1), + )); } #[test] @@ -685,9 +681,9 @@ mod tests { std::fs::write(&meta_path, b"meta contents").unwrap(); assert_warm_hit_reads_no_input_bytes(InputUnit::claude_code( - ClientId::Claude, path, home.path().to_path_buf(), + DecoderSpec::plain(DecoderId::Claude, 1), )); } @@ -696,19 +692,19 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("session.json"); std::fs::write(&path, b"old contents").unwrap(); - let old_unit = InputUnit::plain_file(ClientId::Amp, path.clone()); + let old_unit = plain_unit(path.clone(), DecoderId::Amp); let old_fingerprint = old_unit.input_policy().fingerprint().unwrap(); let mut cache = message_cache::InputMessageCache::default(); cache.insert(message_cache::CachedInputEntry::new_with_version( &path, - old_unit.parser_version, + old_unit.decoder.version(), old_fingerprint, vec![cached_message()], None, )); std::fs::write(&path, b"new and larger contents").unwrap(); - let unit = InputUnit::plain_file(ClientId::Amp, path.clone()) + let unit = plain_unit(path.clone(), DecoderId::Amp) .prepare_snapshot() .unwrap(); let expected_snapshot = unit.input_policy().snapshot().unwrap(); @@ -737,7 +733,7 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("session.json"); std::fs::write(&path, b"input contents").unwrap(); - let unit = InputUnit::plain_file(ClientId::Amp, path.clone()) + let unit = plain_unit(path.clone(), DecoderId::Amp) .prepare_snapshot() .unwrap(); let mut cache = message_cache::InputMessageCache::default(); @@ -747,7 +743,7 @@ mod tests { assert!(miss.prepared_input_snapshot().is_some()); cache.insert(message_cache::CachedInputEntry::new_with_version( &path, - miss.parser_version, + miss.decoder.version(), miss.input_policy().fingerprint().unwrap(), vec![cached_message()], None, @@ -782,10 +778,12 @@ mod tests { }, ); - let health = parsed.input_health(); - assert_eq!(health.client, ClientId::Pi); - assert_eq!(health.path, input_path); - let failure = health.status.failure().expect("input must be unavailable"); + assert_eq!(parsed.unit.path, input_path); + let failure = parsed + .health + .status + .failure() + .expect("input must be unavailable"); assert_eq!(failure.operation, "parse test SQLite"); assert!(failure.message.contains("sqlite root cause")); assert!(matches!( @@ -800,7 +798,7 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("empty.jsonl"); std::fs::write(&path, b"").unwrap(); - let unit = InputUnit::plain_file(ClientId::Amp, path); + let unit = plain_unit(path, DecoderId::Amp); let parsed = load_or_scan_unit_with(unit, &ParseContext { pricing: None }, |_| { Ok(ScannedInput::complete(Vec::new())) @@ -814,7 +812,7 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("all-bad.jsonl"); std::fs::write(&path, b"bad").unwrap(); - let unit = InputUnit::plain_file(ClientId::Amp, path); + let unit = plain_unit(path, DecoderId::Amp); let parsed = load_or_scan_unit_with(unit, &ParseContext { pricing: None }, |_| { let mut scanned = ScannedInput::complete(Vec::new()); @@ -838,7 +836,7 @@ mod tests { let shard_path = message_cache::mark_current_key_shard_as_unsupported_format_for_test( cache_dir.path(), &input_path, - unit.parser_version, + unit.decoder.version(), ); let before = std::fs::read(&shard_path).unwrap(); let cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); @@ -865,7 +863,7 @@ mod tests { let shard_path = message_cache::mark_current_key_shard_as_future_format_for_test( cache_dir.path(), &input_path, - unit.parser_version, + unit.decoder.version(), ); let before = std::fs::read(&shard_path).unwrap(); let cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); @@ -883,11 +881,11 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("session.json"); std::fs::write(&path, b"input contents").unwrap(); - let unit = InputUnit::plain_file(ClientId::Amp, path.clone()); + let unit = plain_unit(path.clone(), DecoderId::Amp); let mut cache = message_cache::InputMessageCache::default(); cache.insert(message_cache::CachedInputEntry::new_with_version( &path, - unit.parser_version, + unit.decoder.version(), unit.input_policy().fingerprint().unwrap(), vec![cached_message()], None, @@ -906,7 +904,7 @@ mod tests { let path = dir.path().join("session.json"); std::fs::write(&path, b"original").unwrap(); let original_mtime = std::fs::metadata(&path).unwrap().modified().unwrap(); - let unit = InputUnit::plain_file(ClientId::Amp, path.clone()) + let unit = plain_unit(path.clone(), DecoderId::Amp) .prepare_snapshot() .unwrap(); let policy = unit.input_policy(); @@ -922,7 +920,7 @@ mod tests { let mut cache = message_cache::InputMessageCache::default(); cache.insert(message_cache::CachedInputEntry::new_with_version( &path, - unit.parser_version, + unit.decoder.version(), fingerprint, vec![cached_message()], None, @@ -971,7 +969,7 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("session.json"); std::fs::write(&path, b"before").unwrap(); - let unit = InputUnit::plain_file(ClientId::Amp, path.clone()); + let unit = plain_unit(path.clone(), DecoderId::Amp); let parsed = load_or_scan_unit_with(unit, &ParseContext { pricing: None }, |_| { std::fs::write(&path, b"after-and-different-size").unwrap(); Ok(ScannedInput::complete(vec![cached_message()])) @@ -989,14 +987,14 @@ mod tests { let related = dir.path().join("metadata.jsonl"); std::fs::write(&primary, b"primary contents").unwrap(); std::fs::write(&related, b"related contents").unwrap(); - let unit = InputUnit::plain_file(ClientId::Kiro, primary.clone()) - .with_optional_dependency(related.clone()); - let parser_version = unit.parser_version; + let unit = + plain_unit(primary.clone(), DecoderId::Kiro).with_optional_dependency(related.clone()); + let decoder_version = unit.decoder.version(); let fingerprint = unit.input_policy().fingerprint().unwrap(); let mut cache = message_cache::InputMessageCache::default(); cache.insert(message_cache::CachedInputEntry::new_with_version( &primary, - unit.parser_version, + unit.decoder.version(), fingerprint, vec![cached_message()], None, @@ -1025,10 +1023,10 @@ mod tests { parsed.messages, UnitMessagePayload::Fresh(ref messages) if messages.len() == 1 )); - let messages = fold_planned_unit(parsed, &mut cache); + let messages = fold_planned_unit(ClientId::Kiro, parsed, &mut cache); assert_eq!(messages.len(), 1); assert!( - cache.get_meta(&primary, parser_version).unwrap().is_none(), + cache.get_meta(&primary, decoder_version).unwrap().is_none(), "the stale shard must be invalidated instead of surviving the partial scan" ); } @@ -1040,8 +1038,7 @@ mod tests { let dependency = dir.path().join("parent.jsonl"); std::fs::write(&primary, b"child contents").unwrap(); std::fs::create_dir(&dependency).unwrap(); - let unit = - InputUnit::plain_file(ClientId::CommandCode, primary).with_dependency(dependency); + let unit = plain_unit(primary, DecoderId::CommandCode).with_dependency(dependency); let scan_called = std::cell::Cell::new(false); let parsed = load_or_scan_unit_with(unit, &ParseContext { pricing: None }, |_| { @@ -1064,7 +1061,7 @@ mod tests { let related = dir.path().join("metadata.jsonl"); std::fs::create_dir(&primary).unwrap(); std::fs::write(&related, b"related contents").unwrap(); - let unit = InputUnit::plain_file(ClientId::Kiro, primary).with_optional_dependency(related); + let unit = plain_unit(primary, DecoderId::Kiro).with_optional_dependency(related); let scan_called = std::cell::Cell::new(false); let parsed = load_or_scan_unit_with(unit, &ParseContext { pricing: None }, |_| { @@ -1087,7 +1084,7 @@ mod tests { let wal_path = dir.path().join("history.db-wal"); std::fs::write(&path, b"database").unwrap(); std::fs::write(&wal_path, b"wal-before").unwrap(); - let unit = InputUnit::sqlite_with_wal(ClientId::Zed, path); + let unit = InputUnit::sqlite_with_wal(path, DecoderSpec::plain(DecoderId::Zed, 1)); let parsed = load_or_scan_unit_with(unit, &ParseContext { pricing: None }, |_| { std::fs::write(&wal_path, b"wal-after-and-larger").unwrap(); Ok(ScannedInput::complete(vec![cached_message()])) @@ -1105,12 +1102,12 @@ mod tests { let input_path = input_dir.path().join("session.jsonl"); std::fs::write(&input_path, PI_INPUT).unwrap(); let unit = pi_unit(&input_path); - let expected_client = unit.client; + let expected_client = ClientId::Pi; let fingerprint = seed_disk_cache(cache_dir.path(), &unit, "stale-cache-session"); message_cache::truncate_shard_after_header_for_test( cache_dir.path(), &input_path, - unit.parser_version, + unit.decoder.version(), ); let mut diagnostic_reader = @@ -1118,7 +1115,7 @@ mod tests { let failure = diagnostic_reader .take_messages(&message_cache::CacheReadPlan::new( &input_path, - unit.parser_version, + unit.decoder.version(), fingerprint, )) .expect_err("truncated cache body must be an explicit read failure"); @@ -1137,9 +1134,9 @@ mod tests { plan_cache_hit(unit.clone().prepare_snapshot().unwrap(), &cache), "valid header must still plan a cache hit", ); - let repaired = fold_planned_unit(parsed, &mut cache); + let repaired = fold_planned_unit(ClientId::Pi, parsed, &mut cache); assert_eq!(repaired.len(), 1); - assert_eq!(repaired[0].client.as_ref(), expected_client.as_str()); + assert_eq!(repaired[0].client, expected_client); assert_eq!(repaired[0].session_id.as_ref(), "input-session"); assert_eq!(repaired[0].tokens.input, 17); @@ -1149,8 +1146,8 @@ mod tests { plan_cache_hit(unit.prepare_snapshot().unwrap(), &warm_cache), "successful recovery must atomically replace the failed shard", ); - let warm_messages = fold_planned_unit(warm, &mut warm_cache); - assert_eq!(warm_messages[0].client.as_ref(), expected_client.as_str()); + let warm_messages = fold_planned_unit(ClientId::Pi, warm, &mut warm_cache); + assert_eq!(warm_messages[0].client, expected_client); assert_eq!(warm_messages[0].session_id.as_ref(), "input-session"); assert_eq!( message_cache::get_input_read_stats(&input_path), @@ -1165,13 +1162,12 @@ mod tests { let cache_dir = tempfile::TempDir::new().unwrap(); let input_path = input_dir.path().join("opencode.db"); std::fs::write(&input_path, b"not-a-database").unwrap(); - let unit = InputUnit::sqlite_with_wal(ClientId::OpenCode, input_path.clone()) - .with_meta(crate::adapters::InputUnitMeta::OpenCodeSqlite); + let unit = InputUnit::sqlite_with_wal(input_path.clone(), DecoderSpec::opencode_sqlite(1)); let fingerprint = unit.input_policy().fingerprint().unwrap(); let mut seed = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); seed.insert(message_cache::CachedInputEntry::new_with_version( &input_path, - unit.parser_version, + unit.decoder.version(), fingerprint, vec![cached_message()], None, @@ -1180,7 +1176,7 @@ mod tests { let shard_path = message_cache::truncate_shard_after_header_for_test( cache_dir.path(), &input_path, - unit.parser_version, + unit.decoder.version(), ); let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); @@ -1190,7 +1186,7 @@ mod tests { ); std::fs::remove_file(&input_path).unwrap(); std::fs::create_dir(&input_path).unwrap(); - let mut ctx = FoldContext::new(&mut cache, None); + let mut ctx = FoldContext::new(binding(ClientId::OpenCode), &mut cache, None); let resolved = resolve_unit(parsed, &mut ctx) .expect("recovery parse failure must isolate the unit, not fail the pipeline"); let failure = resolved @@ -1213,13 +1209,12 @@ mod tests { let cache_dir = tempfile::TempDir::new().unwrap(); let input_path = input_dir.path().join("opencode.db"); std::fs::write(&input_path, b"not-a-database").unwrap(); - let unit = InputUnit::sqlite_with_wal(ClientId::OpenCode, input_path.clone()) - .with_meta(crate::adapters::InputUnitMeta::OpenCodeSqlite); + let unit = InputUnit::sqlite_with_wal(input_path.clone(), DecoderSpec::opencode_sqlite(1)); let fingerprint = unit.input_policy().fingerprint().unwrap(); let mut seed = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); seed.insert(message_cache::CachedInputEntry::new_with_version( &input_path, - unit.parser_version, + unit.decoder.version(), fingerprint, vec![cached_message()], None, @@ -1228,7 +1223,7 @@ mod tests { let shard_path = message_cache::truncate_shard_after_header_for_test( cache_dir.path(), &input_path, - unit.parser_version, + unit.decoder.version(), ); let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); @@ -1238,7 +1233,7 @@ mod tests { ); std::fs::remove_file(&input_path).unwrap(); std::fs::create_dir(&input_path).unwrap(); - let mut ctx = FoldContext::new(&mut cache, None); + let mut ctx = FoldContext::new(binding(ClientId::OpenCode), &mut cache, None); let resolved = resolve_unit(parsed, &mut ctx) .expect("recovery parse failure must isolate the unit, not fail the pipeline"); assert!( @@ -1268,7 +1263,7 @@ mod tests { message_cache::replace_shard_message_count_for_test( cache_dir.path(), &input_path, - unit.parser_version, + unit.decoder.version(), 2, ); @@ -1277,7 +1272,7 @@ mod tests { plan_cache_hit(unit.prepare_snapshot().unwrap(), &cache), "mismatched body count remains a planned header hit", ); - let messages = fold_planned_unit(parsed, &mut cache); + let messages = fold_planned_unit(ClientId::Pi, parsed, &mut cache); assert_eq!(messages.len(), 1); assert_eq!(messages[0].session_id.as_ref(), "input-session"); @@ -1298,11 +1293,14 @@ mod tests { plan_cache_hit(unit.clone().prepare_snapshot().unwrap(), &cache), "seeded shard must plan a cache hit", ); - let shard_path = - message_cache::shard_path_for_test(cache_dir.path(), &input_path, unit.parser_version); + let shard_path = message_cache::shard_path_for_test( + cache_dir.path(), + &input_path, + unit.decoder.version(), + ); std::fs::remove_file(&shard_path).unwrap(); - let messages = fold_planned_unit_result(parsed, &mut cache) + let messages = fold_planned_unit_result(ClientId::Pi, parsed, &mut cache) .expect("a missing derived shard must be rebuilt from the input"); assert_eq!(messages.len(), 1); assert_eq!(messages[0].session_id.as_ref(), "input-session"); @@ -1330,7 +1328,7 @@ mod tests { std::fs::write(&input_path, PI_REPLACEMENT_INPUT).unwrap(); seed_disk_cache(cache_dir.path(), &unit, "replacement-input-session"); - let messages = fold_planned_unit_result(parsed, &mut reader) + let messages = fold_planned_unit_result(ClientId::Pi, parsed, &mut reader) .expect("a stale derived shard plan must reparse the current input"); assert_eq!(messages.len(), 1); assert_eq!(messages[0].session_id.as_ref(), "replacement-input-session"); @@ -1340,7 +1338,7 @@ mod tests { plan_cache_hit(unit.prepare_snapshot().unwrap(), &repaired_cache), "current input fingerprint must have a repaired shard", ); - let cached = fold_planned_unit(repaired, &mut repaired_cache); + let cached = fold_planned_unit(ClientId::Pi, repaired, &mut repaired_cache); assert_eq!(cached[0].session_id.as_ref(), "replacement-input-session"); } @@ -1360,18 +1358,23 @@ mod tests { ); std::fs::write(&input_path, PI_REPLACEMENT_INPUT).unwrap(); seed_disk_cache(cache_dir.path(), &unit, "replacement-cache-session"); - let shard_path = - message_cache::shard_path_for_test(cache_dir.path(), &input_path, unit.parser_version); + let shard_path = message_cache::shard_path_for_test( + cache_dir.path(), + &input_path, + unit.decoder.version(), + ); let replacement_bytes = std::fs::read(&shard_path).unwrap(); std::fs::write(&input_path, b"not a pi jsonl session").unwrap(); - let mut sink = Vec::new(); - let mut ctx = FoldContext::new(&mut reader, None); + let binding = binding(ClientId::Pi); + let mut messages = Vec::new(); + let mut sink = BoundMessageSink::new(binding, &mut messages); + let mut ctx = FoldContext::new(binding, &mut reader, None); fold_units(vec![parsed], &mut ctx, &mut sink) .expect("malformed third-party records must not fail the pipeline"); - assert!(sink.is_empty()); - assert_eq!(ctx.health.rejected_records(), 0); - assert_eq!(ctx.health.failed_inputs(), 1); + assert!(messages.is_empty()); + assert_eq!(ctx.health().rejected_records(), 0); + assert_eq!(ctx.health().failed_inputs(), 1); reader.save_if_dirty().unwrap(); assert_eq!( std::fs::read(&shard_path).unwrap(), @@ -1387,7 +1390,7 @@ mod tests { let input_path = input_dir.path().join("session.jsonl"); std::fs::write(&input_path, PI_INPUT).unwrap(); let unit = pi_unit(&input_path); - let expected_client = unit.client; + let expected_client = ClientId::Pi; seed_disk_cache(cache_dir.path(), &unit, "stale-cache-session"); let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); @@ -1398,17 +1401,19 @@ mod tests { let shard_path = message_cache::truncate_shard_after_header_for_test( cache_dir.path(), &input_path, - unit.parser_version, + unit.decoder.version(), ); std::fs::write(&input_path, b"not a pi jsonl session").unwrap(); - let mut sink = Vec::new(); - let mut ctx = FoldContext::new(&mut cache, None); + let binding = binding(ClientId::Pi); + let mut messages = Vec::new(); + let mut sink = BoundMessageSink::new(binding, &mut messages); + let mut ctx = FoldContext::new(binding, &mut cache, None); fold_units(vec![parsed], &mut ctx, &mut sink) .expect("recovery parse failure must isolate the unit, not fail the fold"); - assert!(sink.is_empty()); - assert_eq!(ctx.health.failed_inputs(), 1); - let health = &ctx.health.inputs()[0]; + assert!(messages.is_empty()); + assert_eq!(ctx.health().failed_inputs(), 1); + let health = &ctx.health().inputs()[0]; assert_eq!(health.path, input_path); assert!(health.status.failure().is_some()); cache.save_if_dirty().unwrap(); @@ -1429,8 +1434,8 @@ mod tests { crate::sessions::pi::parse_pi_file, ); let mut cold_cache = cold_cache; - let cold_messages = fold_planned_unit(cold_parsed, &mut cold_cache); - assert_eq!(cold_messages[0].client.as_ref(), expected_client.as_str()); + let cold_messages = fold_planned_unit(ClientId::Pi, cold_parsed, &mut cold_cache); + assert_eq!(cold_messages[0].client, expected_client); assert_eq!(cold_messages[0].session_id.as_ref(), "input-session"); } @@ -1451,21 +1456,23 @@ mod tests { plan_cache_hit(unit.prepare_snapshot().unwrap(), &cache), "second planned read must hit", ); - let mut sink = Vec::new(); + let binding = binding(ClientId::Pi); + let mut messages = Vec::new(); + let mut sink = BoundMessageSink::new(binding, &mut messages); let error = fold_units( vec![first, second], - &mut FoldContext::new(&mut cache, None), + &mut FoldContext::new(binding, &mut cache, None), &mut sink, ) .expect_err("second consumption must expose a typed pipeline error"); assert!(error.to_string().contains("already consumed")); assert_eq!( - sink.len(), + messages.len(), 1, "pipeline failure must not reparse duplicate output" ); - assert_eq!(sink[0].session_id.as_ref(), "cached-session"); + assert_eq!(messages[0].session_id.as_ref(), "cached-session"); } #[test] diff --git a/crates/tokscale-core/src/adapters/claude.rs b/crates/tokscale-core/src/adapters/claude.rs index 0815b3c7c..e1e3c7099 100644 --- a/crates/tokscale-core/src/adapters/claude.rs +++ b/crates/tokscale-core/src/adapters/claude.rs @@ -8,15 +8,15 @@ use rayon::prelude::*; use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ - AdapterScanContext, FingerprintPolicy, FoldContext, InputDiscoveryError, InputUnit, - LocalInputAdapter, MessageSink, ParseContext, ParsedBatchInput, ParsedUnit, + AdapterScanContext, BoundMessageSink, DecoderSpec, FingerprintPolicy, FoldContext, + InputDiscoveryError, InputUnit, LocalInputAdapter, ParseContext, ParsedBatchInput, ParsedUnit, MODEL_ID_CANONICALIZATION_REVISION, }; use crate::clients::ClientId; -use crate::message_cache::{ParserId, ParserVersion}; +use crate::message_cache::DecoderId; use crate::sessions; -const CLAUDE_PARSER_REVISION: u32 = MODEL_ID_CANONICALIZATION_REVISION + 10; +const CLAUDE_DECODER_REVISION: u32 = MODEL_ID_CANONICALIZATION_REVISION + 10; static CLAUDE_PROJECT_RESOLVERS: LazyLock< Mutex>>, @@ -56,42 +56,34 @@ fn claude_project_resolver( pub(crate) struct ClaudeAdapter; impl LocalInputAdapter for ClaudeAdapter { - fn client(&self) -> ClientId { - ClientId::Claude - } - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { reset_claude_project_resolver(Path::new(ctx.home_dir)); - let def = ClientId::Claude + let def = client .local_def() .expect("Claude adapter must have local scan policy"); let mut roots = vec![def.resolve_path(ctx.home_dir)]; - roots.extend(adapter_discover::extra_roots_for_client( - ClientId::Claude, - ctx, - )?); + roots.extend(adapter_discover::extra_roots_for_client(client, ctx)?); roots.push(PathBuf::from(format!( "{}/.claude/transcripts", ctx.home_dir ))); let units = adapter_discover::input_units_from_paths( - ClientId::Claude, - adapter_discover::scan_roots(ClientId::Claude, roots, def.pattern)?, + client, + adapter_discover::scan_roots(client, roots, def.pattern)?, FingerprintPolicy::ClaudeCodeWithHome { home_dir: PathBuf::from(ctx.home_dir), parent_session_path: None, }, + DecoderSpec::plain(DecoderId::Claude, CLAUDE_DECODER_REVISION), )? .into_iter() .map(configure_claude_parent_dependency) - .map(|unit| { - unit.with_parser_version(ParserVersion::new(ParserId::Claude, CLAUDE_PARSER_REVISION)) - }) .collect(); Ok(units) } @@ -144,7 +136,7 @@ impl LocalInputAdapter for ClaudeAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), crate::adapters::InputPipelineError> { let mut seen_keys = HashSet::new(); fold_claude_units(parsed, ctx, sink, &mut seen_keys) @@ -154,7 +146,7 @@ impl LocalInputAdapter for ClaudeAdapter { &self, batches: &mut ParsedBatchInput<'_>, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), crate::adapters::InputPipelineError> { let mut seen_keys = HashSet::new(); while let Some(parsed) = batches.next(ctx)? { @@ -250,7 +242,7 @@ fn resolve_flat_parent_dependency(path: &std::path::Path) -> FlatParentResolutio fn fold_claude_units( parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, seen_keys: &mut HashSet, ) -> Result<(), crate::adapters::InputPipelineError> { for parsed_unit in parsed { @@ -262,20 +254,14 @@ fn fold_claude_units( status, rejections, } = adapter_cache::resolve_unit(parsed_unit, ctx)?; - ctx.health.record(crate::input_health::InputHealth { - client: unit.client, - path: unit.path.clone(), - status, - rejections, - }); + ctx.record_health(unit.path.clone(), status, rejections); let path = unit.path.clone(); let cache_write_outcome = adapter_cache::write_cache(cache_write, ctx, &messages); if cache_write_outcome.is_err() && invalidate_cache { - ctx.input_cache.remove(&path, unit.parser_version); + ctx.input_cache.remove(&path, unit.decoder.version()); } let cache_write_outcome = cache_write_outcome?; adapter_cache::emit_messages( - unit.client, messages .into_iter() .filter(|message| crate::should_keep_deduped_message(seen_keys, message)), @@ -283,7 +269,7 @@ fn fold_claude_units( ); if cache_write_outcome == adapter_cache::CacheWriteOutcome::NotPlanned && invalidate_cache { - ctx.input_cache.remove(&path, unit.parser_version); + ctx.input_cache.remove(&path, unit.decoder.version()); } } Ok(()) @@ -294,9 +280,10 @@ pub(crate) static CLAUDE_ADAPTER: ClaudeAdapter = ClaudeAdapter; #[cfg(test)] mod tests { use std::collections::BTreeMap; - use std::path::Path; + use std::path::{Path, PathBuf}; use super::*; + use crate::adapters::{adapter_for, AdapterBinding}; use crate::input_health::InputHealth; use crate::message_cache; @@ -318,26 +305,61 @@ mod tests { fn discover_unit(home_dir: &Path, path: &Path) -> InputUnit { let settings = crate::scanner::ScannerSettings::default(); CLAUDE_ADAPTER - .discover_checked(&scan_context(home_dir, &settings)) + .discover_checked(ClientId::Claude, &scan_context(home_dir, &settings)) .unwrap() .into_iter() .find(|unit| unit.path == path) .unwrap_or_else(|| panic!("Claude input was not discovered: {}", path.display())) } - fn scan_and_fold( - unit: InputUnit, + fn binding() -> AdapterBinding<'static> { + adapter_for(ClientId::Claude).expect("Claude must have an adapter binding") + } + + fn decoder() -> DecoderSpec { + DecoderSpec::plain(DecoderId::Claude, CLAUDE_DECODER_REVISION) + } + + fn input_unit(path: PathBuf, home_dir: PathBuf) -> InputUnit { + InputUnit::claude_code(path, home_dir, decoder()) + } + + fn input_health(parsed: &ParsedUnit) -> InputHealth { + InputHealth { + client: ClientId::Claude, + path: parsed.unit.path.clone(), + status: parsed.health.status.clone(), + rejections: parsed.health.rejections.clone(), + } + } + + fn fold_parsed( + parsed: Vec, cache: &mut message_cache::InputMessageCache, - ) -> (Vec, InputHealth) { - let parsed = CLAUDE_ADAPTER.parse_checked(vec![unit], &ParseContext { pricing: None }); - let health = parsed[0].input_health(); + ) -> Vec { + let binding = binding(); let mut messages = Vec::new(); + let mut sink = BoundMessageSink::new(binding, &mut messages); CLAUDE_ADAPTER - .fold(parsed, &mut FoldContext::new(cache, None), &mut messages) + .fold( + parsed, + &mut FoldContext::new(binding, cache, None), + &mut sink, + ) .unwrap(); assert!(messages .iter() - .all(|message| message.client.as_ref() == ClientId::Claude.as_str())); + .all(|message| message.client == ClientId::Claude)); + messages + } + + fn scan_and_fold( + unit: InputUnit, + cache: &mut message_cache::InputMessageCache, + ) -> (Vec, InputHealth) { + let parsed = CLAUDE_ADAPTER.parse_checked(vec![unit], &ParseContext { pricing: None }); + let health = input_health(&parsed[0]); + let messages = fold_parsed(parsed, cache); (messages, health) } @@ -345,18 +367,8 @@ mod tests { parsed: ParsedUnit, cache: &mut message_cache::InputMessageCache, ) -> (Vec, InputHealth) { - let health = parsed.input_health(); - let mut messages = Vec::new(); - CLAUDE_ADAPTER - .fold( - vec![parsed], - &mut FoldContext::new(cache, None), - &mut messages, - ) - .unwrap(); - assert!(messages - .iter() - .all(|message| message.client.as_ref() == ClientId::Claude.as_str())); + let health = input_health(&parsed); + let messages = fold_parsed(vec![parsed], cache); (messages, health) } @@ -404,7 +416,7 @@ mod tests { }; let units = CLAUDE_ADAPTER - .discover_checked(&scan_context(home.path(), &settings)) + .discover_checked(ClientId::Claude, &scan_context(home.path(), &settings)) .unwrap(); let paths: Vec<_> = units.iter().map(|unit| unit.path.clone()).collect(); let mut expected = vec![default_file, workflow_file, transcript_file, extra_file]; @@ -424,11 +436,7 @@ mod tests { .path() .join(".claude/projects/project-a/session-1.jsonl"); write_file(&session_path, ""); - let unit = InputUnit::claude_code( - ClientId::Claude, - session_path.clone(), - home.path().to_path_buf(), - ); + let unit = input_unit(session_path.clone(), home.path().to_path_buf()); let mut digest_paths = unit.digest_paths(); digest_paths.sort_unstable(); @@ -473,16 +481,9 @@ mod tests { ); let mut cache = message_cache::InputMessageCache::default(); - let unit = InputUnit::claude_code( - ClientId::Claude, - session_path.clone(), - home.path().to_path_buf(), - ); + let unit = input_unit(session_path.clone(), home.path().to_path_buf()); let parsed = CLAUDE_ADAPTER.parse_checked(vec![unit], &ParseContext { pricing: None }); - let mut actual = Vec::new(); - CLAUDE_ADAPTER - .fold(parsed, &mut FoldContext::new(&mut cache, None), &mut actual) - .unwrap(); + let actual = fold_parsed(parsed, &mut cache); let expected = finalized( sessions::claudecode::parse_claude_file_with_home(&session_path, Some(home.path())) @@ -491,7 +492,7 @@ mod tests { ); assert!(actual .iter() - .all(|message| message.client.as_ref() == ClientId::Claude.as_str())); + .all(|message| message.client == ClientId::Claude)); assert_eq!(actual, expected); assert_eq!(actual.len(), 1); } @@ -506,17 +507,13 @@ mod tests { {not-json "#, ); - let unit = InputUnit::claude_code( - ClientId::Claude, - session_path.clone(), - home.path().to_path_buf(), - ); - let parser_version = unit.parser_version; + let unit = input_unit(session_path.clone(), home.path().to_path_buf()); + let decoder_version = unit.decoder.version(); let parsed = CLAUDE_ADAPTER.parse_checked(vec![unit], &ParseContext { pricing: None }); assert_eq!(parsed.len(), 1); - let health = parsed[0].input_health(); + let health = input_health(&parsed[0]); assert_eq!(health.client, ClientId::Claude); assert_eq!(health.path, session_path); assert!(matches!( @@ -533,19 +530,12 @@ mod tests { assert!(failure.message.contains("line 2")); let mut cache = message_cache::InputMessageCache::default(); - let mut messages = Vec::new(); - CLAUDE_ADAPTER - .fold( - parsed, - &mut FoldContext::new(&mut cache, None), - &mut messages, - ) - .unwrap(); + let messages = fold_parsed(parsed, &mut cache); assert_eq!(messages.len(), 1); - assert_eq!(messages[0].client.as_ref(), ClientId::Claude.as_str()); + assert_eq!(messages[0].client, ClientId::Claude); assert_eq!(messages[0].tokens.input, 10); assert!(cache - .get_meta(&session_path, parser_version) + .get_meta(&session_path, decoder_version) .unwrap() .is_none()); } @@ -561,23 +551,15 @@ mod tests { {"type":"assistant","timestamp":"2026-07-14T00:00:01Z","message":{"usage":{"input_tokens":99,"output_tokens":9}}} {"type":"assistant","timestamp":"2026-07-14T00:00:02Z","message":{"model":"claude-sonnet-4.6","usage":{"input_tokens":20,"output_tokens":2}}}"#, ); - let unit = InputUnit::claude_code( - ClientId::Claude, - session_path.clone(), - home.path().to_path_buf(), - ) - .with_parser_version(ParserVersion::new(ParserId::Claude, CLAUDE_PARSER_REVISION)) - .prepare_snapshot() - .unwrap(); + let unit = input_unit(session_path.clone(), home.path().to_path_buf()) + .prepare_snapshot() + .unwrap(); let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); let cold = CLAUDE_ADAPTER.parse_checked(vec![unit.clone()], &ParseContext { pricing: None }); - assert_eq!(cold[0].input_health().rejections.total(), 1); - let mut messages = Vec::new(); - CLAUDE_ADAPTER - .fold(cold, &mut FoldContext::new(&mut cache, None), &mut messages) - .unwrap(); + assert_eq!(cold[0].health.rejections.total(), 1); + let messages = fold_parsed(cold, &mut cache); assert_eq!(messages.len(), 2); cache.save_if_dirty().unwrap(); @@ -586,7 +568,7 @@ mod tests { let crate::adapters::CacheHitPlan::Hit(warm) = planned else { panic!("unchanged Claude input must use its complete cached scan"); }; - let health = warm.input_health(); + let health = input_health(&warm); assert!(matches!( health.status, crate::input_health::InputStatus::Complete @@ -610,7 +592,7 @@ mod tests { ); let unit = discover_unit(home.path(), &session_path); - let parser_version = unit.parser_version; + let decoder_version = unit.decoder.version(); let mut cache = message_cache::InputMessageCache::default(); let (unresolved, _) = scan_and_fold(unit, &mut cache); assert_eq!( @@ -618,7 +600,7 @@ mod tests { Some("-home-travis-external-project") ); assert!(cache - .get_meta(&session_path, parser_version) + .get_meta(&session_path, decoder_version) .unwrap() .is_none()); @@ -638,7 +620,7 @@ mod tests { Some("/home/travis/external-project") ); assert!(cache - .get_meta(&session_path, parser_version) + .get_meta(&session_path, decoder_version) .unwrap() .is_none()); } @@ -787,7 +769,7 @@ mod tests { let unit = discover_unit(home.path(), &child_path); assert!(!unit.digest_paths().contains(&parent_path)); - let parser_version = unit.parser_version; + let decoder_version = unit.decoder.version(); let mut cache = message_cache::InputMessageCache::default(); let (unresolved, _) = scan_and_fold(unit, &mut cache); assert_eq!( @@ -795,7 +777,7 @@ mod tests { Some("-home-travis-parent-project") ); assert!(cache - .get_meta(&child_path, parser_version) + .get_meta(&child_path, decoder_version) .unwrap() .is_none()); @@ -815,7 +797,7 @@ mod tests { Some("/home/travis/parent-project") ); assert!(cache - .get_meta(&child_path, parser_version) + .get_meta(&child_path, decoder_version) .unwrap() .is_none()); } @@ -855,7 +837,7 @@ mod tests { let settings = crate::scanner::ScannerSettings::default(); let units = CLAUDE_ADAPTER - .discover_checked(&scan_context(home.path(), &settings)) + .discover_checked(ClientId::Claude, &scan_context(home.path(), &settings)) .expect("one unresolved flat sidechain must not fail Claude discovery"); let unresolved_unit = units.iter().find(|unit| unit.path == unresolved).unwrap(); let regular_unit = units.iter().find(|unit| unit.path == regular).unwrap(); diff --git a/crates/tokscale-core/src/adapters/cline.rs b/crates/tokscale-core/src/adapters/cline.rs index 70e306fc2..90703e8b8 100644 --- a/crates/tokscale-core/src/adapters/cline.rs +++ b/crates/tokscale-core/src/adapters/cline.rs @@ -3,12 +3,14 @@ use rayon::prelude::*; use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ - AdapterScanContext, FingerprintPolicy, FoldContext, InputDiscoveryError, InputPipelineError, - InputUnit, LocalInputAdapter, MessageSink, ParseContext, ParsedUnit, - MODEL_ID_CANONICALIZATION_REVISION, + AdapterScanContext, BoundMessageSink, DecoderSpec, FingerprintPolicy, FoldContext, + InputDiscoveryError, InputPipelineError, InputUnit, LocalInputAdapter, ParseContext, + ParsedUnit, MODEL_ID_CANONICALIZATION_REVISION, }; use crate::clients::ClientId; -use crate::message_cache::{ParserId, ParserVersion}; +use crate::message_cache::DecoderId; +#[cfg(test)] +use crate::message_cache::DecoderVersion; use crate::{local_clients, sessions}; const CLINE_SDK_V1_REVISION: u32 = MODEL_ID_CANONICALIZATION_REVISION + 4; @@ -16,36 +18,30 @@ const CLINE_SDK_V1_REVISION: u32 = MODEL_ID_CANONICALIZATION_REVISION + 4; pub(crate) struct ClineAdapter; impl LocalInputAdapter for ClineAdapter { - fn client(&self) -> ClientId { - ClientId::Cline - } - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { - let def = ClientId::Cline + let def = client .local_def() .expect("Cline adapter must have local scan policy"); let mut roots = vec![local_clients::cline_session_data_dir(ctx.home_dir)]; - roots.extend(adapter_discover::extra_roots_for_client( - ClientId::Cline, - ctx, - )?); + roots.extend(adapter_discover::extra_roots_for_client(client, ctx)?); Ok(adapter_discover::input_units_from_paths( - ClientId::Cline, - adapter_discover::scan_roots(ClientId::Cline, roots, def.pattern)?, + client, + adapter_discover::scan_roots(client, roots, def.pattern)?, FingerprintPolicy::PlainFile, + DecoderSpec::plain(DecoderId::Cline, CLINE_SDK_V1_REVISION), )? .into_iter() - .map(|unit| { - let unit = match sessions::cline::cline_manifest_dependency_path(&unit.path) { + .map( + |unit| match sessions::cline::cline_manifest_dependency_path(&unit.path) { Some(manifest) => unit.with_optional_dependency(manifest), None => unit, - }; - unit.with_parser_version(ParserVersion::new(ParserId::Cline, CLINE_SDK_V1_REVISION)) - }) + }, + ) .collect()) } @@ -70,7 +66,7 @@ impl LocalInputAdapter for ClineAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), InputPipelineError> { adapter_cache::fold_units(parsed, ctx, sink) } @@ -109,14 +105,14 @@ mod tests { let settings = crate::scanner::ScannerSettings::default(); let units = CLINE_ADAPTER - .discover_checked(&scan_context(home.path(), &settings)) + .discover_checked(ClientId::Cline, &scan_context(home.path(), &settings)) .unwrap(); assert_eq!(units.len(), 1); assert_eq!(units[0].path, current); assert_eq!( - units[0].parser_version, - ParserVersion::new(ParserId::Cline, CLINE_SDK_V1_REVISION) + units[0].decoder.version(), + DecoderVersion::new(DecoderId::Cline, CLINE_SDK_V1_REVISION) ); assert_eq!( units[0].fingerprint_policy, @@ -150,7 +146,7 @@ mod tests { ], ); let units = CLINE_ADAPTER - .discover_checked(&scan_context(home.path(), &settings)) + .discover_checked(ClientId::Cline, &scan_context(home.path(), &settings)) .unwrap(); let paths: Vec = units.into_iter().map(|unit| unit.path).collect(); @@ -169,7 +165,7 @@ mod tests { let settings = crate::scanner::ScannerSettings::default(); let first = CLINE_ADAPTER - .discover_checked(&scan_context(home.path(), &settings)) + .discover_checked(ClientId::Cline, &scan_context(home.path(), &settings)) .unwrap() .remove(0) .input_policy() @@ -178,7 +174,7 @@ mod tests { write_file(&manifest, r#"{"workspace_root":"/tmp/project-b"}"#); let second = CLINE_ADAPTER - .discover_checked(&scan_context(home.path(), &settings)) + .discover_checked(ClientId::Cline, &scan_context(home.path(), &settings)) .unwrap() .remove(0) .input_policy() diff --git a/crates/tokscale-core/src/adapters/codebuddy.rs b/crates/tokscale-core/src/adapters/codebuddy.rs index 174825e16..cd4e593f1 100644 --- a/crates/tokscale-core/src/adapters/codebuddy.rs +++ b/crates/tokscale-core/src/adapters/codebuddy.rs @@ -6,12 +6,13 @@ use rayon::prelude::*; use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ - AdapterScanContext, CodeBuddyLogOrigin, FingerprintPolicy, FoldContext, InputDiscoveryError, - InputUnit, InputUnitMeta, LocalInputAdapter, MessageSink, ParseContext, ParsedBatchInput, - ParsedUnit, + AdapterScanContext, BoundMessageSink, CodeBuddyLogOrigin, DecoderRoute, DecoderSpec, + FingerprintPolicy, FoldContext, InputDiscoveryError, InputUnit, LocalInputAdapter, + ParseContext, ParsedBatchInput, ParsedUnit, }; use crate::clients::ClientId; -use crate::message_cache::{ParserId, ParserVersion}; +#[cfg(test)] +use crate::message_cache::DecoderVersion; use crate::sessions::{self, ParsedMessage}; const MIRROR_DEDUP_WINDOW_MS: i64 = 1000; @@ -23,45 +24,33 @@ const CODEBUDDY_EXTENSION_RECORD_REJECTION_REVISION: u32 = pub(crate) struct CodeBuddyAdapter; impl LocalInputAdapter for CodeBuddyAdapter { - fn client(&self) -> ClientId { - ClientId::CodeBuddy - } - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { - let def = ClientId::CodeBuddy + let def = client .local_def() .expect("CodeBuddy adapter must have local scan policy"); let default_root = def.resolve_path(ctx.home_dir); - let extra_roots = adapter_discover::extra_roots_for_client(ClientId::CodeBuddy, ctx)?; + let extra_roots = adapter_discover::extra_roots_for_client(client, ctx)?; - let mut jsonl_paths = - adapter_discover::scan_roots(ClientId::CodeBuddy, [default_root], def.pattern)?; + let mut jsonl_paths = adapter_discover::scan_roots(client, [default_root], def.pattern)?; jsonl_paths.extend(adapter_discover::scan_roots( - ClientId::CodeBuddy, + client, extra_roots.clone(), def.pattern, )?); let mut units = adapter_discover::input_units_from_paths( - ClientId::CodeBuddy, + client, jsonl_paths, FingerprintPolicy::PlainFile, - )? - .into_iter() - .map(|unit| { - unit.with_meta(InputUnitMeta::CodeBuddyJsonl) - .with_parser_version(ParserVersion::new( - ParserId::CodeBuddy, - CODEBUDDY_JSONL_RECORD_REJECTION_REVISION, - )) - }) - .collect::>(); + DecoderSpec::codebuddy_jsonl(CODEBUDDY_JSONL_RECORD_REJECTION_REVISION), + )?; - units.extend(codebuddy_extension_log_units(ctx.home_dir)?); - units.extend(codebuddy_extra_log_units(extra_roots)?); + units.extend(codebuddy_extension_log_units(client, ctx.home_dir)?); + units.extend(codebuddy_extra_log_units(client, extra_roots)?); dedup_units_by_canonical_path(&mut units)?; units.sort_by(|left, right| left.path.cmp(&right.path)); Ok(units) @@ -70,13 +59,13 @@ impl LocalInputAdapter for CodeBuddyAdapter { fn parse_checked(&self, units: Vec, ctx: &ParseContext<'_>) -> Vec { units .into_par_iter() - .map(|unit| match unit.meta { - InputUnitMeta::CodeBuddyJsonl => { + .map(|unit| match unit.decoder.route() { + DecoderRoute::CodeBuddyJsonl => { adapter_cache::load_or_scan_unit_with(unit, ctx, |path| { sessions::codebuddy::parse_codebuddy_jsonl_file(path) }) } - InputUnitMeta::CodeBuddyExtensionLog { .. } => { + DecoderRoute::CodeBuddyExtensionLog { .. } => { adapter_cache::load_or_scan_unit_with(unit, ctx, |path| { sessions::codebuddy::parse_codebuddy_extension_log_file(path) }) @@ -98,7 +87,7 @@ impl LocalInputAdapter for CodeBuddyAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), crate::adapters::InputPipelineError> { let mut deduper = CodeBuddyDeduper::default(); adapter_cache::fold_units_with_filter(parsed, ctx, sink, |unit, messages| { @@ -110,7 +99,7 @@ impl LocalInputAdapter for CodeBuddyAdapter { &self, batches: &mut ParsedBatchInput<'_>, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), crate::adapters::InputPipelineError> { let mut deduper = CodeBuddyDeduper::default(); while let Some(parsed) = batches.next(ctx)? { @@ -122,7 +111,10 @@ impl LocalInputAdapter for CodeBuddyAdapter { } } -fn codebuddy_extension_log_units(home_dir: &str) -> Result, InputDiscoveryError> { +fn codebuddy_extension_log_units( + client: ClientId, + home_dir: &str, +) -> Result, InputDiscoveryError> { let home = PathBuf::from(home_dir); let roots = [ ( @@ -149,43 +141,41 @@ fn codebuddy_extension_log_units(home_dir: &str) -> Result, Input let mut units = Vec::new(); for (root, origin, require_extension_component) in roots { - let paths = adapter_discover::scan_roots(ClientId::CodeBuddy, [root], "*.log")? + let paths = adapter_discover::scan_roots(client, [root], "*.log")? .into_iter() .filter(|path| !require_extension_component || has_codebuddy_extension_component(path)) .collect::>(); - units.extend( - adapter_discover::input_units_from_paths( - ClientId::CodeBuddy, - paths, - FingerprintPolicy::PlainFile, - )? - .into_iter() - .map(|unit| codebuddy_log_unit(unit, origin)), - ); + units.extend(adapter_discover::input_units_from_paths( + client, + paths, + FingerprintPolicy::PlainFile, + DecoderSpec::codebuddy_extension_log( + CODEBUDDY_EXTENSION_RECORD_REJECTION_REVISION, + origin, + ), + )?); } Ok(units) } -fn codebuddy_extra_log_units(roots: Vec) -> Result, InputDiscoveryError> { - let paths = adapter_discover::scan_roots(ClientId::CodeBuddy, roots, "*.log")?; - Ok(adapter_discover::input_units_from_paths( - ClientId::CodeBuddy, - paths, - FingerprintPolicy::PlainFile, - )? - .into_iter() - .filter_map(|unit| { - codebuddy_extra_log_origin(&unit.path).map(|origin| codebuddy_log_unit(unit, origin)) - }) - .collect()) -} - -fn codebuddy_log_unit(unit: InputUnit, origin: CodeBuddyLogOrigin) -> InputUnit { - unit.with_meta(InputUnitMeta::CodeBuddyExtensionLog { origin }) - .with_parser_version(ParserVersion::new( - ParserId::CodeBuddy, - CODEBUDDY_EXTENSION_RECORD_REJECTION_REVISION, - )) +fn codebuddy_extra_log_units( + client: ClientId, + roots: Vec, +) -> Result, InputDiscoveryError> { + let paths = adapter_discover::scan_roots(client, roots, "*.log")?; + Ok(paths + .into_iter() + .filter_map(|unit| { + let origin = codebuddy_extra_log_origin(&unit)?; + Some(InputUnit::plain_file( + unit, + DecoderSpec::codebuddy_extension_log( + CODEBUDDY_EXTENSION_RECORD_REJECTION_REVISION, + origin, + ), + )) + }) + .collect()) } fn codebuddy_extra_log_origin(path: &std::path::Path) -> Option { @@ -206,12 +196,7 @@ fn dedup_units_by_canonical_path(units: &mut Vec) -> Result<(), Input let mut keys = Vec::with_capacity(units.len()); for unit in units.iter() { keys.push(std::fs::canonicalize(&unit.path).map_err(|source| { - InputDiscoveryError::new( - ClientId::CodeBuddy, - &unit.path, - "canonicalize discovered input", - source, - ) + InputDiscoveryError::new(&unit.path, "canonicalize discovered input", source) })?); } let mut index = 0; @@ -251,7 +236,7 @@ impl CodeBuddyDeduper { return self.seen_keys.insert(key); } - let InputUnitMeta::CodeBuddyExtensionLog { origin } = unit.meta else { + let DecoderRoute::CodeBuddyExtensionLog { origin } = unit.decoder.route() else { return true; }; let signature = MirrorSignature::from_message(message); @@ -310,7 +295,7 @@ mod tests { use super::*; use crate::adapters::{FoldContext, ParseContext}; - use crate::message_cache::{self, ParserId}; + use crate::message_cache::{self, DecoderId}; fn scan_context<'a>( home_dir: &'a Path, @@ -335,7 +320,7 @@ mod tests { .collect(); assert!(messages .iter() - .all(|message| message.client.as_ref() == ClientId::CodeBuddy.as_str())); + .all(|message| message.client == ClientId::CodeBuddy)); messages } @@ -343,12 +328,15 @@ mod tests { let mut cache = message_cache::InputMessageCache::default(); let parsed = CODEBUDDY_ADAPTER.parse_checked(units, &ParseContext { pricing: None }); let mut sink = Vec::new(); + let binding = crate::adapters::adapter_for(ClientId::CodeBuddy).unwrap(); + let mut fold_ctx = FoldContext::new(binding, &mut cache, None); + let mut bound_sink = BoundMessageSink::new(binding, &mut sink); CODEBUDDY_ADAPTER - .fold(parsed, &mut FoldContext::new(&mut cache, None), &mut sink) + .fold(parsed, &mut fold_ctx, &mut bound_sink) .unwrap(); assert!(sink .iter() - .all(|message| message.client.as_ref() == ClientId::CodeBuddy.as_str())); + .all(|message| message.client == ClientId::CodeBuddy)); sink } @@ -370,35 +358,37 @@ mod tests { let settings = crate::scanner::ScannerSettings::default(); let ctx = scan_context(home.path(), &settings); - let mut units = CODEBUDDY_ADAPTER.discover_checked(&ctx).unwrap(); + let mut units = CODEBUDDY_ADAPTER + .discover_checked(ClientId::CodeBuddy, &ctx) + .unwrap(); units.sort_by(|left, right| left.path.cmp(&right.path)); let paths: Vec<_> = units.iter().map(|unit| unit.path.clone()).collect(); - let metas: Vec<_> = units.iter().map(|unit| unit.meta).collect(); + let metas: Vec<_> = units.iter().map(|unit| unit.decoder.route()).collect(); let mut expected = vec![project_path, ide_log, vscode_log]; expected.sort_unstable(); assert_eq!(paths, expected); for unit in &units { - let revision = match unit.meta { - InputUnitMeta::CodeBuddyJsonl => CODEBUDDY_JSONL_RECORD_REJECTION_REVISION, - InputUnitMeta::CodeBuddyExtensionLog { .. } => { + let revision = match unit.decoder.route() { + DecoderRoute::CodeBuddyJsonl => CODEBUDDY_JSONL_RECORD_REJECTION_REVISION, + DecoderRoute::CodeBuddyExtensionLog { .. } => { CODEBUDDY_EXTENSION_RECORD_REJECTION_REVISION } _ => unreachable!(), }; assert_eq!( - unit.parser_version, - ParserVersion::new(ParserId::CodeBuddy, revision) + unit.decoder.version(), + DecoderVersion::new(DecoderId::CodeBuddy, revision) ); } assert_eq!( metas, vec![ - InputUnitMeta::CodeBuddyJsonl, - InputUnitMeta::CodeBuddyExtensionLog { + DecoderRoute::CodeBuddyJsonl, + DecoderRoute::CodeBuddyExtensionLog { origin: CodeBuddyLogOrigin::Extension }, - InputUnitMeta::CodeBuddyExtensionLog { + DecoderRoute::CodeBuddyExtensionLog { origin: CodeBuddyLogOrigin::Host }, ] @@ -430,20 +420,22 @@ mod tests { ..Default::default() }; let ctx = scan_context(home.path(), &settings); - let units = CODEBUDDY_ADAPTER.discover_checked(&ctx).unwrap(); + let units = CODEBUDDY_ADAPTER + .discover_checked(ClientId::CodeBuddy, &ctx) + .unwrap(); assert_eq!(units.len(), 3); let mut expected = vec![ - (project_path, InputUnitMeta::CodeBuddyJsonl), + (project_path, DecoderRoute::CodeBuddyJsonl), ( ide_log, - InputUnitMeta::CodeBuddyExtensionLog { + DecoderRoute::CodeBuddyExtensionLog { origin: CodeBuddyLogOrigin::Extension, }, ), ( host_log, - InputUnitMeta::CodeBuddyExtensionLog { + DecoderRoute::CodeBuddyExtensionLog { origin: CodeBuddyLogOrigin::Host, }, ), @@ -452,19 +444,19 @@ mod tests { assert_eq!( units .iter() - .map(|unit| (unit.path.clone(), unit.meta)) + .map(|unit| (unit.path.clone(), unit.decoder.route())) .collect::>(), expected ); assert!(units.iter().all(|unit| { - let expected_revision = match unit.meta { - InputUnitMeta::CodeBuddyJsonl => CODEBUDDY_JSONL_RECORD_REJECTION_REVISION, - InputUnitMeta::CodeBuddyExtensionLog { .. } => { + let expected_revision = match unit.decoder.route() { + DecoderRoute::CodeBuddyJsonl => CODEBUDDY_JSONL_RECORD_REJECTION_REVISION, + DecoderRoute::CodeBuddyExtensionLog { .. } => { CODEBUDDY_EXTENSION_RECORD_REJECTION_REVISION } _ => unreachable!(), }; - unit.parser_version == ParserVersion::new(ParserId::CodeBuddy, expected_revision) + unit.decoder.version() == DecoderVersion::new(DecoderId::CodeBuddy, expected_revision) })); } @@ -483,10 +475,15 @@ mod tests { let settings = crate::scanner::ScannerSettings::default(); let ctx = scan_context(home.path(), &settings); let units = CODEBUDDY_ADAPTER - .discover_checked(&ctx) + .discover_checked(ClientId::CodeBuddy, &ctx) .unwrap() .into_iter() - .filter(|unit| matches!(unit.meta, InputUnitMeta::CodeBuddyExtensionLog { .. })) + .filter(|unit| { + matches!( + unit.decoder.route(), + DecoderRoute::CodeBuddyExtensionLog { .. } + ) + }) .collect::>(); assert_eq!(units.len(), 1); @@ -499,15 +496,19 @@ mod tests { let path = dir.path().join("session.log"); write_file(&path, ""); let mut units = vec![ - InputUnit::plain_file(ClientId::CodeBuddy, path.clone()).with_meta( - InputUnitMeta::CodeBuddyExtensionLog { - origin: CodeBuddyLogOrigin::Extension, - }, + InputUnit::plain_file( + path.clone(), + DecoderSpec::codebuddy_extension_log( + CODEBUDDY_EXTENSION_RECORD_REJECTION_REVISION, + CodeBuddyLogOrigin::Extension, + ), ), - InputUnit::plain_file(ClientId::CodeBuddy, path.clone()).with_meta( - InputUnitMeta::CodeBuddyExtensionLog { - origin: CodeBuddyLogOrigin::Extension, - }, + InputUnit::plain_file( + path.clone(), + DecoderSpec::codebuddy_extension_log( + CODEBUDDY_EXTENSION_RECORD_REJECTION_REVISION, + CodeBuddyLogOrigin::Extension, + ), ), ]; @@ -525,8 +526,10 @@ mod tests { &path, r#"{"id":"assistant-1","timestamp":1780000000100,"type":"message","role":"assistant","status":"completed","sessionId":"session-1","providerData":{"model":"glm-5.2","messageId":"msg-1"},"message":{"usage":{"input_tokens":10,"output_tokens":3}}}"#, ); - let units = vec![InputUnit::plain_file(ClientId::CodeBuddy, path.clone()) - .with_meta(InputUnitMeta::CodeBuddyJsonl)]; + let units = vec![InputUnit::plain_file( + path.clone(), + DecoderSpec::codebuddy_jsonl(CODEBUDDY_JSONL_RECORD_REJECTION_REVISION), + )]; let actual = fold_with_units(units); let expected = finalized( @@ -557,8 +560,13 @@ mod tests { .into_iter() .zip([CodeBuddyLogOrigin::Extension, CodeBuddyLogOrigin::Host]) .map(|(path, origin)| { - InputUnit::plain_file(ClientId::CodeBuddy, path) - .with_meta(InputUnitMeta::CodeBuddyExtensionLog { origin }) + InputUnit::plain_file( + path, + DecoderSpec::codebuddy_extension_log( + CODEBUDDY_EXTENSION_RECORD_REJECTION_REVISION, + origin, + ), + ) }) .collect(); @@ -586,10 +594,12 @@ mod tests { let units = [first, second] .into_iter() .map(|path| { - InputUnit::plain_file(ClientId::CodeBuddy, path).with_meta( - InputUnitMeta::CodeBuddyExtensionLog { - origin: CodeBuddyLogOrigin::Extension, - }, + InputUnit::plain_file( + path, + DecoderSpec::codebuddy_extension_log( + CODEBUDDY_EXTENSION_RECORD_REJECTION_REVISION, + CodeBuddyLogOrigin::Extension, + ), ) }) .collect(); @@ -600,11 +610,13 @@ mod tests { } #[test] - fn codebuddy_parse_cache_uses_codebuddy_parser_id() { + fn codebuddy_parse_cache_uses_codebuddy_decoder_id() { let path = PathBuf::from("/tmp/codebuddy/session.jsonl"); - let unit = InputUnit::plain_file(ClientId::CodeBuddy, path) - .with_meta(InputUnitMeta::CodeBuddyJsonl); + let unit = InputUnit::plain_file( + path, + DecoderSpec::codebuddy_jsonl(CODEBUDDY_JSONL_RECORD_REJECTION_REVISION), + ); - assert_eq!(unit.parser_version.parser_id, ParserId::CodeBuddy); + assert_eq!(unit.decoder.version().decoder_id, DecoderId::CodeBuddy); } } diff --git a/crates/tokscale-core/src/adapters/codebuff.rs b/crates/tokscale-core/src/adapters/codebuff.rs index 2c231d7d4..21a504a98 100644 --- a/crates/tokscale-core/src/adapters/codebuff.rs +++ b/crates/tokscale-core/src/adapters/codebuff.rs @@ -5,11 +5,13 @@ use rayon::prelude::*; use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ - AdapterScanContext, FingerprintPolicy, FoldContext, InputDiscoveryError, InputUnit, - LocalInputAdapter, MessageSink, ParseContext, ParsedUnit, + AdapterScanContext, BoundMessageSink, DecoderSpec, FingerprintPolicy, FoldContext, + InputDiscoveryError, InputUnit, LocalInputAdapter, ParseContext, ParsedUnit, }; use crate::clients::ClientId; -use crate::message_cache::{ParserId, ParserVersion}; +use crate::message_cache::DecoderId; +#[cfg(test)] +use crate::message_cache::DecoderVersion; use crate::sessions; pub(crate) struct CodebuffAdapter; @@ -18,36 +20,23 @@ const CODEBUFF_RECORD_REJECTION_REVISION: u32 = crate::adapters::MODEL_ID_CANONICALIZATION_REVISION + 2; impl LocalInputAdapter for CodebuffAdapter { - fn client(&self) -> ClientId { - ClientId::Codebuff - } - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { - let def = ClientId::Codebuff + let def = client .local_def() .expect("Codebuff adapter must have local scan policy"); let mut roots = codebuff_roots(ctx.home_dir); - roots.extend(adapter_discover::extra_roots_for_client( - ClientId::Codebuff, - ctx, - )?); - - Ok(adapter_discover::input_units_from_paths( - ClientId::Codebuff, - adapter_discover::scan_roots(ClientId::Codebuff, roots, def.pattern)?, + roots.extend(adapter_discover::extra_roots_for_client(client, ctx)?); + + adapter_discover::input_units_from_paths( + client, + adapter_discover::scan_roots(client, roots, def.pattern)?, FingerprintPolicy::PlainFile, - )? - .into_iter() - .map(|unit| { - unit.with_parser_version(ParserVersion::new( - ParserId::Codebuff, - CODEBUFF_RECORD_REJECTION_REVISION, - )) - }) - .collect()) + DecoderSpec::plain(DecoderId::Codebuff, CODEBUFF_RECORD_REJECTION_REVISION), + ) } fn parse_checked(&self, units: Vec, ctx: &ParseContext<'_>) -> Vec { @@ -73,7 +62,7 @@ impl LocalInputAdapter for CodebuffAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), crate::adapters::InputPipelineError> { adapter_cache::fold_units(parsed, ctx, sink) } @@ -123,14 +112,16 @@ mod tests { home_dir: home.path().to_str().unwrap(), scanner_settings: &settings, }; - let units = CODEBUFF_ADAPTER.discover_checked(&ctx).unwrap(); + let units = CODEBUFF_ADAPTER + .discover_checked(ClientId::Codebuff, &ctx) + .unwrap(); let paths: Vec<_> = units.iter().map(|unit| unit.path.clone()).collect(); assert!(paths.contains(&default_file)); assert!(paths.contains(&extra_file)); assert!(units.iter().all(|unit| { - unit.parser_version - == ParserVersion::new(ParserId::Codebuff, CODEBUFF_RECORD_REJECTION_REVISION) + unit.decoder.version() + == DecoderVersion::new(DecoderId::Codebuff, CODEBUFF_RECORD_REJECTION_REVISION) })); } diff --git a/crates/tokscale-core/src/adapters/codex.rs b/crates/tokscale-core/src/adapters/codex.rs index a4ddb5e44..299183c21 100644 --- a/crates/tokscale-core/src/adapters/codex.rs +++ b/crates/tokscale-core/src/adapters/codex.rs @@ -5,10 +5,12 @@ use rayon::prelude::*; use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; +#[cfg(test)] +use crate::adapters::DecoderRoute; use crate::adapters::{ - AdapterScanContext, CacheHitPlan, FingerprintPolicy, FoldContext, InputDiscoveryError, - InputParseError, InputPipelineError, InputPlanningError, InputUnit, InputUnitMeta, - LocalInputAdapter, MessageSink, ParseContext, ParsedBatchInput, ParsedUnit, UnitMessagePayload, + AdapterScanContext, BoundMessageSink, CacheHitPlan, DecoderSpec, FingerprintPolicy, + FoldContext, InputDiscoveryError, InputParseError, InputPipelineError, InputPlanningError, + InputUnit, LocalInputAdapter, ParseContext, ParsedBatchInput, ParsedUnit, UnitMessagePayload, }; use crate::clients::ClientId; use crate::input_health::InputStatus; @@ -21,21 +23,18 @@ pub(crate) struct CodexAdapter; pub(crate) struct CodexAppendInput { path: PathBuf, read_plan: message_cache::CacheReadPlan, - parser_version: message_cache::ParserVersion, + decoder_version: message_cache::DecoderVersion, tail_messages: Vec, cache_write: Option>, } impl LocalInputAdapter for CodexAdapter { - fn client(&self) -> ClientId { - ClientId::Codex - } - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { - let def = ClientId::Codex + let def = client .local_def() .expect("Codex adapter must have local scan policy"); let codex_home = codex_home(ctx.home_dir); @@ -43,19 +42,14 @@ impl LocalInputAdapter for CodexAdapter { def.resolve_path(ctx.home_dir), codex_home.join("archived_sessions"), ]; - roots.extend(adapter_discover::extra_roots_for_client( - ClientId::Codex, - ctx, - )?); + roots.extend(adapter_discover::extra_roots_for_client(client, ctx)?); let units = adapter_discover::input_units_from_paths_preserving_order( - ClientId::Codex, - adapter_discover::scan_roots(ClientId::Codex, roots, def.pattern)?, + client, + adapter_discover::scan_roots(client, roots, def.pattern)?, FingerprintPolicy::PlainFile, - )? - .into_iter() - .map(|unit| unit.with_meta(InputUnitMeta::Codex)) - .collect(); + DecoderSpec::codex(crate::adapters::CODEX_EXEC_IDENTITY_REVISION), + )?; Ok(units) } @@ -87,7 +81,7 @@ impl LocalInputAdapter for CodexAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), InputPipelineError> { let mut seen = HashSet::new(); fold_codex_units(parsed, ctx, sink, &mut seen) @@ -97,7 +91,7 @@ impl LocalInputAdapter for CodexAdapter { &self, batches: &mut ParsedBatchInput<'_>, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), InputPipelineError> { let mut seen = HashSet::new(); while let Some(parsed) = batches.next(ctx)? { @@ -112,7 +106,7 @@ fn plan_exact_codex_cache_hit( input_cache: &message_cache::InputMessageCache, ) -> Result { unit.revalidate_snapshot_for_cache_decision()?; - let cached = match input_cache.get_meta(&unit.path, unit.parser_version) { + let cached = match input_cache.get_meta(&unit.path, unit.decoder.version()) { Ok(Some(cached)) => cached, Ok(None) => { unit.mark_cache_lookup_completed_no_hit(); @@ -137,7 +131,7 @@ fn plan_exact_codex_cache_hit( } let read_plan = - message_cache::CacheReadPlan::new(&unit.path, unit.parser_version, cached.fingerprint); + message_cache::CacheReadPlan::new(&unit.path, unit.decoder.version(), cached.fingerprint); unit.release_prepared_snapshot(); let mut parsed = ParsedUnit::healthy( unit, @@ -152,46 +146,54 @@ fn plan_exact_codex_cache_hit( fn fold_codex_units( parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, seen: &mut HashSet, ) -> Result<(), InputPipelineError> { for parsed in parsed { - let path = parsed.unit.path.clone(); - let client = parsed.unit.client; - let parser_version = parsed.unit.parser_version; - let mut input_health = parsed.input_health(); + let ParsedUnit { + unit, + messages: payload, + cache_write, + invalidate_cache, + health, + } = parsed; + let path = unit.path.clone(); + let decoder_version = unit.decoder.version(); + let crate::adapters::UnitScanHealth { + mut status, + mut rejections, + } = *health; let CodexResolvedMessages { mut messages, cache_write: extra_write, finalization, recovery_requires_removal, health_override, - } = match resolve_codex_messages(parsed.messages, ctx) { + } = match resolve_codex_messages(payload, ctx) { Ok(resolved) => resolved, Err(error) => { let Some(failure) = codex_recovery_input_failure(&error) else { return Err(error); }; - ctx.health.record(crate::input_health::InputHealth { - client, + ctx.record_health( path, - status: crate::input_health::InputStatus::Unavailable { failure }, - rejections: Default::default(), - }); + crate::input_health::InputStatus::Unavailable { failure }, + Default::default(), + ); continue; } }; if let Some(health) = health_override { - input_health.status = health.status; - input_health.rejections = health.rejections; + status = health.status; + rejections = health.rejections; } - ctx.health.record(input_health); + ctx.record_health(path.clone(), status, rejections); write_codex_cache_and_apply_recovery( &path, - parser_version, - parsed.cache_write.or(extra_write), + decoder_version, + cache_write.or(extra_write), &messages, - parsed.invalidate_cache, + invalidate_cache, recovery_requires_removal, ctx, )?; @@ -199,7 +201,6 @@ fn fold_codex_units( finalize_codex_messages(&mut messages, ctx.pricing); } adapter_cache::emit_messages( - client, messages .into_iter() .filter(|message| crate::should_keep_deduped_message(seen, message)), @@ -229,7 +230,7 @@ fn codex_recovery_input_failure( fn write_codex_cache_and_apply_recovery( path: &Path, - parser_version: message_cache::ParserVersion, + decoder_version: message_cache::DecoderVersion, cache_write: Option>, messages: &[ParsedMessage], invalidate_cache: bool, @@ -239,7 +240,7 @@ fn write_codex_cache_and_apply_recovery( let write_result = adapter_cache::write_cache(cache_write, ctx, messages); let should_remove = invalidate_cache || recovery_requires_removal; if should_remove && !matches!(&write_result, Ok(adapter_cache::CacheWriteOutcome::Written)) { - ctx.input_cache.remove(path, parser_version); + ctx.input_cache.remove(path, decoder_version); } write_result.map_err(Into::into) } @@ -290,7 +291,7 @@ fn parse_full_log_input( })?; build_codex_cache_plan( &path, - unit.parser_version, + unit.decoder.version(), CodexCacheMaterial { consumed_offset, state, @@ -344,7 +345,7 @@ struct CodexCacheMaterial { fn build_codex_cache_plan( path: &Path, - parser_version: message_cache::ParserVersion, + decoder_version: message_cache::DecoderVersion, material: CodexCacheMaterial, ) -> Result, message_cache::InputSnapshotError> { let Some((fingerprint, codex_incremental)) = build_codex_cache_metadata(path, material)? else { @@ -353,7 +354,7 @@ fn build_codex_cache_plan( Ok(Some(message_cache::CacheWritePlan::new( path, - parser_version, + decoder_version, fingerprint, Some(codex_incremental), ))) @@ -434,7 +435,7 @@ fn load_or_parse_codex_unit( if message_cache::codex_cache_meta_is_consistent(&cached) { let read_plan = message_cache::CacheReadPlan::new( &path, - unit.parser_version, + unit.decoder.version(), cached.fingerprint, ); let mut parsed = ParsedUnit::healthy( @@ -499,7 +500,7 @@ fn load_or_parse_codex_unit( Some(Box::new( message_cache::CacheWritePlan::new( &path, - unit.parser_version, + unit.decoder.version(), fingerprint, Some(incremental), ) @@ -508,10 +509,10 @@ fn load_or_parse_codex_unit( } else { None }; - let parser_version = unit.parser_version; + let decoder_version = unit.decoder.version(); let read_plan = message_cache::CacheReadPlan::new( &path, - parser_version, + decoder_version, cached.fingerprint.clone(), ); let invalidate_cache = interrupted.is_some(); @@ -520,7 +521,7 @@ fn load_or_parse_codex_unit( UnitMessagePayload::CodexAppend(Box::new(CodexAppendInput { path, read_plan, - parser_version, + decoder_version, tail_messages: parsed.messages, cache_write, })), @@ -577,14 +578,14 @@ fn resolve_codex_messages( let recovery_requires_removal = failure.requires_shard_removal(); if recovery_requires_removal { ctx.input_cache - .remove(&read_plan.path(), read_plan.parser_version()); + .remove(&read_plan.path(), read_plan.decoder_version()); } else { ctx.input_cache - .invalidate_read(&read_plan.path(), read_plan.parser_version()); + .invalidate_read(&read_plan.path(), read_plan.decoder_version()); } reparse_full_codex_messages( &read_plan.path(), - read_plan.parser_version(), + read_plan.decoder_version(), recovery_requires_removal, ) } @@ -594,7 +595,7 @@ fn resolve_codex_messages( let CodexAppendInput { path, read_plan, - parser_version, + decoder_version, tail_messages, cache_write, } = *append; @@ -606,13 +607,13 @@ fn resolve_codex_messages( } let recovery_requires_removal = failure.requires_shard_removal(); if recovery_requires_removal { - ctx.input_cache.remove(&path, parser_version); + ctx.input_cache.remove(&path, decoder_version); } else { - ctx.input_cache.invalidate_read(&path, parser_version); + ctx.input_cache.invalidate_read(&path, decoder_version); } return reparse_full_codex_messages( &path, - parser_version, + decoder_version, recovery_requires_removal, ); } @@ -632,7 +633,7 @@ fn resolve_codex_messages( fn reparse_full_codex_messages( path: &Path, - parser_version: message_cache::ParserVersion, + decoder_version: message_cache::DecoderVersion, recovery_requires_removal: bool, ) -> Result { let input_snapshot = message_cache::InputPolicy::plain(path) @@ -652,9 +653,7 @@ fn reparse_full_codex_messages( 0, sessions::codex::CodexParseState::default(), ) - .map_err(|source| { - InputParseError::from_session(ClientId::Codex, path, parser_version.parser_id, source) - })?; + .map_err(|source| InputParseError::from_session(path, decoder_version.decoder_id, source))?; let cache_write = if interrupted.is_none() { let content_hash = content_hash.ok_or_else(|| { InputPipelineError::contract("Codex full reparse did not produce a content hash") @@ -664,7 +663,7 @@ fn reparse_full_codex_messages( })?; build_codex_cache_plan( path, - parser_version, + decoder_version, CodexCacheMaterial { consumed_offset, state, @@ -780,7 +779,10 @@ mod tests { } fn codex_unit(path: &Path) -> InputUnit { - InputUnit::plain_file(ClientId::Codex, path.to_path_buf()).with_meta(InputUnitMeta::Codex) + InputUnit::plain_file( + path.to_path_buf(), + DecoderSpec::codex(crate::adapters::CODEX_EXEC_IDENTITY_REVISION), + ) } fn prepared_codex_unit(path: &Path) -> InputUnit { @@ -789,6 +791,36 @@ mod tests { .expect("Codex fixture snapshot must succeed") } + fn codex_binding() -> crate::adapters::AdapterBinding<'static> { + crate::adapters::adapter_for(ClientId::Codex) + .expect("Codex must have a registered adapter binding") + } + + fn codex_fold_context<'a>( + cache: &'a mut message_cache::InputMessageCache, + pricing: Option<&'a PricingService>, + ) -> FoldContext<'a> { + FoldContext::new(codex_binding(), cache, pricing) + } + + fn fold_codex_into( + parsed: Vec, + ctx: &mut FoldContext<'_>, + messages: &mut Vec, + ) -> Result<(), InputPipelineError> { + let mut sink = BoundMessageSink::new(codex_binding(), messages); + CODEX_ADAPTER.fold(parsed, ctx, &mut sink) + } + + fn fold_codex_batches_into( + batches: &mut ParsedBatchInput<'_>, + ctx: &mut FoldContext<'_>, + messages: &mut Vec, + ) -> Result<(), InputPipelineError> { + let mut sink = BoundMessageSink::new(codex_binding(), messages); + CODEX_ADAPTER.fold_batches(batches, ctx, &mut sink) + } + fn expect_codex_hit( result: Result, message: &str, @@ -824,13 +856,12 @@ mod tests { ) -> Vec { let parsed = plan_and_parse(units, cache, Some(pricing)); let mut sink = Vec::new(); - CODEX_ADAPTER - .fold( - parsed, - &mut FoldContext::new(cache, Some(pricing)), - &mut sink, - ) - .expect("valid Codex fixture must fold"); + fold_codex_into( + parsed, + &mut codex_fold_context(cache, Some(pricing)), + &mut sink, + ) + .expect("valid Codex fixture must fold"); assert_codex_attribution(&sink); sink } @@ -877,10 +908,13 @@ mod tests { cache: &mut message_cache::InputMessageCache, ) -> Vec { let mut sink = Vec::new(); - let mut batches = crate::adapters::ParsedBatchInput::new(&CODEX_ADAPTER, units); - CODEX_ADAPTER - .fold_batches(&mut batches, &mut FoldContext::new(cache, None), &mut sink) - .unwrap(); + let mut batches = crate::adapters::ParsedBatchInput::new(codex_binding(), units); + fold_codex_batches_into( + &mut batches, + &mut codex_fold_context(cache, None), + &mut sink, + ) + .unwrap(); assert_codex_attribution(&sink); sink } @@ -890,8 +924,7 @@ mod tests { cache: &mut message_cache::InputMessageCache, ) -> Vec { let mut sink = Vec::new(); - CODEX_ADAPTER - .fold(parsed, &mut FoldContext::new(cache, None), &mut sink) + fold_codex_into(parsed, &mut codex_fold_context(cache, None), &mut sink) .expect("valid Codex fixture must fold"); assert_codex_attribution(&sink); sink @@ -909,8 +942,8 @@ mod tests { assert!( messages .iter() - .all(|message| message.client.as_ref() == ClientId::Codex.as_str()), - "Codex fold output must be attributed from its InputUnit client" + .all(|message| message.client == ClientId::Codex), + "Codex fold output must be attributed from its adapter binding" ); } @@ -931,19 +964,19 @@ mod tests { sessions::codex::CodexParseState::default(), ) .unwrap(); - let parser_version = message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + let decoder_version = message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION, ); let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_home); let meta = cache - .get_meta(path, parser_version) + .get_meta(path, decoder_version) .expect("Codex cache lookup must succeed") .expect("Codex fold must persist the raw shard immediately"); let messages = cache .take_messages(&message_cache::CacheReadPlan::new( path, - parser_version, + decoder_version, meta.fingerprint, )) .unwrap(); @@ -973,7 +1006,7 @@ mod tests { }; let units = CODEX_ADAPTER - .discover_checked(&scan_context(home.path(), &settings)) + .discover_checked(ClientId::Codex, &scan_context(home.path(), &settings)) .expect("Codex fixture discovery must succeed"); let paths: Vec<_> = units.iter().map(|unit| unit.path.clone()).collect(); let expected = vec![ @@ -988,7 +1021,7 @@ mod tests { .all(|unit| unit.fingerprint_policy == FingerprintPolicy::PlainFile)); assert!(units .iter() - .all(|unit| matches!(unit.meta, InputUnitMeta::Codex))); + .all(|unit| matches!(unit.decoder.route(), DecoderRoute::Codex))); } #[test] @@ -1011,7 +1044,7 @@ mod tests { let settings = crate::scanner::ScannerSettings::default(); let units = CODEX_ADAPTER - .discover_checked(&scan_context(home.path(), &settings)) + .discover_checked(ClientId::Codex, &scan_context(home.path(), &settings)) .expect("Codex fixture discovery must succeed"); assert_eq!(units.len(), 2); @@ -1055,8 +1088,8 @@ mod tests { assert!(cache .get_meta( &path, - message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION, ), ) @@ -1075,9 +1108,9 @@ mod tests { CODEX_ADAPTER.parse_checked(vec![codex_unit(&path)], &ParseContext { pricing: None }); assert_eq!(parsed.len(), 1); - let health = parsed[0].input_health(); - assert_eq!(health.client, ClientId::Codex); - assert_eq!(health.path, path); + let health = &parsed[0].health; + assert_eq!(codex_binding().client(), ClientId::Codex); + assert_eq!(parsed[0].unit.path, path); assert!(matches!( health.status, crate::input_health::InputStatus::Complete @@ -1114,7 +1147,7 @@ mod tests { let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); let parsed = plan_and_parse(vec![codex_unit(&path)], &cache, None); - let health = parsed[0].input_health(); + let health = &parsed[0].health; assert!(matches!( health.status, crate::input_health::InputStatus::Complete @@ -1122,22 +1155,22 @@ mod tests { assert_eq!(health.rejections.total(), 1); let mut messages = Vec::new(); - let mut ctx = FoldContext::new(&mut cache, None); - CODEX_ADAPTER.fold(parsed, &mut ctx, &mut messages).unwrap(); + let mut ctx = codex_fold_context(&mut cache, None); + fold_codex_into(parsed, &mut ctx, &mut messages).unwrap(); assert_codex_attribution(&messages); assert_eq!(messages.len(), 2); assert_eq!(messages[0].model_id.as_ref(), "gpt-5.4"); assert_eq!(messages[1].model_id.as_ref(), "gpt-5.5"); - assert_eq!(ctx.health.rejected_records(), 1); - assert_eq!(ctx.health.partial_inputs(), 0); - let parser_version = message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + assert_eq!(ctx.health().rejected_records(), 1); + assert_eq!(ctx.health().partial_inputs(), 0); + let decoder_version = message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION, ); let cached = ctx .input_cache - .get_meta(&path, parser_version) + .get_meta(&path, decoder_version) .unwrap() .expect("a complete Codex scan with rejections must remain cacheable"); assert_eq!(cached.rejections.total(), 1); @@ -1158,7 +1191,7 @@ mod tests { let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); let parsed = plan_and_parse(vec![codex_unit(&path)], &cache, None); - let health = parsed[0].input_health(); + let health = &parsed[0].health; assert!(matches!( health.status, crate::input_health::InputStatus::Partial { .. } @@ -1166,20 +1199,20 @@ mod tests { assert_eq!(health.rejections.total(), 1); let mut messages = Vec::new(); - let mut ctx = FoldContext::new(&mut cache, None); - CODEX_ADAPTER.fold(parsed, &mut ctx, &mut messages).unwrap(); + let mut ctx = codex_fold_context(&mut cache, None); + fold_codex_into(parsed, &mut ctx, &mut messages).unwrap(); assert_codex_attribution(&messages); assert_eq!(messages.len(), 1); assert_eq!(messages[0].tokens.input, 8); - assert_eq!(ctx.health.partial_inputs(), 1); - let parser_version = message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + assert_eq!(ctx.health().partial_inputs(), 1); + let decoder_version = message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION, ); assert!(ctx .input_cache - .get_meta(&path, parser_version) + .get_meta(&path, decoder_version) .unwrap() .is_none()); } @@ -1199,7 +1232,7 @@ mod tests { let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); let parsed = plan_and_parse(vec![codex_unit(&path)], &cache, None); - let health = parsed[0].input_health(); + let health = &parsed[0].health; assert!(matches!(health.status, InputStatus::Partial { .. })); assert_eq!(health.rejections.total(), 1); assert_eq!( @@ -1212,19 +1245,19 @@ mod tests { ); let mut messages = Vec::new(); - let mut ctx = FoldContext::new(&mut cache, None); - CODEX_ADAPTER.fold(parsed, &mut ctx, &mut messages).unwrap(); + let mut ctx = codex_fold_context(&mut cache, None); + fold_codex_into(parsed, &mut ctx, &mut messages).unwrap(); assert_codex_attribution(&messages); assert_eq!(messages.len(), 1); assert_eq!(messages[0].model_id.as_ref(), "gpt-5.4"); - let parser_version = message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + let decoder_version = message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION, ); assert!(ctx .input_cache - .get_meta(&path, parser_version) + .get_meta(&path, decoder_version) .unwrap() .is_none()); } @@ -1265,8 +1298,8 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("session.jsonl"); write_file(&path, FIRST_CODEX_ENTRY); - let parser_version = message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + let decoder_version = message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION, ); let mut seed_cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); @@ -1274,7 +1307,7 @@ mod tests { message_cache::truncate_shard_after_header_for_test( cache_home.path(), &path, - parser_version, + decoder_version, ); let mut repair_cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); @@ -1306,23 +1339,26 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("session.jsonl"); write_file(&path, FIRST_CODEX_ENTRY); - let parser_version = message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + let decoder_version = message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION, ); let mut seed_cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); parse_and_fold(vec![codex_unit(&path)], &mut seed_cache); - let meta = seed_cache.get_meta(&path, parser_version).unwrap().unwrap(); + let meta = seed_cache + .get_meta(&path, decoder_version) + .unwrap() + .unwrap(); let raw_messages = seed_cache .take_messages(&message_cache::CacheReadPlan::new( &path, - parser_version, + decoder_version, meta.fingerprint.clone(), )) .unwrap(); let mut entry = message_cache::CachedInputEntry::new_with_version( &path, - parser_version, + decoder_version, meta.fingerprint, raw_messages, meta.codex_incremental, @@ -1339,13 +1375,13 @@ mod tests { "unchanged Codex input must plan an exact warm hit", ); let mut sink = Vec::new(); - let mut ctx = FoldContext::new(&mut warm_cache, None); - CODEX_ADAPTER.fold(vec![hit], &mut ctx, &mut sink).unwrap(); + let mut ctx = codex_fold_context(&mut warm_cache, None); + fold_codex_into(vec![hit], &mut ctx, &mut sink).unwrap(); assert_codex_attribution(&sink); assert!(!sink.is_empty()); - assert_eq!(ctx.health.rejected_records(), 1); - assert_eq!(ctx.health.inputs()[0].path, path); + assert_eq!(ctx.health().rejected_records(), 1); + assert_eq!(ctx.health().inputs()[0].path, path); } #[test] @@ -1354,23 +1390,26 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("session.jsonl"); write_file(&path, FIRST_CODEX_ENTRY); - let parser_version = message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + let decoder_version = message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION, ); let mut seed_cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); parse_and_fold(vec![codex_unit(&path)], &mut seed_cache); - let meta = seed_cache.get_meta(&path, parser_version).unwrap().unwrap(); + let meta = seed_cache + .get_meta(&path, decoder_version) + .unwrap() + .unwrap(); let raw_messages = seed_cache .take_messages(&message_cache::CacheReadPlan::new( &path, - parser_version, + decoder_version, meta.fingerprint.clone(), )) .unwrap(); let mut entry = message_cache::CachedInputEntry::new_with_version( &path, - parser_version, + decoder_version, meta.fingerprint, raw_messages, meta.codex_incremental, @@ -1385,12 +1424,10 @@ mod tests { let mut append_cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); let parsed = plan_and_parse(vec![codex_unit(&path)], &append_cache, None); let mut appended_messages = Vec::new(); - let mut append_ctx = FoldContext::new(&mut append_cache, None); - CODEX_ADAPTER - .fold(parsed, &mut append_ctx, &mut appended_messages) - .unwrap(); + let mut append_ctx = codex_fold_context(&mut append_cache, None); + fold_codex_into(parsed, &mut append_ctx, &mut appended_messages).unwrap(); assert_codex_attribution(&appended_messages); - assert_eq!(append_ctx.health.rejected_records(), 1); + assert_eq!(append_ctx.health().rejected_records(), 1); append_ctx.input_cache.save_if_dirty().unwrap(); let mut warm_cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); @@ -1399,14 +1436,12 @@ mod tests { "appended Codex input must be rewritten as an exact warm shard", ); let mut warm_messages = Vec::new(); - let mut warm_ctx = FoldContext::new(&mut warm_cache, None); - CODEX_ADAPTER - .fold(vec![hit], &mut warm_ctx, &mut warm_messages) - .unwrap(); + let mut warm_ctx = codex_fold_context(&mut warm_cache, None); + fold_codex_into(vec![hit], &mut warm_ctx, &mut warm_messages).unwrap(); assert_codex_attribution(&warm_messages); assert_eq!(warm_messages, appended_messages); - assert_eq!(warm_ctx.health.rejected_records(), 1); + assert_eq!(warm_ctx.health().rejected_records(), 1); } #[test] @@ -1415,8 +1450,8 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("session.jsonl"); write_file(&path, FIRST_CODEX_ENTRY); - let parser_version = message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + let decoder_version = message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION, ); let mut seed_cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); @@ -1429,20 +1464,20 @@ mod tests { let mut append_cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); let parsed = plan_and_parse(vec![codex_unit(&path)], &append_cache, None); - let health = parsed[0].input_health(); + let health = &parsed[0].health; assert!(matches!(health.status, InputStatus::Complete)); assert_eq!(health.rejections.total(), 1); let mut messages = Vec::new(); - let mut ctx = FoldContext::new(&mut append_cache, None); - CODEX_ADAPTER.fold(parsed, &mut ctx, &mut messages).unwrap(); + let mut ctx = codex_fold_context(&mut append_cache, None); + fold_codex_into(parsed, &mut ctx, &mut messages).unwrap(); assert_codex_attribution(&messages); assert_eq!(messages.len(), 2); - assert_eq!(ctx.health.rejected_records(), 1); + assert_eq!(ctx.health().rejected_records(), 1); let cached = ctx .input_cache - .get_meta(&path, parser_version) + .get_meta(&path, decoder_version) .unwrap() .expect("complete appended scan must replace the shard"); assert_eq!(cached.rejections.total(), 1); @@ -1454,8 +1489,8 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("session.jsonl"); write_file(&path, FIRST_CODEX_ENTRY); - let parser_version = message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + let decoder_version = message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION, ); let mut seed_cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); @@ -1471,7 +1506,7 @@ mod tests { let mut append_cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); let parsed = plan_and_parse(vec![codex_unit(&path)], &append_cache, None); - let health = parsed[0].input_health(); + let health = &parsed[0].health; assert!(matches!(health.status, InputStatus::Partial { .. })); assert_eq!(health.rejections.total(), 1); assert_eq!( @@ -1480,16 +1515,16 @@ mod tests { ); let mut messages = Vec::new(); - let mut ctx = FoldContext::new(&mut append_cache, None); - CODEX_ADAPTER.fold(parsed, &mut ctx, &mut messages).unwrap(); + let mut ctx = codex_fold_context(&mut append_cache, None); + fold_codex_into(parsed, &mut ctx, &mut messages).unwrap(); assert_codex_attribution(&messages); assert_eq!(messages.len(), 1); - assert_eq!(ctx.health.partial_inputs(), 1); - assert_eq!(ctx.health.rejected_records(), 1); + assert_eq!(ctx.health().partial_inputs(), 1); + assert_eq!(ctx.health().rejected_records(), 1); assert!(ctx .input_cache - .get_meta(&path, parser_version) + .get_meta(&path, decoder_version) .unwrap() .is_none()); } @@ -1500,8 +1535,8 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("session.jsonl"); write_file(&path, FIRST_CODEX_ENTRY); - let parser_version = message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + let decoder_version = message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION, ); let mut seed_cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); @@ -1517,7 +1552,7 @@ mod tests { let mut append_cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); let parsed = plan_and_parse(vec![codex_unit(&path)], &append_cache, None); - let health = parsed[0].input_health(); + let health = &parsed[0].health; assert!(matches!(health.status, InputStatus::Partial { .. })); assert_eq!(health.rejections.total(), 1); assert_eq!( @@ -1526,15 +1561,15 @@ mod tests { ); let mut messages = Vec::new(); - let mut ctx = FoldContext::new(&mut append_cache, None); - CODEX_ADAPTER.fold(parsed, &mut ctx, &mut messages).unwrap(); + let mut ctx = codex_fold_context(&mut append_cache, None); + fold_codex_into(parsed, &mut ctx, &mut messages).unwrap(); assert_codex_attribution(&messages); assert_eq!(messages.len(), 1); assert_eq!(messages[0].model_id.as_ref(), "gpt-5.4"); assert!(ctx .input_cache - .get_meta(&path, parser_version) + .get_meta(&path, decoder_version) .unwrap() .is_none()); } @@ -1545,8 +1580,8 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("session.jsonl"); write_file(&path, FIRST_CODEX_ENTRY); - let parser_version = message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + let decoder_version = message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION, ); let mut seed_cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); @@ -1554,7 +1589,7 @@ mod tests { message_cache::replace_shard_message_count_for_test( cache_home.path(), &path, - parser_version, + decoder_version, 2, ); @@ -1573,8 +1608,8 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("session.jsonl"); write_file(&path, FIRST_CODEX_ENTRY); - let parser_version = message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + let decoder_version = message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION, ); let mut seed_cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); @@ -1582,7 +1617,7 @@ mod tests { let shard_path = message_cache::truncate_shard_after_header_for_test( cache_home.path(), &path, - parser_version, + decoder_version, ); let mut repair_cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); @@ -1592,15 +1627,14 @@ mod tests { ); write_file(&path, MISSING_TIMESTAMP_CODEX_ENTRY); let mut sink = Vec::new(); - let mut ctx = FoldContext::new(&mut repair_cache, None); - CODEX_ADAPTER - .fold(vec![planned], &mut ctx, &mut sink) + let mut ctx = codex_fold_context(&mut repair_cache, None); + fold_codex_into(vec![planned], &mut ctx, &mut sink) .expect("an interrupted Codex recovery scan must stay inside its input domain"); assert_codex_attribution(&sink); assert!(sink.is_empty()); - assert_eq!(ctx.health.partial_inputs(), 1); - assert_eq!(ctx.health.rejected_records(), 1); - assert!(ctx.health.inputs()[0] + assert_eq!(ctx.health().partial_inputs(), 1); + assert_eq!(ctx.health().rejected_records(), 1); + assert!(ctx.health().inputs()[0] .status .failure() .unwrap() @@ -1620,14 +1654,14 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("session.jsonl"); write_file(&path, FIRST_CODEX_ENTRY); - let parser_version = message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + let decoder_version = message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION, ); let mut seed_cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); parse_and_fold(vec![codex_unit(&path)], &mut seed_cache); let shard_path = - message_cache::shard_path_for_test(cache_home.path(), &path, parser_version); + message_cache::shard_path_for_test(cache_home.path(), &path, decoder_version); let mut repair_cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); let planned = expect_codex_hit( @@ -1638,15 +1672,14 @@ mod tests { std::fs::write(&shard_path, unknown).unwrap(); write_file(&path, MISSING_TIMESTAMP_CODEX_ENTRY); let mut sink = Vec::new(); - let mut ctx = FoldContext::new(&mut repair_cache, None); - CODEX_ADAPTER - .fold(vec![planned], &mut ctx, &mut sink) + let mut ctx = codex_fold_context(&mut repair_cache, None); + fold_codex_into(vec![planned], &mut ctx, &mut sink) .expect("an interrupted Codex recovery scan must stay inside its input domain"); assert_codex_attribution(&sink); assert!(sink.is_empty()); - assert_eq!(ctx.health.partial_inputs(), 1); - assert_eq!(ctx.health.rejected_records(), 1); - assert!(ctx.health.inputs()[0] + assert_eq!(ctx.health().partial_inputs(), 1); + assert_eq!(ctx.health().rejected_records(), 1); + assert!(ctx.health().inputs()[0] .status .failure() .unwrap() @@ -1667,8 +1700,8 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("session.jsonl"); write_file(&path, FIRST_CODEX_ENTRY); - let parser_version = message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + let decoder_version = message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION, ); let mut seed_cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); @@ -1676,7 +1709,7 @@ mod tests { let shard_path = message_cache::truncate_shard_after_header_for_test( cache_home.path(), &path, - parser_version, + decoder_version, ); let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); let planned = expect_codex_hit( @@ -1684,7 +1717,7 @@ mod tests { "valid header must still plan a Codex hit", ); let resolved = - resolve_codex_messages(planned.messages, &mut FoldContext::new(&mut cache, None)) + resolve_codex_messages(planned.messages, &mut codex_fold_context(&mut cache, None)) .expect("valid input must repair the corrupt cache body"); assert!(resolved.recovery_requires_removal); assert!(resolved.cache_write.is_some()); @@ -1695,12 +1728,12 @@ mod tests { std::fs::write(&cache_path, b"block cache directory recreation").unwrap(); let write_result = write_codex_cache_and_apply_recovery( &path, - parser_version, + decoder_version, resolved.cache_write, &resolved.messages, false, resolved.recovery_requires_removal, - &mut FoldContext::new(&mut cache, None), + &mut codex_fold_context(&mut cache, None), ); std::fs::remove_file(&cache_path).unwrap(); std::fs::rename(&backup_path, &cache_path).unwrap(); @@ -1755,12 +1788,12 @@ mod tests { let mut cold_cache = message_cache::InputMessageCache::with_cache_dir(cache_home.path()); assert!(parse_and_fold(vec![codex_unit(&path)], &mut cold_cache).is_empty()); - let parser_version = message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + let decoder_version = message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION, ); let meta = cold_cache - .get_meta(&path, parser_version) + .get_meta(&path, decoder_version) .expect("Codex cache lookup must succeed") .expect("an empty Codex parse is still a valid cache result"); assert!(meta.codex_incremental.is_some()); @@ -1998,8 +2031,8 @@ mod tests { let mut remover = message_cache::InputMessageCache::load().unwrap(); remover.remove( &path, - message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION, ), ); diff --git a/crates/tokscale-core/src/adapters/decoder.rs b/crates/tokscale-core/src/adapters/decoder.rs new file mode 100644 index 000000000..c921aae3b --- /dev/null +++ b/crates/tokscale-core/src/adapters/decoder.rs @@ -0,0 +1,123 @@ +use crate::message_cache::{DecoderId, DecoderRevision, DecoderVersion}; + +use super::CodeBuddyLogOrigin; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub(crate) enum DecoderRoute { + #[default] + None, + OpenCodeSqlite, + AntigravityCliSqlite, + KiroFile, + KiroSqlite, + KiroGlobalStorage, + CodeBuddyJsonl, + CodeBuddyExtensionLog { + origin: CodeBuddyLogOrigin, + }, + Codex, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct DecoderSpec { + version: DecoderVersion, + route: DecoderRoute, +} + +impl DecoderSpec { + const fn new(decoder_id: DecoderId, revision: DecoderRevision, route: DecoderRoute) -> Self { + Self { + version: DecoderVersion::new(decoder_id, revision), + route, + } + } + + pub(crate) const fn plain(decoder_id: DecoderId, revision: DecoderRevision) -> Self { + assert!( + decoder_id.supports_plain_route(), + "routed decoder requires its specialized DecoderSpec constructor" + ); + Self::new(decoder_id, revision, DecoderRoute::None) + } + + pub(crate) const fn opencode_sqlite(revision: DecoderRevision) -> Self { + Self::new( + DecoderId::OpenCodeSqlite, + revision, + DecoderRoute::OpenCodeSqlite, + ) + } + + pub(crate) const fn antigravity_cli_sqlite(revision: DecoderRevision) -> Self { + Self::new( + DecoderId::AntigravityCliSqlite, + revision, + DecoderRoute::AntigravityCliSqlite, + ) + } + + pub(crate) const fn kiro_file(revision: DecoderRevision) -> Self { + Self::new(DecoderId::KiroFile, revision, DecoderRoute::KiroFile) + } + + pub(crate) const fn kiro_sqlite(revision: DecoderRevision) -> Self { + Self::new(DecoderId::KiroSqlite, revision, DecoderRoute::KiroSqlite) + } + + pub(crate) const fn kiro_global_storage(revision: DecoderRevision) -> Self { + Self::new( + DecoderId::KiroGlobalStorage, + revision, + DecoderRoute::KiroGlobalStorage, + ) + } + + pub(crate) const fn codebuddy_jsonl(revision: DecoderRevision) -> Self { + Self::new(DecoderId::CodeBuddy, revision, DecoderRoute::CodeBuddyJsonl) + } + + pub(crate) const fn codebuddy_extension_log( + revision: DecoderRevision, + origin: CodeBuddyLogOrigin, + ) -> Self { + Self::new( + DecoderId::CodeBuddy, + revision, + DecoderRoute::CodeBuddyExtensionLog { origin }, + ) + } + + pub(crate) const fn codex(revision: DecoderRevision) -> Self { + Self::new(DecoderId::Codex, revision, DecoderRoute::Codex) + } + + pub(crate) const fn version(self) -> DecoderVersion { + self.version + } + + pub(crate) const fn route(self) -> DecoderRoute { + self.route + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + #[should_panic(expected = "routed decoder requires its specialized DecoderSpec constructor")] + fn routed_decoder_cannot_be_constructed_as_plain() { + let _ = DecoderSpec::plain(DecoderId::Codex, 1); + } + + #[test] + fn specialized_constructor_binds_route_and_cache_identity() { + let spec = DecoderSpec::kiro_sqlite(7); + + assert_eq!(spec.route(), DecoderRoute::KiroSqlite); + assert_eq!( + spec.version(), + DecoderVersion::new(DecoderId::KiroSqlite, 7) + ); + } +} diff --git a/crates/tokscale-core/src/adapters/discover.rs b/crates/tokscale-core/src/adapters/discover.rs index 50b2e1df4..34a8143d8 100644 --- a/crates/tokscale-core/src/adapters/discover.rs +++ b/crates/tokscale-core/src/adapters/discover.rs @@ -1,7 +1,9 @@ use std::collections::HashSet; use std::path::{Path, PathBuf}; -use crate::adapters::{AdapterScanContext, FingerprintPolicy, InputDiscoveryError, InputUnit}; +use crate::adapters::{ + AdapterScanContext, DecoderSpec, FingerprintPolicy, InputDiscoveryError, InputUnit, +}; use crate::clients::ClientId; use crate::scanner; @@ -9,6 +11,7 @@ pub(crate) fn discover_default_scanned_units( client: ClientId, ctx: &AdapterScanContext<'_>, fingerprint_policy: FingerprintPolicy, + decoder: DecoderSpec, ) -> Result, InputDiscoveryError> { let def = client .local_def() @@ -21,7 +24,7 @@ pub(crate) fn discover_default_scanned_units( extra_roots_for_client(client, ctx)?, def.pattern, )?); - input_units_from_paths(client, paths, fingerprint_policy) + input_units_from_paths(client, paths, fingerprint_policy, decoder) } pub(crate) fn extra_roots_for_client( @@ -43,7 +46,7 @@ pub(crate) fn extra_roots_for_client( } pub(crate) fn scan_roots( - client: ClientId, + _client: ClientId, roots: I, pattern: &str, ) -> Result, InputDiscoveryError> @@ -53,9 +56,8 @@ where let mut paths = Vec::new(); for root in roots { paths.extend( - scanner::scan_directory(&root, pattern).map_err(|source| { - InputDiscoveryError::new(client, &root, "walk directory", source) - })?, + scanner::scan_directory(&root, pattern) + .map_err(|source| InputDiscoveryError::new(&root, "walk directory", source))?, ); } Ok(paths) @@ -65,6 +67,7 @@ pub(crate) fn input_units_from_paths( client: ClientId, paths: Vec, fingerprint_policy: FingerprintPolicy, + decoder: DecoderSpec, ) -> Result, InputDiscoveryError> { let mut seen = HashSet::new(); let mut units = Vec::new(); @@ -72,7 +75,7 @@ pub(crate) fn input_units_from_paths( for path in paths { let key = canonical_key(client, &path)?; if seen.insert(key) { - units.push(input_unit_for_policy(client, path, &fingerprint_policy)?); + units.push(input_unit_for_policy(path, &fingerprint_policy, decoder)); } } @@ -84,6 +87,7 @@ pub(crate) fn input_units_from_paths_preserving_order( client: ClientId, paths: Vec, fingerprint_policy: FingerprintPolicy, + decoder: DecoderSpec, ) -> Result, InputDiscoveryError> { let mut seen = HashSet::new(); let mut units = Vec::new(); @@ -91,7 +95,7 @@ pub(crate) fn input_units_from_paths_preserving_order( for path in paths { let key = canonical_key(client, &path)?; if seen.insert(key) { - units.push(input_unit_for_policy(client, path, &fingerprint_policy)?); + units.push(input_unit_for_policy(path, &fingerprint_policy, decoder)); } } @@ -99,7 +103,7 @@ pub(crate) fn input_units_from_paths_preserving_order( } pub(crate) fn push_existing_file( - client: ClientId, + _client: ClientId, path: PathBuf, paths: &mut Vec, ) -> Result<(), InputDiscoveryError> { @@ -107,7 +111,6 @@ pub(crate) fn push_existing_file( Ok(metadata) if metadata.is_file() => paths.push(path), Ok(_) => { return Err(InputDiscoveryError::new( - client, &path, "validate file candidate", std::io::Error::new( @@ -119,7 +122,6 @@ pub(crate) fn push_existing_file( Err(source) if source.kind() == std::io::ErrorKind::NotFound => {} Err(source) => { return Err(InputDiscoveryError::new( - client, &path, "read file candidate metadata", source, @@ -129,28 +131,27 @@ pub(crate) fn push_existing_file( Ok(()) } -fn canonical_key(client: ClientId, path: &Path) -> Result { - std::fs::canonicalize(path).map_err(|source| { - InputDiscoveryError::new(client, path, "canonicalize discovered input", source) - }) +fn canonical_key(_client: ClientId, path: &Path) -> Result { + std::fs::canonicalize(path) + .map_err(|source| InputDiscoveryError::new(path, "canonicalize discovered input", source)) } fn input_unit_for_policy( - client: ClientId, path: PathBuf, fingerprint_policy: &FingerprintPolicy, -) -> Result { - let unit = match fingerprint_policy { - FingerprintPolicy::PlainFile => InputUnit::plain_file(client, path), - FingerprintPolicy::SqliteWithWal => InputUnit::sqlite_with_wal(client, path), + decoder: DecoderSpec, +) -> InputUnit { + match fingerprint_policy { + FingerprintPolicy::PlainFile => InputUnit::plain_file(path, decoder), + FingerprintPolicy::SqliteWithWal => InputUnit::sqlite_with_wal(path, decoder), FingerprintPolicy::ClaudeCodeWithHome { home_dir, .. } => { - InputUnit::claude_code(client, path, home_dir.clone()) + InputUnit::claude_code(path, home_dir.clone(), decoder) } FingerprintPolicy::PrimaryWithSiblings { sibling_names, related_failure_policy, } => { - let mut unit = InputUnit::plain_file(client, path); + let mut unit = InputUnit::plain_file(path, decoder); unit.fingerprint_policy = FingerprintPolicy::PrimaryWithSiblings { sibling_names, related_failure_policy: *related_failure_policy, @@ -162,14 +163,13 @@ fn input_unit_for_policy( related_failure_policy, } => match related_failure_policy { crate::message_cache::RelatedInputFailurePolicy::FailInput => { - InputUnit::plain_file(client, path).with_dependency(dependency_path.clone()) + InputUnit::plain_file(path, decoder).with_dependency(dependency_path.clone()) } crate::message_cache::RelatedInputFailurePolicy::PreservePrimary => { - InputUnit::plain_file(client, path) + InputUnit::plain_file(path, decoder) .with_optional_dependency(dependency_path.clone()) } }, - FingerprintPolicy::NoMessageCache => InputUnit::no_message_cache(client, path), - }; - Ok(unit) + FingerprintPolicy::NoMessageCache => InputUnit::no_message_cache(path, decoder), + } } diff --git a/crates/tokscale-core/src/adapters/error.rs b/crates/tokscale-core/src/adapters/error.rs index 80edd79a5..5f6f6fd44 100644 --- a/crates/tokscale-core/src/adapters/error.rs +++ b/crates/tokscale-core/src/adapters/error.rs @@ -2,8 +2,7 @@ use std::error::Error; use std::fmt; use std::path::{Path, PathBuf}; -use crate::clients::ClientId; -use crate::message_cache::ParserId; +use crate::message_cache::DecoderId; use crate::message_cache::{CacheLookupFailure, CacheReadFailure, InputCacheError}; use crate::sessions::error::SessionParseError; @@ -11,7 +10,6 @@ type BoxInputError = Box; #[derive(Debug)] pub(crate) struct InputDiscoveryError { - pub(crate) client: ClientId, pub(crate) path: PathBuf, pub(crate) operation: &'static str, source: BoxInputError, @@ -19,13 +17,11 @@ pub(crate) struct InputDiscoveryError { impl InputDiscoveryError { pub(crate) fn new( - client: ClientId, path: impl Into, operation: &'static str, source: impl Error + Send + Sync + 'static, ) -> Self { Self { - client, path: path.into(), operation, source: Box::new(source), @@ -37,8 +33,7 @@ impl fmt::Display for InputDiscoveryError { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { write!( formatter, - "{} input discovery failed to {} `{}`: {}", - self.client.as_str(), + "input discovery failed to {} `{}`: {}", self.operation, self.path.display(), self.source @@ -54,38 +49,30 @@ impl Error for InputDiscoveryError { #[derive(Debug)] pub(crate) struct InputParseError { - pub(crate) client: ClientId, pub(crate) path: PathBuf, - pub(crate) parser: ParserId, + pub(crate) decoder: DecoderId, pub(crate) operation: &'static str, source: BoxInputError, } impl InputParseError { pub(crate) fn new( - client: ClientId, path: impl Into, - parser: ParserId, + decoder: DecoderId, operation: &'static str, source: impl Error + Send + Sync + 'static, ) -> Self { Self { - client, path: path.into(), - parser, + decoder, operation, source: Box::new(source), } } - pub(crate) fn from_session( - client: ClientId, - path: &Path, - parser: ParserId, - source: SessionParseError, - ) -> Self { + pub(crate) fn from_session(path: &Path, decoder: DecoderId, source: SessionParseError) -> Self { let input_path = source.path().unwrap_or(path).to_path_buf(); - Self::new(client, input_path, parser, source.operation(), source) + Self::new(input_path, decoder, source.operation(), source) } } @@ -93,9 +80,8 @@ impl fmt::Display for InputParseError { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { write!( formatter, - "{} parser `{}` failed to {} input `{}`: {}", - self.client.as_str(), - self.parser.stable_name(), + "decoder `{}` failed to {} input `{}`: {}", + self.decoder.stable_name(), self.operation, self.path.display(), self.source @@ -156,3 +142,38 @@ impl InputPipelineError { } } } + +#[cfg(test)] +mod tests { + use super::*; + use std::io; + + #[test] + fn discovery_diagnostic_is_source_neutral() { + let error = InputDiscoveryError::new( + "/inputs/history", + "walk directory", + io::Error::new(io::ErrorKind::PermissionDenied, "permission denied"), + ); + + assert_eq!( + error.to_string(), + "input discovery failed to walk directory `/inputs/history`: permission denied" + ); + } + + #[test] + fn parse_diagnostic_identifies_the_decoder_not_a_client() { + let error = InputParseError::new( + "/inputs/session.jsonl", + DecoderId::Codex, + "decode session", + io::Error::new(io::ErrorKind::InvalidData, "invalid record"), + ); + + assert_eq!( + error.to_string(), + "decoder `codex` failed to decode session input `/inputs/session.jsonl`: invalid record" + ); + } +} diff --git a/crates/tokscale-core/src/adapters/file.rs b/crates/tokscale-core/src/adapters/file.rs index 546bfff56..bb19760d3 100644 --- a/crates/tokscale-core/src/adapters/file.rs +++ b/crates/tokscale-core/src/adapters/file.rs @@ -5,13 +5,15 @@ use rayon::prelude::*; use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ - AdapterScanContext, FingerprintPolicy, FoldContext, InputDiscoveryError, InputPipelineError, - InputUnit, LocalInputAdapter, MessageSink, ParseContext, ParsedUnit, - MODEL_ID_CANONICALIZATION_REVISION, + AdapterScanContext, BoundMessageSink, DecoderSpec, FingerprintPolicy, FoldContext, + InputDiscoveryError, InputPipelineError, InputUnit, LocalInputAdapter, ParseContext, + ParsedUnit, MODEL_ID_CANONICALIZATION_REVISION, }; use crate::clients::ClientId; use crate::input_health::ScannedInput; -use crate::message_cache::{ParserId, ParserVersion, RelatedInputFailurePolicy}; +#[cfg(test)] +use crate::message_cache::DecoderVersion; +use crate::message_cache::{DecoderId, RelatedInputFailurePolicy}; use crate::sessions; use crate::sessions::error::SessionParseResult; use crate::sessions::{ParsedMessage, WorkspaceMetadata}; @@ -37,8 +39,7 @@ const DROID_AGENT_ATTRIBUTION_REVISION: u32 = DROID_RECORD_REJECTION_REVISION + const GROK_RELATED_METADATA_SIBLINGS: &[&str] = &["summary.json", "events.jsonl"]; pub(crate) struct CachedFileAdapter { - client: ClientId, - parser_version: ParserVersion, + decoder: DecoderSpec, fingerprint_policy: FingerprintPolicy, dependency_failure_policy: RelatedInputFailurePolicy, dependency_path: Option Option>, @@ -48,14 +49,12 @@ pub(crate) struct CachedFileAdapter { impl CachedFileAdapter { pub(crate) const fn new( - client: ClientId, - parser_id: ParserId, + decoder_id: DecoderId, revision: u32, parse: fn(&Path) -> SessionParseResult, ) -> Self { Self { - client, - parser_version: ParserVersion::new(parser_id, revision), + decoder: DecoderSpec::plain(decoder_id, revision), fingerprint_policy: FingerprintPolicy::PlainFile, dependency_failure_policy: RelatedInputFailurePolicy::FailInput, dependency_path: None, @@ -65,15 +64,13 @@ impl CachedFileAdapter { } pub(crate) const fn new_with_required_dependency( - client: ClientId, - parser_id: ParserId, + decoder_id: DecoderId, revision: u32, dependency_path: fn(&Path) -> Option, parse: fn(&Path) -> SessionParseResult, ) -> Self { Self { - client, - parser_version: ParserVersion::new(parser_id, revision), + decoder: DecoderSpec::plain(decoder_id, revision), fingerprint_policy: FingerprintPolicy::PlainFile, dependency_failure_policy: RelatedInputFailurePolicy::FailInput, dependency_path: Some(dependency_path), @@ -83,15 +80,13 @@ impl CachedFileAdapter { } pub(crate) const fn new_with_optional_dependency( - client: ClientId, - parser_id: ParserId, + decoder_id: DecoderId, revision: u32, dependency_path: fn(&Path) -> Option, parse: fn(&Path) -> SessionParseResult, ) -> Self { Self { - client, - parser_version: ParserVersion::new(parser_id, revision), + decoder: DecoderSpec::plain(decoder_id, revision), fingerprint_policy: FingerprintPolicy::PlainFile, dependency_failure_policy: RelatedInputFailurePolicy::PreservePrimary, dependency_path: Some(dependency_path), @@ -101,15 +96,13 @@ impl CachedFileAdapter { } pub(crate) const fn new_with_optional_siblings( - client: ClientId, - parser_id: ParserId, + decoder_id: DecoderId, revision: u32, sibling_names: &'static [&'static str], parse: fn(&Path) -> SessionParseResult, ) -> Self { Self { - client, - parser_version: ParserVersion::new(parser_id, revision), + decoder: DecoderSpec::plain(decoder_id, revision), fingerprint_policy: FingerprintPolicy::PrimaryWithSiblings { sibling_names, related_failure_policy: RelatedInputFailurePolicy::PreservePrimary, @@ -128,25 +121,23 @@ impl CachedFileAdapter { } impl LocalInputAdapter for CachedFileAdapter { - fn client(&self) -> ClientId { - self.client - } - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { Ok(adapter_discover::discover_default_scanned_units( - self.client, + client, ctx, self.fingerprint_policy.clone(), + self.decoder, )? .into_iter() .map(|unit| { let dependency_path = self .dependency_path .and_then(|dependency_path| dependency_path(&unit.path)); - let unit = match dependency_path { + match dependency_path { Some(dependency_path) if self.dependency_failure_policy == RelatedInputFailurePolicy::PreservePrimary => @@ -155,8 +146,7 @@ impl LocalInputAdapter for CachedFileAdapter { } Some(dependency_path) => unit.with_dependency(dependency_path), None => unit, - }; - unit.with_parser_version(self.parser_version) + } }) .collect()) } @@ -181,7 +171,7 @@ impl LocalInputAdapter for CachedFileAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), InputPipelineError> { let workspace_enrichment = self.workspace_enrichment; // Companion workspace metadata is projected after cache resolution so @@ -220,40 +210,29 @@ fn enrich_gemini_workspace(path: &Path, messages: &mut [ParsedMessage]) { pub(crate) struct CopilotAdapter; impl LocalInputAdapter for CopilotAdapter { - fn client(&self) -> ClientId { - ClientId::Copilot - } - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { - let def = ClientId::Copilot + let def = client .local_def() .expect("Copilot adapter must have local scan policy"); let default_root = def.resolve_path(ctx.home_dir); - let mut paths = - adapter_discover::scan_roots(ClientId::Copilot, [default_root], def.pattern)?; + let mut paths = adapter_discover::scan_roots(client, [default_root], def.pattern)?; paths.extend(adapter_discover::scan_roots( - ClientId::Copilot, - adapter_discover::extra_roots_for_client(ClientId::Copilot, ctx)?, + client, + adapter_discover::extra_roots_for_client(client, ctx)?, def.pattern, )?); - Ok(adapter_discover::input_units_from_paths( - ClientId::Copilot, + adapter_discover::input_units_from_paths( + client, paths, FingerprintPolicy::PlainFile, - )? - .into_iter() - .map(|unit| { - unit.with_parser_version(ParserVersion::new( - ParserId::Copilot, - COPILOT_AGENT_IDENTITY_REVISION, - )) - }) - .collect()) + DecoderSpec::plain(DecoderId::Copilot, COPILOT_AGENT_IDENTITY_REVISION), + ) } fn parse_checked(&self, units: Vec, ctx: &ParseContext<'_>) -> Vec { @@ -285,7 +264,7 @@ impl LocalInputAdapter for CopilotAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), InputPipelineError> { adapter_cache::fold_units(parsed, ctx, sink) } @@ -293,29 +272,25 @@ impl LocalInputAdapter for CopilotAdapter { pub(crate) static COPILOT_ADAPTER: CopilotAdapter = CopilotAdapter; pub(crate) static GEMINI_ADAPTER: CachedFileAdapter = CachedFileAdapter::new( - ClientId::Gemini, - ParserId::Gemini, + DecoderId::Gemini, GEMINI_RECORD_REJECTION_REVISION, sessions::gemini::parse_gemini_file, ) .with_workspace_enrichment(enrich_gemini_workspace); pub(crate) static GROK_ADAPTER: CachedFileAdapter = CachedFileAdapter::new_with_optional_siblings( - ClientId::Grok, - ParserId::Grok, + DecoderId::Grok, GROK_RELATED_METADATA_REVISION, GROK_RELATED_METADATA_SIBLINGS, sessions::grok::parse_grok_updates_file, ); pub(crate) static AMP_ADAPTER: CachedFileAdapter = CachedFileAdapter::new( - ClientId::Amp, - ParserId::Amp, + DecoderId::Amp, AMP_RECORD_REJECTION_REVISION, sessions::amp::parse_amp_file, ); pub(crate) static DROID_ADAPTER: CachedFileAdapter = CachedFileAdapter::new_with_optional_dependency( - ClientId::Droid, - ParserId::Droid, + DecoderId::Droid, DROID_AGENT_ATTRIBUTION_REVISION, sessions::droid::droid_agent_dependency_path, sessions::droid::parse_droid_file, @@ -323,36 +298,31 @@ pub(crate) static DROID_ADAPTER: CachedFileAdapter = .with_workspace_enrichment(enrich_droid_metadata); pub(crate) static KIMI_ADAPTER: CachedFileAdapter = CachedFileAdapter::new_with_optional_dependency( - ClientId::Kimi, - ParserId::Kimi, + DecoderId::Kimi, KIMI_RECORD_REJECTION_REVISION, sessions::kimi::kimi_config_dependency_path, sessions::kimi::parse_kimi_file, ) .with_workspace_enrichment(enrich_kimi_workspace); pub(crate) static QWEN_ADAPTER: CachedFileAdapter = CachedFileAdapter::new( - ClientId::Qwen, - ParserId::Qwen, + DecoderId::Qwen, QWEN_RECORD_REJECTION_REVISION, sessions::qwen::parse_qwen_file, ); pub(crate) static MUX_ADAPTER: CachedFileAdapter = CachedFileAdapter::new( - ClientId::Mux, - ParserId::Mux, + DecoderId::Mux, MUX_RECORD_REJECTION_REVISION, sessions::mux::parse_mux_file, ); pub(crate) static COMMANDCODE_ADAPTER: CachedFileAdapter = CachedFileAdapter::new_with_required_dependency( - ClientId::CommandCode, - ParserId::CommandCode, + DecoderId::CommandCode, COMMANDCODE_WORKSPACE_REVISION, sessions::commandcode::commandcode_config_dependency_path, sessions::commandcode::parse_commandcode_file, ); pub(crate) static ZCODE_ADAPTER: CachedFileAdapter = CachedFileAdapter::new( - ClientId::Zcode, - ParserId::Zcode, + DecoderId::Zcode, ZCODE_RECORD_REJECTION_REVISION, sessions::zcode::parse_zcode_file, ); @@ -396,26 +366,36 @@ not-json .into_iter() .map(|message| message.attribute(client)) .collect(); - assert!(messages - .iter() - .all(|message| message.client.as_ref() == client.as_str())); + assert!(messages.iter().all(|message| message.client == client)); messages } fn fold_with_adapter( + client: ClientId, adapter: &'static dyn LocalInputAdapter, units: Vec, cache: &mut message_cache::InputMessageCache, ) -> Vec { let parsed = adapter.parse_checked(units, &ParseContext { pricing: None }); + let (sink, _) = fold_parsed(client, adapter, parsed, cache); + assert!(sink.iter().all(|message| message.client == client)); + sink + } + + fn fold_parsed( + client: ClientId, + adapter: &'static dyn LocalInputAdapter, + parsed: Vec, + cache: &mut message_cache::InputMessageCache, + ) -> (Vec, crate::input_health::DataHealth) { let mut sink = Vec::new(); + let binding = crate::adapters::AdapterBinding::new(client, adapter); + let mut fold_ctx = FoldContext::new(binding, cache, None); + let mut bound_sink = BoundMessageSink::new(binding, &mut sink); adapter - .fold(parsed, &mut FoldContext::new(cache, None), &mut sink) + .fold(parsed, &mut fold_ctx, &mut bound_sink) .unwrap(); - assert!(sink - .iter() - .all(|message| message.client.as_ref() == adapter.client().as_str())); - sink + (sink, fold_ctx.take_health()) } #[test] @@ -436,7 +416,7 @@ not-json }; let ctx = scan_context(home.path(), &settings); - let units = AMP_ADAPTER.discover_checked(&ctx).unwrap(); + let units = AMP_ADAPTER.discover_checked(ClientId::Amp, &ctx).unwrap(); let paths: Vec<_> = units.iter().map(|unit| unit.path.clone()).collect(); let mut expected = vec![default_path, extra_path]; expected.sort_unstable(); @@ -452,10 +432,13 @@ not-json let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("T-test.json"); write_file(&path, AMP_CONTENT); - let units = vec![InputUnit::plain_file(ClientId::Amp, path.clone())]; + let units = vec![InputUnit::plain_file( + path.clone(), + DecoderSpec::plain(DecoderId::Amp, AMP_RECORD_REJECTION_REVISION), + )]; let mut cache = message_cache::InputMessageCache::default(); - let actual = fold_with_adapter(&_ADAPTER, units, &mut cache); + let actual = fold_with_adapter(ClientId::Amp, &_ADAPTER, units, &mut cache); let expected = finalized( ClientId::Amp, sessions::amp::parse_amp_file(&path).unwrap().messages, @@ -486,7 +469,11 @@ model = "gpt-5" let ctx = scan_context(home.path(), &settings); let mut cache = message_cache::InputMessageCache::default(); - let cold_unit = KIMI_ADAPTER.discover_checked(&ctx).unwrap().pop().unwrap(); + let cold_unit = KIMI_ADAPTER + .discover_checked(ClientId::Kimi, &ctx) + .unwrap() + .pop() + .unwrap(); assert_eq!( cold_unit.fingerprint_policy, FingerprintPolicy::PrimaryWithDependency { @@ -494,7 +481,8 @@ model = "gpt-5" related_failure_policy: RelatedInputFailurePolicy::PreservePrimary, } ); - let cold_messages = fold_with_adapter(&KIMI_ADAPTER, vec![cold_unit], &mut cache); + let cold_messages = + fold_with_adapter(ClientId::Kimi, &KIMI_ADAPTER, vec![cold_unit], &mut cache); assert_eq!(cold_messages.len(), 1); assert_eq!(cold_messages[0].model_id.as_ref(), "gpt-5"); assert_eq!(cold_messages[0].provider_id.as_ref(), "openai"); @@ -506,14 +494,23 @@ provider = "anthropic" model = "claude-sonnet-4" "#, ); - let changed_unit = KIMI_ADAPTER.discover_checked(&ctx).unwrap().pop().unwrap(); + let changed_unit = KIMI_ADAPTER + .discover_checked(ClientId::Kimi, &ctx) + .unwrap() + .pop() + .unwrap(); let changed_unit = match KIMI_ADAPTER.plan_cache_hit(changed_unit, &cache).unwrap() { crate::adapters::CacheHitPlan::Miss(unit) => unit, crate::adapters::CacheHitPlan::Hit(_) => { panic!("changed Kimi config must invalidate the usage cache") } }; - let changed_messages = fold_with_adapter(&KIMI_ADAPTER, vec![changed_unit], &mut cache); + let changed_messages = fold_with_adapter( + ClientId::Kimi, + &KIMI_ADAPTER, + vec![changed_unit], + &mut cache, + ); assert_eq!(changed_messages.len(), 1); assert_eq!(changed_messages[0].model_id.as_ref(), "claude-sonnet-4"); assert_eq!(changed_messages[0].provider_id.as_ref(), "anthropic"); @@ -533,30 +530,30 @@ model = "claude-sonnet-4" std::fs::create_dir(&config_path).unwrap(); let settings = crate::scanner::ScannerSettings::default(); let ctx = scan_context(home.path(), &settings); - let unit = KIMI_ADAPTER.discover_checked(&ctx).unwrap().pop().unwrap(); + let unit = KIMI_ADAPTER + .discover_checked(ClientId::Kimi, &ctx) + .unwrap() + .pop() + .unwrap(); let mut cache = message_cache::InputMessageCache::default(); let parsed = KIMI_ADAPTER.parse_checked(vec![unit], &ParseContext { pricing: None }); assert_eq!(parsed.len(), 1); assert!(matches!( - parsed[0].input_health().status, + parsed[0].health.status, crate::input_health::InputStatus::Partial { .. } )); assert!(parsed[0].cache_write.is_none()); - let mut messages = Vec::new(); - let mut fold_ctx = FoldContext::new(&mut cache, None); - KIMI_ADAPTER - .fold(parsed, &mut fold_ctx, &mut messages) - .unwrap(); + let (messages, health) = fold_parsed(ClientId::Kimi, &KIMI_ADAPTER, parsed, &mut cache); assert_eq!(messages.len(), 1); assert_eq!(messages[0].model_id.as_ref(), "active-model"); assert_eq!(messages[0].tokens.input, 10); assert_eq!(messages[0].tokens.output, 2); - assert_eq!(fold_ctx.health.partial_inputs(), 1); + assert_eq!(health.partial_inputs(), 1); assert!(cache - .get_meta(&wire_path, KIMI_ADAPTER.parser_version) + .get_meta(&wire_path, KIMI_ADAPTER.decoder.version()) .unwrap() .is_none()); } @@ -575,17 +572,21 @@ model = "claude-sonnet-4" assert!(!config_path.exists()); let settings = crate::scanner::ScannerSettings::default(); let ctx = scan_context(home.path(), &settings); - let unit = KIMI_ADAPTER.discover_checked(&ctx).unwrap().pop().unwrap(); + let unit = KIMI_ADAPTER + .discover_checked(ClientId::Kimi, &ctx) + .unwrap() + .pop() + .unwrap(); let mut cache = message_cache::InputMessageCache::default(); - let messages = fold_with_adapter(&KIMI_ADAPTER, vec![unit], &mut cache); + let messages = fold_with_adapter(ClientId::Kimi, &KIMI_ADAPTER, vec![unit], &mut cache); assert_eq!(messages.len(), 1); assert_eq!(messages[0].model_id.as_ref(), "active-model"); assert_eq!(messages[0].tokens.input, 10); assert_eq!(messages[0].tokens.output, 2); assert!(cache - .get_meta(&wire_path, KIMI_ADAPTER.parser_version) + .get_meta(&wire_path, KIMI_ADAPTER.decoder.version()) .unwrap() .is_some()); } @@ -611,7 +612,7 @@ model = "claude-sonnet-4" let mut cache = message_cache::InputMessageCache::default(); let cold_unit = COMMANDCODE_ADAPTER - .discover_checked(&ctx) + .discover_checked(ClientId::CommandCode, &ctx) .unwrap() .pop() .unwrap(); @@ -622,14 +623,19 @@ model = "claude-sonnet-4" related_failure_policy: RelatedInputFailurePolicy::FailInput, } ); - let cold_messages = fold_with_adapter(&COMMANDCODE_ADAPTER, vec![cold_unit], &mut cache); + let cold_messages = fold_with_adapter( + ClientId::CommandCode, + &COMMANDCODE_ADAPTER, + vec![cold_unit], + &mut cache, + ); assert_eq!(cold_messages.len(), 1); assert_eq!(cold_messages[0].model_id.as_ref(), "gpt-5"); assert_eq!(cold_messages[0].provider_id.as_ref(), "openai"); write_file(&config_path, r#"{"model":"private-preview"}"#); let changed_unit = COMMANDCODE_ADAPTER - .discover_checked(&ctx) + .discover_checked(ClientId::CommandCode, &ctx) .unwrap() .pop() .unwrap(); @@ -642,8 +648,12 @@ model = "claude-sonnet-4" panic!("changed Command Code config must invalidate the usage cache") } }; - let changed_messages = - fold_with_adapter(&COMMANDCODE_ADAPTER, vec![changed_unit], &mut cache); + let changed_messages = fold_with_adapter( + ClientId::CommandCode, + &COMMANDCODE_ADAPTER, + vec![changed_unit], + &mut cache, + ); assert_eq!(changed_messages.len(), 1); assert_eq!(changed_messages[0].model_id.as_ref(), "private-preview"); assert_eq!(changed_messages[0].provider_id.as_ref(), "unknown"); @@ -663,7 +673,9 @@ model = "claude-sonnet-4" let settings = crate::scanner::ScannerSettings::default(); let ctx = scan_context(home.path(), &settings); - let units = COMMANDCODE_ADAPTER.discover_checked(&ctx).unwrap(); + let units = COMMANDCODE_ADAPTER + .discover_checked(ClientId::CommandCode, &ctx) + .unwrap(); assert!(units.is_empty()); } @@ -687,7 +699,7 @@ model = "claude-sonnet-4" let settings = crate::scanner::ScannerSettings::default(); let ctx = scan_context(home.path(), &settings); let unit = COMMANDCODE_ADAPTER - .discover_checked(&ctx) + .discover_checked(ClientId::CommandCode, &ctx) .unwrap() .pop() .unwrap(); @@ -696,22 +708,20 @@ model = "claude-sonnet-4" let parsed = COMMANDCODE_ADAPTER.parse_checked(vec![unit], &ParseContext { pricing: None }); assert_eq!(parsed.len(), 1); assert!(matches!( - parsed[0].input_health().status, + parsed[0].health.status, crate::input_health::InputStatus::Unavailable { .. } )); assert!(parsed[0].cache_write.is_none()); - let mut messages = Vec::new(); - COMMANDCODE_ADAPTER - .fold( - parsed, - &mut FoldContext::new(&mut cache, None), - &mut messages, - ) - .unwrap(); + let (messages, _) = fold_parsed( + ClientId::CommandCode, + &COMMANDCODE_ADAPTER, + parsed, + &mut cache, + ); assert!(messages.is_empty()); assert!(cache - .get_meta(&session_path, COMMANDCODE_ADAPTER.parser_version) + .get_meta(&session_path, COMMANDCODE_ADAPTER.decoder.version()) .unwrap() .is_none()); } @@ -722,19 +732,18 @@ model = "claude-sonnet-4" let cache_dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("session.jsonl"); write_file(&path, QWEN_MIXED_CONTENT); - let unit = InputUnit::plain_file(ClientId::Qwen, path.clone()) - .with_parser_version(ParserVersion::new( - ParserId::Qwen, - QWEN_RECORD_REJECTION_REVISION, - )) - .prepare_snapshot() - .unwrap(); + let unit = InputUnit::plain_file( + path.clone(), + DecoderSpec::plain(DecoderId::Qwen, QWEN_RECORD_REJECTION_REVISION), + ) + .prepare_snapshot() + .unwrap(); let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); let parsed = QWEN_ADAPTER.parse_checked(vec![unit.clone()], &ParseContext { pricing: None }); assert_eq!(parsed.len(), 1); - let health = parsed[0].input_health(); + let health = &parsed[0].health; assert!(matches!( health.status, crate::input_health::InputStatus::Complete @@ -745,11 +754,9 @@ model = "claude-sonnet-4" "malformed-record" ); - let mut sink = Vec::new(); - let mut fold_ctx = FoldContext::new(&mut cache, None); - QWEN_ADAPTER.fold(parsed, &mut fold_ctx, &mut sink).unwrap(); + let (sink, health) = fold_parsed(ClientId::Qwen, &QWEN_ADAPTER, parsed, &mut cache); assert_eq!(sink.len(), 2); - assert_eq!(fold_ctx.health.rejected_records(), 1); + assert_eq!(health.rejected_records(), 1); cache.save_if_dirty().unwrap(); let warm_cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); @@ -757,7 +764,7 @@ model = "claude-sonnet-4" let crate::adapters::CacheHitPlan::Hit(hit) = planned else { panic!("unchanged Qwen input must use its cached complete scan"); }; - let warm_health = hit.input_health(); + let warm_health = &hit.health; assert_eq!(warm_health.rejections.total(), 1); assert_eq!( warm_health.rejections.entries().next().unwrap().key, @@ -777,12 +784,16 @@ model = "claude-sonnet-4" write_file(&summary_path, "not-json"); let settings = crate::scanner::ScannerSettings::default(); let ctx = scan_context(home.path(), &settings); - let unit = GROK_ADAPTER.discover_checked(&ctx).unwrap().pop().unwrap(); + let unit = GROK_ADAPTER + .discover_checked(ClientId::Grok, &ctx) + .unwrap() + .pop() + .unwrap(); let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); let parsed = GROK_ADAPTER.parse_checked(vec![unit.clone()], &ParseContext { pricing: None }); - let health = parsed[0].input_health(); + let health = &parsed[0].health; assert!(matches!( health.status, crate::input_health::InputStatus::Complete @@ -793,12 +804,9 @@ model = "claude-sonnet-4" "malformed-record" ); - let mut sink = Vec::new(); - let mut fold_ctx = FoldContext::new(&mut cache, None); - GROK_ADAPTER.fold(parsed, &mut fold_ctx, &mut sink).unwrap(); + let (sink, health) = fold_parsed(ClientId::Grok, &GROK_ADAPTER, parsed, &mut cache); assert_eq!(sink.len(), 1); - assert_eq!(fold_ctx.health.rejected_records(), 1); - drop(fold_ctx); + assert_eq!(health.rejected_records(), 1); cache.save_if_dirty().unwrap(); let warm_cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); @@ -806,7 +814,7 @@ model = "claude-sonnet-4" let crate::adapters::CacheHitPlan::Hit(hit) = planned else { panic!("unchanged Grok siblings must restore the complete cached scan"); }; - let warm_health = hit.input_health(); + let warm_health = &hit.health; assert_eq!(warm_health.rejections.total(), 1); assert_eq!( warm_health.rejections.entries().next().unwrap().key, @@ -828,14 +836,15 @@ model = "claude-sonnet-4" ); let settings = crate::scanner::ScannerSettings::default(); let ctx = scan_context(home.path(), &settings); - let unit = GROK_ADAPTER.discover_checked(&ctx).unwrap().pop().unwrap(); + let unit = GROK_ADAPTER + .discover_checked(ClientId::Grok, &ctx) + .unwrap() + .pop() + .unwrap(); let mut cache = message_cache::InputMessageCache::default(); let parsed = GROK_ADAPTER.parse_checked(vec![unit.clone()], &ParseContext { pricing: None }); - let mut sink = Vec::new(); - GROK_ADAPTER - .fold(parsed, &mut FoldContext::new(&mut cache, None), &mut sink) - .unwrap(); + let (sink, _) = fold_parsed(ClientId::Grok, &GROK_ADAPTER, parsed, &mut cache); assert_eq!(sink.len(), 1); write_file( @@ -846,7 +855,11 @@ model = "claude-sonnet-4" std::fs::read_to_string(&updates_path).unwrap(), GROK_SELF_CONTAINED_UPDATES ); - let changed_unit = GROK_ADAPTER.discover_checked(&ctx).unwrap().pop().unwrap(); + let changed_unit = GROK_ADAPTER + .discover_checked(ClientId::Grok, &ctx) + .unwrap() + .pop() + .unwrap(); assert!(matches!( GROK_ADAPTER.plan_cache_hit(changed_unit, &cache).unwrap(), @@ -865,11 +878,15 @@ model = "claude-sonnet-4" std::fs::create_dir(&events_path).unwrap(); let settings = crate::scanner::ScannerSettings::default(); let ctx = scan_context(home.path(), &settings); - let unit = GROK_ADAPTER.discover_checked(&ctx).unwrap().pop().unwrap(); + let unit = GROK_ADAPTER + .discover_checked(ClientId::Grok, &ctx) + .unwrap() + .pop() + .unwrap(); let mut cache = message_cache::InputMessageCache::default(); let parsed = GROK_ADAPTER.parse_checked(vec![unit], &ParseContext { pricing: None }); - let health = parsed[0].input_health(); + let health = &parsed[0].health; assert!(matches!( health.status, crate::input_health::InputStatus::Partial { .. } @@ -878,15 +895,13 @@ model = "claude-sonnet-4" assert_eq!(failure.operation, "read related events line"); assert!(failure.message.contains(&events_path.display().to_string())); - let mut sink = Vec::new(); - let mut fold_ctx = FoldContext::new(&mut cache, None); - GROK_ADAPTER.fold(parsed, &mut fold_ctx, &mut sink).unwrap(); + let (sink, health) = fold_parsed(ClientId::Grok, &GROK_ADAPTER, parsed, &mut cache); assert_eq!(sink.len(), 1); - assert_eq!(fold_ctx.health.partial_inputs(), 1); - assert_eq!(fold_ctx.health.failed_inputs(), 0); - assert_eq!(fold_ctx.health.rejected_records(), 1); + assert_eq!(health.partial_inputs(), 1); + assert_eq!(health.failed_inputs(), 0); + assert_eq!(health.rejected_records(), 1); assert!(cache - .get_meta(&updates_path, GROK_ADAPTER.parser_version) + .get_meta(&updates_path, GROK_ADAPTER.decoder.version()) .unwrap() .is_none()); } @@ -903,15 +918,16 @@ model = "claude-sonnet-4" "tokenUsage": {"inputTokens": 10} }"#, ); - let unit = InputUnit::plain_file(ClientId::Droid, path.clone()).with_parser_version( - ParserVersion::new(ParserId::Droid, DROID_RECORD_REJECTION_REVISION), + let unit = InputUnit::plain_file( + path.clone(), + DecoderSpec::plain(DecoderId::Droid, DROID_RECORD_REJECTION_REVISION), ); let mut cache = message_cache::InputMessageCache::default(); let parsed = DROID_ADAPTER.parse_checked(vec![unit], &ParseContext { pricing: None }); assert_eq!(parsed.len(), 1); - let health = parsed[0].input_health(); + let health = &parsed[0].health; assert!(matches!( health.status, crate::input_health::InputStatus::Complete @@ -920,14 +936,10 @@ model = "claude-sonnet-4" assert_eq!(rejection.key, "missing-model"); assert_eq!(rejection.count, 1); - let mut sink = Vec::new(); - let mut fold_ctx = FoldContext::new(&mut cache, None); - DROID_ADAPTER - .fold(parsed, &mut fold_ctx, &mut sink) - .unwrap(); + let (sink, health) = fold_parsed(ClientId::Droid, &DROID_ADAPTER, parsed, &mut cache); assert!(sink.is_empty()); - assert_eq!(fold_ctx.health.rejected_records(), 1); - assert_eq!(fold_ctx.health.failed_inputs(), 0); + assert_eq!(health.rejected_records(), 1); + assert_eq!(health.failed_inputs(), 0); } #[test] @@ -954,8 +966,13 @@ model = "claude-sonnet-4" let ctx = scan_context(home.path(), &settings); let mut cache = message_cache::InputMessageCache::default(); - let cold_unit = DROID_ADAPTER.discover_checked(&ctx).unwrap().pop().unwrap(); - let cold_messages = fold_with_adapter(&DROID_ADAPTER, vec![cold_unit], &mut cache); + let cold_unit = DROID_ADAPTER + .discover_checked(ClientId::Droid, &ctx) + .unwrap() + .pop() + .unwrap(); + let cold_messages = + fold_with_adapter(ClientId::Droid, &DROID_ADAPTER, vec![cold_unit], &mut cache); assert_eq!( cold_messages[0].workspace_key.as_deref(), Some("/home/travis/01-workspace/tokscale") @@ -971,20 +988,18 @@ model = "claude-sonnet-4" r#"{"type":"session_start","cwd":"/home/travis/02-workspace/oh-my-openagent"} "#, ); - let warm_unit = DROID_ADAPTER.discover_checked(&ctx).unwrap().pop().unwrap(); + let warm_unit = DROID_ADAPTER + .discover_checked(ClientId::Droid, &ctx) + .unwrap() + .pop() + .unwrap(); let crate::adapters::CacheHitPlan::Hit(warm_hit) = DROID_ADAPTER.plan_cache_hit(warm_unit, &cache).unwrap() else { panic!("unchanged Droid settings must retain the usage-cache hit"); }; - let mut warm_messages = Vec::new(); - DROID_ADAPTER - .fold( - vec![warm_hit], - &mut FoldContext::new(&mut cache, None), - &mut warm_messages, - ) - .unwrap(); + let (warm_messages, _) = + fold_parsed(ClientId::Droid, &DROID_ADAPTER, vec![warm_hit], &mut cache); assert_eq!( warm_messages[0].workspace_key.as_deref(), @@ -1027,7 +1042,11 @@ model = "claude-sonnet-4" ); let settings = crate::scanner::ScannerSettings::default(); let ctx = scan_context(home.path(), &settings); - let unit = DROID_ADAPTER.discover_checked(&ctx).unwrap().pop().unwrap(); + let unit = DROID_ADAPTER + .discover_checked(ClientId::Droid, &ctx) + .unwrap() + .pop() + .unwrap(); assert_eq!( unit.fingerprint_policy, FingerprintPolicy::PrimaryWithDependency { @@ -1037,7 +1056,8 @@ model = "claude-sonnet-4" ); let mut cache = message_cache::InputMessageCache::default(); - let worker_messages = fold_with_adapter(&DROID_ADAPTER, vec![unit], &mut cache); + let worker_messages = + fold_with_adapter(ClientId::Droid, &DROID_ADAPTER, vec![unit], &mut cache); assert_eq!(worker_messages.len(), 1); assert_eq!(worker_messages[0].agent.as_deref(), Some("Droid Worker")); @@ -1045,14 +1065,23 @@ model = "claude-sonnet-4" &features_path, r#"{"features":[{"id":"scrutiny","skillName":"scrutiny-validator","workerSessionIds":["mission-worker"]}]}"#, ); - let changed_unit = DROID_ADAPTER.discover_checked(&ctx).unwrap().pop().unwrap(); + let changed_unit = DROID_ADAPTER + .discover_checked(ClientId::Droid, &ctx) + .unwrap() + .pop() + .unwrap(); let changed_unit = match DROID_ADAPTER.plan_cache_hit(changed_unit, &cache).unwrap() { crate::adapters::CacheHitPlan::Miss(unit) => unit, crate::adapters::CacheHitPlan::Hit(_) => { panic!("changed Mission feature must invalidate the Droid input cache") } }; - let validator_messages = fold_with_adapter(&DROID_ADAPTER, vec![changed_unit], &mut cache); + let validator_messages = fold_with_adapter( + ClientId::Droid, + &DROID_ADAPTER, + vec![changed_unit], + &mut cache, + ); assert_eq!(validator_messages.len(), 1); assert_eq!( validator_messages[0].agent.as_deref(), @@ -1089,84 +1118,84 @@ model = "claude-sonnet-4" std::fs::create_dir_all(&features_path).unwrap(); let settings = crate::scanner::ScannerSettings::default(); let ctx = scan_context(home.path(), &settings); - let unit = DROID_ADAPTER.discover_checked(&ctx).unwrap().pop().unwrap(); + let unit = DROID_ADAPTER + .discover_checked(ClientId::Droid, &ctx) + .unwrap() + .pop() + .unwrap(); let mut cache = message_cache::InputMessageCache::default(); let parsed = DROID_ADAPTER.parse_checked(vec![unit], &ParseContext { pricing: None }); assert_eq!(parsed.len(), 1); assert!(matches!( - parsed[0].input_health().status, + parsed[0].health.status, crate::input_health::InputStatus::Partial { .. } )); assert!(parsed[0].cache_write.is_none()); - let mut messages = Vec::new(); - let mut fold_ctx = FoldContext::new(&mut cache, None); - DROID_ADAPTER - .fold(parsed, &mut fold_ctx, &mut messages) - .unwrap(); + let (messages, health) = fold_parsed(ClientId::Droid, &DROID_ADAPTER, parsed, &mut cache); assert_eq!(messages.len(), 1); assert_eq!(messages[0].tokens.input, 10); assert_eq!(messages[0].tokens.output, 5); assert_eq!(messages[0].agent.as_deref(), Some("Droid Worker")); - assert_eq!(fold_ctx.health.partial_inputs(), 1); + assert_eq!(health.partial_inputs(), 1); assert!(cache - .get_meta(&settings_path, DROID_ADAPTER.parser_version) + .get_meta(&settings_path, DROID_ADAPTER.decoder.version()) .unwrap() .is_none()); } #[test] fn cached_file_adapters_use_their_actual_record_rejection_revisions() { - for (actual, parser_id, revision) in [ + for (actual, decoder_id, revision) in [ ( - GEMINI_ADAPTER.parser_version, - ParserId::Gemini, + GEMINI_ADAPTER.decoder.version(), + DecoderId::Gemini, GEMINI_RECORD_REJECTION_REVISION, ), ( - GROK_ADAPTER.parser_version, - ParserId::Grok, + GROK_ADAPTER.decoder.version(), + DecoderId::Grok, GROK_RELATED_METADATA_REVISION, ), ( - AMP_ADAPTER.parser_version, - ParserId::Amp, + AMP_ADAPTER.decoder.version(), + DecoderId::Amp, AMP_RECORD_REJECTION_REVISION, ), ( - DROID_ADAPTER.parser_version, - ParserId::Droid, + DROID_ADAPTER.decoder.version(), + DecoderId::Droid, DROID_AGENT_ATTRIBUTION_REVISION, ), ( - KIMI_ADAPTER.parser_version, - ParserId::Kimi, + KIMI_ADAPTER.decoder.version(), + DecoderId::Kimi, KIMI_RECORD_REJECTION_REVISION, ), ( - QWEN_ADAPTER.parser_version, - ParserId::Qwen, + QWEN_ADAPTER.decoder.version(), + DecoderId::Qwen, QWEN_RECORD_REJECTION_REVISION, ), ( - MUX_ADAPTER.parser_version, - ParserId::Mux, + MUX_ADAPTER.decoder.version(), + DecoderId::Mux, MUX_RECORD_REJECTION_REVISION, ), ( - COMMANDCODE_ADAPTER.parser_version, - ParserId::CommandCode, + COMMANDCODE_ADAPTER.decoder.version(), + DecoderId::CommandCode, COMMANDCODE_WORKSPACE_REVISION, ), ( - ZCODE_ADAPTER.parser_version, - ParserId::Zcode, + ZCODE_ADAPTER.decoder.version(), + DecoderId::Zcode, ZCODE_RECORD_REJECTION_REVISION, ), ] { - assert_eq!(actual, ParserVersion::new(parser_id, revision)); + assert_eq!(actual, DecoderVersion::new(decoder_id, revision)); } } @@ -1187,7 +1216,9 @@ model = "claude-sonnet-4" }; let ctx = scan_context(home.path(), &settings); - let units = COPILOT_ADAPTER.discover_checked(&ctx).unwrap(); + let units = COPILOT_ADAPTER + .discover_checked(ClientId::Copilot, &ctx) + .unwrap(); assert_eq!( units .iter() @@ -1195,8 +1226,8 @@ model = "claude-sonnet-4" .collect::>(), vec![default_path, extra_path] ); - assert!(units.iter().all(|unit| unit.parser_version - == ParserVersion::new(ParserId::Copilot, COPILOT_AGENT_IDENTITY_REVISION))); + assert!(units.iter().all(|unit| unit.decoder.version() + == DecoderVersion::new(DecoderId::Copilot, COPILOT_AGENT_IDENTITY_REVISION))); } #[test] @@ -1207,12 +1238,14 @@ model = "claude-sonnet-4" let settings = crate::scanner::ScannerSettings::default(); let ctx = scan_context(home.path(), &settings); - let units = ZCODE_ADAPTER.discover_checked(&ctx).unwrap(); + let units = ZCODE_ADAPTER + .discover_checked(ClientId::Zcode, &ctx) + .unwrap(); let paths: Vec<_> = units.iter().map(|unit| unit.path.clone()).collect(); assert_eq!(paths, vec![default_path]); - assert!(units.iter().all(|unit| unit.parser_version - == ParserVersion::new(ParserId::Zcode, ZCODE_RECORD_REJECTION_REVISION))); + assert!(units.iter().all(|unit| unit.decoder.version() + == DecoderVersion::new(DecoderId::Zcode, ZCODE_RECORD_REJECTION_REVISION))); } #[test] @@ -1225,16 +1258,16 @@ model = "claude-sonnet-4" let settings = crate::scanner::ScannerSettings::default(); let ctx = scan_context(home.path(), &settings); - let units = GROK_ADAPTER.discover_checked(&ctx).unwrap(); + let units = GROK_ADAPTER.discover_checked(ClientId::Grok, &ctx).unwrap(); assert_eq!(units.len(), 1); assert_eq!( - units[0].parser_version, - ParserVersion::new(ParserId::Grok, GROK_RELATED_METADATA_REVISION) + units[0].decoder.version(), + DecoderVersion::new(DecoderId::Grok, GROK_RELATED_METADATA_REVISION) ); assert_ne!( - units[0].parser_version, - ParserVersion::new(ParserId::Grok, GROK_RECORD_REJECTION_REVISION) + units[0].decoder.version(), + DecoderVersion::new(DecoderId::Grok, GROK_RECORD_REJECTION_REVISION) ); assert_eq!( units[0].fingerprint_policy, @@ -1250,14 +1283,13 @@ model = "claude-sonnet-4" let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("session.jsonl"); write_file(&path, ZCODE_CONTENT); - let units = vec![ - InputUnit::plain_file(ClientId::Zcode, path.clone()).with_parser_version( - ParserVersion::new(ParserId::Zcode, ZCODE_RECORD_REJECTION_REVISION), - ), - ]; + let units = vec![InputUnit::plain_file( + path.clone(), + DecoderSpec::plain(DecoderId::Zcode, ZCODE_RECORD_REJECTION_REVISION), + )]; let mut cache = message_cache::InputMessageCache::default(); - let actual = fold_with_adapter(&ZCODE_ADAPTER, units, &mut cache); + let actual = fold_with_adapter(ClientId::Zcode, &ZCODE_ADAPTER, units, &mut cache); let expected = finalized( ClientId::Zcode, sessions::zcode::parse_zcode_file(&path).unwrap().messages, @@ -1274,13 +1306,15 @@ model = "claude-sonnet-4" &path, "{\"type\":\"init\",\"model\":\"gemini-2.5-pro\",\"session_id\":\"session-1\"}\nnot-json\n{\"type\":\"result\",\"stats\":{\"input_tokens\":10,\"output_tokens\":20}}\n", ); - let units = vec![InputUnit::plain_file(ClientId::Gemini, path.clone())]; + let units = vec![InputUnit::plain_file( + path.clone(), + DecoderSpec::plain(DecoderId::Gemini, GEMINI_RECORD_REJECTION_REVISION), + )]; let parsed = GEMINI_ADAPTER.parse_checked(units, &ParseContext { pricing: None }); assert_eq!(parsed.len(), 1); - let health = parsed[0].input_health(); - assert_eq!(health.client, ClientId::Gemini); - assert_eq!(health.path, path); + let health = &parsed[0].health; + assert_eq!(parsed[0].unit.path, path); let failure = health.status.failure().expect("input must be partial"); assert_eq!(failure.operation, "decode JSONL line"); assert!(matches!( diff --git a/crates/tokscale-core/src/adapters/goose.rs b/crates/tokscale-core/src/adapters/goose.rs index 5256a5001..300d86e33 100644 --- a/crates/tokscale-core/src/adapters/goose.rs +++ b/crates/tokscale-core/src/adapters/goose.rs @@ -5,11 +5,13 @@ use rayon::prelude::*; use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ - AdapterScanContext, FoldContext, InputDiscoveryError, InputUnit, LocalInputAdapter, - MessageSink, ParseContext, ParsedUnit, + AdapterScanContext, BoundMessageSink, DecoderSpec, FoldContext, InputDiscoveryError, InputUnit, + LocalInputAdapter, ParseContext, ParsedUnit, }; use crate::clients::ClientId; -use crate::message_cache::{ParserId, ParserVersion}; +use crate::message_cache::DecoderId; +#[cfg(test)] +use crate::message_cache::DecoderVersion; use crate::sessions; const GOOSE_RECORD_REJECTION_REVISION: u32 = 5; @@ -17,27 +19,17 @@ const GOOSE_RECORD_REJECTION_REVISION: u32 = 5; pub(crate) struct GooseAdapter; impl LocalInputAdapter for GooseAdapter { - fn client(&self) -> ClientId { - ClientId::Goose - } - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { - Ok(adapter_discover::input_units_from_paths_preserving_order( - ClientId::Goose, - goose_db_paths(ctx)?, + adapter_discover::input_units_from_paths_preserving_order( + client, + goose_db_paths(client, ctx)?, crate::adapters::FingerprintPolicy::SqliteWithWal, - )? - .into_iter() - .map(|path| { - path.with_parser_version(ParserVersion::new( - ParserId::Goose, - GOOSE_RECORD_REJECTION_REVISION, - )) - }) - .collect()) + DecoderSpec::plain(DecoderId::Goose, GOOSE_RECORD_REJECTION_REVISION), + ) } fn parse_checked(&self, units: Vec, ctx: &ParseContext<'_>) -> Vec { @@ -53,14 +45,17 @@ impl LocalInputAdapter for GooseAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), crate::adapters::InputPipelineError> { adapter_cache::fold_units(parsed, ctx, sink) } } -fn goose_db_paths(ctx: &AdapterScanContext<'_>) -> Result, InputDiscoveryError> { - let def = ClientId::Goose +fn goose_db_paths( + client: ClientId, + ctx: &AdapterScanContext<'_>, +) -> Result, InputDiscoveryError> { + let def = client .local_def() .expect("Goose adapter must have local scan policy"); let default_candidates = [ @@ -70,13 +65,13 @@ fn goose_db_paths(ctx: &AdapterScanContext<'_>) -> Result, InputDis let mut existing_defaults = Vec::new(); for candidate in default_candidates { - adapter_discover::push_existing_file(ClientId::Goose, candidate, &mut existing_defaults)?; + adapter_discover::push_existing_file(client, candidate, &mut existing_defaults)?; } let mut paths: Vec<_> = existing_defaults.into_iter().take(1).collect(); paths.extend(adapter_discover::scan_roots( - ClientId::Goose, - adapter_discover::extra_roots_for_client(ClientId::Goose, ctx)?, + client, + adapter_discover::extra_roots_for_client(client, ctx)?, "sessions.db", )?); Ok(paths) @@ -105,13 +100,15 @@ mod tests { scanner_settings: &settings, }; - let units = GOOSE_ADAPTER.discover_checked(&ctx).unwrap(); + let units = GOOSE_ADAPTER + .discover_checked(ClientId::Goose, &ctx) + .unwrap(); assert_eq!(units.len(), 1); assert_eq!(units[0].path, xdg_db); assert_eq!( - units[0].parser_version, - ParserVersion::new(ParserId::Goose, GOOSE_RECORD_REJECTION_REVISION) + units[0].decoder.version(), + DecoderVersion::new(DecoderId::Goose, GOOSE_RECORD_REJECTION_REVISION) ); } @@ -147,7 +144,9 @@ mod tests { scanner_settings: &settings, }; - let units = GOOSE_ADAPTER.discover_checked(&ctx).unwrap(); + let units = GOOSE_ADAPTER + .discover_checked(ClientId::Goose, &ctx) + .unwrap(); assert_eq!( units.into_iter().map(|unit| unit.path).collect::>(), diff --git a/crates/tokscale-core/src/adapters/hermes.rs b/crates/tokscale-core/src/adapters/hermes.rs index e51aea913..38552b5bf 100644 --- a/crates/tokscale-core/src/adapters/hermes.rs +++ b/crates/tokscale-core/src/adapters/hermes.rs @@ -5,11 +5,13 @@ use rayon::prelude::*; use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ - AdapterScanContext, FingerprintPolicy, FoldContext, InputDiscoveryError, InputUnit, - LocalInputAdapter, MessageSink, ParseContext, ParsedBatchInput, ParsedUnit, + AdapterScanContext, BoundMessageSink, DecoderSpec, FingerprintPolicy, FoldContext, + InputDiscoveryError, InputUnit, LocalInputAdapter, ParseContext, ParsedBatchInput, ParsedUnit, }; use crate::clients::ClientId; -use crate::message_cache::{ParserId, ParserVersion}; +use crate::message_cache::DecoderId; +#[cfg(test)] +use crate::message_cache::DecoderVersion; use crate::sessions; const HERMES_RECORD_REJECTION_REVISION: u32 = 5; @@ -17,44 +19,30 @@ const HERMES_RECORD_REJECTION_REVISION: u32 = 5; pub(crate) struct HermesAdapter; impl LocalInputAdapter for HermesAdapter { - fn client(&self) -> ClientId { - ClientId::Hermes - } - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { - let def = ClientId::Hermes + let def = client .local_def() .expect("Hermes adapter must have local scan policy"); let mut paths = Vec::new(); - adapter_discover::push_existing_file( - ClientId::Hermes, - def.resolve_path(ctx.home_dir), - &mut paths, - )?; + adapter_discover::push_existing_file(client, def.resolve_path(ctx.home_dir), &mut paths)?; paths.extend(adapter_discover::scan_roots( - ClientId::Hermes, - adapter_discover::extra_roots_for_client(ClientId::Hermes, ctx)?, + client, + adapter_discover::extra_roots_for_client(client, ctx)?, def.pattern, )?); let units = adapter_discover::input_units_from_paths_preserving_order( - ClientId::Hermes, + client, paths, FingerprintPolicy::SqliteWithWal, + DecoderSpec::plain(DecoderId::Hermes, HERMES_RECORD_REJECTION_REVISION), )?; - Ok(units - .into_iter() - .map(|unit| { - unit.with_parser_version(ParserVersion::new( - ParserId::Hermes, - HERMES_RECORD_REJECTION_REVISION, - )) - }) - .collect()) + Ok(units) } fn parse_checked(&self, units: Vec, ctx: &ParseContext<'_>) -> Vec { @@ -70,7 +58,7 @@ impl LocalInputAdapter for HermesAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), crate::adapters::InputPipelineError> { let mut seen = HashSet::new(); fold_hermes_units(parsed, ctx, sink, &mut seen) @@ -80,7 +68,7 @@ impl LocalInputAdapter for HermesAdapter { &self, batches: &mut ParsedBatchInput<'_>, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), crate::adapters::InputPipelineError> { let mut seen = HashSet::new(); while let Some(parsed) = batches.next(ctx)? { @@ -93,7 +81,7 @@ impl LocalInputAdapter for HermesAdapter { fn fold_hermes_units( parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, seen: &mut HashSet, ) -> Result<(), crate::adapters::InputPipelineError> { adapter_cache::fold_units_with_filter(parsed, ctx, sink, |_, messages| { @@ -115,13 +103,16 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("state.db"); std::fs::write(&path, b"direct parser input").unwrap(); - let unit = InputUnit::sqlite_with_wal(ClientId::Hermes, path.clone()) - .prepare_snapshot() - .unwrap(); + let unit = InputUnit::sqlite_with_wal( + path.clone(), + DecoderSpec::plain(DecoderId::Hermes, HERMES_RECORD_REJECTION_REVISION), + ) + .prepare_snapshot() + .unwrap(); let mut cache = crate::message_cache::InputMessageCache::default(); cache.insert(crate::message_cache::CachedInputEntry::new_with_version( &path, - unit.parser_version, + unit.decoder.version(), unit.input_policy().fingerprint().unwrap(), vec![crate::sessions::ParsedMessage::new( "model", @@ -164,13 +155,15 @@ mod tests { scanner_settings: &settings, }; - let units = HERMES_ADAPTER.discover_checked(&ctx).unwrap(); + let units = HERMES_ADAPTER + .discover_checked(ClientId::Hermes, &ctx) + .unwrap(); let paths: Vec<_> = units.iter().map(|unit| unit.path.clone()).collect(); assert_eq!(paths, vec![default_db, profile_db]); assert!(units.iter().all(|unit| { - unit.parser_version - == ParserVersion::new(ParserId::Hermes, HERMES_RECORD_REJECTION_REVISION) + unit.decoder.version() + == DecoderVersion::new(DecoderId::Hermes, HERMES_RECORD_REJECTION_REVISION) })); } } diff --git a/crates/tokscale-core/src/adapters/junie.rs b/crates/tokscale-core/src/adapters/junie.rs index e186c97a8..eb3a81141 100644 --- a/crates/tokscale-core/src/adapters/junie.rs +++ b/crates/tokscale-core/src/adapters/junie.rs @@ -3,10 +3,13 @@ use rayon::prelude::*; use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ - AdapterScanContext, FingerprintPolicy, FoldContext, InputDiscoveryError, InputUnit, - LocalInputAdapter, MessageSink, ParseContext, ParsedUnit, EXPLICIT_TOKEN_OVERFLOW_REVISION, + AdapterScanContext, BoundMessageSink, DecoderSpec, FingerprintPolicy, FoldContext, + InputDiscoveryError, InputUnit, LocalInputAdapter, ParseContext, ParsedUnit, + EXPLICIT_TOKEN_OVERFLOW_REVISION, }; -use crate::message_cache::{ParserId, ParserVersion}; +use crate::message_cache::DecoderId; +#[cfg(test)] +use crate::message_cache::DecoderVersion; use crate::{sessions, ClientId}; pub(crate) struct JunieAdapter; @@ -14,27 +17,17 @@ pub(crate) struct JunieAdapter; const JUNIE_RECORD_REJECTION_REVISION: u32 = EXPLICIT_TOKEN_OVERFLOW_REVISION + 3; impl LocalInputAdapter for JunieAdapter { - fn client(&self) -> ClientId { - ClientId::Junie - } - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { let units = adapter_discover::discover_default_scanned_units( - ClientId::Junie, + client, ctx, FingerprintPolicy::PlainFile, - )? - .into_iter() - .map(|unit| { - unit.with_parser_version(ParserVersion::new( - ParserId::Junie, - JUNIE_RECORD_REJECTION_REVISION, - )) - }) - .collect(); + DecoderSpec::plain(DecoderId::Junie, JUNIE_RECORD_REJECTION_REVISION), + )?; Ok(units) } @@ -61,7 +54,7 @@ impl LocalInputAdapter for JunieAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), crate::adapters::InputPipelineError> { adapter_cache::fold_units(parsed, ctx, sink) } @@ -116,12 +109,15 @@ mod tests { ) -> Vec { let parsed = JUNIE_ADAPTER.parse_checked(units, &ParseContext { pricing }); let mut messages = Vec::new(); + let binding = crate::adapters::adapter_for(ClientId::Junie).unwrap(); + let mut fold_ctx = FoldContext::new(binding, cache, pricing); + let mut sink = BoundMessageSink::new(binding, &mut messages); JUNIE_ADAPTER - .fold(parsed, &mut FoldContext::new(cache, pricing), &mut messages) + .fold(parsed, &mut fold_ctx, &mut sink) .unwrap(); assert!(messages .iter() - .all(|message| message.client.as_ref() == ClientId::Junie.as_str())); + .all(|message| message.client == ClientId::Junie)); messages } @@ -155,16 +151,15 @@ mod tests { let settings = ScannerSettings::default(); let units = JUNIE_ADAPTER - .discover_checked(&scan_context(home.path(), &settings)) + .discover_checked(ClientId::Junie, &scan_context(home.path(), &settings)) .unwrap(); assert_eq!(units.len(), 1); - assert_eq!(units[0].client, ClientId::Junie); assert_eq!(units[0].path, path); assert_eq!(units[0].fingerprint_policy, FingerprintPolicy::PlainFile); assert_eq!( - units[0].parser_version, - ParserVersion::new(ParserId::Junie, JUNIE_RECORD_REJECTION_REVISION) + units[0].decoder.version(), + DecoderVersion::new(DecoderId::Junie, JUNIE_RECORD_REJECTION_REVISION) ); } @@ -175,7 +170,10 @@ mod tests { let mut cache = message_cache::InputMessageCache::default(); let actual = fold_with_adapter( - vec![InputUnit::plain_file(ClientId::Junie, path.clone())], + vec![InputUnit::plain_file( + path.clone(), + DecoderSpec::plain(DecoderId::Junie, JUNIE_RECORD_REJECTION_REVISION), + )], &mut cache, None, ); @@ -183,7 +181,7 @@ mod tests { assert!(actual .iter() - .all(|message| message.client.as_ref() == ClientId::Junie.as_str())); + .all(|message| message.client == ClientId::Junie)); assert_eq!(actual, expected); } @@ -197,7 +195,10 @@ mod tests { let path = write_session(home.path()); let mut cache = message_cache::InputMessageCache::load().unwrap(); - let units = vec![InputUnit::plain_file(ClientId::Junie, path.clone())]; + let units = vec![InputUnit::plain_file( + path.clone(), + DecoderSpec::plain(DecoderId::Junie, JUNIE_RECORD_REJECTION_REVISION), + )]; let fresh = fold_with_adapter(units.clone(), &mut cache, None); let planned = JUNIE_ADAPTER @@ -209,13 +210,16 @@ mod tests { }; let mut cached = Vec::new(); + let binding = crate::adapters::adapter_for(ClientId::Junie).unwrap(); + let mut fold_ctx = FoldContext::new(binding, &mut cache, None); + let mut sink = BoundMessageSink::new(binding, &mut cached); JUNIE_ADAPTER - .fold(parsed, &mut FoldContext::new(&mut cache, None), &mut cached) + .fold(parsed, &mut fold_ctx, &mut sink) .unwrap(); assert!(cached .iter() - .all(|message| message.client.as_ref() == ClientId::Junie.as_str())); + .all(|message| message.client == ClientId::Junie)); assert_eq!(cached, fresh); restore_env_var("TOKSCALE_CONFIG_DIR", previous_config_dir); } @@ -237,7 +241,10 @@ mod tests { let mut cache = message_cache::InputMessageCache::default(); let messages = fold_with_adapter( - vec![InputUnit::plain_file(ClientId::Junie, path)], + vec![InputUnit::plain_file( + path, + DecoderSpec::plain(DecoderId::Junie, JUNIE_RECORD_REJECTION_REVISION), + )], &mut cache, Some(&pricing), ); diff --git a/crates/tokscale-core/src/adapters/kilo.rs b/crates/tokscale-core/src/adapters/kilo.rs index 374e54d57..fef9bffcd 100644 --- a/crates/tokscale-core/src/adapters/kilo.rs +++ b/crates/tokscale-core/src/adapters/kilo.rs @@ -3,11 +3,13 @@ use rayon::prelude::*; use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ - AdapterScanContext, FoldContext, InputDiscoveryError, InputUnit, LocalInputAdapter, - MessageSink, ParseContext, ParsedUnit, + AdapterScanContext, BoundMessageSink, DecoderSpec, FoldContext, InputDiscoveryError, InputUnit, + LocalInputAdapter, ParseContext, ParsedUnit, }; use crate::clients::ClientId; -use crate::message_cache::{ParserId, ParserVersion}; +use crate::message_cache::DecoderId; +#[cfg(test)] +use crate::message_cache::DecoderVersion; use crate::sessions; const KILO_RECORD_REJECTION_REVISION: u32 = 6; @@ -15,42 +17,28 @@ const KILO_RECORD_REJECTION_REVISION: u32 = 6; pub(crate) struct KiloAdapter; impl LocalInputAdapter for KiloAdapter { - fn client(&self) -> ClientId { - ClientId::Kilo - } - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { - let def = ClientId::Kilo + let def = client .local_def() .expect("Kilo adapter must have local scan policy"); let mut paths = Vec::new(); - adapter_discover::push_existing_file( - ClientId::Kilo, - def.resolve_path(ctx.home_dir), - &mut paths, - )?; + adapter_discover::push_existing_file(client, def.resolve_path(ctx.home_dir), &mut paths)?; paths.extend(adapter_discover::scan_roots( - ClientId::Kilo, - adapter_discover::extra_roots_for_client(ClientId::Kilo, ctx)?, + client, + adapter_discover::extra_roots_for_client(client, ctx)?, def.pattern, )?); - Ok(adapter_discover::input_units_from_paths_preserving_order( - ClientId::Kilo, + adapter_discover::input_units_from_paths_preserving_order( + client, paths, crate::adapters::FingerprintPolicy::SqliteWithWal, - )? - .into_iter() - .map(|unit| { - unit.with_parser_version(ParserVersion::new( - ParserId::Kilo, - KILO_RECORD_REJECTION_REVISION, - )) - }) - .collect()) + DecoderSpec::plain(DecoderId::Kilo, KILO_RECORD_REJECTION_REVISION), + ) } fn parse_checked(&self, units: Vec, ctx: &ParseContext<'_>) -> Vec { @@ -66,7 +54,7 @@ impl LocalInputAdapter for KiloAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), crate::adapters::InputPipelineError> { adapter_cache::fold_units(parsed, ctx, sink) } @@ -110,7 +98,7 @@ mod tests { scanner_settings: &settings, }; - let units = KILO_ADAPTER.discover_checked(&ctx).unwrap(); + let units = KILO_ADAPTER.discover_checked(ClientId::Kilo, &ctx).unwrap(); assert_eq!( units @@ -120,8 +108,8 @@ mod tests { vec![default_db, first_extra_db, second_extra_db] ); assert!(units.iter().all(|unit| { - unit.parser_version - == ParserVersion::new(ParserId::Kilo, KILO_RECORD_REJECTION_REVISION) + unit.decoder.version() + == DecoderVersion::new(DecoderId::Kilo, KILO_RECORD_REJECTION_REVISION) && unit.fingerprint_policy == crate::adapters::FingerprintPolicy::SqliteWithWal })); } diff --git a/crates/tokscale-core/src/adapters/kiro.rs b/crates/tokscale-core/src/adapters/kiro.rs index 984229193..80ab51b27 100644 --- a/crates/tokscale-core/src/adapters/kiro.rs +++ b/crates/tokscale-core/src/adapters/kiro.rs @@ -7,11 +7,12 @@ use walkdir::WalkDir; use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ - AdapterScanContext, FingerprintPolicy, FoldContext, InputDiscoveryError, InputUnit, - InputUnitMeta, LocalInputAdapter, MessageSink, ParseContext, ParsedUnit, + AdapterScanContext, BoundMessageSink, DecoderRoute, DecoderSpec, FingerprintPolicy, + FoldContext, InputDiscoveryError, InputUnit, LocalInputAdapter, ParseContext, ParsedUnit, }; use crate::clients::ClientId; -use crate::message_cache::{ParserId, ParserVersion}; +#[cfg(test)] +use crate::message_cache::{DecoderId, DecoderVersion}; use crate::sessions; const KIRO_RECORD_REJECTION_REVISION: u32 = 5; @@ -19,63 +20,44 @@ const KIRO_RECORD_REJECTION_REVISION: u32 = 5; pub(crate) struct KiroAdapter; impl LocalInputAdapter for KiroAdapter { - fn client(&self) -> ClientId { - ClientId::Kiro - } - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { - let def = ClientId::Kiro + let def = client .local_def() .expect("Kiro adapter must have local scan policy"); let mut units = adapter_discover::input_units_from_paths( - ClientId::Kiro, - adapter_discover::scan_roots( - ClientId::Kiro, - [def.resolve_path(ctx.home_dir)], - def.pattern, - )?, + client, + adapter_discover::scan_roots(client, [def.resolve_path(ctx.home_dir)], def.pattern)?, FingerprintPolicy::PlainFile, + DecoderSpec::kiro_file(KIRO_RECORD_REJECTION_REVISION), )? .into_iter() - .map(kiro_file_unit) + .map(configure_kiro_file_unit) .collect::>(); - if let Some(db_path) = kiro_db_path(ctx.home_dir)? { - units.push( - InputUnit::sqlite_with_wal(ClientId::Kiro, db_path) - .with_meta(InputUnitMeta::KiroSqlite) - .with_parser_version(ParserVersion::new( - ParserId::KiroSqlite, - KIRO_RECORD_REJECTION_REVISION, - )), - ); + if let Some(db_path) = kiro_db_path(client, ctx.home_dir)? { + units.push(InputUnit::sqlite_with_wal( + db_path, + DecoderSpec::kiro_sqlite(KIRO_RECORD_REJECTION_REVISION), + )); } - units.extend( - adapter_discover::input_units_from_paths( - ClientId::Kiro, - adapter_discover::scan_roots( - ClientId::Kiro, - kiro_global_storage_roots(ctx.home_dir), - "kiro-globalstorage", - )?, - FingerprintPolicy::PlainFile, - )? - .into_iter() - .map(|unit| { - unit.with_meta(InputUnitMeta::KiroGlobalStorage) - .with_parser_version(ParserVersion::new( - ParserId::KiroGlobalStorage, - KIRO_RECORD_REJECTION_REVISION, - )) - }), - ); + units.extend(adapter_discover::input_units_from_paths( + client, + adapter_discover::scan_roots( + client, + kiro_global_storage_roots(ctx.home_dir), + "kiro-globalstorage", + )?, + FingerprintPolicy::PlainFile, + DecoderSpec::kiro_global_storage(KIRO_RECORD_REJECTION_REVISION), + )?); - units.extend(kiro_extra_units(ctx)?); - dedup_units_by_canonical_path(&mut units)?; + units.extend(kiro_extra_units(client, ctx)?); + dedup_units_by_canonical_path(client, &mut units)?; units.sort_by(|left, right| left.path.cmp(&right.path)); Ok(units) } @@ -83,26 +65,26 @@ impl LocalInputAdapter for KiroAdapter { fn parse_checked(&self, units: Vec, ctx: &ParseContext<'_>) -> Vec { units .into_par_iter() - .map(|unit| match unit.meta { - InputUnitMeta::KiroFile => adapter_cache::load_or_scan_unit_with( + .map(|unit| match unit.decoder.route() { + DecoderRoute::KiroFile => adapter_cache::load_or_scan_unit_with( unit, ctx, sessions::kiro::parse_kiro_file, ), - InputUnitMeta::KiroSqlite => { + DecoderRoute::KiroSqlite => { adapter_cache::parse_uncached_unit(unit, ctx, sessions::kiro::parse_kiro_sqlite) } - InputUnitMeta::KiroGlobalStorage => adapter_cache::load_or_scan_unit_with( + DecoderRoute::KiroGlobalStorage => adapter_cache::load_or_scan_unit_with( unit, ctx, sessions::kiro::parse_kiro_file, ), - InputUnitMeta::None - | InputUnitMeta::AntigravityCliSqlite - | InputUnitMeta::OpenCodeSqlite - | InputUnitMeta::CodeBuddyJsonl - | InputUnitMeta::CodeBuddyExtensionLog { .. } - | InputUnitMeta::Codex => unreachable!("unexpected Kiro input unit meta"), + DecoderRoute::None + | DecoderRoute::AntigravityCliSqlite + | DecoderRoute::OpenCodeSqlite + | DecoderRoute::CodeBuddyJsonl + | DecoderRoute::CodeBuddyExtensionLog { .. } + | DecoderRoute::Codex => unreachable!("unexpected Kiro input unit meta"), }) .collect() } @@ -112,11 +94,11 @@ impl LocalInputAdapter for KiroAdapter { unit: InputUnit, input_cache: &crate::message_cache::InputMessageCache, ) -> Result { - match unit.meta { - InputUnitMeta::KiroFile | InputUnitMeta::KiroGlobalStorage => { + match unit.decoder.route() { + DecoderRoute::KiroFile | DecoderRoute::KiroGlobalStorage => { adapter_cache::plan_cache_hit(unit, input_cache) } - InputUnitMeta::KiroSqlite => Ok(crate::adapters::CacheHitPlan::Miss(unit)), + DecoderRoute::KiroSqlite => Ok(crate::adapters::CacheHitPlan::Miss(unit)), _ => unreachable!("unexpected Kiro input unit meta"), } } @@ -125,16 +107,16 @@ impl LocalInputAdapter for KiroAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), crate::adapters::InputPipelineError> { adapter_cache::fold_units(parsed, ctx, sink) } } -fn kiro_db_path(home_dir: &str) -> Result, InputDiscoveryError> { +fn kiro_db_path(client: ClientId, home_dir: &str) -> Result, InputDiscoveryError> { let xdg_path = PathBuf::from(format!("{}/.local/share/kiro-cli/data.sqlite3", home_dir)); let mut paths = Vec::new(); - adapter_discover::push_existing_file(ClientId::Kiro, xdg_path, &mut paths)?; + adapter_discover::push_existing_file(client, xdg_path, &mut paths)?; if let Some(path) = paths.pop() { return Ok(Some(path)); } @@ -143,7 +125,7 @@ fn kiro_db_path(home_dir: &str) -> Result, InputDiscoveryError> "{}/Library/Application Support/kiro-cli/data.sqlite3", home_dir )); - adapter_discover::push_existing_file(ClientId::Kiro, macos_path, &mut paths)?; + adapter_discover::push_existing_file(client, macos_path, &mut paths)?; Ok(paths.pop()) } @@ -176,28 +158,25 @@ fn kiro_global_storage_roots(home_dir: &str) -> Vec { ] } -fn kiro_file_unit(unit: InputUnit) -> InputUnit { +fn configure_kiro_file_unit(unit: InputUnit) -> InputUnit { let sidecar = unit.path.with_extension("jsonl"); unit.with_optional_dependency(sidecar) - .with_meta(InputUnitMeta::KiroFile) - .with_parser_version(ParserVersion::new( - ParserId::KiroFile, - KIRO_RECORD_REJECTION_REVISION, - )) } -fn kiro_extra_units(ctx: &AdapterScanContext<'_>) -> Result, InputDiscoveryError> { +fn kiro_extra_units( + client: ClientId, + ctx: &AdapterScanContext<'_>, +) -> Result, InputDiscoveryError> { let mut cli_paths = Vec::new(); let mut sqlite_paths = Vec::new(); let mut global_storage_paths = Vec::new(); - for root in adapter_discover::extra_roots_for_client(ClientId::Kiro, ctx)? { + for root in adapter_discover::extra_roots_for_client(client, ctx)? { match std::fs::metadata(&root) { Ok(_) => {} Err(source) if source.kind() == std::io::ErrorKind::NotFound => continue, Err(source) => { return Err(InputDiscoveryError::new( - ClientId::Kiro, &root, "read extra scan root metadata", source, @@ -207,7 +186,7 @@ fn kiro_extra_units(ctx: &AdapterScanContext<'_>) -> Result, Inpu for entry in WalkDir::new(&root) { let entry = entry.map_err(|source| { - InputDiscoveryError::new(ClientId::Kiro, &root, "walk extra scan root", source) + InputDiscoveryError::new(&root, "walk extra scan root", source) })?; if !entry.file_type().is_file() { continue; @@ -228,43 +207,26 @@ fn kiro_extra_units(ctx: &AdapterScanContext<'_>) -> Result, Inpu } let mut units = adapter_discover::input_units_from_paths( - ClientId::Kiro, + client, cli_paths, FingerprintPolicy::PlainFile, + DecoderSpec::kiro_file(KIRO_RECORD_REJECTION_REVISION), )? .into_iter() - .map(kiro_file_unit) + .map(configure_kiro_file_unit) .collect::>(); - units.extend( - adapter_discover::input_units_from_paths( - ClientId::Kiro, - sqlite_paths, - FingerprintPolicy::SqliteWithWal, - )? - .into_iter() - .map(|unit| { - unit.with_meta(InputUnitMeta::KiroSqlite) - .with_parser_version(ParserVersion::new( - ParserId::KiroSqlite, - KIRO_RECORD_REJECTION_REVISION, - )) - }), - ); - units.extend( - adapter_discover::input_units_from_paths( - ClientId::Kiro, - global_storage_paths, - FingerprintPolicy::PlainFile, - )? - .into_iter() - .map(|unit| { - unit.with_meta(InputUnitMeta::KiroGlobalStorage) - .with_parser_version(ParserVersion::new( - ParserId::KiroGlobalStorage, - KIRO_RECORD_REJECTION_REVISION, - )) - }), - ); + units.extend(adapter_discover::input_units_from_paths( + client, + sqlite_paths, + FingerprintPolicy::SqliteWithWal, + DecoderSpec::kiro_sqlite(KIRO_RECORD_REJECTION_REVISION), + )?); + units.extend(adapter_discover::input_units_from_paths( + client, + global_storage_paths, + FingerprintPolicy::PlainFile, + DecoderSpec::kiro_global_storage(KIRO_RECORD_REJECTION_REVISION), + )?); Ok(units) } @@ -287,17 +249,15 @@ fn is_kiro_global_storage_input(path: &Path) -> bool { .is_some_and(|extension| extension == "chat" || extension == "json") } -fn dedup_units_by_canonical_path(units: &mut Vec) -> Result<(), InputDiscoveryError> { +fn dedup_units_by_canonical_path( + _client: ClientId, + units: &mut Vec, +) -> Result<(), InputDiscoveryError> { let mut seen = HashSet::new(); let mut keys = Vec::with_capacity(units.len()); for unit in units.iter() { keys.push(std::fs::canonicalize(&unit.path).map_err(|source| { - InputDiscoveryError::new( - ClientId::Kiro, - &unit.path, - "canonicalize discovered input", - source, - ) + InputDiscoveryError::new(&unit.path, "canonicalize discovered input", source) })?); } let mut index = 0; @@ -316,23 +276,16 @@ mod tests { use super::*; fn kiro_global_unit(path: PathBuf) -> InputUnit { - InputUnit::plain_file(ClientId::Kiro, path) - .with_meta(InputUnitMeta::KiroGlobalStorage) - .with_parser_version(ParserVersion::new( - ParserId::KiroGlobalStorage, - KIRO_RECORD_REJECTION_REVISION, - )) + InputUnit::plain_file( + path, + DecoderSpec::kiro_global_storage(KIRO_RECORD_REJECTION_REVISION), + ) } fn kiro_file_unit(path: PathBuf) -> InputUnit { let sidecar = path.with_extension("jsonl"); - InputUnit::plain_file(ClientId::Kiro, path) + InputUnit::plain_file(path, DecoderSpec::kiro_file(KIRO_RECORD_REJECTION_REVISION)) .with_optional_dependency(sidecar) - .with_meta(InputUnitMeta::KiroFile) - .with_parser_version(ParserVersion::new( - ParserId::KiroFile, - KIRO_RECORD_REJECTION_REVISION, - )) } #[test] @@ -353,21 +306,24 @@ mod tests { scanner_settings: &settings, }; - let units = KIRO_ADAPTER.discover_checked(&ctx).unwrap(); + let units = KIRO_ADAPTER.discover_checked(ClientId::Kiro, &ctx).unwrap(); assert_eq!(units.len(), 3); assert!(units .iter() - .any(|unit| unit.path == file_path && matches!(unit.meta, InputUnitMeta::KiroFile))); + .any(|unit| unit.path == file_path + && matches!(unit.decoder.route(), DecoderRoute::KiroFile))); assert!(units .iter() - .any(|unit| unit.path == db_path && matches!(unit.meta, InputUnitMeta::KiroSqlite))); + .any(|unit| unit.path == db_path + && matches!(unit.decoder.route(), DecoderRoute::KiroSqlite))); assert!(units.iter().any(|unit| { - unit.path == global_path && matches!(unit.meta, InputUnitMeta::KiroGlobalStorage) + unit.path == global_path + && matches!(unit.decoder.route(), DecoderRoute::KiroGlobalStorage) })); let file_unit = units .iter() - .find(|unit| matches!(unit.meta, InputUnitMeta::KiroFile)) + .find(|unit| matches!(unit.decoder.route(), DecoderRoute::KiroFile)) .unwrap(); assert_eq!( file_unit.fingerprint_policy, @@ -379,19 +335,19 @@ mod tests { ); let global_unit = units .iter() - .find(|unit| matches!(unit.meta, InputUnitMeta::KiroGlobalStorage)) + .find(|unit| matches!(unit.decoder.route(), DecoderRoute::KiroGlobalStorage)) .unwrap(); assert_eq!(global_unit.fingerprint_policy, FingerprintPolicy::PlainFile); for unit in &units { - let parser_id = match unit.meta { - InputUnitMeta::KiroFile => ParserId::KiroFile, - InputUnitMeta::KiroSqlite => ParserId::KiroSqlite, - InputUnitMeta::KiroGlobalStorage => ParserId::KiroGlobalStorage, + let decoder_id = match unit.decoder.route() { + DecoderRoute::KiroFile => DecoderId::KiroFile, + DecoderRoute::KiroSqlite => DecoderId::KiroSqlite, + DecoderRoute::KiroGlobalStorage => DecoderId::KiroGlobalStorage, _ => unreachable!(), }; assert_eq!( - unit.parser_version, - ParserVersion::new(parser_id, KIRO_RECORD_REJECTION_REVISION) + unit.decoder.version(), + DecoderVersion::new(decoder_id, KIRO_RECORD_REJECTION_REVISION) ); } } @@ -422,11 +378,11 @@ mod tests { scanner_settings: &settings, }; - let units = KIRO_ADAPTER.discover_checked(&ctx).unwrap(); + let units = KIRO_ADAPTER.discover_checked(ClientId::Kiro, &ctx).unwrap(); assert_eq!(units.len(), 3); let cli_unit = units.iter().find(|unit| unit.path == cli_path).unwrap(); - assert_eq!(cli_unit.meta, InputUnitMeta::KiroFile); + assert_eq!(cli_unit.decoder.route(), DecoderRoute::KiroFile); assert_eq!( cli_unit.fingerprint_policy, FingerprintPolicy::PrimaryWithDependency { @@ -436,26 +392,26 @@ mod tests { } ); assert_eq!( - cli_unit.parser_version, - ParserVersion::new(ParserId::KiroFile, KIRO_RECORD_REJECTION_REVISION) + cli_unit.decoder.version(), + DecoderVersion::new(DecoderId::KiroFile, KIRO_RECORD_REJECTION_REVISION) ); let sqlite_unit = units.iter().find(|unit| unit.path == sqlite_path).unwrap(); - assert_eq!(sqlite_unit.meta, InputUnitMeta::KiroSqlite); + assert_eq!(sqlite_unit.decoder.route(), DecoderRoute::KiroSqlite); assert_eq!( sqlite_unit.fingerprint_policy, FingerprintPolicy::SqliteWithWal ); assert_eq!( - sqlite_unit.parser_version, - ParserVersion::new(ParserId::KiroSqlite, KIRO_RECORD_REJECTION_REVISION) + sqlite_unit.decoder.version(), + DecoderVersion::new(DecoderId::KiroSqlite, KIRO_RECORD_REJECTION_REVISION) ); let global_unit = units.iter().find(|unit| unit.path == global_path).unwrap(); - assert_eq!(global_unit.meta, InputUnitMeta::KiroGlobalStorage); + assert_eq!(global_unit.decoder.route(), DecoderRoute::KiroGlobalStorage); assert_eq!( - global_unit.parser_version, - ParserVersion::new(ParserId::KiroGlobalStorage, KIRO_RECORD_REJECTION_REVISION) + global_unit.decoder.version(), + DecoderVersion::new(DecoderId::KiroGlobalStorage, KIRO_RECORD_REJECTION_REVISION) ); } @@ -484,18 +440,20 @@ mod tests { let units = paths.into_iter().map(kiro_global_unit).collect(); let parsed = KIRO_ADAPTER.parse_checked(units, &ParseContext { pricing: None }); let mut cache = crate::message_cache::InputMessageCache::default(); - let mut ctx = FoldContext::new(&mut cache, None); let mut messages = Vec::new(); + let binding = crate::adapters::adapter_for(ClientId::Kiro).unwrap(); + let mut ctx = FoldContext::new(binding, &mut cache, None); + let mut sink = BoundMessageSink::new(binding, &mut messages); - KIRO_ADAPTER.fold(parsed, &mut ctx, &mut messages).unwrap(); + KIRO_ADAPTER.fold(parsed, &mut ctx, &mut sink).unwrap(); assert_eq!(messages.len(), 2); assert!(messages .iter() - .all(|message| message.client.as_ref() == ClientId::Kiro.as_str())); - assert_eq!(ctx.health.rejected_records(), 1); - assert_eq!(ctx.health.failed_inputs(), 0); - assert_eq!(ctx.health.partial_inputs(), 0); + .all(|message| message.client == ClientId::Kiro)); + assert_eq!(ctx.health().rejected_records(), 1); + assert_eq!(ctx.health().failed_inputs(), 0); + assert_eq!(ctx.health().partial_inputs(), 0); } #[test] @@ -507,7 +465,7 @@ mod tests { let parsed = KIRO_ADAPTER.parse_checked(vec![kiro_file_unit(path)], &ParseContext { pricing: None }); - let health = parsed[0].input_health(); + let health = &parsed[0].health; let failure = health.status.failure().unwrap(); assert!(matches!( health.status, @@ -537,11 +495,11 @@ mod tests { .unwrap(); std::fs::create_dir(&sidecar).unwrap(); let unit = kiro_file_unit(path.clone()); - let parser_version = unit.parser_version; + let decoder_version = unit.decoder.version(); let parsed = KIRO_ADAPTER.parse_checked(vec![unit], &ParseContext { pricing: None }); - let health = parsed[0].input_health(); + let health = &parsed[0].health; let failure = health.status.failure().unwrap(); assert!(matches!( health.status, @@ -555,16 +513,18 @@ mod tests { assert!(parsed[0].cache_write.is_none()); let mut cache = crate::message_cache::InputMessageCache::default(); - let mut ctx = FoldContext::new(&mut cache, None); let mut messages = Vec::new(); - KIRO_ADAPTER.fold(parsed, &mut ctx, &mut messages).unwrap(); + let binding = crate::adapters::adapter_for(ClientId::Kiro).unwrap(); + let mut ctx = FoldContext::new(binding, &mut cache, None); + let mut sink = BoundMessageSink::new(binding, &mut messages); + KIRO_ADAPTER.fold(parsed, &mut ctx, &mut sink).unwrap(); assert_eq!(messages.len(), 1); - assert_eq!(messages[0].client.as_ref(), ClientId::Kiro.as_str()); + assert_eq!(messages[0].client, ClientId::Kiro); assert_eq!(messages[0].tokens.input, 13); assert_eq!(messages[0].tokens.output, 5); - assert_eq!(ctx.health.partial_inputs(), 1); - assert!(cache.get_meta(&path, parser_version).unwrap().is_none()); + assert_eq!(ctx.health().partial_inputs(), 1); + assert!(cache.get_meta(&path, decoder_version).unwrap().is_none()); } #[test] @@ -578,7 +538,7 @@ mod tests { let mut cache = crate::message_cache::InputMessageCache::default(); cache.insert(crate::message_cache::CachedInputEntry::new_with_version( &path, - unit.parser_version, + unit.decoder.version(), unit.input_policy().fingerprint().unwrap(), vec![crate::sessions::ParsedMessage::new( "cached-model", diff --git a/crates/tokscale-core/src/adapters/mod.rs b/crates/tokscale-core/src/adapters/mod.rs index 1d54be280..eead28ce0 100644 --- a/crates/tokscale-core/src/adapters/mod.rs +++ b/crates/tokscale-core/src/adapters/mod.rs @@ -5,6 +5,7 @@ mod cline; mod codebuddy; mod codebuff; mod codex; +mod decoder; pub(crate) mod discover; pub(crate) mod error; pub(crate) mod file; @@ -19,6 +20,7 @@ mod opencode; mod omp; mod pi; mod roocode; +mod runtime; mod warp; mod zed; @@ -29,31 +31,34 @@ use rayon::prelude::*; use crate::clients::ClientId; use crate::input_health::{DataHealth, InputFailure, InputHealth, InputStatus, RejectionSummary}; -use crate::message_cache::{InputFileIdentity, ParserId, ParserRevision, ParserVersion}; +#[cfg(test)] +use crate::message_cache::DecoderId; +use crate::message_cache::{DecoderRevision, InputFileIdentity}; use crate::{message_cache, pricing, scanner, sessions::ParsedMessage, UnifiedMessage}; +pub(crate) use decoder::{DecoderRoute, DecoderSpec}; pub(crate) use error::{ InputDiscoveryError, InputParseError, InputPipelineError, InputPlanningError, }; +pub(crate) use runtime::{BoundMessageSink, FoldContext}; -pub(crate) const MODEL_ID_CANONICALIZATION_REVISION: ParserRevision = 3; +pub(crate) const MODEL_ID_CANONICALIZATION_REVISION: DecoderRevision = 3; // Record-level rejection and parsed Agent identity changes alter the cached // scan outcome, so old OpenCode shards must be rebuilt. -pub(crate) const OPENCODE_CURRENT_SQLITE_REVISION: ParserRevision = +pub(crate) const OPENCODE_CURRENT_SQLITE_REVISION: DecoderRevision = MODEL_ID_CANONICALIZATION_REVISION + 6; -pub(crate) const EXPLICIT_TOKEN_OVERFLOW_REVISION: ParserRevision = +pub(crate) const EXPLICIT_TOKEN_OVERFLOW_REVISION: DecoderRevision = MODEL_ID_CANONICALIZATION_REVISION + 1; -pub(crate) const ZED_RECORD_FILTER_REVISION: ParserRevision = EXPLICIT_TOKEN_OVERFLOW_REVISION + 2; +pub(crate) const ZED_RECORD_FILTER_REVISION: DecoderRevision = EXPLICIT_TOKEN_OVERFLOW_REVISION + 2; // Codex Agent roles use the shared neutral text normalizer. Keep cached // messages aligned when that parser-owned identity changes. -pub(crate) const CODEX_EXEC_IDENTITY_REVISION: ParserRevision = +pub(crate) const CODEX_EXEC_IDENTITY_REVISION: DecoderRevision = MODEL_ID_CANONICALIZATION_REVISION + 4; pub(crate) trait LocalInputAdapter: Sync { - fn client(&self) -> ClientId; - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError>; @@ -71,14 +76,14 @@ pub(crate) trait LocalInputAdapter: Sync { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), InputPipelineError>; fn fold_batches( &self, batches: &mut ParsedBatchInput<'_>, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), InputPipelineError> { while let Some(parsed) = batches.next(ctx)? { self.fold(parsed, ctx, sink)?; @@ -96,30 +101,11 @@ pub(crate) struct ParseContext<'a> { pub pricing: Option<&'a pricing::PricingService>, } -pub(crate) struct FoldContext<'a> { - pub input_cache: &'a mut message_cache::InputMessageCache, - pub pricing: Option<&'a pricing::PricingService>, - pub health: DataHealth, -} - -impl<'a> FoldContext<'a> { - pub(crate) fn new( - input_cache: &'a mut message_cache::InputMessageCache, - pricing: Option<&'a pricing::PricingService>, - ) -> Self { - Self { - input_cache, - pricing, - health: DataHealth::default(), - } - } -} - -pub(crate) trait MessageSink { +pub(crate) trait UnifiedMessageSink { fn push_message(&mut self, message: UnifiedMessage); } -impl MessageSink for Vec { +impl UnifiedMessageSink for Vec { fn push_message(&mut self, message: UnifiedMessage) { self.push(message); } @@ -127,11 +113,9 @@ impl MessageSink for Vec { #[derive(Debug, Clone)] pub(crate) struct InputUnit { - client: ClientId, pub path: PathBuf, pub fingerprint_policy: FingerprintPolicy, - pub meta: InputUnitMeta, - pub parser_version: ParserVersion, + pub decoder: DecoderSpec, prepared_snapshot: Option, snapshot_confirmed_for_execution: bool, planned_cache_meta: Option, @@ -139,18 +123,11 @@ pub(crate) struct InputUnit { } impl InputUnit { - /// Adapter-issued source provenance for this input. - pub(crate) const fn client(&self) -> ClientId { - self.client - } - - pub(crate) fn plain_file(client: ClientId, path: PathBuf) -> Self { + pub(crate) fn plain_file(path: PathBuf, decoder: DecoderSpec) -> Self { Self { - client, path, fingerprint_policy: FingerprintPolicy::PlainFile, - meta: InputUnitMeta::None, - parser_version: InputUnitMeta::None.parser_version(client), + decoder, prepared_snapshot: None, snapshot_confirmed_for_execution: false, planned_cache_meta: None, @@ -158,13 +135,11 @@ impl InputUnit { } } - pub(crate) fn sqlite_with_wal(client: ClientId, path: PathBuf) -> Self { + pub(crate) fn sqlite_with_wal(path: PathBuf, decoder: DecoderSpec) -> Self { Self { - client, path, fingerprint_policy: FingerprintPolicy::SqliteWithWal, - meta: InputUnitMeta::None, - parser_version: InputUnitMeta::None.parser_version(client), + decoder, prepared_snapshot: None, snapshot_confirmed_for_execution: false, planned_cache_meta: None, @@ -172,13 +147,11 @@ impl InputUnit { } } - pub(crate) fn no_message_cache(client: ClientId, path: PathBuf) -> Self { + pub(crate) fn no_message_cache(path: PathBuf, decoder: DecoderSpec) -> Self { Self { - client, path, fingerprint_policy: FingerprintPolicy::NoMessageCache, - meta: InputUnitMeta::None, - parser_version: InputUnitMeta::None.parser_version(client), + decoder, prepared_snapshot: None, snapshot_confirmed_for_execution: false, planned_cache_meta: None, @@ -186,16 +159,14 @@ impl InputUnit { } } - pub(crate) fn claude_code(client: ClientId, path: PathBuf, home_dir: PathBuf) -> Self { + pub(crate) fn claude_code(path: PathBuf, home_dir: PathBuf, decoder: DecoderSpec) -> Self { Self { - client, path, fingerprint_policy: FingerprintPolicy::ClaudeCodeWithHome { home_dir, parent_session_path: None, }, - meta: InputUnitMeta::None, - parser_version: InputUnitMeta::None.parser_version(client), + decoder, prepared_snapshot: None, snapshot_confirmed_for_execution: false, planned_cache_meta: None, @@ -203,17 +174,6 @@ impl InputUnit { } } - pub(crate) fn with_meta(mut self, meta: InputUnitMeta) -> Self { - self.parser_version = meta.parser_version(self.client); - self.meta = meta; - self - } - - pub(crate) fn with_parser_version(mut self, parser_version: ParserVersion) -> Self { - self.parser_version = parser_version; - self - } - pub(crate) fn with_dependency(mut self, dependency_path: PathBuf) -> Self { self.fingerprint_policy = FingerprintPolicy::PrimaryWithDependency { dependency_path, @@ -334,13 +294,12 @@ impl InputUnit { .prepared_snapshot .as_ref() .expect("inventory units must carry a prepared input snapshot"); - message_cache::hash_inventory_bytes(hasher, self.client.as_str().as_bytes()); message_cache::hash_inventory_bytes( hasher, - self.parser_version.parser_id.stable_name().as_bytes(), + self.decoder.version().decoder_id.stable_name().as_bytes(), ); - hasher.update(self.parser_version.revision.to_le_bytes()); - self.update_meta_inventory_signature(hasher); + hasher.update(self.decoder.version().revision.to_le_bytes()); + self.update_decoder_inventory_signature(hasher); self.update_policy_inventory_signature(hasher); self.input_policy() .update_inventory_signature(snapshot, hasher); @@ -355,23 +314,23 @@ impl InputUnit { hasher.finalize().into() } - fn update_meta_inventory_signature(&self, hasher: &mut sha2::Sha256) { - let (name, detail) = match self.meta { - InputUnitMeta::None => ("none", None), - InputUnitMeta::OpenCodeSqlite => ("opencode-sqlite", None), - InputUnitMeta::AntigravityCliSqlite => ("antigravity-cli-sqlite", None), - InputUnitMeta::KiroFile => ("kiro-file", None), - InputUnitMeta::KiroSqlite => ("kiro-sqlite", None), - InputUnitMeta::KiroGlobalStorage => ("kiro-global-storage", None), - InputUnitMeta::CodeBuddyJsonl => ("codebuddy-jsonl", None), - InputUnitMeta::CodeBuddyExtensionLog { origin } => ( + fn update_decoder_inventory_signature(&self, hasher: &mut sha2::Sha256) { + let (name, detail) = match self.decoder.route() { + DecoderRoute::None => ("none", None), + DecoderRoute::OpenCodeSqlite => ("opencode-sqlite", None), + DecoderRoute::AntigravityCliSqlite => ("antigravity-cli-sqlite", None), + DecoderRoute::KiroFile => ("kiro-file", None), + DecoderRoute::KiroSqlite => ("kiro-sqlite", None), + DecoderRoute::KiroGlobalStorage => ("kiro-global-storage", None), + DecoderRoute::CodeBuddyJsonl => ("codebuddy-jsonl", None), + DecoderRoute::CodeBuddyExtensionLog { origin } => ( "codebuddy-extension-log", Some(match origin { CodeBuddyLogOrigin::Extension => "extension", CodeBuddyLogOrigin::Host => "host", }), ), - InputUnitMeta::Codex => ("codex", None), + DecoderRoute::Codex => ("codex", None), }; message_cache::hash_inventory_bytes(hasher, name.as_bytes()); message_cache::hash_inventory_bytes(hasher, detail.unwrap_or("").as_bytes()); @@ -473,93 +432,12 @@ impl InputUnit { } } -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] -pub(crate) enum InputUnitMeta { - #[default] - None, - OpenCodeSqlite, - AntigravityCliSqlite, - KiroFile, - KiroSqlite, - KiroGlobalStorage, - CodeBuddyJsonl, - CodeBuddyExtensionLog { - origin: CodeBuddyLogOrigin, - }, - Codex, -} - #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum CodeBuddyLogOrigin { Extension, Host, } -impl InputUnitMeta { - fn parser_version(&self, client: ClientId) -> ParserVersion { - match self { - Self::None => ParserVersion::new( - default_parser_id(client), - MODEL_ID_CANONICALIZATION_REVISION, - ), - Self::OpenCodeSqlite => { - ParserVersion::new(ParserId::OpenCodeSqlite, OPENCODE_CURRENT_SQLITE_REVISION) - } - Self::AntigravityCliSqlite => ParserVersion::new( - ParserId::AntigravityCliSqlite, - EXPLICIT_TOKEN_OVERFLOW_REVISION, - ), - Self::KiroFile => { - ParserVersion::new(ParserId::KiroFile, MODEL_ID_CANONICALIZATION_REVISION) - } - Self::KiroSqlite => { - ParserVersion::new(ParserId::KiroSqlite, MODEL_ID_CANONICALIZATION_REVISION) - } - Self::KiroGlobalStorage => ParserVersion::new( - ParserId::KiroGlobalStorage, - MODEL_ID_CANONICALIZATION_REVISION, - ), - Self::CodeBuddyJsonl | Self::CodeBuddyExtensionLog { .. } => { - ParserVersion::new(ParserId::CodeBuddy, MODEL_ID_CANONICALIZATION_REVISION) - } - Self::Codex => ParserVersion::new(ParserId::Codex, CODEX_EXEC_IDENTITY_REVISION), - } - } -} - -fn default_parser_id(client: ClientId) -> ParserId { - match client { - ClientId::OpenCode => ParserId::OpenCodeSqlite, - ClientId::Claude => ParserId::Claude, - ClientId::Codex => ParserId::Codex, - ClientId::Gemini => ParserId::Gemini, - ClientId::Amp => ParserId::Amp, - ClientId::Droid => ParserId::Droid, - ClientId::OpenClaw => ParserId::OpenClaw, - ClientId::Pi => ParserId::Pi, - ClientId::Omp => ParserId::Omp, - ClientId::Kimi => ParserId::Kimi, - ClientId::Qwen => ParserId::Qwen, - ClientId::RooCode => ParserId::RooCode, - ClientId::Mux => ParserId::Mux, - ClientId::Kilo => ParserId::Kilo, - ClientId::Hermes => ParserId::Hermes, - ClientId::Copilot => ParserId::Copilot, - ClientId::Goose => ParserId::Goose, - ClientId::Codebuff => ParserId::Codebuff, - ClientId::CodeBuddy => ParserId::CodeBuddy, - ClientId::Antigravity => ParserId::AntigravityCliSqlite, - ClientId::Zed => ParserId::Zed, - ClientId::Zcode => ParserId::Zcode, - ClientId::Kiro => ParserId::Kiro, - ClientId::Junie => ParserId::Junie, - ClientId::Cline => ParserId::Cline, - ClientId::CommandCode => ParserId::CommandCode, - ClientId::Grok => ParserId::Grok, - ClientId::Warp => ParserId::Warp, - } -} - #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) enum FingerprintPolicy { PlainFile, @@ -640,64 +518,75 @@ impl ParsedUnit { }), } } +} - pub(crate) fn input_health(&self) -> InputHealth { - InputHealth { - client: self.unit.client, - path: self.unit.path.clone(), - status: self.health.status.clone(), - rejections: self.health.rejections.clone(), - } +#[derive(Clone, Copy)] +pub(crate) struct AdapterBinding<'a> { + client: ClientId, + adapter: &'a dyn LocalInputAdapter, +} + +impl<'a> AdapterBinding<'a> { + const fn new(client: ClientId, adapter: &'a dyn LocalInputAdapter) -> Self { + Self { client, adapter } + } + + pub(crate) const fn client(self) -> ClientId { + self.client + } + + pub(crate) const fn adapter(self) -> &'a dyn LocalInputAdapter { + self.adapter } } -static LOCAL_INPUT_ADAPTERS: [&dyn LocalInputAdapter; 28] = [ - &zed::ZED_ADAPTER, - &pi::PI_ADAPTER, - &omp::OMP_ADAPTER, - &claude::CLAUDE_ADAPTER, - &codex::CODEX_ADAPTER, - &opencode::OPENCODE_ADAPTER, - &file::COPILOT_ADAPTER, - &file::GEMINI_ADAPTER, - &file::GROK_ADAPTER, - &file::AMP_ADAPTER, - &file::DROID_ADAPTER, - &file::KIMI_ADAPTER, - &file::QWEN_ADAPTER, - &file::MUX_ADAPTER, - &codebuff::CODEBUFF_ADAPTER, - &codebuddy::CODEBUDDY_ADAPTER, - &openclaw::OPENCLAW_ADAPTER, - &roocode::ROOCODE_ADAPTER, - &cline::CLINE_ADAPTER, - &antigravity::ANTIGRAVITY_ADAPTER, - &kilo::KILO_ADAPTER, - &hermes::HERMES_ADAPTER, - &goose::GOOSE_ADAPTER, - &file::ZCODE_ADAPTER, - &kiro::KIRO_ADAPTER, - &junie::JUNIE_ADAPTER, - &file::COMMANDCODE_ADAPTER, - &warp::WARP_ADAPTER, +static LOCAL_INPUT_ADAPTERS: &[AdapterBinding<'static>] = &[ + AdapterBinding::new(ClientId::Zed, &zed::ZED_ADAPTER), + AdapterBinding::new(ClientId::Pi, &pi::PI_ADAPTER), + AdapterBinding::new(ClientId::Omp, &omp::OMP_ADAPTER), + AdapterBinding::new(ClientId::Claude, &claude::CLAUDE_ADAPTER), + AdapterBinding::new(ClientId::Codex, &codex::CODEX_ADAPTER), + AdapterBinding::new(ClientId::OpenCode, &opencode::OPENCODE_ADAPTER), + AdapterBinding::new(ClientId::Copilot, &file::COPILOT_ADAPTER), + AdapterBinding::new(ClientId::Gemini, &file::GEMINI_ADAPTER), + AdapterBinding::new(ClientId::Grok, &file::GROK_ADAPTER), + AdapterBinding::new(ClientId::Amp, &file::AMP_ADAPTER), + AdapterBinding::new(ClientId::Droid, &file::DROID_ADAPTER), + AdapterBinding::new(ClientId::Kimi, &file::KIMI_ADAPTER), + AdapterBinding::new(ClientId::Qwen, &file::QWEN_ADAPTER), + AdapterBinding::new(ClientId::Mux, &file::MUX_ADAPTER), + AdapterBinding::new(ClientId::Codebuff, &codebuff::CODEBUFF_ADAPTER), + AdapterBinding::new(ClientId::CodeBuddy, &codebuddy::CODEBUDDY_ADAPTER), + AdapterBinding::new(ClientId::OpenClaw, &openclaw::OPENCLAW_ADAPTER), + AdapterBinding::new(ClientId::RooCode, &roocode::ROOCODE_ADAPTER), + AdapterBinding::new(ClientId::Cline, &cline::CLINE_ADAPTER), + AdapterBinding::new(ClientId::Antigravity, &antigravity::ANTIGRAVITY_ADAPTER), + AdapterBinding::new(ClientId::Kilo, &kilo::KILO_ADAPTER), + AdapterBinding::new(ClientId::Hermes, &hermes::HERMES_ADAPTER), + AdapterBinding::new(ClientId::Goose, &goose::GOOSE_ADAPTER), + AdapterBinding::new(ClientId::Zcode, &file::ZCODE_ADAPTER), + AdapterBinding::new(ClientId::Kiro, &kiro::KIRO_ADAPTER), + AdapterBinding::new(ClientId::Junie, &junie::JUNIE_ADAPTER), + AdapterBinding::new(ClientId::CommandCode, &file::COMMANDCODE_ADAPTER), + AdapterBinding::new(ClientId::Warp, &warp::WARP_ADAPTER), ]; -pub(crate) fn local_input_adapters() -> &'static [&'static dyn LocalInputAdapter] { - &LOCAL_INPUT_ADAPTERS +pub(crate) fn local_input_adapters() -> &'static [AdapterBinding<'static>] { + LOCAL_INPUT_ADAPTERS } -pub(crate) fn adapter_for(client: ClientId) -> Option<&'static dyn LocalInputAdapter> { +pub(crate) fn adapter_for(client: ClientId) -> Option> { local_input_adapters() .iter() .copied() - .find(|adapter| adapter.client() == client) + .find(|binding| binding.client == client) } pub(crate) fn selected_adapters( - clients: &[String], -) -> Result, String> { + clients: &[ClientId], +) -> Result>, String> { let include_all = clients.is_empty(); - let requested = requested_client_ids(clients)?; + let requested: HashSet = clients.iter().copied().collect(); let missing = ClientId::iter().find(|client| { (include_all || requested.contains(client)) && adapter_for(*client).is_none() @@ -712,12 +601,12 @@ pub(crate) fn selected_adapters( Ok(local_input_adapters() .iter() .copied() - .filter(|adapter| include_all || requested.contains(&adapter.client())) + .filter(|binding| include_all || requested.contains(&binding.client)) .collect()) } pub(crate) struct PreparedAdapterInputs { - pub adapter: &'static dyn LocalInputAdapter, + pub binding: AdapterBinding<'static>, pub units: Vec, } @@ -727,24 +616,8 @@ pub(crate) struct ConfirmedAdapterInputs { pub present_files: Vec<(InputFileIdentity, u64)>, } -pub(crate) fn validate_discovered_unit_sources( - adapter: &dyn LocalInputAdapter, - units: &[InputUnit], -) -> Result<(), InputPipelineError> { - let expected = adapter.client(); - if let Some(unit) = units.iter().find(|unit| unit.client != expected) { - return Err(InputPipelineError::contract(format!( - "local input adapter `{}` discovered `{}` input `{}`", - expected.as_str(), - unit.client.as_str(), - unit.path.display() - ))); - } - Ok(()) -} - pub(crate) struct ParsedBatchInput<'a> { - adapter: &'a dyn LocalInputAdapter, + binding: AdapterBinding<'a>, units: Option>, planned: VecDeque, confirmed_inventory_digests: Vec<[u8; 32]>, @@ -770,9 +643,9 @@ enum BatchSlot { } impl<'a> ParsedBatchInput<'a> { - fn new(adapter: &'a dyn LocalInputAdapter, units: Vec) -> Self { + fn new(binding: AdapterBinding<'a>, units: Vec) -> Self { Self { - adapter, + binding, units: Some(units), planned: VecDeque::new(), confirmed_inventory_digests: Vec::new(), @@ -817,7 +690,7 @@ impl<'a> ParsedBatchInput<'a> { let parsed_misses = if miss_units.is_empty() { Vec::new() } else { - self.adapter.parse_checked( + self.binding.adapter().parse_checked( miss_units, &ParseContext { pricing: ctx.pricing, @@ -892,11 +765,10 @@ impl<'a> ParsedBatchInput<'a> { let planned: Vec = units .into_par_iter() .map(|mut unit| { - let client = unit.client; let path = unit.path.clone(); if let Err(source) = unit.revalidate_snapshot_for_cache_decision() { return PlannedOrFailed::Failed(InputHealth { - client, + client: self.binding.client(), path, status: InputStatus::Unavailable { failure: InputFailure::new( @@ -912,7 +784,11 @@ impl<'a> ParsedBatchInput<'a> { unit.prepared_input_snapshot() .expect("revalidated input unit must retain its confirmed snapshot") .visit_present_files(|identity, size| present_files.push((identity, size))); - match self.adapter.plan_cache_hit(unit, &*ctx.input_cache) { + match self + .binding + .adapter() + .plan_cache_hit(unit, &*ctx.input_cache) + { Ok(plan) => { let planned = match plan { CacheHitPlan::Hit(parsed) => PlannedInputUnit::Hit(parsed), @@ -945,24 +821,23 @@ pub(crate) fn run_prepared_local_input_adapters( prepared: Vec, input_cache: &mut message_cache::InputMessageCache, pricing: Option<&pricing::PricingService>, - sink: &mut dyn MessageSink, + sink: &mut dyn UnifiedMessageSink, health: &mut DataHealth, ) -> Result, InputPipelineError> { - for group in &prepared { - validate_discovered_unit_sources(group.adapter, &group.units)?; - } - let mut confirmed = Vec::with_capacity(prepared.len()); - for PreparedAdapterInputs { adapter, units } in prepared { - let mut batches = ParsedBatchInput::new(adapter, units); - let mut fold_ctx = FoldContext::new(input_cache, pricing); - adapter.fold_batches(&mut batches, &mut fold_ctx, sink)?; - health.merge(std::mem::take(&mut fold_ctx.health)); + for PreparedAdapterInputs { binding, units } in prepared { + let mut batches = ParsedBatchInput::new(binding, units); + let mut fold_ctx = FoldContext::new(binding, input_cache, pricing); + let mut bound_sink = BoundMessageSink::new(binding, sink); + binding + .adapter() + .fold_batches(&mut batches, &mut fold_ctx, &mut bound_sink)?; + health.merge(fold_ctx.take_health()); for failed in batches.take_failed_health() { health.record(failed); } confirmed.push(ConfirmedAdapterInputs { - client: adapter.client(), + client: binding.client(), unit_digests: batches.confirmed_inventory_digests, present_files: batches.confirmed_present_files.into_iter().collect(), }); @@ -970,15 +845,6 @@ pub(crate) fn run_prepared_local_input_adapters( Ok(confirmed) } -fn requested_client_ids(clients: &[String]) -> Result, String> { - clients - .iter() - .map(|client| { - ClientId::from_str(client).ok_or_else(|| format!("unknown local client `{client}`")) - }) - .collect() -} - #[cfg(test)] mod tests { use std::collections::HashSet; @@ -1000,6 +866,10 @@ mod tests { planner_calls: AtomicUsize, } + fn amp_decoder() -> DecoderSpec { + DecoderSpec::plain(DecoderId::Amp, 0) + } + #[test] fn local_adapter_registry_is_unique_and_covers_catalog() { let adapters: Vec = local_input_adapters() @@ -1021,32 +891,28 @@ mod tests { } #[test] - fn adapter_rejects_input_unit_with_foreign_source_identity() { - let prepared = PreparedAdapterInputs { - adapter: adapter_for(ClientId::Amp).unwrap(), - units: vec![InputUnit::plain_file( - ClientId::Codex, - PathBuf::from("foreign-input.jsonl"), - )], - }; - - let error = match validate_discovered_unit_sources(prepared.adapter, &prepared.units) { - Ok(_) => panic!("foreign source identity must be rejected"), - Err(error) => error, - }; + fn bound_sink_attributes_messages_with_the_registry_client() { + let binding = adapter_for(ClientId::Amp).unwrap(); + let mut messages = Vec::new(); + let mut sink = BoundMessageSink::new(binding, &mut messages); + + sink.emit(ParsedMessage::new( + "model", + "provider", + "session", + 1, + crate::TokenBreakdown::default(), + 0.0, + )); - assert!(error - .to_string() - .contains("local input adapter `amp` discovered `codex` input `foreign-input.jsonl`")); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].client, ClientId::Amp); } impl LocalInputAdapter for RecordingAdapter { - fn client(&self) -> ClientId { - ClientId::Amp - } - fn discover_checked( &self, + _client: ClientId, _ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { unreachable!("test adapter does not discover inputs") @@ -1075,19 +941,16 @@ mod tests { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), InputPipelineError> { cache::fold_units(parsed, ctx, sink) } } impl LocalInputAdapter for BatchLifetimeAdapter { - fn client(&self) -> ClientId { - ClientId::Amp - } - fn discover_checked( &self, + _client: ClientId, _ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { unreachable!("test adapter does not discover inputs") @@ -1129,19 +992,16 @@ mod tests { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), InputPipelineError> { cache::fold_units(parsed, ctx, sink) } } impl LocalInputAdapter for PlannedWeaveAdapter { - fn client(&self) -> ClientId { - ClientId::Amp - } - fn discover_checked( &self, + _client: ClientId, _ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { unreachable!("test adapter does not discover inputs") @@ -1181,7 +1041,7 @@ mod tests { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), InputPipelineError> { cache::fold_units(parsed, ctx, sink) } @@ -1189,7 +1049,7 @@ mod tests { struct DroppingSink; - impl MessageSink for DroppingSink { + impl UnifiedMessageSink for DroppingSink { fn push_message(&mut self, _message: UnifiedMessage) {} } @@ -1214,19 +1074,18 @@ mod tests { let units = input_paths .iter() .cloned() - .map(|path| InputUnit::plain_file(ClientId::Amp, path)) + .map(|path| InputUnit::plain_file(path, amp_decoder())) .collect(); let sessions = pool.install(|| { let mut cache = message_cache::InputMessageCache::default(); let mut sink = Vec::new(); - let mut batches = ParsedBatchInput::new(&adapter, units); + let binding = AdapterBinding::new(ClientId::Amp, &adapter); + let mut batches = ParsedBatchInput::new(binding, units); + let mut fold_ctx = FoldContext::new(binding, &mut cache, None); + let mut bound_sink = BoundMessageSink::new(binding, &mut sink); adapter - .fold_batches( - &mut batches, - &mut FoldContext::new(&mut cache, None), - &mut sink, - ) + .fold_batches(&mut batches, &mut fold_ctx, &mut bound_sink) .unwrap(); sink.into_iter() .map(|message| message.session_id.to_string()) @@ -1259,17 +1118,17 @@ mod tests { .map(|index| { let path = dir.path().join(index.to_string()); std::fs::write(&path, format!("input {index}")).unwrap(); - InputUnit::plain_file(ClientId::Amp, path) + InputUnit::plain_file(path, amp_decoder()) }) .collect(); let mut cache = message_cache::InputMessageCache::default(); - let mut batches = ParsedBatchInput::new(&adapter, units); + let binding = AdapterBinding::new(ClientId::Amp, &adapter); + let mut batches = ParsedBatchInput::new(binding, units); + let mut fold_ctx = FoldContext::new(binding, &mut cache, None); + let mut dropping_sink = DroppingSink; + let mut bound_sink = BoundMessageSink::new(binding, &mut dropping_sink); adapter - .fold_batches( - &mut batches, - &mut FoldContext::new(&mut cache, None), - &mut DroppingSink, - ) + .fold_batches(&mut batches, &mut fold_ctx, &mut bound_sink) .unwrap(); }); } @@ -1286,7 +1145,7 @@ mod tests { .map(|index| { let path = dir.path().join(index.to_string()); std::fs::write(&path, format!("input {index}")).unwrap(); - InputUnit::plain_file(ClientId::Amp, path) + InputUnit::plain_file(path, amp_decoder()) .prepare_snapshot() .unwrap() }) @@ -1300,7 +1159,7 @@ mod tests { let unit = &units[index]; cache.insert(message_cache::CachedInputEntry::new_with_version( &unit.path, - unit.parser_version, + unit.decoder.version(), unit.input_policy().fingerprint().unwrap(), vec![ParsedMessage::new( "model", @@ -1317,13 +1176,12 @@ mod tests { )); } let mut sink = Vec::new(); - let mut batches = ParsedBatchInput::new(&adapter, units); + let binding = AdapterBinding::new(ClientId::Amp, &adapter); + let mut batches = ParsedBatchInput::new(binding, units); + let mut fold_ctx = FoldContext::new(binding, &mut cache, None); + let mut bound_sink = BoundMessageSink::new(binding, &mut sink); adapter - .fold_batches( - &mut batches, - &mut FoldContext::new(&mut cache, None), - &mut sink, - ) + .fold_batches(&mut batches, &mut fold_ctx, &mut bound_sink) .unwrap(); assert_eq!(adapter.planner_calls.load(Ordering::Relaxed), 6); @@ -1349,7 +1207,7 @@ mod tests { .map(|index| { let path = dir.path().join(index.to_string()); std::fs::write(&path, format!("input {index}")).unwrap(); - InputUnit::plain_file(ClientId::Amp, path) + InputUnit::plain_file(path, amp_decoder()) .prepare_snapshot() .unwrap() }) @@ -1362,7 +1220,7 @@ mod tests { for (index, unit) in units.iter().enumerate() { cache.insert(message_cache::CachedInputEntry::new_with_version( &unit.path, - unit.parser_version, + unit.decoder.version(), unit.input_policy().fingerprint().unwrap(), vec![ParsedMessage::new( "model", @@ -1380,13 +1238,12 @@ mod tests { message_cache::reset_input_read_stats(&unit.path); } let mut sink = Vec::new(); - let mut batches = ParsedBatchInput::new(&adapter, units); + let binding = AdapterBinding::new(ClientId::Amp, &adapter); + let mut batches = ParsedBatchInput::new(binding, units); + let mut fold_ctx = FoldContext::new(binding, &mut cache, None); + let mut bound_sink = BoundMessageSink::new(binding, &mut sink); adapter - .fold_batches( - &mut batches, - &mut FoldContext::new(&mut cache, None), - &mut sink, - ) + .fold_batches(&mut batches, &mut fold_ctx, &mut bound_sink) .unwrap(); assert_eq!(adapter.planner_calls.load(Ordering::Relaxed), 4); @@ -1406,7 +1263,7 @@ mod tests { let cache_dir = tempfile::TempDir::new().unwrap(); let path = input_dir.path().join("future-cache-input"); std::fs::write(&path, b"input").unwrap(); - let unit = InputUnit::plain_file(ClientId::Amp, path.clone()) + let unit = InputUnit::plain_file(path.clone(), amp_decoder()) .prepare_snapshot() .unwrap(); let adapter = PlannedWeaveAdapter { @@ -1416,7 +1273,7 @@ mod tests { let mut seeded = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); seeded.insert(message_cache::CachedInputEntry::new_with_version( &path, - unit.parser_version, + unit.decoder.version(), unit.input_policy().fingerprint().unwrap(), vec![ParsedMessage::new( "model", @@ -1432,17 +1289,17 @@ mod tests { message_cache::mark_current_key_shard_as_future_format_for_test( cache_dir.path(), &path, - unit.parser_version, + unit.decoder.version(), ); let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); - let mut batches = ParsedBatchInput::new(&adapter, vec![unit]); + let binding = AdapterBinding::new(ClientId::Amp, &adapter); + let mut batches = ParsedBatchInput::new(binding, vec![unit]); + let mut fold_ctx = FoldContext::new(binding, &mut cache, None); + let mut dropping_sink = DroppingSink; + let mut bound_sink = BoundMessageSink::new(binding, &mut dropping_sink); adapter - .fold_batches( - &mut batches, - &mut FoldContext::new(&mut cache, None), - &mut DroppingSink, - ) + .fold_batches(&mut batches, &mut fold_ctx, &mut bound_sink) .expect("a newer cache shard must be discarded and reparsed"); } @@ -1451,13 +1308,13 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("direct-input"); std::fs::write(&path, b"direct input").unwrap(); - let unit = InputUnit::plain_file(ClientId::Amp, path.clone()) + let unit = InputUnit::plain_file(path.clone(), amp_decoder()) .prepare_snapshot() .unwrap(); let mut cache = message_cache::InputMessageCache::default(); cache.insert(message_cache::CachedInputEntry::new_with_version( &path, - unit.parser_version, + unit.decoder.version(), unit.input_policy().fingerprint().unwrap(), vec![ParsedMessage::new( "model", @@ -1476,14 +1333,13 @@ mod tests { batch_sizes: Mutex::new(Vec::new()), }; let mut sink = Vec::new(); - let mut batches = ParsedBatchInput::new(&adapter, vec![unit]); + let binding = AdapterBinding::new(ClientId::Amp, &adapter); + let mut batches = ParsedBatchInput::new(binding, vec![unit]); + let mut fold_ctx = FoldContext::new(binding, &mut cache, None); + let mut bound_sink = BoundMessageSink::new(binding, &mut sink); adapter - .fold_batches( - &mut batches, - &mut FoldContext::new(&mut cache, None), - &mut sink, - ) + .fold_batches(&mut batches, &mut fold_ctx, &mut bound_sink) .unwrap(); assert_eq!(*adapter.batch_sizes.lock().unwrap(), [1]); @@ -1499,10 +1355,16 @@ mod tests { let adapter = RecordingAdapter { batch_sizes: Mutex::new(Vec::new()), }; - let mut batches = - ParsedBatchInput::new(&adapter, vec![InputUnit::plain_file(ClientId::Amp, path)]); + let mut batches = ParsedBatchInput::new( + AdapterBinding::new(ClientId::Amp, &adapter), + vec![InputUnit::plain_file(path, amp_decoder())], + ); let mut cache = message_cache::InputMessageCache::default(); - let fold_context = FoldContext::new(&mut cache, None); + let fold_context = FoldContext::new( + AdapterBinding::new(ClientId::Amp, &adapter), + &mut cache, + None, + ); let planned = batches.take_all_planned_units(&fold_context).unwrap(); @@ -1520,7 +1382,7 @@ mod tests { #[test] fn plain_db_input_does_not_guess_a_wal_input() { let path = PathBuf::from("/tmp/plain-history.db"); - let unit = InputUnit::plain_file(ClientId::Amp, path.clone()); + let unit = InputUnit::plain_file(path.clone(), amp_decoder()); assert_eq!(unit.digest_paths(), vec![path]); } @@ -1532,8 +1394,9 @@ mod tests { let dependency = dir.path().join("parent.jsonl"); std::fs::write(&primary, b"child").unwrap(); - let mut unit = InputUnit::plain_file(ClientId::Omp, primary.clone()) - .with_dependency(dependency.clone()); + let mut unit = + InputUnit::plain_file(primary.clone(), DecoderSpec::plain(DecoderId::Omp, 0)) + .with_dependency(dependency.clone()); assert_eq!( unit.digest_paths(), vec![primary.clone(), dependency.clone()] @@ -1560,12 +1423,12 @@ mod tests { adapter_for(ClientId::Warp).map(|adapter| adapter.client()), Some(ClientId::Warp) ); - assert_eq!(selected_adapters(&["warp".to_string()]).unwrap().len(), 1); + assert_eq!(selected_adapters(&[ClientId::Warp]).unwrap().len(), 1); } #[test] fn antigravity_uses_one_adapter_for_all_local_inputs() { - let adapters = selected_adapters(&["antigravity".to_string()]).unwrap(); + let adapters = selected_adapters(&[ClientId::Antigravity]).unwrap(); assert_eq!(adapters.len(), 1); assert_eq!(adapters[0].client(), ClientId::Antigravity); @@ -1573,7 +1436,7 @@ mod tests { #[test] fn kilo_uses_one_current_runtime_adapter() { - let adapters = selected_adapters(&["kilo".to_string()]).unwrap(); + let adapters = selected_adapters(&[ClientId::Kilo]).unwrap(); assert_eq!(adapters.len(), 1); assert_eq!(adapters[0].client(), ClientId::Kilo); diff --git a/crates/tokscale-core/src/adapters/omp.rs b/crates/tokscale-core/src/adapters/omp.rs index 3a01a9a8d..c789a8b78 100644 --- a/crates/tokscale-core/src/adapters/omp.rs +++ b/crates/tokscale-core/src/adapters/omp.rs @@ -6,12 +6,14 @@ use rayon::prelude::*; use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ - AdapterScanContext, CacheHitPlan, FingerprintPolicy, FoldContext, InputDiscoveryError, - InputPipelineError, InputPlanningError, InputUnit, LocalInputAdapter, MessageSink, - ParseContext, ParsedBatchInput, ParsedUnit, + AdapterScanContext, BoundMessageSink, CacheHitPlan, DecoderSpec, FingerprintPolicy, + FoldContext, InputDiscoveryError, InputPipelineError, InputPlanningError, InputUnit, + LocalInputAdapter, ParseContext, ParsedBatchInput, ParsedUnit, }; use crate::clients::ClientId; -use crate::message_cache::{ParserId, ParserVersion}; +use crate::message_cache::DecoderId; +#[cfg(test)] +use crate::message_cache::DecoderVersion; use crate::sessions; pub(crate) struct OmpAdapter; @@ -25,28 +27,22 @@ const OMP_RECORD_REJECTION_REVISION: u32 = crate::adapters::MODEL_ID_CANONICALIZ const OMP_PARENT_HEALTH_REVISION: u32 = 2; impl LocalInputAdapter for OmpAdapter { - fn client(&self) -> ClientId { - ClientId::Omp - } - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { let units = adapter_discover::discover_default_scanned_units( - ClientId::Omp, + client, ctx, FingerprintPolicy::PlainFile, + DecoderSpec::plain(DecoderId::Omp, OMP_RECORD_REJECTION_REVISION), )? .into_iter() .map(|unit| { let dependency_path = sessions::pi::omp_parent_candidate_path(&unit.path) .expect("discovered OMP input must have a parent directory"); unit.with_optional_dependency(dependency_path) - .with_parser_version(ParserVersion::new( - ParserId::Omp, - OMP_RECORD_REJECTION_REVISION, - )) }) .collect(); Ok(units) @@ -76,7 +72,7 @@ impl LocalInputAdapter for OmpAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), InputPipelineError> { let mut hit_units = Vec::new(); let mut parsed_misses = Vec::new(); @@ -91,7 +87,7 @@ impl LocalInputAdapter for OmpAdapter { } } - let failed_hits = fold_omp_cache_hits(hit_units, ctx, sink)?; + let (failed_hits, _) = fold_omp_cache_hits(hit_units, ctx, Some(sink))?; if !failed_hits.is_empty() { let failed_hit_count = failed_hits.len(); let recovery_invalidations: Vec<_> = failed_hits @@ -134,7 +130,7 @@ impl LocalInputAdapter for OmpAdapter { &self, batches: &mut ParsedBatchInput<'_>, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), InputPipelineError> { let mut hit_units = Vec::new(); let mut miss_units = Vec::new(); @@ -149,7 +145,7 @@ impl LocalInputAdapter for OmpAdapter { plan_parent_health_cache(parent_health_candidates, ctx.input_cache)?; let batch_width = batches.batch_width(); - let failed_hits = fold_omp_cache_hits(hit_units, ctx, sink)?; + let (failed_hits, _) = fold_omp_cache_hits(hit_units, ctx, Some(sink))?; let mut remaining_failed_hits = failed_hits.len(); let recovery_invalidations: Vec<_> = failed_hits .iter() @@ -167,10 +163,9 @@ impl LocalInputAdapter for OmpAdapter { hit.parsed }) .collect(); - let mut parent_health_messages = Vec::new(); - let failed_parent_hits = - fold_omp_cache_hits(parent_hit_units, ctx, &mut parent_health_messages)?; - if !parent_health_messages.is_empty() { + let (failed_parent_hits, parent_health_message_count) = + fold_omp_cache_hits(parent_hit_units, ctx, None)?; + if parent_health_message_count != 0 { return Err(InputPipelineError::contract( "OMP parent-health cache contained usage messages", )); @@ -309,15 +304,12 @@ fn child_only_parent_health_candidates( } fn omp_parent_health_unit(path: PathBuf, cacheable: bool) -> InputUnit { - let unit = if cacheable { - InputUnit::plain_file(ClientId::Omp, path) + let decoder = DecoderSpec::plain(DecoderId::OmpParentHealth, OMP_PARENT_HEALTH_REVISION); + if cacheable { + InputUnit::plain_file(path, decoder) } else { - InputUnit::no_message_cache(ClientId::Omp, path) - }; - unit.with_parser_version(ParserVersion::new( - ParserId::OmpParentHealth, - OMP_PARENT_HEALTH_REVISION, - )) + InputUnit::no_message_cache(path, decoder) + } } fn plan_parent_health_cache( @@ -436,9 +428,10 @@ fn parse_parent_health_cache_misses( fn fold_omp_cache_hits( hit_units: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, -) -> Result, InputPipelineError> { + mut sink: Option<&mut BoundMessageSink<'_>>, +) -> Result<(Vec, usize), InputPipelineError> { let mut failed_units = Vec::new(); + let mut message_count = 0; for parsed in hit_units { let ParsedUnit { mut unit, @@ -455,26 +448,24 @@ fn fold_omp_cache_hits( match adapter_cache::resolve_messages(messages, ctx) { Ok(messages) => { let crate::adapters::UnitScanHealth { status, rejections } = *health; - ctx.health.record(crate::input_health::InputHealth { - client: unit.client, - path: unit.path.clone(), - status, - rejections, - }); - adapter_cache::emit_messages(unit.client, messages, sink); + ctx.record_health(unit.path.clone(), status, rejections); + message_count += messages.len(); + if let Some(sink) = sink.as_deref_mut() { + adapter_cache::emit_messages(messages, sink); + } } Err(failure) => { if !failure.can_reparse_input() { return Err(failure.into()); } debug_assert_eq!(failure.input_path, unit.path); - debug_assert_eq!(failure.parser_version, unit.parser_version); + debug_assert_eq!(failure.decoder_version, unit.decoder.version()); let remove_failed_shard = failure.requires_shard_removal(); if remove_failed_shard { - ctx.input_cache.remove(&unit.path, unit.parser_version); + ctx.input_cache.remove(&unit.path, unit.decoder.version()); } else { ctx.input_cache - .invalidate_read(&unit.path, unit.parser_version); + .invalidate_read(&unit.path, unit.decoder.version()); } unit.mark_cache_lookup_completed_no_hit(); failed_units.push(OmpFailedCacheHit { @@ -484,7 +475,7 @@ fn fold_omp_cache_hits( } } } - Ok(failed_units) + Ok((failed_units, message_count)) } fn parse_omp_miss_units( @@ -536,8 +527,9 @@ fn omp_parent_health_units( .into_iter() .filter(|health| !owned_paths.contains(&health.path)) .map(|health| { - let unit = InputUnit::no_message_cache(ClientId::Omp, health.path).with_parser_version( - ParserVersion::new(ParserId::OmpParentHealth, OMP_PARENT_HEALTH_REVISION), + let unit = InputUnit::no_message_cache( + health.path, + DecoderSpec::plain(DecoderId::OmpParentHealth, OMP_PARENT_HEALTH_REVISION), ); let mut parsed = ParsedUnit::healthy( unit, @@ -584,6 +576,39 @@ mod tests { std::fs::write(path, content).unwrap(); } + fn binding() -> crate::adapters::AdapterBinding<'static> { + crate::adapters::adapter_for(ClientId::Omp).expect("OMP adapter must be registered") + } + + fn decoder() -> DecoderSpec { + DecoderSpec::plain(DecoderId::Omp, OMP_RECORD_REJECTION_REVISION) + } + + fn unit(path: PathBuf) -> InputUnit { + InputUnit::plain_file(path, decoder()) + } + + fn unit_with_parent(path: PathBuf, parent_path: PathBuf) -> InputUnit { + unit(path).with_optional_dependency(parent_path) + } + + fn fold_parsed( + parsed: Vec, + cache: &mut message_cache::InputMessageCache, + ) -> Vec { + let binding = binding(); + let mut messages = Vec::new(); + let mut sink = crate::adapters::BoundMessageSink::new(binding, &mut messages); + OMP_ADAPTER + .fold( + parsed, + &mut FoldContext::new(binding, cache, None), + &mut sink, + ) + .unwrap(); + messages + } + fn finalized(mut messages: Vec) -> Vec { crate::finalize_token_priced_messages(&mut messages, None); messages @@ -597,30 +622,38 @@ mod tests { cache: &mut message_cache::InputMessageCache, ) -> Vec { let parsed = OMP_ADAPTER.parse_checked(units, &ParseContext { pricing: None }); - let mut sink = Vec::new(); + let mut messages = Vec::new(); + let binding = binding(); + let mut sink = crate::adapters::BoundMessageSink::new(binding, &mut messages); OMP_ADAPTER - .fold(parsed, &mut FoldContext::new(cache, None), &mut sink) + .fold( + parsed, + &mut FoldContext::new(binding, cache, None), + &mut sink, + ) .unwrap(); - assert!(sink + assert!(messages .iter() - .all(|message| message.client.as_ref() == ClientId::Omp.as_str())); - sink + .all(|message| message.client == ClientId::Omp)); + messages } fn fold_batches_with_omp_adapter( units: Vec, cache: &mut message_cache::InputMessageCache, ) -> (Vec, crate::input_health::DataHealth) { - let mut sink = Vec::new(); - let mut batches = crate::adapters::ParsedBatchInput::new(&OMP_ADAPTER, units); - let mut ctx = FoldContext::new(cache, None); + let binding = binding(); + let mut messages = Vec::new(); + let mut sink = crate::adapters::BoundMessageSink::new(binding, &mut messages); + let mut batches = crate::adapters::ParsedBatchInput::new(binding, units); + let mut ctx = FoldContext::new(binding, cache, None); OMP_ADAPTER .fold_batches(&mut batches, &mut ctx, &mut sink) .unwrap(); - assert!(sink + assert!(messages .iter() - .all(|message| message.client.as_ref() == ClientId::Omp.as_str())); - (sink, std::mem::take(&mut ctx.health)) + .all(|message| message.client == ClientId::Omp)); + (messages, ctx.take_health()) } fn omp_content(session_id: &str) -> String { @@ -631,7 +664,7 @@ mod tests { let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_dir); cache.insert(message_cache::CachedInputEntry::new_with_version( &unit.path, - unit.parser_version, + unit.decoder.version(), unit.input_policy().fingerprint().unwrap(), vec![crate::sessions::ParsedMessage::new( "gpt-5.5", @@ -669,7 +702,7 @@ mod tests { }; let ctx = scan_context(home.path(), &settings); - let units = OMP_ADAPTER.discover_checked(&ctx).unwrap(); + let units = OMP_ADAPTER.discover_checked(ClientId::Omp, &ctx).unwrap(); let paths: Vec<_> = units.iter().map(|unit| unit.path.clone()).collect(); let mut expected = vec![default_path, extra_path]; expected.sort_unstable(); @@ -686,7 +719,8 @@ mod tests { ) })); assert!(units.iter().all(|unit| { - unit.parser_version == ParserVersion::new(ParserId::Omp, OMP_RECORD_REJECTION_REVISION) + unit.decoder.version() + == DecoderVersion::new(DecoderId::Omp, OMP_RECORD_REJECTION_REVISION) })); } @@ -708,7 +742,7 @@ mod tests { }; let ctx = scan_context(home.path(), &settings); - let units = OMP_ADAPTER.discover_checked(&ctx).unwrap(); + let units = OMP_ADAPTER.discover_checked(ClientId::Omp, &ctx).unwrap(); assert_eq!(units.len(), 1); assert_eq!(units[0].path, artifact_path); @@ -732,7 +766,7 @@ mod tests { write_file(&parent_path, OMP_PARENT_CONTENT); write_file(&child_path, OMP_CHILD_CONTENT); - let units = vec![InputUnit::plain_file(ClientId::Omp, child_path.clone())]; + let units = vec![unit(child_path.clone())]; let mut cache = message_cache::InputMessageCache::default(); let actual = fold_with_omp_adapter(units, &mut cache); @@ -760,7 +794,7 @@ mod tests { let settings = crate::scanner::ScannerSettings::default(); let ctx = scan_context(home.path(), &settings); let child_unit = OMP_ADAPTER - .discover_checked(&ctx) + .discover_checked(ClientId::Omp, &ctx) .unwrap() .into_iter() .find(|unit| unit.path == child_path) @@ -768,10 +802,7 @@ mod tests { let mut cache = message_cache::InputMessageCache::default(); let parsed = OMP_ADAPTER.parse_checked(vec![child_unit.clone()], &ParseContext { pricing: None }); - let mut first = Vec::new(); - OMP_ADAPTER - .fold(parsed, &mut FoldContext::new(&mut cache, None), &mut first) - .unwrap(); + let first = fold_parsed(parsed, &mut cache); assert_eq!(first[0].session_id.as_ref(), "child-session"); assert_eq!(first[0].agent.as_deref(), Some("OMP Reviewer")); @@ -779,14 +810,7 @@ mod tests { CacheHitPlan::Hit(hit) => hit, CacheHitPlan::Miss(_) => panic!("unchanged OMP child and parent must use warm cache"), }; - let mut warm_messages = Vec::new(); - OMP_ADAPTER - .fold( - vec![warm], - &mut FoldContext::new(&mut cache, None), - &mut warm_messages, - ) - .unwrap(); + let warm_messages = fold_parsed(vec![warm], &mut cache); assert_eq!(warm_messages[0].agent.as_deref(), Some("OMP Reviewer")); write_file( @@ -794,7 +818,7 @@ mod tests { &OMP_PARENT_CONTENT.replace(r#""agent":"reviewer""#, r#""agent":"oracle""#), ); let refreshed_child_unit = OMP_ADAPTER - .discover_checked(&ctx) + .discover_checked(ClientId::Omp, &ctx) .unwrap() .into_iter() .find(|unit| unit.path == child_path) @@ -808,10 +832,7 @@ mod tests { CacheHitPlan::Hit(_) => panic!("parent-only change must invalidate child cache"), }; let parsed = OMP_ADAPTER.parse_checked(vec![miss], &ParseContext { pricing: None }); - let mut second = Vec::new(); - OMP_ADAPTER - .fold(parsed, &mut FoldContext::new(&mut cache, None), &mut second) - .unwrap(); + let second = fold_parsed(parsed, &mut cache); assert_eq!(second[0].session_id.as_ref(), "child-session"); assert_eq!(second[0].agent.as_deref(), Some("OMP Oracle")); } @@ -827,14 +848,9 @@ mod tests { write_file(&first_child_path, &omp_content("first-child")); write_file(&second_child_path, &omp_content("second-child")); - let parser_version = ParserVersion::new(ParserId::Omp, OMP_RECORD_REJECTION_REVISION); let units = [first_child_path, second_child_path] .into_iter() - .map(|path| { - InputUnit::plain_file(ClientId::Omp, path) - .with_optional_dependency(parent_path.clone()) - .with_parser_version(parser_version) - }) + .map(|path| unit_with_parent(path, parent_path.clone())) .collect(); let mut cache = message_cache::InputMessageCache::default(); message_cache::reset_input_read_stats(&parent_path); @@ -865,12 +881,7 @@ mod tests { &parent_path, &OMP_PARENT_CONTENT.replace(r#""agent":"reviewer""#, r#""agent":"new-reviewer""#), ); - let unit = InputUnit::plain_file(ClientId::Omp, child_path) - .with_optional_dependency(parent_path) - .with_parser_version(ParserVersion::new( - ParserId::Omp, - OMP_RECORD_REJECTION_REVISION, - )); + let unit = unit_with_parent(child_path, parent_path); let parsed = parse_omp_miss_units(vec![unit], &ParseContext { pricing: None }, &parent_index); @@ -929,19 +940,14 @@ mod tests { write_file(&child_path, OMP_CHILD_CONTENT); write_file(&cached_path, OMP_CHILD_CONTENT); - let parser_version = ParserVersion::new(ParserId::Omp, OMP_RECORD_REJECTION_REVISION); - let child_unit = - InputUnit::plain_file(ClientId::Omp, child_path).with_parser_version(parser_version); - let cached_unit = InputUnit::plain_file(ClientId::Omp, cached_path.clone()) - .with_parser_version(parser_version) - .prepare_snapshot() - .unwrap(); - let parent_unit = - InputUnit::plain_file(ClientId::Omp, parent_path).with_parser_version(parser_version); + let decoder_version = decoder().version(); + let child_unit = unit(child_path); + let cached_unit = unit(cached_path.clone()).prepare_snapshot().unwrap(); + let parent_unit = unit(parent_path); let mut cache = message_cache::InputMessageCache::default(); cache.insert(message_cache::CachedInputEntry::new_with_version( &cached_path, - parser_version, + decoder_version, cached_unit.input_policy().fingerprint().unwrap(), vec![crate::sessions::ParsedMessage::new( "gpt-5.5", @@ -962,19 +968,21 @@ mod tests { .build() .unwrap() .install(|| { - let mut sink = Vec::new(); + let binding = binding(); + let mut messages = Vec::new(); + let mut sink = crate::adapters::BoundMessageSink::new(binding, &mut messages); let mut batches = crate::adapters::ParsedBatchInput::new( - &OMP_ADAPTER, + binding, vec![child_unit, cached_unit, parent_unit], ); OMP_ADAPTER .fold_batches( &mut batches, - &mut FoldContext::new(&mut cache, None), + &mut FoldContext::new(binding, &mut cache, None), &mut sink, ) .unwrap(); - sink + messages }); let sessions: Vec<_> = messages @@ -997,30 +1005,27 @@ mod tests { write_file(&parent_path, &format!("{{not-json\n{OMP_PARENT_CONTENT}")); write_file(&child_path, OMP_CHILD_CONTENT); - let unit = InputUnit::plain_file(ClientId::Omp, child_path) - .with_optional_dependency(parent_path) - .with_parser_version(ParserVersion::new( - ParserId::Omp, - OMP_RECORD_REJECTION_REVISION, - )); + let unit = unit_with_parent(child_path, parent_path); let mut cache = message_cache::InputMessageCache::default(); - let mut sink = Vec::new(); - let mut batches = crate::adapters::ParsedBatchInput::new(&OMP_ADAPTER, vec![unit]); - let mut ctx = FoldContext::new(&mut cache, None); + let binding = binding(); + let mut messages = Vec::new(); + let mut sink = crate::adapters::BoundMessageSink::new(binding, &mut messages); + let mut batches = crate::adapters::ParsedBatchInput::new(binding, vec![unit]); + let mut ctx = FoldContext::new(binding, &mut cache, None); OMP_ADAPTER .fold_batches(&mut batches, &mut ctx, &mut sink) .expect("a malformed parent record must stay inside the OMP input health domain"); - assert_eq!(sink.len(), 1); - assert_eq!(sink[0].agent.as_deref(), Some("OMP Reviewer")); - assert_eq!(ctx.health.rejected_records(), 1); - assert_eq!(ctx.health.partial_inputs(), 0); - assert_eq!(ctx.health.failed_inputs(), 0); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].agent.as_deref(), Some("OMP Reviewer")); + assert_eq!(ctx.health().rejected_records(), 1); + assert_eq!(ctx.health().partial_inputs(), 0); + assert_eq!(ctx.health().failed_inputs(), 0); } #[test] - fn unowned_parent_health_uses_parent_health_parser_version() { + fn unowned_parent_health_uses_parent_health_decoder_version() { let dir = tempfile::TempDir::new().unwrap(); let session_root = dir.path().join(".omp/agent/sessions/project/root-session"); let parent_path = session_root.with_extension("jsonl"); @@ -1033,8 +1038,8 @@ mod tests { assert_eq!(units.len(), 1); assert_eq!( - units[0].unit.parser_version, - ParserVersion::new(ParserId::OmpParentHealth, OMP_PARENT_HEALTH_REVISION) + units[0].unit.decoder.version(), + DecoderVersion::new(DecoderId::OmpParentHealth, OMP_PARENT_HEALTH_REVISION) ); } @@ -1056,7 +1061,9 @@ mod tests { ..Default::default() }; let scan_ctx = scan_context(home.path(), &settings); - let cold_units = OMP_ADAPTER.discover_checked(&scan_ctx).unwrap(); + let cold_units = OMP_ADAPTER + .discover_checked(ClientId::Omp, &scan_ctx) + .unwrap(); assert_eq!( cold_units .iter() @@ -1073,7 +1080,9 @@ mod tests { assert_eq!(sessions::pi::omp_parent_scan_count(&parent_path), 1); sessions::pi::reset_omp_parent_scan_count(&parent_path); - let warm_units = OMP_ADAPTER.discover_checked(&scan_ctx).unwrap(); + let warm_units = OMP_ADAPTER + .discover_checked(ClientId::Omp, &scan_ctx) + .unwrap(); assert!(warm_units.iter().cloned().all(|unit| matches!( OMP_ADAPTER.plan_cache_hit(unit, &cache).unwrap(), CacheHitPlan::Hit(_) @@ -1097,7 +1106,9 @@ mod tests { &OMP_PARENT_CONTENT.replace(r#""agent":"reviewer""#, r#""agent":"oracle""#), ); sessions::pi::reset_omp_parent_scan_count(&parent_path); - let changed_units = OMP_ADAPTER.discover_checked(&scan_ctx).unwrap(); + let changed_units = OMP_ADAPTER + .discover_checked(ClientId::Omp, &scan_ctx) + .unwrap(); assert!(changed_units.iter().cloned().all(|unit| matches!( OMP_ADAPTER.plan_cache_hit(unit, &cache).unwrap(), CacheHitPlan::Miss(_) @@ -1132,11 +1143,13 @@ mod tests { ..Default::default() }; let scan_ctx = scan_context(home.path(), &settings); - let parent_parser_version = - ParserVersion::new(ParserId::OmpParentHealth, OMP_PARENT_HEALTH_REVISION); + let parent_decoder_version = + DecoderVersion::new(DecoderId::OmpParentHealth, OMP_PARENT_HEALTH_REVISION); let mut cold_cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); - let cold_units = OMP_ADAPTER.discover_checked(&scan_ctx).unwrap(); + let cold_units = OMP_ADAPTER + .discover_checked(ClientId::Omp, &scan_ctx) + .unwrap(); let (_, cold_health) = fold_batches_with_omp_adapter(cold_units, &mut cold_cache); assert_eq!(cold_health.issue_count(), 1); drop(cold_cache); @@ -1144,11 +1157,13 @@ mod tests { message_cache::truncate_shard_after_header_for_test( cache_dir.path(), &parent_path, - parent_parser_version, + parent_decoder_version, ); sessions::pi::reset_omp_parent_scan_count(&parent_path); let mut repair_cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); - let repair_units = OMP_ADAPTER.discover_checked(&scan_ctx).unwrap(); + let repair_units = OMP_ADAPTER + .discover_checked(ClientId::Omp, &scan_ctx) + .unwrap(); let (_, repair_health) = fold_batches_with_omp_adapter(repair_units, &mut repair_cache); assert_eq!(repair_health.issue_count(), 1); assert_eq!(sessions::pi::omp_parent_scan_count(&parent_path), 1); @@ -1156,7 +1171,9 @@ mod tests { sessions::pi::reset_omp_parent_scan_count(&parent_path); let mut warm_cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); - let warm_units = OMP_ADAPTER.discover_checked(&scan_ctx).unwrap(); + let warm_units = OMP_ADAPTER + .discover_checked(ClientId::Omp, &scan_ctx) + .unwrap(); let (_, warm_health) = fold_batches_with_omp_adapter(warm_units, &mut warm_cache); assert_eq!(warm_health.issue_count(), 1); assert_eq!(sessions::pi::omp_parent_scan_count(&parent_path), 0); @@ -1184,7 +1201,7 @@ mod tests { message_cache::mark_current_key_shard_as_unsupported_format_for_test( unsupported_cache_dir.path(), &parent_path, - parent_unit.parser_version, + parent_unit.decoder.version(), ); let unsupported_cache = message_cache::InputMessageCache::with_cache_dir(unsupported_cache_dir.path()); @@ -1202,7 +1219,7 @@ mod tests { message_cache::mark_current_key_shard_as_future_format_for_test( future_cache_dir.path(), &parent_path, - parent_unit.parser_version, + parent_unit.decoder.version(), ); let future_cache = message_cache::InputMessageCache::with_cache_dir(future_cache_dir.path()); @@ -1229,20 +1246,24 @@ mod tests { let scan_ctx = scan_context(home.path(), &settings); let mut cache = message_cache::InputMessageCache::default(); - let cold_units = OMP_ADAPTER.discover_checked(&scan_ctx).unwrap(); + let cold_units = OMP_ADAPTER + .discover_checked(ClientId::Omp, &scan_ctx) + .unwrap(); let (cold_messages, cold_health) = fold_batches_with_omp_adapter(cold_units, &mut cache); assert_eq!(cold_health.issue_count(), 0); assert_eq!(sessions::pi::omp_parent_scan_count(&parent_path), 1); assert!(cache .get_meta( &parent_path, - ParserVersion::new(ParserId::OmpParentHealth, OMP_PARENT_HEALTH_REVISION), + DecoderVersion::new(DecoderId::OmpParentHealth, OMP_PARENT_HEALTH_REVISION), ) .unwrap() .is_some()); sessions::pi::reset_omp_parent_scan_count(&parent_path); - let warm_units = OMP_ADAPTER.discover_checked(&scan_ctx).unwrap(); + let warm_units = OMP_ADAPTER + .discover_checked(ClientId::Omp, &scan_ctx) + .unwrap(); assert!(warm_units.iter().cloned().all(|unit| matches!( OMP_ADAPTER.plan_cache_hit(unit, &cache).unwrap(), CacheHitPlan::Hit(_) @@ -1271,12 +1292,16 @@ mod tests { }; let scan_ctx = scan_context(home.path(), &settings); let mut cache = message_cache::InputMessageCache::default(); - let cold_units = OMP_ADAPTER.discover_checked(&scan_ctx).unwrap(); + let cold_units = OMP_ADAPTER + .discover_checked(ClientId::Omp, &scan_ctx) + .unwrap(); fold_batches_with_omp_adapter(cold_units, &mut cache); write_file(&child_path, &omp_content("changed-child-with-longer-id")); sessions::pi::reset_omp_parent_scan_count(&parent_path); - let changed_units = OMP_ADAPTER.discover_checked(&scan_ctx).unwrap(); + let changed_units = OMP_ADAPTER + .discover_checked(ClientId::Omp, &scan_ctx) + .unwrap(); let (messages, health) = fold_batches_with_omp_adapter(changed_units, &mut cache); assert_eq!(messages.len(), 1); @@ -1298,14 +1323,7 @@ mod tests { std::fs::create_dir_all(&parent_path).unwrap(); write_file(&child_path, OMP_CHILD_CONTENT); - let make_unit = || { - InputUnit::plain_file(ClientId::Omp, child_path.clone()) - .with_optional_dependency(parent_path.clone()) - .with_parser_version(ParserVersion::new( - ParserId::Omp, - OMP_RECORD_REJECTION_REVISION, - )) - }; + let make_unit = || unit_with_parent(child_path.clone(), parent_path.clone()); let mut cache = message_cache::InputMessageCache::default(); let (cold_messages, cold_health) = fold_batches_with_omp_adapter(vec![make_unit()], &mut cache); @@ -1317,7 +1335,7 @@ mod tests { assert!(cache .get_meta( &parent_path, - ParserVersion::new(ParserId::OmpParentHealth, OMP_PARENT_HEALTH_REVISION), + DecoderVersion::new(DecoderId::OmpParentHealth, OMP_PARENT_HEALTH_REVISION), ) .unwrap() .is_none()); @@ -1347,13 +1365,15 @@ mod tests { ..Default::default() }; let scan_ctx = scan_context(home.path(), &settings); - let parent_parser_version = - ParserVersion::new(ParserId::OmpParentHealth, OMP_PARENT_HEALTH_REVISION); + let parent_decoder_version = + DecoderVersion::new(DecoderId::OmpParentHealth, OMP_PARENT_HEALTH_REVISION); let mut cache = message_cache::InputMessageCache::default(); - let cold_units = OMP_ADAPTER.discover_checked(&scan_ctx).unwrap(); + let cold_units = OMP_ADAPTER + .discover_checked(ClientId::Omp, &scan_ctx) + .unwrap(); fold_batches_with_omp_adapter(cold_units, &mut cache); let old_fingerprint = cache - .get_meta(&parent_path, parent_parser_version) + .get_meta(&parent_path, parent_decoder_version) .unwrap() .unwrap() .fingerprint; @@ -1364,7 +1384,9 @@ mod tests { ); sessions::pi::force_omp_parent_open_failure(&parent_path, true); sessions::pi::reset_omp_parent_scan_count(&parent_path); - let unavailable_units = OMP_ADAPTER.discover_checked(&scan_ctx).unwrap(); + let unavailable_units = OMP_ADAPTER + .discover_checked(ClientId::Omp, &scan_ctx) + .unwrap(); let (unavailable_messages, unavailable_health) = fold_batches_with_omp_adapter(unavailable_units, &mut cache); @@ -1375,7 +1397,7 @@ mod tests { assert_eq!(sessions::pi::omp_parent_scan_count(&parent_path), 1); assert_eq!( cache - .get_meta(&parent_path, parent_parser_version) + .get_meta(&parent_path, parent_decoder_version) .unwrap() .unwrap() .fingerprint, @@ -1393,7 +1415,9 @@ mod tests { sessions::pi::force_omp_parent_open_failure(&parent_path, false); sessions::pi::reset_omp_parent_scan_count(&parent_path); - let retry_units = OMP_ADAPTER.discover_checked(&scan_ctx).unwrap(); + let retry_units = OMP_ADAPTER + .discover_checked(ClientId::Omp, &scan_ctx) + .unwrap(); let (retry_messages, retry_health) = fold_batches_with_omp_adapter(retry_units, &mut cache); assert_eq!(retry_health.issue_count(), 0); assert_eq!(retry_messages[0].agent.as_deref(), Some("OMP Oracle")); @@ -1415,7 +1439,6 @@ mod tests { write_file(&first_child, &omp_content("first-child")); write_file(&second_child, &omp_content("second-child")); - let parser_version = ParserVersion::new(ParserId::Omp, OMP_RECORD_REJECTION_REVISION); let make_units = || { vec![ parent_path.clone(), @@ -1425,9 +1448,7 @@ mod tests { .into_iter() .map(|path| { let dependency_path = sessions::pi::omp_parent_candidate_path(&path).unwrap(); - InputUnit::plain_file(ClientId::Omp, path) - .with_optional_dependency(dependency_path) - .with_parser_version(parser_version) + unit_with_parent(path, dependency_path) }) .collect() }; @@ -1451,7 +1472,7 @@ mod tests { assert!(cache .get_meta( &parent_path, - ParserVersion::new(ParserId::OmpParentHealth, OMP_PARENT_HEALTH_REVISION), + DecoderVersion::new(DecoderId::OmpParentHealth, OMP_PARENT_HEALTH_REVISION), ) .unwrap() .is_none()); @@ -1476,27 +1497,24 @@ mod tests { write_file(&first_child, &omp_content("first-child")); write_file(&second_child, &omp_content("second-child")); - let parser_version = ParserVersion::new(ParserId::Omp, OMP_RECORD_REJECTION_REVISION); let units = vec![first_child, second_child] .into_iter() - .map(|path| { - InputUnit::plain_file(ClientId::Omp, path) - .with_optional_dependency(parent_path.clone()) - .with_parser_version(parser_version) - }) + .map(|path| unit_with_parent(path, parent_path.clone())) .collect(); let mut cache = message_cache::InputMessageCache::default(); let parsed = OMP_ADAPTER.parse_checked(units, &ParseContext { pricing: None }); - let mut sink = Vec::new(); - let mut ctx = FoldContext::new(&mut cache, None); + let binding = binding(); + let mut messages = Vec::new(); + let mut sink = crate::adapters::BoundMessageSink::new(binding, &mut messages); + let mut ctx = FoldContext::new(binding, &mut cache, None); OMP_ADAPTER.fold(parsed, &mut ctx, &mut sink).unwrap(); - assert_eq!(sink.len(), 2); - assert_eq!(ctx.health.partial_inputs(), 1); - assert_eq!(ctx.health.failed_inputs(), 0); - assert_eq!(ctx.health.inputs()[0].path, parent_path); - let failure = ctx.health.inputs()[0].status.failure().unwrap(); + assert_eq!(messages.len(), 2); + assert_eq!(ctx.health().partial_inputs(), 1); + assert_eq!(ctx.health().failed_inputs(), 0); + assert_eq!(ctx.health().inputs()[0].path, parent_path); + let failure = ctx.health().inputs()[0].status.failure().unwrap(); assert_eq!(failure.operation, "read OMP parent JSONL line"); } @@ -1505,16 +1523,14 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("cached.jsonl"); write_file(&path, OMP_CHILD_CONTENT); - let parser_version = ParserVersion::new(ParserId::Omp, OMP_RECORD_REJECTION_REVISION); + let decoder_version = DecoderVersion::new(DecoderId::Omp, OMP_RECORD_REJECTION_REVISION); let dependency_path = path.parent().unwrap().with_extension("jsonl"); - let unit = InputUnit::plain_file(ClientId::Omp, path.clone()) - .with_optional_dependency(dependency_path) - .with_parser_version(parser_version) + let unit = unit_with_parent(path.clone(), dependency_path) .prepare_snapshot() .unwrap(); let mut entry = message_cache::CachedInputEntry::new_with_version( &path, - parser_version, + decoder_version, unit.input_policy().fingerprint().unwrap(), vec![crate::sessions::ParsedMessage::new( "gpt-5.5", @@ -1538,14 +1554,16 @@ mod tests { CacheHitPlan::Hit(hit) => hit, CacheHitPlan::Miss(_) => panic!("unchanged OMP dependency cache must hit"), }; - let mut sink = Vec::new(); - let mut ctx = FoldContext::new(&mut cache, None); + let binding = binding(); + let mut messages = Vec::new(); + let mut sink = crate::adapters::BoundMessageSink::new(binding, &mut messages); + let mut ctx = FoldContext::new(binding, &mut cache, None); OMP_ADAPTER.fold(vec![hit], &mut ctx, &mut sink).unwrap(); - assert_eq!(sink.len(), 1); - assert_eq!(sink[0].session_id.as_ref(), "cached-session"); - assert_eq!(ctx.health.rejected_records(), 1); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].session_id.as_ref(), "cached-session"); + assert_eq!(ctx.health().rejected_records(), 1); } #[test] @@ -1572,12 +1590,7 @@ mod tests { write_file(&ordinary_a, &omp_content("ordinary-a")); write_file(&ordinary_b, &omp_content("ordinary-b")); - let make_unit = |path: PathBuf| { - InputUnit::plain_file(ClientId::Omp, path).with_parser_version(ParserVersion::new( - ParserId::Omp, - OMP_RECORD_REJECTION_REVISION, - )) - }; + let make_unit = |path: PathBuf| unit(path); let first_child_unit = make_unit(first_child.clone()); let second_child_unit = make_unit(second_child.clone()); let successful_unit = make_unit(successful_path.clone()); @@ -1587,12 +1600,12 @@ mod tests { message_cache::truncate_shard_after_header_for_test( cache_dir.path(), &first_child, - first_child_unit.parser_version, + first_child_unit.decoder.version(), ); message_cache::truncate_shard_after_header_for_test( cache_dir.path(), &second_child, - second_child_unit.parser_version, + second_child_unit.decoder.version(), ); let units = vec![ @@ -1608,16 +1621,18 @@ mod tests { .build() .unwrap() .install(|| { - let mut sink = Vec::new(); - let mut batches = crate::adapters::ParsedBatchInput::new(&OMP_ADAPTER, units); + let binding = binding(); + let mut messages = Vec::new(); + let mut sink = crate::adapters::BoundMessageSink::new(binding, &mut messages); + let mut batches = crate::adapters::ParsedBatchInput::new(binding, units); OMP_ADAPTER .fold_batches( &mut batches, - &mut FoldContext::new(&mut cache, None), + &mut FoldContext::new(binding, &mut cache, None), &mut sink, ) .unwrap(); - sink + messages }); let sessions: Vec<_> = messages diff --git a/crates/tokscale-core/src/adapters/openclaw.rs b/crates/tokscale-core/src/adapters/openclaw.rs index 184e0160a..639315ad5 100644 --- a/crates/tokscale-core/src/adapters/openclaw.rs +++ b/crates/tokscale-core/src/adapters/openclaw.rs @@ -3,11 +3,13 @@ use rayon::prelude::*; use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ - AdapterScanContext, FingerprintPolicy, FoldContext, InputDiscoveryError, InputUnit, - LocalInputAdapter, MessageSink, ParseContext, ParsedUnit, + AdapterScanContext, BoundMessageSink, DecoderSpec, FingerprintPolicy, FoldContext, + InputDiscoveryError, InputUnit, LocalInputAdapter, ParseContext, ParsedUnit, }; use crate::clients::ClientId; -use crate::message_cache::{ParserId, ParserVersion}; +use crate::message_cache::DecoderId; +#[cfg(test)] +use crate::message_cache::DecoderVersion; use crate::sessions; pub(crate) struct OpenClawAdapter; @@ -16,36 +18,23 @@ const OPENCLAW_RECORD_REJECTION_REVISION: u32 = crate::adapters::MODEL_ID_CANONICALIZATION_REVISION + 2; impl LocalInputAdapter for OpenClawAdapter { - fn client(&self) -> ClientId { - ClientId::OpenClaw - } - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { - let def = ClientId::OpenClaw + let def = client .local_def() .expect("OpenClaw adapter must have local scan policy"); let mut roots = vec![def.resolve_path(ctx.home_dir)]; - roots.extend(adapter_discover::extra_roots_for_client( - ClientId::OpenClaw, - ctx, - )?); + roots.extend(adapter_discover::extra_roots_for_client(client, ctx)?); - Ok(adapter_discover::input_units_from_paths( - ClientId::OpenClaw, - adapter_discover::scan_roots(ClientId::OpenClaw, roots, def.pattern)?, + adapter_discover::input_units_from_paths( + client, + adapter_discover::scan_roots(client, roots, def.pattern)?, FingerprintPolicy::PlainFile, - )? - .into_iter() - .map(|unit| { - unit.with_parser_version(ParserVersion::new( - ParserId::OpenClaw, - OPENCLAW_RECORD_REJECTION_REVISION, - )) - }) - .collect()) + DecoderSpec::plain(DecoderId::OpenClaw, OPENCLAW_RECORD_REJECTION_REVISION), + ) } fn parse_checked(&self, units: Vec, ctx: &ParseContext<'_>) -> Vec { @@ -71,7 +60,7 @@ impl LocalInputAdapter for OpenClawAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), crate::adapters::InputPipelineError> { adapter_cache::fold_units(parsed, ctx, sink) } @@ -111,15 +100,17 @@ mod tests { scanner_settings: &settings, }; - let units = OPENCLAW_ADAPTER.discover_checked(&ctx).unwrap(); + let units = OPENCLAW_ADAPTER + .discover_checked(ClientId::OpenClaw, &ctx) + .unwrap(); let paths: Vec<_> = units.iter().map(|unit| unit.path.clone()).collect(); let mut expected = vec![default_path, extra_path]; expected.sort_unstable(); assert_eq!(paths, expected); assert!(units.iter().all(|unit| { - unit.parser_version - == ParserVersion::new(ParserId::OpenClaw, OPENCLAW_RECORD_REJECTION_REVISION) + unit.decoder.version() + == DecoderVersion::new(DecoderId::OpenClaw, OPENCLAW_RECORD_REJECTION_REVISION) })); } } diff --git a/crates/tokscale-core/src/adapters/opencode.rs b/crates/tokscale-core/src/adapters/opencode.rs index 7b73b1e6e..4db2d075f 100644 --- a/crates/tokscale-core/src/adapters/opencode.rs +++ b/crates/tokscale-core/src/adapters/opencode.rs @@ -4,8 +4,9 @@ use rayon::prelude::*; use crate::adapters::cache as adapter_cache; use crate::adapters::{ - AdapterScanContext, FoldContext, InputDiscoveryError, InputPipelineError, InputUnit, - InputUnitMeta, LocalInputAdapter, MessageSink, ParseContext, ParsedBatchInput, ParsedUnit, + AdapterScanContext, BoundMessageSink, DecoderRoute, DecoderSpec, FoldContext, + InputDiscoveryError, InputPipelineError, InputUnit, LocalInputAdapter, ParseContext, + ParsedBatchInput, ParsedUnit, OPENCODE_CURRENT_SQLITE_REVISION, }; use crate::clients::ClientId; use crate::{scanner, sessions}; @@ -13,22 +14,14 @@ use crate::{scanner, sessions}; pub(crate) struct OpenCodeAdapter; impl LocalInputAdapter for OpenCodeAdapter { - fn client(&self) -> ClientId { - ClientId::OpenCode - } - fn discover_checked( &self, + _client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { let data_dir = scanner::opencode_data_dir(ctx.home_dir); let mut db_paths = scanner::discover_opencode_dbs(&data_dir).map_err(|source| { - InputDiscoveryError::new( - ClientId::OpenCode, - &data_dir, - "discover OpenCode databases", - source, - ) + InputDiscoveryError::new(&data_dir, "discover OpenCode databases", source) })?; scanner::merge_user_opencode_db_paths( &mut db_paths, @@ -39,8 +32,10 @@ impl LocalInputAdapter for OpenCodeAdapter { Ok(db_paths .into_iter() .map(|path| { - InputUnit::sqlite_with_wal(ClientId::OpenCode, path) - .with_meta(InputUnitMeta::OpenCodeSqlite) + InputUnit::sqlite_with_wal( + path, + DecoderSpec::opencode_sqlite(OPENCODE_CURRENT_SQLITE_REVISION), + ) }) .collect()) } @@ -48,8 +43,8 @@ impl LocalInputAdapter for OpenCodeAdapter { fn parse_checked(&self, units: Vec, ctx: &ParseContext<'_>) -> Vec { units .into_par_iter() - .map(|unit| match unit.meta { - InputUnitMeta::OpenCodeSqlite => { + .map(|unit| match unit.decoder.route() { + DecoderRoute::OpenCodeSqlite => { adapter_cache::load_or_scan_unit_with(unit, ctx, |path| { sessions::opencode::parse_opencode_sqlite(path).map_err(|error| { crate::sessions::error::SessionParseError::new( @@ -59,14 +54,14 @@ impl LocalInputAdapter for OpenCodeAdapter { }) }) } - InputUnitMeta::None - | InputUnitMeta::AntigravityCliSqlite - | InputUnitMeta::KiroFile - | InputUnitMeta::KiroSqlite - | InputUnitMeta::KiroGlobalStorage - | InputUnitMeta::CodeBuddyJsonl - | InputUnitMeta::CodeBuddyExtensionLog { .. } - | InputUnitMeta::Codex => { + DecoderRoute::None + | DecoderRoute::AntigravityCliSqlite + | DecoderRoute::KiroFile + | DecoderRoute::KiroSqlite + | DecoderRoute::KiroGlobalStorage + | DecoderRoute::CodeBuddyJsonl + | DecoderRoute::CodeBuddyExtensionLog { .. } + | DecoderRoute::Codex => { unreachable!("unexpected OpenCode input unit meta") } }) @@ -85,7 +80,7 @@ impl LocalInputAdapter for OpenCodeAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), InputPipelineError> { let mut seen = HashSet::new(); for unit in parsed { @@ -98,7 +93,7 @@ impl LocalInputAdapter for OpenCodeAdapter { &self, batches: &mut ParsedBatchInput<'_>, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), InputPipelineError> { let mut seen = HashSet::new(); while let Some(parsed) = batches.next(ctx)? { @@ -113,7 +108,7 @@ impl LocalInputAdapter for OpenCodeAdapter { fn fold_opencode_unit( parsed: ParsedUnit, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, seen: &mut HashSet, ) -> Result<(), InputPipelineError> { let adapter_cache::ResolvedUnit { @@ -124,20 +119,14 @@ fn fold_opencode_unit( status, rejections, } = adapter_cache::resolve_unit(parsed, ctx)?; - ctx.health.record(crate::input_health::InputHealth { - client: unit.client, - path: unit.path.clone(), - status, - rejections, - }); + ctx.record_health(unit.path.clone(), status, rejections); let path = unit.path.clone(); let cache_write_outcome = adapter_cache::write_cache(cache_write, ctx, &messages); if cache_write_outcome.is_err() && invalidate_cache { - ctx.input_cache.remove(&path, unit.parser_version); + ctx.input_cache.remove(&path, unit.decoder.version()); } let cache_write_outcome = cache_write_outcome?; adapter_cache::emit_messages( - unit.client, messages .into_iter() .filter(|message| message.dedup_key.is_none_or(|key| seen.insert(key))), @@ -145,7 +134,7 @@ fn fold_opencode_unit( ); if cache_write_outcome == adapter_cache::CacheWriteOutcome::NotPlanned && invalidate_cache { - ctx.input_cache.remove(&path, unit.parser_version); + ctx.input_cache.remove(&path, unit.decoder.version()); } Ok(()) } @@ -209,7 +198,9 @@ mod tests { scanner_settings: &settings, }; - let units = OPENCODE_ADAPTER.discover_checked(&ctx).unwrap(); + let units = OPENCODE_ADAPTER + .discover_checked(ClientId::OpenCode, &ctx) + .unwrap(); assert_eq!( units .iter() @@ -219,7 +210,7 @@ mod tests { ); assert!(units .iter() - .all(|unit| matches!(unit.meta, InputUnitMeta::OpenCodeSqlite))); + .all(|unit| matches!(unit.decoder.route(), DecoderRoute::OpenCodeSqlite))); assert!(units.iter().all(|unit| unit.digest_paths().len() == 2)); } @@ -245,8 +236,10 @@ mod tests { Some(key), ); ParsedUnit::healthy( - InputUnit::sqlite_with_wal(ClientId::OpenCode, dir.path().join(name)) - .with_meta(InputUnitMeta::OpenCodeSqlite), + InputUnit::sqlite_with_wal( + dir.path().join(name), + DecoderSpec::opencode_sqlite(OPENCODE_CURRENT_SQLITE_REVISION), + ), UnitMessagePayload::Fresh(vec![message]), None, false, @@ -255,12 +248,15 @@ mod tests { .collect(); let mut cache = message_cache::InputMessageCache::default(); let mut sink = Vec::new(); + let binding = crate::adapters::adapter_for(ClientId::OpenCode).unwrap(); + let mut fold_ctx = FoldContext::new(binding, &mut cache, None); + let mut bound_sink = BoundMessageSink::new(binding, &mut sink); OPENCODE_ADAPTER - .fold(parsed, &mut FoldContext::new(&mut cache, None), &mut sink) + .fold(parsed, &mut fold_ctx, &mut bound_sink) .unwrap(); assert_eq!(sink.len(), 1); - assert_eq!(sink[0].client.as_ref(), ClientId::OpenCode.as_str()); + assert_eq!(sink[0].client, ClientId::OpenCode); assert_eq!(sink[0].session_id.as_ref(), "session-0"); } @@ -274,10 +270,12 @@ mod tests { let units = vec![first, second] .into_iter() .map(|path| { - InputUnit::sqlite_with_wal(ClientId::OpenCode, path) - .with_meta(InputUnitMeta::OpenCodeSqlite) - .prepare_snapshot() - .unwrap() + InputUnit::sqlite_with_wal( + path, + DecoderSpec::opencode_sqlite(OPENCODE_CURRENT_SQLITE_REVISION), + ) + .prepare_snapshot() + .unwrap() }) .collect(); @@ -288,18 +286,17 @@ mod tests { .install(|| { let mut cache = message_cache::InputMessageCache::default(); let mut sink = Vec::new(); - let mut batches = ParsedBatchInput::new(&OPENCODE_ADAPTER, units); + let binding = crate::adapters::adapter_for(ClientId::OpenCode).unwrap(); + let mut batches = ParsedBatchInput::new(binding, units); + let mut fold_ctx = FoldContext::new(binding, &mut cache, None); + let mut bound_sink = BoundMessageSink::new(binding, &mut sink); OPENCODE_ADAPTER - .fold_batches( - &mut batches, - &mut FoldContext::new(&mut cache, None), - &mut sink, - ) + .fold_batches(&mut batches, &mut fold_ctx, &mut bound_sink) .unwrap(); sink }); assert_eq!(messages.len(), 1); - assert_eq!(messages[0].client.as_ref(), ClientId::OpenCode.as_str()); + assert_eq!(messages[0].client, ClientId::OpenCode); } #[test] @@ -310,17 +307,19 @@ mod tests { conn.execute_batch("CREATE TABLE message (id TEXT, session_id TEXT, data TEXT);") .unwrap(); drop(conn); - let unit = InputUnit::sqlite_with_wal(ClientId::OpenCode, path.clone()) - .with_meta(InputUnitMeta::OpenCodeSqlite) - .prepare_snapshot() - .unwrap(); - let parser_version = unit.parser_version; + let unit = InputUnit::sqlite_with_wal( + path.clone(), + DecoderSpec::opencode_sqlite(OPENCODE_CURRENT_SQLITE_REVISION), + ) + .prepare_snapshot() + .unwrap(); + let decoder_version = unit.decoder.version(); let fingerprint = unit.input_policy().fingerprint().unwrap(); let mut cache = message_cache::InputMessageCache::default(); cache.insert(message_cache::CachedInputEntry::new_with_version( &path, - message_cache::ParserVersion::new( - message_cache::ParserId::OpenCodeSqlite, + message_cache::DecoderVersion::new( + message_cache::DecoderId::OpenCodeSqlite, crate::adapters::MODEL_ID_CANONICALIZATION_REVISION, ), fingerprint, @@ -340,12 +339,12 @@ mod tests { let parsed = OPENCODE_ADAPTER.parse_checked(vec![unit], &ParseContext { pricing: None }); assert_eq!(parsed.len(), 1); - let health = parsed[0].input_health(); - assert_eq!(health.path, path); + let health = &parsed[0].health; + assert_eq!(parsed[0].unit.path, path); let failure = health.status.failure().expect("input must be unavailable"); assert_eq!(failure.operation, "parse OpenCode SQLite"); assert!(failure.message.contains("current session schema")); - assert!(cache.get_meta(&path, parser_version).unwrap().is_none()); + assert!(cache.get_meta(&path, decoder_version).unwrap().is_none()); } #[test] @@ -374,18 +373,20 @@ mod tests { .unwrap(); drop(conn); - let unit = InputUnit::sqlite_with_wal(ClientId::OpenCode, path.clone()) - .with_meta(InputUnitMeta::OpenCodeSqlite) - .prepare_snapshot() - .unwrap(); - let parser_version = unit.parser_version; + let unit = InputUnit::sqlite_with_wal( + path.clone(), + DecoderSpec::opencode_sqlite(OPENCODE_CURRENT_SQLITE_REVISION), + ) + .prepare_snapshot() + .unwrap(); + let decoder_version = unit.decoder.version(); let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); let parsed = OPENCODE_ADAPTER.parse_checked(vec![unit.clone()], &ParseContext { pricing: None }); assert_eq!(parsed.len(), 1); - let health = parsed[0].input_health(); - assert_eq!(health.path, path); + let health = &parsed[0].health; + assert_eq!(parsed[0].unit.path, path); assert!(matches!( health.status, crate::input_health::InputStatus::Complete @@ -395,12 +396,15 @@ mod tests { assert_eq!(rejection.key, "malformed-record"); let mut sink = Vec::new(); + let binding = crate::adapters::adapter_for(ClientId::OpenCode).unwrap(); + let mut fold_ctx = FoldContext::new(binding, &mut cache, None); + let mut bound_sink = BoundMessageSink::new(binding, &mut sink); OPENCODE_ADAPTER - .fold(parsed, &mut FoldContext::new(&mut cache, None), &mut sink) + .fold(parsed, &mut fold_ctx, &mut bound_sink) .unwrap(); assert!(sink.is_empty()); cache.save_if_dirty().unwrap(); - let cached = cache.get_meta(&path, parser_version).unwrap().unwrap(); + let cached = cache.get_meta(&path, decoder_version).unwrap().unwrap(); assert_eq!(cached.rejections.total(), 1); let warm_cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); @@ -408,7 +412,7 @@ mod tests { let crate::adapters::CacheHitPlan::Hit(warm) = warm else { panic!("unchanged all-bad input must use the complete cached scan"); }; - let warm_health = warm.input_health(); + let warm_health = &warm.health; assert!(matches!( warm_health.status, crate::input_health::InputStatus::Complete @@ -435,14 +439,16 @@ mod tests { .unwrap(); drop(conn); - let unit = InputUnit::sqlite_with_wal(ClientId::OpenCode, path.clone()) - .with_meta(InputUnitMeta::OpenCodeSqlite) - .prepare_snapshot() - .unwrap(); - let parser_version = unit.parser_version; + let unit = InputUnit::sqlite_with_wal( + path.clone(), + DecoderSpec::opencode_sqlite(OPENCODE_CURRENT_SQLITE_REVISION), + ) + .prepare_snapshot() + .unwrap(); + let decoder_version = unit.decoder.version(); let parsed = OPENCODE_ADAPTER.parse_checked(vec![unit], &ParseContext { pricing: None }); assert_eq!(parsed.len(), 1); - let health = parsed[0].input_health(); + let health = &parsed[0].health; assert!(matches!( health.status, crate::input_health::InputStatus::Complete @@ -455,16 +461,18 @@ mod tests { let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); let mut sink = Vec::new(); - let mut fold_ctx = FoldContext::new(&mut cache, None); + let binding = crate::adapters::adapter_for(ClientId::OpenCode).unwrap(); + let mut fold_ctx = FoldContext::new(binding, &mut cache, None); + let mut bound_sink = BoundMessageSink::new(binding, &mut sink); OPENCODE_ADAPTER - .fold(parsed, &mut fold_ctx, &mut sink) + .fold(parsed, &mut fold_ctx, &mut bound_sink) .unwrap(); assert_eq!(sink.len(), 1); assert_eq!(sink[0].session_id.as_ref(), "session-1"); - assert_eq!(fold_ctx.health.partial_inputs(), 0); - assert_eq!(fold_ctx.health.rejected_records(), 1); - let cached = cache.get_meta(&path, parser_version).unwrap().unwrap(); + assert_eq!(fold_ctx.health().partial_inputs(), 0); + assert_eq!(fold_ctx.health().rejected_records(), 1); + let cached = cache.get_meta(&path, decoder_version).unwrap().unwrap(); assert_eq!(cached.rejections.total(), 1); } } diff --git a/crates/tokscale-core/src/adapters/pi.rs b/crates/tokscale-core/src/adapters/pi.rs index a0e8a753d..7e0965ad5 100644 --- a/crates/tokscale-core/src/adapters/pi.rs +++ b/crates/tokscale-core/src/adapters/pi.rs @@ -3,11 +3,13 @@ use rayon::prelude::*; use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ - AdapterScanContext, FingerprintPolicy, FoldContext, InputDiscoveryError, InputUnit, - LocalInputAdapter, MessageSink, ParseContext, ParsedUnit, + AdapterScanContext, BoundMessageSink, DecoderSpec, FingerprintPolicy, FoldContext, + InputDiscoveryError, InputUnit, LocalInputAdapter, ParseContext, ParsedUnit, }; use crate::clients::ClientId; -use crate::message_cache::{ParserId, ParserVersion}; +use crate::message_cache::DecoderId; +#[cfg(test)] +use crate::message_cache::DecoderVersion; use crate::sessions; pub(crate) struct PiAdapter; @@ -19,27 +21,17 @@ pub(crate) static PI_ADAPTER: PiAdapter = PiAdapter; const PI_RECORD_REJECTION_REVISION: u32 = crate::adapters::MODEL_ID_CANONICALIZATION_REVISION + 5; impl LocalInputAdapter for PiAdapter { - fn client(&self) -> ClientId { - ClientId::Pi - } - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { let units = adapter_discover::discover_default_scanned_units( - ClientId::Pi, + client, ctx, FingerprintPolicy::PlainFile, - )? - .into_iter() - .map(|unit| { - unit.with_parser_version(ParserVersion::new( - ParserId::Pi, - PI_RECORD_REJECTION_REVISION, - )) - }) - .collect(); + DecoderSpec::plain(DecoderId::Pi, PI_RECORD_REJECTION_REVISION), + )?; Ok(units) } @@ -66,7 +58,7 @@ impl LocalInputAdapter for PiAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), crate::adapters::InputPipelineError> { adapter_cache::fold_units(parsed, ctx, sink) } @@ -124,12 +116,13 @@ mod tests { ) -> Vec { let parsed = adapter.parse_checked(units, &ParseContext { pricing: None }); let mut sink = Vec::new(); + let binding = crate::adapters::AdapterBinding::new(ClientId::Pi, adapter); + let mut fold_ctx = FoldContext::new(binding, cache, None); + let mut bound_sink = BoundMessageSink::new(binding, &mut sink); adapter - .fold(parsed, &mut FoldContext::new(cache, None), &mut sink) + .fold(parsed, &mut fold_ctx, &mut bound_sink) .unwrap(); - assert!(sink - .iter() - .all(|message| message.client.as_ref() == adapter.client().as_str())); + assert!(sink.iter().all(|message| message.client == ClientId::Pi)); sink } @@ -151,7 +144,7 @@ mod tests { }; let ctx = scan_context(home.path(), &settings); - let units = PI_ADAPTER.discover_checked(&ctx).unwrap(); + let units = PI_ADAPTER.discover_checked(ClientId::Pi, &ctx).unwrap(); let paths: Vec<_> = units.iter().map(|unit| unit.path.clone()).collect(); let mut expected = vec![default_path, extra_path]; expected.sort_unstable(); @@ -161,7 +154,8 @@ mod tests { .iter() .all(|unit| unit.fingerprint_policy == FingerprintPolicy::PlainFile)); assert!(units.iter().all(|unit| { - unit.parser_version == ParserVersion::new(ParserId::Pi, PI_RECORD_REJECTION_REVISION) + unit.decoder.version() + == DecoderVersion::new(DecoderId::Pi, PI_RECORD_REJECTION_REVISION) })); } @@ -170,7 +164,10 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("pi.jsonl"); write_file(&path, PI_CONTENT); - let units = vec![InputUnit::plain_file(ClientId::Pi, path.clone())]; + let units = vec![InputUnit::plain_file( + path.clone(), + DecoderSpec::plain(DecoderId::Pi, PI_RECORD_REJECTION_REVISION), + )]; let mut cache = message_cache::InputMessageCache::default(); let actual = fold_with_adapter(&PI_ADAPTER, units, &mut cache); @@ -185,14 +182,16 @@ mod tests { let path = dir.path().join("missing.jsonl"); let parsed = PI_ADAPTER.parse_checked( - vec![InputUnit::plain_file(ClientId::Pi, path.clone())], + vec![InputUnit::plain_file( + path.clone(), + DecoderSpec::plain(DecoderId::Pi, PI_RECORD_REJECTION_REVISION), + )], &ParseContext { pricing: None }, ); assert_eq!(parsed.len(), 1); - let health = parsed[0].input_health(); - assert_eq!(health.client, ClientId::Pi); - assert_eq!(health.path, path); + let health = &parsed[0].health; + assert_eq!(parsed[0].unit.path, path); let failure = health.status.failure().expect("input must be unavailable"); assert_eq!(failure.operation, "snapshot input metadata and content"); assert!(failure.message.contains(path.to_str().unwrap())); @@ -211,14 +210,16 @@ mod tests { ); let parsed = PI_ADAPTER.parse_checked( - vec![InputUnit::plain_file(ClientId::Pi, path.clone())], + vec![InputUnit::plain_file( + path.clone(), + DecoderSpec::plain(DecoderId::Pi, PI_RECORD_REJECTION_REVISION), + )], &ParseContext { pricing: None }, ); assert_eq!(parsed.len(), 1); - let health = parsed[0].input_health(); - assert_eq!(health.client, ClientId::Pi); - assert_eq!(health.path, path); + let health = &parsed[0].health; + assert_eq!(parsed[0].unit.path, path); assert!(matches!( &health.status, crate::input_health::InputStatus::Complete @@ -231,8 +232,9 @@ mod tests { let dir = tempfile::TempDir::new().unwrap(); let path = dir.path().join("partial.jsonl"); write_file(&path, PI_CONTENT); - let unit = InputUnit::plain_file(ClientId::Pi, path.clone()).with_parser_version( - ParserVersion::new(ParserId::Pi, PI_RECORD_REJECTION_REVISION), + let unit = InputUnit::plain_file( + path.clone(), + DecoderSpec::plain(DecoderId::Pi, PI_RECORD_REJECTION_REVISION), ); let parsed = adapter_cache::load_or_scan_unit_with( @@ -248,9 +250,9 @@ mod tests { }, ); - assert_eq!(parsed.input_health().path, path); + assert_eq!(parsed.unit.path, path); assert!(matches!( - &parsed.input_health().status, + &parsed.health.status, crate::input_health::InputStatus::Partial { .. } )); assert!(matches!( @@ -271,7 +273,10 @@ mod tests { let path = dir.path().join("pi.jsonl"); write_file(&path, PI_CONTENT); - let units = vec![InputUnit::plain_file(ClientId::Pi, path.clone())]; + let units = vec![InputUnit::plain_file( + path.clone(), + DecoderSpec::plain(DecoderId::Pi, PI_RECORD_REJECTION_REVISION), + )]; let mut cache = message_cache::InputMessageCache::load().unwrap(); let first = fold_with_adapter(&PI_ADAPTER, units.clone(), &mut cache); @@ -288,13 +293,14 @@ mod tests { )); let mut second = Vec::new(); + let binding = crate::adapters::adapter_for(ClientId::Pi).unwrap(); + let mut fold_ctx = FoldContext::new(binding, &mut cache, None); + let mut bound_sink = BoundMessageSink::new(binding, &mut second); PI_ADAPTER - .fold(parsed, &mut FoldContext::new(&mut cache, None), &mut second) + .fold(parsed, &mut fold_ctx, &mut bound_sink) .unwrap(); - assert!(second - .iter() - .all(|message| message.client.as_ref() == ClientId::Pi.as_str())); + assert!(second.iter().all(|message| message.client == ClientId::Pi)); assert_eq!(second, first); restore_env_var("TOKSCALE_CONFIG_DIR", previous_config_dir); } @@ -302,7 +308,10 @@ mod tests { #[test] fn input_unit_plain_file_digest_is_just_path() { let path = PathBuf::from("/tmp/pi.jsonl"); - let unit = InputUnit::plain_file(ClientId::Pi, path.clone()); + let unit = InputUnit::plain_file( + path.clone(), + DecoderSpec::plain(DecoderId::Pi, PI_RECORD_REJECTION_REVISION), + ); assert_eq!(unit.digest_paths(), vec![path]); } diff --git a/crates/tokscale-core/src/adapters/roocode.rs b/crates/tokscale-core/src/adapters/roocode.rs index 14d30368e..bd417b565 100644 --- a/crates/tokscale-core/src/adapters/roocode.rs +++ b/crates/tokscale-core/src/adapters/roocode.rs @@ -5,51 +5,42 @@ use rayon::prelude::*; use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ - AdapterScanContext, FingerprintPolicy, FoldContext, InputDiscoveryError, InputPipelineError, - InputUnit, LocalInputAdapter, MessageSink, ParseContext, ParsedUnit, - MODEL_ID_CANONICALIZATION_REVISION, + AdapterScanContext, BoundMessageSink, DecoderSpec, FingerprintPolicy, FoldContext, + InputDiscoveryError, InputPipelineError, InputUnit, LocalInputAdapter, ParseContext, + ParsedUnit, MODEL_ID_CANONICALIZATION_REVISION, }; use crate::clients::ClientId; -use crate::message_cache::{ParserId, ParserVersion, RelatedInputFailurePolicy}; +#[cfg(test)] +use crate::message_cache::DecoderVersion; +use crate::message_cache::{DecoderId, RelatedInputFailurePolicy}; use crate::sessions; const ROOCODE_SIBLINGS: &[&str] = &["api_conversation_history.json"]; const ROOCODE_RECORD_REJECTION_REVISION: u32 = MODEL_ID_CANONICALIZATION_REVISION + 4; -const ROOCODE_PARSER_VERSION: ParserVersion = - ParserVersion::new(ParserId::RooCode, ROOCODE_RECORD_REJECTION_REVISION); - pub(crate) struct RooCodeAdapter; impl LocalInputAdapter for RooCodeAdapter { - fn client(&self) -> ClientId { - ClientId::RooCode - } - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { - let def = ClientId::RooCode + let def = client .local_def() .expect("Roo Code adapter must have local scan policy"); let mut roots = vec![def.resolve_path(ctx.home_dir)]; roots.extend(roocode_additional_roots(ctx.home_dir)); - roots.extend(adapter_discover::extra_roots_for_client( - ClientId::RooCode, - ctx, - )?); - - Ok(adapter_discover::input_units_from_paths( - ClientId::RooCode, - adapter_discover::scan_roots(ClientId::RooCode, roots, def.pattern)?, + roots.extend(adapter_discover::extra_roots_for_client(client, ctx)?); + + adapter_discover::input_units_from_paths( + client, + adapter_discover::scan_roots(client, roots, def.pattern)?, FingerprintPolicy::PrimaryWithSiblings { sibling_names: ROOCODE_SIBLINGS, related_failure_policy: RelatedInputFailurePolicy::FailInput, }, - )? - .into_iter() - .map(|unit| unit.with_parser_version(ROOCODE_PARSER_VERSION)) - .collect()) + DecoderSpec::plain(DecoderId::RooCode, ROOCODE_RECORD_REJECTION_REVISION), + ) } fn parse_checked(&self, units: Vec, ctx: &ParseContext<'_>) -> Vec { @@ -77,7 +68,7 @@ impl LocalInputAdapter for RooCodeAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), InputPipelineError> { adapter_cache::fold_units(parsed, ctx, sink) } @@ -126,7 +117,7 @@ mod tests { let settings = crate::scanner::ScannerSettings::default(); let ctx = scan_context(home.path(), &settings); let paths: Vec<_> = ROOCODE_ADAPTER - .discover_checked(&ctx) + .discover_checked(ClientId::RooCode, &ctx) .unwrap() .into_iter() .map(|unit| unit.path) @@ -145,7 +136,9 @@ mod tests { let settings = crate::scanner::ScannerSettings::default(); let ctx = scan_context(home.path(), &settings); - let units = ROOCODE_ADAPTER.discover_checked(&ctx).unwrap(); + let units = ROOCODE_ADAPTER + .discover_checked(ClientId::RooCode, &ctx) + .unwrap(); assert_eq!(units.len(), 1); assert_eq!( @@ -183,20 +176,24 @@ mod tests { let settings = crate::scanner::ScannerSettings::default(); let scan_ctx = scan_context(home.path(), &settings); - let mut units = ROOCODE_ADAPTER.discover_checked(&scan_ctx).unwrap(); + let mut units = ROOCODE_ADAPTER + .discover_checked(ClientId::RooCode, &scan_ctx) + .unwrap(); assert_eq!(units.len(), 1); let unit = units.pop().unwrap(); assert_eq!(unit.path, path); - assert_eq!(unit.parser_version, ROOCODE_PARSER_VERSION); + assert_eq!( + unit.decoder.version(), + DecoderVersion::new(DecoderId::RooCode, ROOCODE_RECORD_REJECTION_REVISION) + ); let cache_dir = tempfile::TempDir::new().unwrap(); let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); let parsed = ROOCODE_ADAPTER.parse_checked(vec![unit.clone()], &ParseContext { pricing: None }); assert_eq!(parsed.len(), 1); - let health = parsed[0].input_health(); - assert_eq!(health.client, ClientId::RooCode); - assert_eq!(health.path, path); + let health = &parsed[0].health; + assert_eq!(parsed[0].unit.path, path); assert!(matches!( health.status, crate::input_health::InputStatus::Complete @@ -208,12 +205,14 @@ mod tests { ); let mut sink = Vec::new(); - let mut fold_ctx = FoldContext::new(&mut cache, None); + let binding = crate::adapters::adapter_for(ClientId::RooCode).unwrap(); + let mut fold_ctx = FoldContext::new(binding, &mut cache, None); + let mut bound_sink = BoundMessageSink::new(binding, &mut sink); ROOCODE_ADAPTER - .fold(parsed, &mut fold_ctx, &mut sink) + .fold(parsed, &mut fold_ctx, &mut bound_sink) .unwrap(); assert!(sink.is_empty()); - assert_eq!(fold_ctx.health.rejected_records(), 1); + assert_eq!(fold_ctx.health().rejected_records(), 1); cache.save_if_dirty().unwrap(); let warm_cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); @@ -221,8 +220,7 @@ mod tests { let crate::adapters::CacheHitPlan::Hit(warm) = warm else { panic!("unchanged Roo Code all-bad scan must use cached health"); }; - let warm_health = warm.input_health(); - assert_eq!(warm_health.client, ClientId::RooCode); + let warm_health = &warm.health; assert_eq!(warm_health.rejections.total(), 1); assert_eq!( warm_health.rejections.entries().next().unwrap().key, diff --git a/crates/tokscale-core/src/adapters/runtime.rs b/crates/tokscale-core/src/adapters/runtime.rs new file mode 100644 index 000000000..0f7d698c6 --- /dev/null +++ b/crates/tokscale-core/src/adapters/runtime.rs @@ -0,0 +1,100 @@ +use std::path::PathBuf; + +use crate::input_health::{DataHealth, InputHealth, InputStatus, RejectionSummary}; +use crate::message_cache; +use crate::pricing; +use crate::sessions::ParsedMessage; + +use super::{ + AdapterBinding, InputPipelineError, InputUnit, ParseContext, ParsedUnit, UnifiedMessageSink, +}; + +pub(crate) struct FoldContext<'a> { + binding: AdapterBinding<'a>, + pub input_cache: &'a mut message_cache::InputMessageCache, + pub pricing: Option<&'a pricing::PricingService>, + health: DataHealth, +} + +impl<'a> FoldContext<'a> { + pub(crate) fn new( + binding: AdapterBinding<'a>, + input_cache: &'a mut message_cache::InputMessageCache, + pricing: Option<&'a pricing::PricingService>, + ) -> Self { + Self { + binding, + input_cache, + pricing, + health: DataHealth::default(), + } + } + + pub(crate) fn record_health( + &mut self, + path: PathBuf, + status: InputStatus, + rejections: RejectionSummary, + ) { + self.health.record(InputHealth { + client: self.binding.client(), + path, + status, + rejections, + }); + } + + #[cfg(test)] + pub(crate) fn health(&self) -> &DataHealth { + &self.health + } + + pub(crate) fn take_health(&mut self) -> DataHealth { + std::mem::take(&mut self.health) + } + + pub(crate) fn reparse_one(&self, unit: InputUnit) -> Result { + let mut reparsed = self.binding.adapter().parse_checked( + vec![unit], + &ParseContext { + pricing: self.pricing, + }, + ); + if reparsed.len() != 1 { + return Err(InputPipelineError::contract(format!( + "single-input cache recovery returned {} parsed units instead of one", + reparsed.len() + ))); + } + reparsed.pop().ok_or_else(|| { + InputPipelineError::contract("single-input cache recovery result disappeared") + }) + } +} + +pub(crate) struct BoundMessageSink<'a> { + client: crate::clients::ClientId, + downstream: &'a mut dyn UnifiedMessageSink, +} + +impl<'a> BoundMessageSink<'a> { + pub(crate) fn new( + binding: AdapterBinding<'_>, + downstream: &'a mut dyn UnifiedMessageSink, + ) -> Self { + Self { + client: binding.client(), + downstream, + } + } + + pub(crate) fn emit(&mut self, message: ParsedMessage) { + self.downstream.push_message(message.attribute(self.client)); + } + + pub(crate) fn emit_all(&mut self, messages: impl IntoIterator) { + for message in messages { + self.emit(message); + } + } +} diff --git a/crates/tokscale-core/src/adapters/warp.rs b/crates/tokscale-core/src/adapters/warp.rs index 34985f20f..4a8f18aac 100644 --- a/crates/tokscale-core/src/adapters/warp.rs +++ b/crates/tokscale-core/src/adapters/warp.rs @@ -3,12 +3,14 @@ use rayon::prelude::*; use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ - AdapterScanContext, FingerprintPolicy, FoldContext, InputDiscoveryError, InputUnit, - LocalInputAdapter, MessageSink, ParseContext, ParsedUnit, + AdapterScanContext, BoundMessageSink, DecoderSpec, FingerprintPolicy, FoldContext, + InputDiscoveryError, InputUnit, LocalInputAdapter, ParseContext, ParsedUnit, }; use crate::clients::ClientId; use crate::local_clients; -use crate::message_cache::{ParserId, ParserVersion}; +use crate::message_cache::DecoderId; +#[cfg(test)] +use crate::message_cache::DecoderVersion; use crate::sessions; const WARP_RECORD_REJECTION_REVISION: u32 = 5; @@ -16,43 +18,33 @@ const WARP_RECORD_REJECTION_REVISION: u32 = 5; pub(crate) struct WarpAdapter; impl LocalInputAdapter for WarpAdapter { - fn client(&self) -> ClientId { - ClientId::Warp - } - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { - let def = ClientId::Warp + let def = client .local_def() .expect("Warp adapter must have local scan policy"); let mut paths = adapter_discover::scan_roots( - ClientId::Warp, + client, local_clients::warp_sqlite_roots(ctx.home_dir), def.pattern, )?; paths.extend(adapter_discover::scan_roots( - ClientId::Warp, - adapter_discover::extra_roots_for_client(ClientId::Warp, ctx)?, + client, + adapter_discover::extra_roots_for_client(client, ctx)?, def.pattern, )?); let units = adapter_discover::input_units_from_paths( - ClientId::Warp, + client, paths, FingerprintPolicy::SqliteWithWal, + DecoderSpec::plain(DecoderId::Warp, WARP_RECORD_REJECTION_REVISION), )?; - Ok(units - .into_iter() - .map(|unit| { - unit.with_parser_version(ParserVersion::new( - ParserId::Warp, - WARP_RECORD_REJECTION_REVISION, - )) - }) - .collect()) + Ok(units) } fn parse_checked(&self, units: Vec, ctx: &ParseContext<'_>) -> Vec { @@ -68,7 +60,7 @@ impl LocalInputAdapter for WarpAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), crate::adapters::InputPipelineError> { adapter_cache::fold_units(parsed, ctx, sink) } @@ -102,7 +94,7 @@ mod tests { scanner_settings: &settings, }; - let units = WARP_ADAPTER.discover_checked(&ctx).unwrap(); + let units = WARP_ADAPTER.discover_checked(ClientId::Warp, &ctx).unwrap(); let paths: Vec<_> = units.iter().map(|unit| unit.path.clone()).collect(); assert_eq!(paths, vec![default_db, extra_db]); @@ -110,8 +102,8 @@ mod tests { .iter() .all(|unit| unit.fingerprint_policy == FingerprintPolicy::SqliteWithWal)); assert!(units.iter().all(|unit| { - unit.parser_version - == ParserVersion::new(ParserId::Warp, WARP_RECORD_REJECTION_REVISION) + unit.decoder.version() + == DecoderVersion::new(DecoderId::Warp, WARP_RECORD_REJECTION_REVISION) })); } } diff --git a/crates/tokscale-core/src/adapters/zed.rs b/crates/tokscale-core/src/adapters/zed.rs index 592405b3f..25b7d2f4c 100644 --- a/crates/tokscale-core/src/adapters/zed.rs +++ b/crates/tokscale-core/src/adapters/zed.rs @@ -3,11 +3,12 @@ use std::path::PathBuf; use crate::adapters::cache as adapter_cache; use crate::adapters::discover as adapter_discover; use crate::adapters::{ - AdapterScanContext, FingerprintPolicy, FoldContext, InputDiscoveryError, InputUnit, - LocalInputAdapter, MessageSink, ParseContext, ParsedUnit, ZED_RECORD_FILTER_REVISION, + AdapterScanContext, BoundMessageSink, DecoderSpec, FingerprintPolicy, FoldContext, + InputDiscoveryError, InputUnit, LocalInputAdapter, ParseContext, ParsedUnit, + ZED_RECORD_FILTER_REVISION, }; use crate::clients::ClientId; -use crate::message_cache::{ParserId, ParserVersion}; +use crate::message_cache::DecoderId; use crate::sessions; pub(crate) struct ZedAdapter; @@ -15,44 +16,34 @@ pub(crate) struct ZedAdapter; pub(crate) static ZED_ADAPTER: ZedAdapter = ZedAdapter; impl LocalInputAdapter for ZedAdapter { - fn client(&self) -> ClientId { - ClientId::Zed - } - fn discover_checked( &self, + client: ClientId, ctx: &AdapterScanContext<'_>, ) -> Result, InputDiscoveryError> { - let def = ClientId::Zed + let def = client .local_def() .expect("Zed adapter must have local scan policy"); let mut paths = Vec::new(); adapter_discover::push_existing_file( - ClientId::Zed, + client, zed_default_db_path(ctx.home_dir), &mut paths, )?; paths.extend(adapter_discover::scan_roots( - ClientId::Zed, - adapter_discover::extra_roots_for_client(ClientId::Zed, ctx)?, + client, + adapter_discover::extra_roots_for_client(client, ctx)?, def.pattern, )?); let units = adapter_discover::input_units_from_paths( - ClientId::Zed, + client, paths, FingerprintPolicy::SqliteWithWal, - )? - .into_iter() - .map(|unit| { - unit.with_parser_version(ParserVersion::new( - ParserId::Zed, - ZED_RECORD_FILTER_REVISION, - )) - }) - .collect(); + DecoderSpec::plain(DecoderId::Zed, ZED_RECORD_FILTER_REVISION), + )?; Ok(units) } @@ -81,7 +72,7 @@ impl LocalInputAdapter for ZedAdapter { &self, parsed: Vec, ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, + sink: &mut BoundMessageSink<'_>, ) -> Result<(), crate::adapters::InputPipelineError> { adapter_cache::fold_units(parsed, ctx, sink) } @@ -204,7 +195,7 @@ mod tests { }; let ctx = scan_context(home.path(), &settings); - let units = ZED_ADAPTER.discover_checked(&ctx).unwrap(); + let units = ZED_ADAPTER.discover_checked(ClientId::Zed, &ctx).unwrap(); let paths: Vec<_> = units.iter().map(|unit| unit.path.clone()).collect(); let mut expected = vec![default_db, extra_db]; expected.sort_unstable(); @@ -234,20 +225,27 @@ mod tests { .unwrap(); drop(conn); - let units = vec![InputUnit::sqlite_with_wal(ClientId::Zed, db_path.clone())]; + let units = vec![InputUnit::sqlite_with_wal( + db_path.clone(), + DecoderSpec::plain(DecoderId::Zed, ZED_RECORD_FILTER_REVISION), + )]; let mut cache = message_cache::InputMessageCache::default(); let parsed = ZED_ADAPTER.parse_checked(units, &ParseContext { pricing: None }); let mut sink = Vec::new(); - let mut fold_ctx = FoldContext::new(&mut cache, None); - ZED_ADAPTER.fold(parsed, &mut fold_ctx, &mut sink).unwrap(); + let binding = crate::adapters::adapter_for(ClientId::Zed).unwrap(); + let mut fold_ctx = FoldContext::new(binding, &mut cache, None); + let mut bound_sink = BoundMessageSink::new(binding, &mut sink); + ZED_ADAPTER + .fold(parsed, &mut fold_ctx, &mut bound_sink) + .unwrap(); assert_eq!(sink.len(), 1); - assert_eq!(sink[0].client.as_ref(), ClientId::Zed.as_str()); + assert_eq!(sink[0].client, ClientId::Zed); assert_eq!(sink[0].session_id.as_ref(), "zed-thread-good"); - assert_eq!(fold_ctx.health.rejected_records(), 1); - assert_eq!(fold_ctx.health.failed_inputs(), 0); - assert_eq!(fold_ctx.health.partial_inputs(), 0); - let input = &fold_ctx.health.inputs()[0]; + assert_eq!(fold_ctx.health().rejected_records(), 1); + assert_eq!(fold_ctx.health().failed_inputs(), 0); + assert_eq!(fold_ctx.health().partial_inputs(), 0); + let input = &fold_ctx.health().inputs()[0]; assert_eq!(input.client, ClientId::Zed); assert_eq!(input.path, db_path); let entries: Vec<_> = input.rejections.entries().collect(); @@ -275,18 +273,21 @@ mod tests { drop(conn); let mut cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); - let unit = InputUnit::sqlite_with_wal(ClientId::Zed, db_path.clone()) - .with_parser_version(ParserVersion::new( - ParserId::Zed, - ZED_RECORD_FILTER_REVISION, - )) - .prepare_snapshot() - .unwrap(); + let unit = InputUnit::sqlite_with_wal( + db_path.clone(), + DecoderSpec::plain(DecoderId::Zed, ZED_RECORD_FILTER_REVISION), + ) + .prepare_snapshot() + .unwrap(); let cold = ZED_ADAPTER.parse_checked(vec![unit.clone()], &ParseContext { pricing: None }); let mut sink = Vec::new(); - let mut fold_ctx = FoldContext::new(&mut cache, None); - ZED_ADAPTER.fold(cold, &mut fold_ctx, &mut sink).unwrap(); - assert_eq!(fold_ctx.health.rejected_records(), 1); + let binding = crate::adapters::adapter_for(ClientId::Zed).unwrap(); + let mut fold_ctx = FoldContext::new(binding, &mut cache, None); + let mut bound_sink = BoundMessageSink::new(binding, &mut sink); + ZED_ADAPTER + .fold(cold, &mut fold_ctx, &mut bound_sink) + .unwrap(); + assert_eq!(fold_ctx.health().rejected_records(), 1); cache.save_if_dirty().unwrap(); let warm_cache = message_cache::InputMessageCache::with_cache_dir(cache_dir.path()); @@ -299,7 +300,7 @@ mod tests { panic!("unchanged input with cached scan must plan a warm hit") } }; - let health = hit.input_health(); + let health = &hit.health; assert_eq!(health.rejections.total(), 1); let entries: Vec<_> = health.rejections.entries().collect(); assert_eq!(entries[0].key, "missing-model"); @@ -313,18 +314,22 @@ mod tests { insert_thread(&conn, "zed-thread-1", "claude-sonnet-4-5"); drop(conn); - let units = vec![InputUnit::sqlite_with_wal(ClientId::Zed, db_path.clone())]; + let units = vec![InputUnit::sqlite_with_wal( + db_path.clone(), + DecoderSpec::plain(DecoderId::Zed, ZED_RECORD_FILTER_REVISION), + )]; let mut cache = message_cache::InputMessageCache::default(); let parsed = ZED_ADAPTER.parse_checked(units, &ParseContext { pricing: None }); let mut actual = Vec::new(); + let binding = crate::adapters::adapter_for(ClientId::Zed).unwrap(); + let mut fold_ctx = FoldContext::new(binding, &mut cache, None); + let mut bound_sink = BoundMessageSink::new(binding, &mut actual); ZED_ADAPTER - .fold(parsed, &mut FoldContext::new(&mut cache, None), &mut actual) + .fold(parsed, &mut fold_ctx, &mut bound_sink) .unwrap(); let expected = finalized(sessions::zed::parse_zed_sqlite(&db_path).unwrap().messages); - assert!(actual - .iter() - .all(|message| message.client.as_ref() == ClientId::Zed.as_str())); + assert!(actual.iter().all(|message| message.client == ClientId::Zed)); assert_eq!(actual, expected); } } diff --git a/crates/tokscale-core/src/aggregate/keys.rs b/crates/tokscale-core/src/aggregate/keys.rs index e49d53c73..2c5e0b115 100644 --- a/crates/tokscale-core/src/aggregate/keys.rs +++ b/crates/tokscale-core/src/aggregate/keys.rs @@ -4,7 +4,7 @@ use std::{collections::HashSet, hash::Hash, sync::Arc}; use serde::{Deserialize, Serialize}; -use crate::{sessions, GroupBy, UnifiedMessage}; +use crate::{sessions, ClientId, GroupBy, UnifiedMessage}; pub const UNKNOWN_WORKSPACE_LABEL: &str = "Unknown workspace"; @@ -129,11 +129,11 @@ impl WorkspaceKey { pub(crate) enum GroupedModelKey { Model(Arc), ClientModel { - client: Arc, + client: ClientId, model: Arc, }, ClientProviderModel { - client: Arc, + client: ClientId, provider: Arc, model: Arc, }, @@ -150,11 +150,11 @@ impl GroupedModelKey { match group_by { GroupBy::Model => Self::Model(Arc::clone(&msg.model_id)), GroupBy::ClientModel => Self::ClientModel { - client: Arc::clone(&msg.client), + client: msg.client, model: Arc::clone(&msg.model_id), }, GroupBy::ClientProviderModel => Self::ClientProviderModel { - client: Arc::clone(&msg.client), + client: msg.client, provider: Arc::clone(&msg.provider_id), model: Arc::clone(&msg.model_id), }, @@ -182,7 +182,7 @@ impl GroupedModelKey { } Self::ClientModel { client, model } => { output.push_str("cm|"); - push_len_prefixed(&mut output, client); + push_len_prefixed(&mut output, client.as_str()); push_len_prefixed(&mut output, model); } Self::ClientProviderModel { @@ -191,7 +191,7 @@ impl GroupedModelKey { model, } => { output.push_str("cpm|"); - push_len_prefixed(&mut output, client); + push_len_prefixed(&mut output, client.as_str()); push_len_prefixed(&mut output, provider); push_len_prefixed(&mut output, model); } @@ -240,7 +240,7 @@ impl HourlyModelKey { /// switching groupings is an in-memory re-fold instead of a rescan. #[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)] pub(crate) struct FineModelKey { - pub(crate) client: Arc, + pub(crate) client: ClientId, pub(crate) provider: Arc, pub(crate) workspace: WorkspaceKey, pub(crate) session: Arc, @@ -250,7 +250,7 @@ pub(crate) struct FineModelKey { impl FineModelKey { pub(crate) fn from_message(msg: &UnifiedMessage) -> Self { Self { - client: Arc::clone(&msg.client), + client: msg.client, provider: Arc::clone(&msg.provider_id), workspace: WorkspaceKey::from_message(msg), session: Arc::clone(&msg.session_id), @@ -263,11 +263,11 @@ impl FineModelKey { match group_by { GroupBy::Model => GroupedModelKey::Model(Arc::clone(&self.model)), GroupBy::ClientModel => GroupedModelKey::ClientModel { - client: Arc::clone(&self.client), + client: self.client, model: Arc::clone(&self.model), }, GroupBy::ClientProviderModel => GroupedModelKey::ClientProviderModel { - client: Arc::clone(&self.client), + client: self.client, provider: Arc::clone(&self.provider), model: Arc::clone(&self.model), }, @@ -328,26 +328,20 @@ mod tests { use std::sync::Arc; use super::*; - use crate::TokenBreakdown; + use crate::{ClientId, TokenBreakdown}; fn message() -> UnifiedMessage { - UnifiedMessage { - client: Arc::from("client"), - model_id: Arc::from("model"), - provider_id: Arc::from("provider"), - session_id: Arc::from("session"), - is_main_session: true, - workspace_key: Some(Arc::from("workspace")), - workspace_label: Some(Arc::from("Workspace")), - timestamp: 0, - tokens: TokenBreakdown::default(), - cost: 0.0, - message_count: 1, - agent: None, - agent_instance: None, - dedup_key: None, - is_turn_start: false, - } + let mut message = UnifiedMessage::new( + ClientId::Codex, + "model", + "provider", + "session", + 0, + TokenBreakdown::default(), + 0.0, + ); + message.set_workspace(Some("workspace".to_string()), Some("Workspace".to_string())); + message } #[test] @@ -372,31 +366,28 @@ mod tests { } #[test] - fn structured_keys_do_not_alias_separator_collisions() { + fn structured_keys_include_typed_client_identity() { let mut left = message(); - left.client = Arc::from("a:b"); + left.client = ClientId::Amp; left.model_id = Arc::from("c"); let mut right = message(); - right.client = Arc::from("a"); - right.model_id = Arc::from("b:c"); + right.client = ClientId::Codex; + right.model_id = Arc::from("c"); let left = GroupedModelKey::from_message(&GroupBy::ClientModel, &left); let right = GroupedModelKey::from_message(&GroupBy::ClientModel, &right); assert_ne!(left, right); assert_ne!(left.map_key(), right.map_key()); - assert_eq!(left.map_key(), "v1|cm|3:a:b1:c"); - assert_eq!(right.map_key(), "v1|cm|1:a3:b:c"); + assert_eq!(left.map_key(), "v1|cm|3:amp1:c"); + assert_eq!(right.map_key(), "v1|cm|5:codex1:c"); } #[test] - fn keys_compare_string_values_instead_of_arc_addresses() { + fn model_keys_compare_string_values_instead_of_arc_addresses() { let mut left = message(); - left.client = Arc::from(String::from("same-client")); left.model_id = Arc::from(String::from("same-model")); let mut right = message(); - right.client = Arc::from(String::from("same-client")); right.model_id = Arc::from(String::from("same-model")); - assert!(!Arc::ptr_eq(&left.client, &right.client)); assert!(!Arc::ptr_eq(&left.model_id, &right.model_id)); assert_eq!( @@ -405,6 +396,16 @@ mod tests { ); } + #[test] + fn fine_model_key_persists_client_as_stable_catalog_id() { + let key = FineModelKey::from_message(&message()); + let encoded = serde_json::to_value(&key).unwrap(); + assert_eq!(encoded["client"], "codex"); + + let restored: FineModelKey = serde_json::from_value(encoded).unwrap(); + assert_eq!(restored.client, ClientId::Codex); + } + #[test] fn hourly_structured_keys_do_not_alias_separator_collisions() { let mut left = message(); @@ -466,7 +467,9 @@ mod tests { #[test] fn unrelated_grouping_does_not_clone_workspace_or_session() { - let msg = message(); + let mut msg = message(); + msg.workspace_key = Some(Arc::from(String::from("unique-workspace"))); + msg.session_id = Arc::from(String::from("unique-session")); let workspace = Arc::clone(msg.workspace_key.as_ref().expect("workspace")); let session = Arc::clone(&msg.session_id); let workspace_before = Arc::strong_count(&workspace); diff --git a/crates/tokscale-core/src/aggregate/tui.rs b/crates/tokscale-core/src/aggregate/tui.rs index ffc046a99..c3795037b 100644 --- a/crates/tokscale-core/src/aggregate/tui.rs +++ b/crates/tokscale-core/src/aggregate/tui.rs @@ -493,7 +493,7 @@ pub fn find_peak_hour(hourly: &[HourlyUsage]) -> Option<(u32, u64, f64)> { #[derive(Default, Serialize, Deserialize)] pub struct TuiAcc { #[serde(with = "map_as_vec")] - usage_totals_by_client: HashMap, UsageTotalsBucket>, + usage_totals_by_client: HashMap, #[serde(with = "map_as_vec")] model_map: HashMap, #[serde(with = "map_as_vec")] @@ -675,27 +675,27 @@ struct FineModelBucket { struct TuiModelBucket { model: Arc, providers: IdentitySet>, - client: Arc, + client: ClientId, workspace_key: Option>, workspace_label: Option>, tokens: UsageTokenBreakdown, cost: f64, - sessions: IdentitySet<(Arc, Arc)>, + sessions: IdentitySet<(ClientId, Arc)>, // Boxed only for grouping modes that merge clients; keeps client-scoped // high-cardinality buckets free of an inline HashMap. #[allow(clippy::box_collection)] - client_totals: Option, ClientContributionOrder>>>, + client_totals: Option>>, } #[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)] enum AgentInstanceKey { Explicit(Arc), - Derived { client: Arc, session: Arc }, + Derived { client: ClientId, session: Arc }, } #[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)] struct AgentKey { - client: Arc, + client: ClientId, agent: Arc, } @@ -710,7 +710,7 @@ struct AgentBucket { #[derive(Serialize, Deserialize)] struct DailyBucket { date: NaiveDate, - clients: HashMap, DailyClientBucket>, + clients: HashMap, } #[derive(Serialize, Deserialize)] @@ -748,7 +748,7 @@ struct DailyModelBucket { #[derive(Serialize, Deserialize)] struct HourlyBucket { datetime: NaiveDateTime, - clients: HashMap, HourlyClientBucket>, + clients: HashMap, } #[derive(Serialize, Deserialize)] @@ -791,7 +791,7 @@ fn materialize_tui_model(bucket: TuiModelBucket) -> UsageModelEntry { }); clients .iter() - .map(|(client, _)| client.as_ref()) + .map(|(client, _)| client.as_str()) .collect::>() .join(crate::usage_views::MODEL_CLIENT_SEPARATOR) } else { @@ -833,8 +833,8 @@ fn materialize_daily_model(model: DailyModelBucket) -> DailyModelInfo { /// floating-point sums stay deterministic; the first-created bucket in the group /// attributes the provider and workspace label (its first message is the /// group's first message, matching a direct grouped fold). -fn client_is_selected(client: &str, selected: Option<&HashSet>) -> bool { - selected.is_none_or(|selected| crate::selected_client_ids_include(client, selected)) +fn client_is_selected(client: ClientId, selected: Option<&HashSet>) -> bool { + selected.is_none_or(|selected| selected.contains(&client)) } fn materialize_daily_common( @@ -850,9 +850,9 @@ fn materialize_daily_common( let mut selected_clients: Vec<_> = bucket .clients .iter() - .filter(|(client, _)| client_is_selected(client, selected)) + .filter(|(client, _)| client_is_selected(**client, selected)) .collect(); - selected_clients.sort_by_key(|(client, _)| *client); + selected_clients.sort_by_key(|(client, _)| **client); for (client, client_bucket) in selected_clients { clients.insert( @@ -938,7 +938,7 @@ fn materialize_daily_models( let client_models: BTreeMap<_, _> = bucket .clients .iter() - .filter(|(client, _)| client_is_selected(client, selected)) + .filter(|(client, _)| client_is_selected(**client, selected)) .map(|(client, client_bucket)| { ( client.to_string(), @@ -980,9 +980,9 @@ fn materialize_hourly_common( let mut selected_clients: Vec<_> = bucket .clients .iter() - .filter(|(client, _)| client_is_selected(client, selected)) + .filter(|(client, _)| client_is_selected(**client, selected)) .collect(); - selected_clients.sort_by_key(|(client, _)| *client); + selected_clients.sort_by_key(|(client, _)| **client); for (client, client_bucket) in selected_clients { clients.insert(client.to_string()); @@ -1011,7 +1011,7 @@ fn materialize_hourly_models( ) -> Option { let mut fine_models = Vec::new(); for (client, client_bucket) in &bucket.clients { - if client_is_selected(client, selected) { + if client_is_selected(*client, selected) { fine_models.extend(client_bucket.models.iter()); } } @@ -1055,9 +1055,10 @@ impl TuiAcc { .expect("TUI aggregation sequence exceeds usize::MAX"); let msg_cost = sane_cost(msg.cost); + let client = msg.client; self.usage_totals_by_client - .entry(Arc::clone(&msg.client)) + .entry(client) .or_default() .push(msg, msg_cost); @@ -1082,7 +1083,7 @@ impl TuiAcc { let agent_entry = self .agent_map .entry(AgentKey { - client: Arc::clone(&msg.client), + client, agent: Arc::clone(agent), }) .or_insert_with(|| AgentBucket { @@ -1098,7 +1099,7 @@ impl TuiAcc { .saturating_add(msg.message_count.max(0) as u32); let instance_key = msg.agent_instance.as_ref().map_or_else( || AgentInstanceKey::Derived { - client: Arc::clone(&msg.client), + client, session: Arc::clone(&msg.session_id), }, |instance| AgentInstanceKey::Explicit(Arc::clone(instance)), @@ -1112,16 +1113,17 @@ impl TuiAcc { clients: HashMap::new(), }); - let client_entry = daily_entry - .clients - .entry(Arc::clone(&msg.client)) - .or_insert_with(|| DailyClientBucket { - tokens: UsageTokenBreakdown::default(), - cost: 0.0, - message_count: 0, - turn_count: 0, - models: HashMap::new(), - }); + let client_entry = + daily_entry + .clients + .entry(client) + .or_insert_with(|| DailyClientBucket { + tokens: UsageTokenBreakdown::default(), + cost: 0.0, + message_count: 0, + turn_count: 0, + models: HashMap::new(), + }); add_unified_tokens(&mut client_entry.tokens, &msg.tokens); client_entry.cost += msg_cost; client_entry.message_count = client_entry @@ -1156,16 +1158,17 @@ impl TuiAcc { datetime: bucket, clients: HashMap::new(), }); - let client_entry = hourly_entry - .clients - .entry(Arc::clone(&msg.client)) - .or_insert_with(|| HourlyClientBucket { - tokens: UsageTokenBreakdown::default(), - cost: 0.0, - models: HashMap::new(), - message_count: 0, - turn_count: 0, - }); + let client_entry = + hourly_entry + .clients + .entry(client) + .or_insert_with(|| HourlyClientBucket { + tokens: UsageTokenBreakdown::default(), + cost: 0.0, + models: HashMap::new(), + message_count: 0, + turn_count: 0, + }); add_unified_tokens(&mut client_entry.tokens, &msg.tokens); client_entry.cost += msg_cost; client_entry.message_count = client_entry @@ -1203,7 +1206,7 @@ impl TuiAcc { OneOrMany<(&FineModelKey, &FineModelBucket)>, > = HashMap::new(); for (fine_key, fine_model) in &self.model_map { - if !client_is_selected(&fine_key.client, selected) { + if !client_is_selected(fine_key.client, selected) { continue; } let fine_bucket = (fine_key, fine_model); @@ -1232,7 +1235,7 @@ impl TuiAcc { TuiModelBucket { model: Arc::clone(&fine_key.model), providers: IdentitySet::default(), - client: Arc::clone(&fine_key.client), + client: fine_key.client, workspace_key, workspace_label, tokens: UsageTokenBreakdown::default(), @@ -1247,7 +1250,7 @@ impl TuiAcc { .client_totals .as_mut() .expect("merge-client TUI grouping has client totals") - .entry(Arc::clone(&fine_key.client)) + .entry(fine_key.client) .or_insert_with(|| ClientContributionOrder { first_seen: fine_model.first_seen, total_tokens: 0, @@ -1265,7 +1268,7 @@ impl TuiAcc { model_entry .sessions - .insert((Arc::clone(&fine_key.client), Arc::clone(&fine_key.session))); + .insert((fine_key.client, Arc::clone(&fine_key.session))); } model_buckets.push(( key, @@ -1314,7 +1317,7 @@ impl TuiAcc { .agent_map .iter() .filter_map(|(key, agent)| { - if !client_is_selected(&key.client, selected) { + if !client_is_selected(key.client, selected) { return None; } Some(AgentEntry { @@ -1360,9 +1363,9 @@ impl TuiAcc { let mut client_totals: Vec<_> = self .usage_totals_by_client .iter() - .filter(|(client, _)| client_is_selected(client, selected)) + .filter(|(client, _)| client_is_selected(**client, selected)) .collect(); - client_totals.sort_by_key(|(client, _)| *client); + client_totals.sort_by_key(|(client, _)| **client); let mut total_token_breakdown = UsageTokenBreakdown::default(); let mut total_cost = 0.0; @@ -1475,7 +1478,7 @@ mod tests { } fn make_workspace_message( - client: &str, + client: ClientId, model_id: &str, provider_id: &str, session_id: &str, @@ -1507,7 +1510,7 @@ mod tests { #[allow(clippy::too_many_arguments)] fn make_message_with_tokens( - client: &str, + client: ClientId, model_id: &str, provider_id: &str, session_id: &str, @@ -1775,7 +1778,7 @@ mod tests { .aggregate_messages( vec![ make_workspace_message( - "opencode", + ClientId::OpenCode, "mimo-v2.5-pro", "xiaomi", "session-1", @@ -1784,7 +1787,7 @@ mod tests { None, ), make_workspace_message( - "opencode", + ClientId::OpenCode, "mimo-v2.5-pro", "xiaomi", "session-2", @@ -1811,7 +1814,7 @@ mod tests { .aggregate_messages( vec![ make_workspace_message( - "opencode", + ClientId::OpenCode, "mimo-v2.5-pro", "xiaomi", "session-1", @@ -1820,7 +1823,7 @@ mod tests { None, ), make_workspace_message( - "opencode", + ClientId::OpenCode, "mimo-v2.5-pro", "xiaomi", "session-2", @@ -1852,7 +1855,7 @@ mod tests { let usage = loader .aggregate_messages( vec![make_workspace_message( - "opencode", + ClientId::OpenCode, "gpt-5.5", "openai", "session-1", @@ -1885,7 +1888,7 @@ mod tests { .aggregate_messages( vec![ make_workspace_message( - "opencode", + ClientId::OpenCode, "gpt-5.5", "openai", "session-1", @@ -1894,7 +1897,7 @@ mod tests { None, ), make_workspace_message( - "opencode", + ClientId::OpenCode, "gpt-5.5", "microsoft", "session-2", @@ -1925,7 +1928,7 @@ mod tests { .aggregate_messages( vec![ make_workspace_message( - "claude", + ClientId::Claude, "kimi-for-coding", "kimi", "session-1", @@ -1934,7 +1937,7 @@ mod tests { None, ), make_workspace_message( - "claude", + ClientId::Claude, "kimi-for-coding", "kimi", "session-2", @@ -1957,7 +1960,7 @@ mod tests { let loader = TuiUsageHarness; let messages = vec![ UnifiedMessage::new_with_agent( - "opencode", + ClientId::OpenCode, "claude-sonnet-4", "anthropic", "session-1", @@ -1973,7 +1976,7 @@ mod tests { Some("Builder".to_string()), ), UnifiedMessage::new_with_agent( - "roocode", + ClientId::RooCode, "claude-sonnet-4", "anthropic", "session-2", @@ -2023,7 +2026,7 @@ mod tests { .aggregate_messages( vec![ make_message_with_tokens( - "opencode", + ClientId::OpenCode, "gpt-5.5", "openai", "session-opencode", @@ -2034,7 +2037,7 @@ mod tests { 0, ), make_message_with_tokens( - "codex", + ClientId::Codex, "gpt-5.5", "openai", "session-codex", @@ -2045,7 +2048,7 @@ mod tests { 0, ), make_message_with_tokens( - "pi", + ClientId::Pi, "gpt-5.5", "openai", "session-pi", @@ -2071,7 +2074,7 @@ mod tests { .aggregate_messages( vec![ make_workspace_message( - "claude", + ClientId::Claude, "claude-sonnet-4.5", "anthropic", "session-1", @@ -2080,7 +2083,7 @@ mod tests { Some("repo-a"), ), make_workspace_message( - "qwen", + ClientId::Qwen, "claude-sonnet-4.5", "anthropic", "session-2", @@ -2110,7 +2113,7 @@ mod tests { .aggregate_messages( vec![ make_workspace_message( - "claude", + ClientId::Claude, "claude-sonnet-4.5", "anthropic", "session-1", @@ -2119,7 +2122,7 @@ mod tests { None, ), make_workspace_message( - "claude", + ClientId::Claude, "claude-sonnet-4.5", "anthropic", "session-2", @@ -2149,7 +2152,7 @@ mod tests { .aggregate_messages( vec![ make_workspace_message( - "claude", + ClientId::Claude, "claude-sonnet-4.5", "anthropic", "session-1", @@ -2158,7 +2161,7 @@ mod tests { Some("unknown-workspace"), ), make_workspace_message( - "claude", + ClientId::Claude, "claude-sonnet-4.5", "anthropic", "session-2", @@ -2191,7 +2194,7 @@ mod tests { .aggregate_messages( vec![ make_workspace_message( - "claude", + ClientId::Claude, "claude-sonnet-4.5", "anthropic", "session-1", @@ -2200,7 +2203,7 @@ mod tests { Some("repo-a"), ), make_workspace_message( - "claude", + ClientId::Claude, "claude-sonnet-4.5", "anthropic", "session-2", @@ -2259,7 +2262,7 @@ mod tests { .aggregate_messages( vec![ make_workspace_message( - "claude", + ClientId::Claude, "claude-sonnet-4.5", "anthropic", "session-1", @@ -2268,7 +2271,7 @@ mod tests { Some("demo"), ), make_workspace_message( - "claude", + ClientId::Claude, "claude-sonnet-4.5", "anthropic", "session-2", @@ -2328,7 +2331,7 @@ mod tests { let usage = loader .aggregate_messages( vec![make_workspace_message( - "claude", + ClientId::Claude, "claude-sonnet-4.5", "anthropic", "session-1", @@ -2369,7 +2372,7 @@ mod tests { .aggregate_messages( vec![ make_workspace_message( - "claude", + ClientId::Claude, "c", "anthropic", "session-1", @@ -2378,7 +2381,7 @@ mod tests { Some("workspace-ab"), ), make_workspace_message( - "claude", + ClientId::Claude, "b:c", "anthropic", "session-2", @@ -2414,7 +2417,7 @@ mod tests { .aggregate_messages( vec![ UnifiedMessage::new( - "claude", + ClientId::Claude, "claude-sonnet-4.5", "anthropic", "session-1", @@ -2429,7 +2432,7 @@ mod tests { 1.0, ), UnifiedMessage::new( - "claude", + ClientId::Claude, "claude-sonnet-4.5", "microsoft", "session-2", @@ -2474,7 +2477,7 @@ mod tests { .aggregate_messages( vec![ UnifiedMessage::new( - "claude", + ClientId::Claude, "claude-sonnet-4.5", "anthropic", "session-1", @@ -2489,7 +2492,7 @@ mod tests { 1.0, ), UnifiedMessage::new( - "gemini", + ClientId::Gemini, "claude-sonnet-4.5", "anthropic", "session-2", @@ -2531,7 +2534,7 @@ mod tests { let loader = TuiUsageHarness; let messages = vec![ UnifiedMessage::new_with_agent( - "opencode", + ClientId::OpenCode, "claude-opus-4.6", "anthropic", "session-1", @@ -2547,7 +2550,7 @@ mod tests { Some("Sisyphus".to_string()), ), UnifiedMessage::new_with_agent( - "opencode", + ClientId::OpenCode, "claude-opus-4.6", "anthropic", "session-2", @@ -2584,7 +2587,7 @@ mod tests { let loader = TuiUsageHarness; let messages = vec![ UnifiedMessage::new_with_agent( - "opencode", + ClientId::OpenCode, "claude-opus-4.6", "anthropic", "session-1", @@ -2600,7 +2603,7 @@ mod tests { Some("Hephaestus".to_string()), ), UnifiedMessage::new_with_agent( - "opencode", + ClientId::OpenCode, "claude-opus-4.6", "anthropic", "session-2", @@ -2631,7 +2634,7 @@ mod tests { let loader = TuiUsageHarness; let messages = vec![ UnifiedMessage::new_with_agent( - "claude", + ClientId::Claude, "claude-opus-4.6", "anthropic", "session-1", @@ -2647,7 +2650,7 @@ mod tests { Some("Sisyphus".to_string()), ), UnifiedMessage::new_with_agent( - "claude", + ClientId::Claude, "claude-opus-4.6", "anthropic", "session-2", @@ -2677,7 +2680,7 @@ mod tests { } fn collision_message( - client: &str, + client: ClientId, provider: &str, session: &str, model: &str, @@ -2704,13 +2707,13 @@ mod tests { let cases = [ ( GroupBy::ClientModel, - collision_message("a:b", "first", "same", "c", 10, timestamp), - collision_message("a", "second", "same", "b:c", 20, timestamp), + collision_message(ClientId::Codex, "first", "same", "c", 10, timestamp), + collision_message(ClientId::Amp, "second", "same", "b:c", 20, timestamp), ), ( GroupBy::ClientProviderModel, - collision_message("a", "b:c", "same", "d", 10, timestamp), - collision_message("a", "b", "same", "c:d", 20, timestamp), + collision_message(ClientId::Amp, "b:c", "same", "d", 10, timestamp), + collision_message(ClientId::Amp, "b", "same", "c:d", 20, timestamp), ), ]; @@ -2730,23 +2733,23 @@ mod tests { #[test] fn daily_and_hourly_maps_use_collision_free_structured_keys() { let timestamp = 1_735_689_600_000; - let first = collision_message("a", "b:c", "same", "d", 10, timestamp); - let second = collision_message("a", "b", "same", "c:d", 20, timestamp); + let first = collision_message(ClientId::Amp, "b:c", "same", "d", 10, timestamp); + let second = collision_message(ClientId::Amp, "b", "same", "c:d", 20, timestamp); let mut acc = TuiAcc::new(); acc.push(&first); acc.push(&second); let usage = acc.project(&GroupBy::ClientProviderModel); - let daily = &usage.daily[0].client_breakdown["a"].models; + let daily = &usage.daily[0].client_breakdown["amp"].models; assert_eq!(daily.len(), 2); - let first_daily = &daily["v1|cpm|1:a3:b:c1:d"]; + let first_daily = &daily["v1|cpm|3:amp3:b:c1:d"]; assert_eq!(first_daily.provider, "b:c"); assert_eq!(first_daily.model_id, "d"); assert_eq!(first_daily.display_name, "d"); assert_eq!(first_daily.tokens.total(), 10); assert_eq!(first_daily.cost, 10.0); assert_eq!(first_daily.messages, 1); - let second_daily = &daily["v1|cpm|1:a1:b3:c:d"]; + let second_daily = &daily["v1|cpm|3:amp1:b3:c:d"]; assert_eq!(second_daily.provider, "b"); assert_eq!(second_daily.model_id, "c:d"); assert_eq!(second_daily.display_name, "c:d"); @@ -2775,7 +2778,7 @@ mod tests { let timestamp = 1_735_689_600_000; let mut acc = TuiAcc::new(); acc.push(&collision_message( - "z-client", + ClientId::Zed, "provider", "session-z", "model", @@ -2783,7 +2786,7 @@ mod tests { timestamp, )); acc.push(&collision_message( - "a-client", + ClientId::Amp, "provider", "session-a", "model", @@ -2799,18 +2802,31 @@ mod tests { .iter() .map(String::as_str) .collect::>(), - ["a-client", "z-client"] + ["amp", "zed"] ); } #[test] fn workspace_maps_tag_unknown_and_known_keys_separately() { let timestamp = 1_735_689_600_000; - let mut unknown = - collision_message("client", "provider", "session", "model", 10, timestamp); + let mut unknown = collision_message( + ClientId::Codex, + "provider", + "session", + "model", + 10, + timestamp, + ); unknown.workspace_key = None; unknown.workspace_label = None; - let mut known = collision_message("client", "provider", "session", "model", 20, timestamp); + let mut known = collision_message( + ClientId::Codex, + "provider", + "session", + "model", + 20, + timestamp, + ); known.workspace_key = Some(Arc::from("")); known.workspace_label = Some(Arc::from("Empty workspace key")); @@ -2820,7 +2836,7 @@ mod tests { let usage = acc.project(&GroupBy::WorkspaceModel); assert_eq!(usage.models.len(), 2); - let daily_models = &usage.daily[0].client_breakdown["client"].models; + let daily_models = &usage.daily[0].client_breakdown["codex"].models; assert_eq!(daily_models.len(), 2); assert!(daily_models.contains_key("v1|wmu|5:model")); assert!(daily_models.contains_key("v1|wmk|0:5:model")); @@ -2830,8 +2846,8 @@ mod tests { fn structured_session_and_agent_instance_identities_do_not_alias_delimiters() { let timestamp = 1_735_689_600_000; let model_messages = [ - collision_message("a:b", "provider", "c", "model", 10, timestamp), - collision_message("a", "provider", "b:c", "model", 20, timestamp), + collision_message(ClientId::Codex, "provider", "c", "model", 10, timestamp), + collision_message(ClientId::Amp, "provider", "b:c", "model", 20, timestamp), ]; let mut model_acc = TuiAcc::new(); for message in &model_messages { @@ -2843,7 +2859,7 @@ mod tests { ); let mut explicit = UnifiedMessage::new_with_agent( - "a", + ClientId::Amp, "model", "provider", "session", @@ -2854,7 +2870,7 @@ mod tests { ); explicit.set_agent_instance(Some("a:b:c".to_string())); let derived_left = UnifiedMessage::new_with_agent( - "a", + ClientId::Amp, "model", "provider", "c", @@ -2864,7 +2880,7 @@ mod tests { Some("builder".to_string()), ); let derived_right = UnifiedMessage::new_with_agent( - "a", + ClientId::Amp, "model", "provider", "b:c", @@ -2887,7 +2903,7 @@ mod tests { #[allow(clippy::too_many_arguments)] fn reprojection_message( - client: &str, + client: ClientId, model: &str, provider: &str, session: &str, @@ -2937,7 +2953,7 @@ mod tests { fn reprojection_corpus() -> Vec { vec![ reprojection_message( - "claude", + ClientId::Claude, "gpt-5.5", "openai", "s1", @@ -2951,7 +2967,7 @@ mod tests { Some("builder"), ), reprojection_message( - "codex", + ClientId::Codex, "gpt-5.5", "openai", "s2", @@ -2965,7 +2981,7 @@ mod tests { None, ), reprojection_message( - "claude", + ClientId::Claude, "gpt-5.5", "azure", "s1", @@ -2979,7 +2995,7 @@ mod tests { None, ), reprojection_message( - "claude", + ClientId::Claude, "claude-sonnet-4.5", "anthropic", "s3", @@ -2993,7 +3009,7 @@ mod tests { None, ), reprojection_message( - "qwen", + ClientId::Qwen, "gpt-5.5", "openai", "s4", @@ -3263,7 +3279,7 @@ mod tests { ] { let mut expected = TuiAcc::new(); for message in &corpus { - if crate::selected_client_ids_include(&message.client, &selected) { + if selected.contains(&message.client) { expected.push(message); } } diff --git a/crates/tokscale-core/src/aggregate/tui_sessions.rs b/crates/tokscale-core/src/aggregate/tui_sessions.rs index 10797fca8..ebabdff96 100644 --- a/crates/tokscale-core/src/aggregate/tui_sessions.rs +++ b/crates/tokscale-core/src/aggregate/tui_sessions.rs @@ -9,7 +9,7 @@ use std::{ sync::Arc, }; -use crate::UnifiedMessage; +use crate::{ClientId, UnifiedMessage}; #[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)] #[serde(rename_all = "camelCase")] @@ -68,7 +68,7 @@ pub struct TuiSessionEntry { #[derive(Default)] pub(crate) struct TuiSessionAcc { - sessions: HashMap<(Arc, Arc), TuiSessionBucket>, + sessions: HashMap<(ClientId, Arc), TuiSessionBucket>, } struct TuiSessionBucket { @@ -93,7 +93,7 @@ impl TuiSessionAcc { let timestamp = timestamp_seconds(message.timestamp); let entry = self .sessions - .entry((Arc::clone(&message.client), Arc::clone(&message.session_id))) + .entry((message.client, Arc::clone(&message.session_id))) .or_insert_with(|| TuiSessionBucket { is_main_session: false, workspace_key: message.workspace_key.as_ref().map(Arc::clone), @@ -174,9 +174,9 @@ fn timestamp_seconds(timestamp: i64) -> i64 { #[cfg(test)] mod tests { use super::*; - use crate::{AggregationConfig, DateRange, GroupBy, TokenBreakdown, ViewSet}; + use crate::{AggregationConfig, ClientId, DateRange, GroupBy, TokenBreakdown, ViewSet}; - fn message(client: &str, session_id: &str, timestamp: i64) -> UnifiedMessage { + fn message(client: ClientId, session_id: &str, timestamp: i64) -> UnifiedMessage { UnifiedMessage::new( client, "gpt-5.6", @@ -196,13 +196,13 @@ mod tests { #[test] fn aggregates_sessions_with_incomplete_workspace_and_counter_values() { - let mut first = message("codex", "session-a", 1_700_000_000_000); + let mut first = message(ClientId::Codex, "session-a", 1_700_000_000_000); first.is_main_session = false; first.workspace_key = Some(Arc::from("")); first.workspace_label = None; first.message_count = -3; - let mut second = message("codex", "session-a", 1_700_000_010); + let mut second = message(ClientId::Codex, "session-a", 1_700_000_010); second.model_id = Arc::from("o3"); second.is_main_session = true; second.workspace_key = Some(Arc::from("later-key")); @@ -238,10 +238,10 @@ mod tests { #[test] fn sorts_by_recent_then_client_then_session() { let mut acc = TuiSessionAcc::new(); - acc.push(&message("zed", "b", 9)); - acc.push(&message("codex", "z", 9)); - acc.push(&message("codex", "a", 9)); - acc.push(&message("amp", "old", 8)); + acc.push(&message(ClientId::Zed, "b", 9)); + acc.push(&message(ClientId::Codex, "z", 9)); + acc.push(&message(ClientId::Codex, "a", 9)); + acc.push(&message(ClientId::Amp, "old", 8)); let keys = acc .finish() @@ -269,8 +269,8 @@ mod tests { }, views: ViewSet::TUI | ViewSet::TUI_SESSIONS, }); - engine.push(&message("codex", "kept", 1_704_110_400_000)); - engine.push(&message("codex", "filtered", 1_735_732_800_000)); + engine.push(&message(ClientId::Codex, "kept", 1_704_110_400_000)); + engine.push(&message(ClientId::Codex, "filtered", 1_735_732_800_000)); let (usage, sessions) = engine.into_tui_bundle(); let sessions = sessions.expect("tui sessions view requested"); diff --git a/crates/tokscale-core/src/input_footprint.rs b/crates/tokscale-core/src/input_footprint.rs new file mode 100644 index 000000000..4fb7d495e --- /dev/null +++ b/crates/tokscale-core/src/input_footprint.rs @@ -0,0 +1,173 @@ +use std::collections::BTreeMap; + +use serde::{Deserialize, Deserializer, Serialize}; + +use crate::ClientId; + +/// Confirmed local input bytes, partitioned by the client that owns each +/// input. This is the sole authority for both per-client space and the +/// Overview total; the total is always derived. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize)] +#[serde(transparent)] +pub struct InputFootprint(BTreeMap); + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct InputFootprintOverflow; + +impl std::fmt::Display for InputFootprintOverflow { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("input data size exceeds u64::MAX") + } +} + +impl std::error::Error for InputFootprintOverflow {} + +impl InputFootprint { + pub fn for_clients(clients: impl IntoIterator) -> Self { + Self(clients.into_iter().map(|client| (client, 0)).collect()) + } + + pub fn from_client_bytes( + entries: impl IntoIterator, + ) -> Result { + let mut footprint = Self::default(); + for (client, bytes) in entries { + footprint.set_bytes(client, bytes)?; + } + Ok(footprint) + } + + pub fn set_bytes( + &mut self, + client: ClientId, + bytes: u64, + ) -> Result<(), InputFootprintOverflow> { + let previous = self.bytes_for(client); + let remaining = self + .total_bytes()? + .checked_sub(previous) + .expect("client bytes must be included in the footprint total"); + remaining.checked_add(bytes).ok_or(InputFootprintOverflow)?; + self.0.insert(client, bytes); + Ok(()) + } + + pub fn add_bytes( + &mut self, + client: ClientId, + bytes: u64, + ) -> Result<(), InputFootprintOverflow> { + let total = self + .bytes_for(client) + .checked_add(bytes) + .ok_or(InputFootprintOverflow)?; + self.set_bytes(client, total) + } + + pub fn bytes_for(&self, client: ClientId) -> u64 { + self.0.get(&client).copied().unwrap_or(0) + } + + pub fn contains_client(&self, client: ClientId) -> bool { + self.0.contains_key(&client) + } + + pub fn len(&self) -> usize { + self.0.len() + } + + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } + + pub fn iter(&self) -> impl ExactSizeIterator + '_ { + self.0.iter().map(|(&client, &bytes)| (client, bytes)) + } + + pub fn total_bytes(&self) -> Result { + self.0.values().copied().try_fold(0_u64, |total, bytes| { + total.checked_add(bytes).ok_or(InputFootprintOverflow) + }) + } +} + +impl<'de> Deserialize<'de> for InputFootprint { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let footprint = Self(BTreeMap::deserialize(deserializer)?); + footprint.total_bytes().map_err(serde::de::Error::custom)?; + Ok(footprint) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn total_is_derived_from_typed_client_entries() { + let footprint = + InputFootprint::from_client_bytes([(ClientId::Amp, 13), (ClientId::CodeBuddy, 8)]) + .unwrap(); + + assert_eq!(footprint.bytes_for(ClientId::Amp), 13); + assert_eq!(footprint.bytes_for(ClientId::Claude), 0); + assert_eq!(footprint.total_bytes().unwrap(), 21); + } + + #[test] + fn serde_uses_canonical_client_ids_as_map_keys() { + let footprint = + InputFootprint::from_client_bytes([(ClientId::Claude, 4_096), (ClientId::Codex, 7)]) + .unwrap(); + + assert_eq!( + footprint + .iter() + .map(|(client, _)| client.as_str()) + .collect::>(), + ["claude", "codex"] + ); + let value = serde_json::to_value(&footprint).unwrap(); + assert_eq!(value, serde_json::json!({"claude": 4096, "codex": 7})); + assert_eq!( + serde_json::from_value::(value).unwrap(), + footprint + ); + } + + #[test] + fn serde_rejects_unknown_client_keys() { + let error = serde_json::from_value::(serde_json::json!({ + "not-a-client": 1 + })) + .unwrap_err(); + + assert!(error.to_string().contains("unknown local client")); + } + + #[test] + fn serde_rejects_an_overflowing_total() { + let error = serde_json::from_value::(serde_json::json!({ + "amp": u64::MAX, + "codex": 1 + })) + .unwrap_err(); + + assert!(error.to_string().contains("exceeds u64::MAX")); + } + + #[test] + fn failed_add_does_not_create_an_invalid_footprint() { + let mut footprint = InputFootprint::from_client_bytes([(ClientId::Amp, u64::MAX)]).unwrap(); + + assert_eq!( + footprint.add_bytes(ClientId::Codex, 1), + Err(InputFootprintOverflow) + ); + assert_eq!(footprint.bytes_for(ClientId::Codex), 0); + assert_eq!(footprint.total_bytes().unwrap(), u64::MAX); + } +} diff --git a/crates/tokscale-core/src/input_health.rs b/crates/tokscale-core/src/input_health.rs index 22629c1ac..464a91983 100644 --- a/crates/tokscale-core/src/input_health.rs +++ b/crates/tokscale-core/src/input_health.rs @@ -180,7 +180,6 @@ impl InputHealth { pub struct DataHealth { inputs: Vec, examined_inputs: usize, - input_data_bytes: u64, } impl DataHealth { @@ -195,18 +194,6 @@ impl DataHealth { pub fn merge(&mut self, other: DataHealth) { self.inputs.extend(other.inputs); self.examined_inputs += other.examined_inputs; - self.input_data_bytes = self - .input_data_bytes - .checked_add(other.input_data_bytes) - .expect("input data size must fit in u64"); - } - - pub fn set_input_data_bytes(&mut self, input_data_bytes: u64) { - self.input_data_bytes = input_data_bytes; - } - - pub fn input_data_bytes(&self) -> u64 { - self.input_data_bytes } pub fn inputs(&self) -> &[InputHealth] { @@ -325,7 +312,6 @@ impl DataHealth { rejected_records: self.rejected_records(), partial_inputs: self.partial_inputs(), failed_inputs: self.failed_inputs(), - input_data_bytes: self.input_data_bytes, issues: grouped.into_values().collect(), } } @@ -342,10 +328,6 @@ pub struct HealthReport { pub rejected_records: u64, pub partial_inputs: usize, pub failed_inputs: usize, - /// Sum of per-client input bytes at the latest inventory snapshot, - /// deduplicated within each client. Tokscale's own cache files are not - /// included. - pub input_data_bytes: u64, #[serde(default)] pub issues: Vec, } @@ -359,7 +341,6 @@ impl Default for HealthReport { rejected_records: 0, partial_inputs: 0, failed_inputs: 0, - input_data_bytes: 0, issues: Vec::new(), } } @@ -491,18 +472,9 @@ mod tests { assert_eq!(report.rejected_records, 0); assert_eq!(report.partial_inputs, 0); assert_eq!(report.failed_inputs, 0); - assert_eq!(report.input_data_bytes, 0); assert!(report.issues.is_empty()); } - #[test] - fn report_projection_preserves_nonzero_input_data_bytes() { - let mut health = DataHealth::default(); - health.set_input_data_bytes(3_072); - - assert_eq!(health.to_report().input_data_bytes, 3_072); - } - #[test] fn empty_health_json_deserializes_as_a_complete_load() { let report: HealthReport = serde_json::from_str("{}").unwrap(); @@ -521,7 +493,7 @@ mod tests { assert_eq!(value["rejectedRecords"], 0); assert_eq!(value["partialInputs"], 0); assert_eq!(value["failedInputs"], 0); - assert_eq!(value["inputDataBytes"], 0); + assert!(value.get("inputDataBytes").is_none()); assert_eq!(value["issues"], serde_json::json!([])); } @@ -541,7 +513,6 @@ mod tests { rejected_records: 3, partial_inputs: 2, failed_inputs: 1, - input_data_bytes: 1_024, issues: Vec::new(), }; @@ -594,14 +565,12 @@ mod tests { #[test] fn merging_data_health_preserves_clean_and_degraded_input_counts() { let mut left = DataHealth::default(); - left.set_input_data_bytes(1_024); left.record(health(InputStatus::Complete, RejectionSummary::default())); let mut rejected = RejectionSummary::default(); rejected.record(RecordRejectionReason::MissingModel); left.record(health(InputStatus::Complete, rejected)); let mut right = DataHealth::default(); - right.set_input_data_bytes(2_048); right.record(health(InputStatus::Complete, RejectionSummary::default())); right.record(health( InputStatus::Unavailable { @@ -617,7 +586,6 @@ mod tests { assert_eq!(left.inputs().len(), 2); assert_eq!(left.failed_inputs(), 1); assert_eq!(left.rejected_records(), 1); - assert_eq!(left.input_data_bytes(), 3_072); } #[test] diff --git a/crates/tokscale-core/src/lib.rs b/crates/tokscale-core/src/lib.rs index f7629ed33..1e6612a35 100644 --- a/crates/tokscale-core/src/lib.rs +++ b/crates/tokscale-core/src/lib.rs @@ -4,6 +4,7 @@ mod adapters; mod client_catalog; pub mod clients; pub mod fs_atomic; +mod input_footprint; pub mod input_health; mod local_clients; mod local_report_error; @@ -28,6 +29,7 @@ pub use aggregate::{ pub use clients::{ cline_session_data_dir, warp_sqlite_roots, ClientId, ClientIdentity, LocalClientDef, PathRoot, }; +pub use input_footprint::{InputFootprint, InputFootprintOverflow}; pub use input_health::{ DataHealth, InputFailure, InputHealth, InputStatus, RecordRejectionReason, RejectionEntry, RejectionSummary, ScannedInput, @@ -37,7 +39,7 @@ pub use message_cache::{prune_input_message_cache, InputCachePruneError, InputCa pub use provider_identity::{inferred_provider_from_model, normalize_provider_for_grouping}; pub use sessions::UnifiedMessage; -use std::collections::{BTreeMap, HashMap, HashSet}; +use std::collections::{HashMap, HashSet}; use std::hash::Hasher; use std::sync::Arc; @@ -82,23 +84,6 @@ pub fn build_tui_accumulator(messages: &[UnifiedMessage], date_range: DateRange) .expect("TUI view must create a TUI accumulator") } -fn retain_for_requested_clients( - client: &str, - _model_id: &str, - _provider_id: &str, - requested: &HashSet<&str>, -) -> bool { - requested_clients_include(client, requested) -} - -pub(crate) fn requested_clients_include(client: &str, requested: &HashSet<&str>) -> bool { - requested.contains(client) -} - -pub(crate) fn selected_client_ids_include(client: &str, selected: &HashSet) -> bool { - ClientId::from_str(client).is_some_and(|client_id| selected.contains(&client_id)) -} - #[derive(Debug, Clone, Default, PartialEq, serde::Serialize)] pub enum GroupBy { #[default] @@ -217,6 +202,8 @@ pub(crate) fn checked_token_sum(values: impl IntoIterator) -> i64 { #[derive(Debug, Clone)] pub struct LocalLoadMetadata { pub input_inventory_signature: InputInventorySignature, + /// Confirmed input bytes keyed by the canonical client that owns them. + pub input_footprint: InputFootprint, } #[derive(Debug)] @@ -271,9 +258,10 @@ impl InputInventorySignature { /// the exact units whose signature was compared by the caller. pub struct PreparedLocalInputs { options: LocalParseOptions, - clients: Vec, + clients: Vec, groups: Vec, signature: InputInventorySignature, + input_footprint: InputFootprint, health: DataHealth, #[cfg(test)] input_cache_dir: std::path::PathBuf, @@ -288,14 +276,18 @@ impl PreparedLocalInputs { self.signature.process_digest() } + pub fn input_footprint(&self) -> &InputFootprint { + &self.input_footprint + } + /// Refresh metadata and stable identity without rediscovering inputs or /// reading their bodies. This is the narrow probe used by TUI auto-refresh /// before it decides that a prepared inventory is unchanged. pub fn refresh_input_inventory_signature(&mut self) -> Result { let mut unavailable = Vec::new(); for group in &mut self.groups { + let client = group.binding.client(); group.units.retain_mut(|unit| { - let client = unit.client(); let path = unit.path.clone(); match unit.refresh_prepared_snapshot_for_inventory_probe() { Ok(()) => true, @@ -319,62 +311,53 @@ impl PreparedLocalInputs { for health in unavailable { self.health.record(health); } - self.health.set_input_data_bytes(input_data_bytes( - self.groups.iter().flat_map(|group| group.units.iter()), - )); + self.input_footprint = prepared_input_footprint(&self.clients, &self.groups); self.signature = input_inventory_signature(&self.clients, &self.groups); Ok(self.signature) } } -fn input_data_bytes<'a>(units: impl IntoIterator) -> u64 { - let mut seen_by_client = HashMap::>::new(); - let mut total = 0_u64; - for unit in units { - let seen = seen_by_client.entry(unit.client()).or_default(); - let snapshot = unit - .prepared_input_snapshot() - .expect("prepared input unit must carry an inventory snapshot"); - snapshot.visit_present_files(|identity, size| { - if seen.insert(identity) { - total = total - .checked_add(size) - .expect("input data size must fit in u64"); - } - }); +fn selected_client_footprint(clients: &[ClientId]) -> InputFootprint { + InputFootprint::for_clients(clients.iter().copied()) +} + +fn prepared_input_footprint( + clients: &[ClientId], + groups: &[adapters::PreparedAdapterInputs], +) -> InputFootprint { + let mut footprint = selected_client_footprint(clients); + for group in groups { + let client = group.binding.client(); + let mut seen = HashSet::new(); + for unit in &group.units { + let snapshot = unit + .prepared_input_snapshot() + .expect("prepared input unit must carry an inventory snapshot"); + snapshot.visit_present_files(|identity, size| { + if seen.insert(identity) { + footprint + .add_bytes(client, size) + .expect("input data size must fit in u64"); + } + }); + } } - total + footprint } -fn confirmed_client_space( - clients: &[String], +fn confirmed_input_footprint( + clients: &[ClientId], groups: &[adapters::ConfirmedAdapterInputs], -) -> BTreeMap { - let mut totals = clients - .iter() - .cloned() - .map(|client| (client, 0_u64)) - .collect::>(); - +) -> InputFootprint { + let mut footprint = selected_client_footprint(clients); for group in groups { - let client = group.client.as_str().to_string(); - let total = totals.entry(client).or_default(); for &(_, size) in &group.present_files { - *total = total - .checked_add(size) - .expect("per-client input data size must fit in u64"); + footprint + .add_bytes(group.client, size) + .expect("input data size must fit in u64"); } } - - totals -} - -fn client_space_total(client_space: &BTreeMap) -> u64 { - client_space.values().copied().fold(0_u64, |total, size| { - total - .checked_add(size) - .expect("input data size must fit in u64") - }) + footprint } #[derive(Debug, Clone, Default)] @@ -461,7 +444,7 @@ fn parse_all_messages_with_health_with_settings( struct FoldOutcome { input_inventory_signature: InputInventorySignature, - client_space: BTreeMap, + input_footprint: InputFootprint, health: DataHealth, } @@ -473,7 +456,7 @@ fn input_cache_dir_for_test_home(home_dir: &str) -> std::path::PathBuf { fn fold_prepared_local_inputs_with_pricing( prepared: PreparedLocalInputs, pricing: Option<&pricing::PricingService>, - sink: &mut dyn adapters::MessageSink, + sink: &mut dyn adapters::UnifiedMessageSink, ) -> Result { #[cfg(test)] let input_cache_dir = prepared.input_cache_dir.clone(); @@ -499,7 +482,7 @@ fn fold_prepared_local_inputs_with_pricing( &mut health, ) } else { - let requested: HashSet<&str> = clients.iter().map(String::as_str).collect(); + let requested: HashSet = clients.iter().copied().collect(); let mut filtered_sink = RequestedClientFilterSink { requested: &requested, inner: sink, @@ -526,11 +509,10 @@ fn fold_prepared_local_inputs_with_pricing( .map(|confirmed| { let input_inventory_signature = confirmed_input_inventory_signature(&clients, &confirmed); - let client_space = confirmed_client_space(&clients, &confirmed); - health.set_input_data_bytes(client_space_total(&client_space)); + let input_footprint = confirmed_input_footprint(&clients, &confirmed); FoldOutcome { input_inventory_signature, - client_space, + input_footprint, health, } }) @@ -538,18 +520,13 @@ fn fold_prepared_local_inputs_with_pricing( } struct RequestedClientFilterSink<'a> { - requested: &'a HashSet<&'a str>, - inner: &'a mut dyn adapters::MessageSink, + requested: &'a HashSet, + inner: &'a mut dyn adapters::UnifiedMessageSink, } -impl adapters::MessageSink for RequestedClientFilterSink<'_> { +impl adapters::UnifiedMessageSink for RequestedClientFilterSink<'_> { fn push_message(&mut self, message: UnifiedMessage) { - if retain_for_requested_clients( - &message.client, - &message.model_id, - &message.provider_id, - self.requested, - ) { + if self.requested.contains(&message.client) { self.inner.push_message(message); } } @@ -557,7 +534,7 @@ impl adapters::MessageSink for RequestedClientFilterSink<'_> { struct AggregationSink<'a>(&'a mut crate::aggregate::AggregationEngine); -impl adapters::MessageSink for AggregationSink<'_> { +impl adapters::UnifiedMessageSink for AggregationSink<'_> { fn push_message(&mut self, message: UnifiedMessage) { self.0.push(&message); } @@ -589,16 +566,19 @@ pub fn prepare_local_inputs( let mut health = DataHealth::default(); let groups: Vec<_> = selected_adapters .into_iter() - .map(|adapter| -> Result<_, LocalReportError> { + .map(|binding| -> Result<_, LocalReportError> { #[cfg(test)] PREPARE_DISCOVERY_COUNT.with(|count| count.set(count.get() + 1)); // Third-party input and snapshot failures stay inside their // input's failure domain. - let units = match adapter.discover_checked(&scan_ctx) { + let units = match binding + .adapter() + .discover_checked(binding.client(), &scan_ctx) + { Ok(units) => units, Err(error) => { health.record(InputHealth { - client: error.client, + client: binding.client(), path: error.path.clone(), status: InputStatus::Unavailable { failure: InputFailure::new(error.operation, error.to_string()), @@ -606,17 +586,15 @@ pub fn prepare_local_inputs( rejections: RejectionSummary::default(), }); return Ok(adapters::PreparedAdapterInputs { - adapter, + binding, units: Vec::new(), }); } }; - adapters::validate_discovered_unit_sources(adapter, &units) - .map_err(LocalReportError::operational)?; let units = units .into_iter() .filter_map(|unit| { - let client = unit.client(); + let client = binding.client(); let path = unit.path.clone(); match unit.prepare_snapshot() { Ok(unit) => Some(unit), @@ -637,18 +615,17 @@ pub fn prepare_local_inputs( } }) .collect(); - Ok(adapters::PreparedAdapterInputs { adapter, units }) + Ok(adapters::PreparedAdapterInputs { binding, units }) }) .collect::>()?; - health.set_input_data_bytes(input_data_bytes( - groups.iter().flat_map(|group| group.units.iter()), - )); + let input_footprint = prepared_input_footprint(&clients, &groups); let signature = input_inventory_signature(&clients, &groups); Ok(PreparedLocalInputs { options, clients, groups, signature, + input_footprint, health, #[cfg(test)] input_cache_dir: input_cache_dir_for_test_home(&home_dir), @@ -671,24 +648,24 @@ fn prepare_discovery_count() -> usize { } fn input_inventory_signature( - clients: &[String], + clients: &[ClientId], groups: &[adapters::PreparedAdapterInputs], ) -> InputInventorySignature { let mut hasher = Sha256::new(); message_cache::hash_inventory_bytes(&mut hasher, b"tokscale/local-input-inventory"); - hasher.update(2_u32.to_le_bytes()); - let mut sorted_clients: Vec<&str> = clients.iter().map(String::as_str).collect(); + hasher.update(3_u32.to_le_bytes()); + let mut sorted_clients = clients.to_vec(); sorted_clients.sort_unstable(); sorted_clients.dedup(); message_cache::hash_inventory_len(&mut hasher, sorted_clients.len()); for client in sorted_clients { - message_cache::hash_inventory_bytes(&mut hasher, client.as_bytes()); + message_cache::hash_inventory_bytes(&mut hasher, client.as_str().as_bytes()); } message_cache::hash_inventory_len(&mut hasher, groups.len()); for group in groups { message_cache::hash_inventory_bytes( &mut hasher, - group.adapter.client().as_str().as_bytes(), + group.binding.client().as_str().as_bytes(), ); message_cache::hash_inventory_len(&mut hasher, group.units.len()); for unit in &group.units { @@ -699,18 +676,18 @@ fn input_inventory_signature( } fn confirmed_input_inventory_signature( - clients: &[String], + clients: &[ClientId], groups: &[adapters::ConfirmedAdapterInputs], ) -> InputInventorySignature { let mut hasher = Sha256::new(); message_cache::hash_inventory_bytes(&mut hasher, b"tokscale/local-input-inventory"); - hasher.update(2_u32.to_le_bytes()); - let mut sorted_clients: Vec<&str> = clients.iter().map(String::as_str).collect(); + hasher.update(3_u32.to_le_bytes()); + let mut sorted_clients = clients.to_vec(); sorted_clients.sort_unstable(); sorted_clients.dedup(); message_cache::hash_inventory_len(&mut hasher, sorted_clients.len()); for client in sorted_clients { - message_cache::hash_inventory_bytes(&mut hasher, client.as_bytes()); + message_cache::hash_inventory_bytes(&mut hasher, client.as_str().as_bytes()); } message_cache::hash_inventory_len(&mut hasher, groups.len()); for group in groups { @@ -813,7 +790,7 @@ struct ResolvedAggregationRequest<'a> { fn load_aggregated_views_resolved( request: ResolvedAggregationRequest<'_>, -) -> Result { +) -> Result<(AggregatedViews, LocalLoadMetadata), LocalReportError> { let prepared = prepare_local_inputs(LocalParseOptions { home_dir: Some(request.home_dir.to_string()), clients: Some(request.clients.to_vec()), @@ -829,7 +806,6 @@ fn load_aggregated_views_resolved( request.views, request.pricing, ) - .map(|(views, _)| views) } fn load_prepared_aggregated_views( @@ -838,7 +814,7 @@ fn load_prepared_aggregated_views( date_range: DateRange, views: ViewSet, pricing: Option<&pricing::PricingService>, -) -> Result<(AggregatedViews, InputInventorySignature), LocalReportError> { +) -> Result<(AggregatedViews, LocalLoadMetadata), LocalReportError> { let mut engine = crate::aggregate::AggregationEngine::new(AggregationConfig { group_by, date_range, @@ -846,8 +822,8 @@ fn load_prepared_aggregated_views( }); let FoldOutcome { input_inventory_signature, + input_footprint, health, - .. } = match stream_local_inputs_into_engine(prepared, pricing, &mut engine) { Ok(outcome) => outcome, Err(error) => { @@ -863,7 +839,13 @@ fn load_prepared_aggregated_views( // is the narrow lifecycle seam where dead weak identity indices can be // reclaimed without sweeping caller-owned generic message slices. sessions::intern::prune_dead(); - Ok((views, input_inventory_signature)) + Ok(( + views, + LocalLoadMetadata { + input_inventory_signature, + input_footprint, + }, + )) } fn load_aggregated_views_for_resolved_report( @@ -882,6 +864,7 @@ fn load_aggregated_views_for_resolved_report( views, pricing, }) + .map(|(views, _)| views) } /// Build any requested union of aggregation views with one adapter fold. @@ -904,116 +887,89 @@ pub fn load_aggregated_views_with_pricing( /// /// Headless renderers select fields from this value instead of maintaining /// command-specific aggregation rules. -pub async fn get_usage_data( +pub async fn get_usage_report( options: ReportOptions, -) -> Result { +) -> Result, LocalReportError> { let (home_dir, clients) = resolve_report_request(&options)?; let pricing = load_pricing_for_local_parse().await; - let mut views = load_aggregated_views_for_resolved_report( - &options, - &home_dir, - &clients, - ViewSet::TUI, - pricing.as_deref(), - )?; + let (mut views, metadata) = load_aggregated_views_resolved(ResolvedAggregationRequest { + home_dir: &home_dir, + clients: &clients, + group_by: options.group_by.clone(), + date_range: DateRange::from_options(&options), + scanner_settings: &options.scanner_settings, + views: ViewSet::TUI, + pricing: pricing.as_deref(), + })?; + let health = views.health.to_report(); let mut data = views.tui_usage.take().expect("tui view requested"); - data.health = views.health.to_report(); - Ok(data) -} - -trait FinalizableMessage { - fn model_id(&self) -> &Arc; - fn set_model_id(&mut self, model_id: Arc); - fn provider_id(&self) -> &Arc; - fn set_provider_id(&mut self, provider_id: Arc); - fn tokens(&self) -> &TokenBreakdown; - fn tokens_mut(&mut self) -> &mut TokenBreakdown; - fn set_cost(&mut self, cost: f64); + // Preserve the established UsageData envelope while exposing the same + // report health beside load metadata for generic report consumers. + data.health = health.clone(); + Ok(LocalReport { + data, + health, + metadata, + }) } -macro_rules! impl_finalizable_message { - ($message:ty) => { - impl FinalizableMessage for $message { - fn model_id(&self) -> &Arc { - &self.model_id - } - - fn set_model_id(&mut self, model_id: Arc) { - self.model_id = model_id; - } - - fn provider_id(&self) -> &Arc { - &self.provider_id - } - - fn set_provider_id(&mut self, provider_id: Arc) { - self.provider_id = provider_id; - } - - fn tokens(&self) -> &TokenBreakdown { - &self.tokens - } - - fn tokens_mut(&mut self) -> &mut TokenBreakdown { - &mut self.tokens - } - - fn set_cost(&mut self, cost: f64) { - self.cost = cost; - } - } - }; +/// Build the canonical TUI usage projection without its generic report +/// envelope. +/// +/// New callers that also need confirmed input attribution should prefer +/// [`get_usage_report`]. +pub async fn get_usage_data( + options: ReportOptions, +) -> Result { + get_usage_report(options).await.map(|report| report.data) } -impl_finalizable_message!(ParsedMessage); -impl_finalizable_message!(UnifiedMessage); - -fn apply_token_pricing( - message: &mut M, +fn apply_token_pricing( + message: &mut sessions::UsageRecord, pricing: Option<&pricing::PricingService>, ) { - message.set_cost(0.0); + message.cost = 0.0; let Some(pricing) = pricing else { return; }; let calculated_cost = pricing.calculate_cost_with_provider( - message.model_id(), - Some(message.provider_id().as_ref()), - message.tokens(), + &message.model_id, + Some(message.provider_id.as_ref()), + &message.tokens, ); if calculated_cost > 0.0 { - message.set_cost(calculated_cost); + message.cost = calculated_cost; } } -fn refresh_derived_message_fields(message: &mut M) { +fn refresh_derived_message_fields(message: &mut sessions::UsageRecord) { if let Some(provider) = provider_identity::provider_override_from_model_and_provider( - message.model_id(), - message.provider_id(), + &message.model_id, + &message.provider_id, ) { - message.set_provider_id(sessions::intern::intern(provider)); + message.provider_id = sessions::intern::intern(provider); } } -fn canonicalize_message_provider(message: &mut M) { +fn canonicalize_message_provider(message: &mut sessions::UsageRecord) { let provider = - provider_identity::finalized_provider_id(message.provider_id(), message.model_id()); - message.set_provider_id(sessions::intern::intern(&provider)); + provider_identity::finalized_provider_id(&message.provider_id, &message.model_id); + message.provider_id = sessions::intern::intern(&provider); } -fn canonicalize_message_model( - message: &mut M, +fn canonicalize_message_model( + message: &mut sessions::UsageRecord, model_cache: &mut HashMap, Arc>, ) { - if let Some(canonical) = model_cache.get(message.model_id()) { - message.set_model_id(Arc::clone(canonical)); + if let Some(canonical) = model_cache.get(&message.model_id) { + message.model_id = Arc::clone(canonical); return; } - let raw = Arc::clone(message.model_id()); + let raw = Arc::clone(&message.model_id); let canonical = model_aliases::canonicalize_model_id(raw.as_ref()); let canonical = if canonical == raw.as_ref() { Arc::clone(&raw) @@ -1022,21 +978,22 @@ fn canonicalize_message_model( }; model_cache.insert(raw, Arc::clone(&canonical)); - message.set_model_id(canonical); + message.model_id = canonical; } -fn finalize_token_priced_messages( +fn finalize_token_priced_messages>( messages: &mut Vec, pricing: Option<&pricing::PricingService>, ) { let mut model_cache = HashMap::new(); messages.retain_mut(|message| { - normalize_token_breakdown(message.tokens_mut()); + let message = message.as_mut(); + normalize_token_breakdown(&mut message.tokens); canonicalize_message_model(message, &mut model_cache); refresh_derived_message_fields(message); canonicalize_message_provider(message); - if !has_positive_tokens(message.tokens()) { + if !has_positive_tokens(&message.tokens) { return false; } apply_token_pricing(message, pricing); @@ -1118,18 +1075,24 @@ async fn load_pricing_for_local_parse_with_diagnostics( fn resolve_local_parse_request( options: &LocalParseOptions, -) -> Result<(String, Vec), LocalReportError> { +) -> Result<(String, Vec), LocalReportError> { let home_dir = get_home_dir_string(&options.home_dir) .map_err(LocalReportError::invalid_environment_message)?; - let clients = options - .clients - .clone() - .unwrap_or_else(|| ClientId::iter().map(|c| c.as_str().to_string()).collect()); - for client in &clients { - ClientId::from_str(client).ok_or_else(|| { - LocalReportError::invalid_request_message(format!("unknown local client `{client}`")) - })?; - } + let mut clients: Vec = match options.clients.as_ref() { + Some(clients) => clients + .iter() + .map(|client| { + ClientId::from_str(client).ok_or_else(|| { + LocalReportError::invalid_request_message(format!( + "unknown local client `{client}`" + )) + }) + }) + .collect::>()?, + None => ClientId::iter().collect(), + }; + clients.sort_unstable(); + clients.dedup(); Ok((home_dir, clients)) } @@ -1141,14 +1104,15 @@ fn parse_prepared_local_unified_messages( let mut messages = Vec::new(); let FoldOutcome { input_inventory_signature, + input_footprint, health, - .. } = fold_prepared_local_inputs_with_pricing(prepared, pricing, &mut messages)?; Ok(LocalReport { data: filter_unified_messages(messages, &filters), health: health.to_report(), metadata: LocalLoadMetadata { input_inventory_signature, + input_footprint, }, }) } @@ -1218,7 +1182,7 @@ pub struct TuiBundleWithDiagnostics { pub accumulator: TuiAcc, pub sessions: Vec, /// Confirmed input bytes keyed by canonical local client id. - pub client_space: BTreeMap, + pub input_footprint: InputFootprint, pub pricing_diagnostics: pricing::PricingDiagnostics, pub input_inventory_signature: InputInventorySignature, pub health: DataHealth, @@ -1243,7 +1207,7 @@ pub async fn load_prepared_tui_bundle_with_diagnostics( }); let FoldOutcome { input_inventory_signature, - client_space, + input_footprint, health, } = match stream_local_inputs_into_engine(prepared, pricing.as_deref(), &mut engine) { Ok(outcome) => outcome, @@ -1258,7 +1222,7 @@ pub async fn load_prepared_tui_bundle_with_diagnostics( Ok(TuiBundleWithDiagnostics { accumulator: accumulator.expect("tui usage view requested"), sessions: sessions.expect("tui sessions view requested"), - client_space, + input_footprint, pricing_diagnostics, input_inventory_signature, health, diff --git a/crates/tokscale-core/src/lib_tests.rs b/crates/tokscale-core/src/lib_tests.rs index ebaf40b2c..e15d878a6 100644 --- a/crates/tokscale-core/src/lib_tests.rs +++ b/crates/tokscale-core/src/lib_tests.rs @@ -3,10 +3,10 @@ use super::{ load_cache_only_pricing_with_diagnostics, load_usage_data_with_pricing, message_cache, normalize_model_for_grouping, parse_all_messages_with_health, parse_all_messages_with_health_with_settings, parse_all_messages_with_pricing, - parse_all_messages_with_pricing_with_settings, positive_token_total, pricing, - retain_for_requested_clients, scanner, select_local_parse_pricing, AggregatedViews, - AggregationConfig, ClientId, DateRange, GroupBy, LocalParseOptions, ReportOptions, - TokenBreakdown, UnifiedMessage, ViewSet, UNKNOWN_WORKSPACE_LABEL, + parse_all_messages_with_pricing_with_settings, positive_token_total, pricing, scanner, + select_local_parse_pricing, AggregatedViews, AggregationConfig, ClientId, DateRange, GroupBy, + LocalParseOptions, ReportOptions, TokenBreakdown, UnifiedMessage, ViewSet, + UNKNOWN_WORKSPACE_LABEL, }; use std::collections::{BTreeMap, HashMap, HashSet}; use std::ffi::OsString; @@ -35,10 +35,8 @@ impl TestClientCounts { fn from_messages(messages: &[UnifiedMessage]) -> Self { let mut counts = Self::default(); for message in messages { - if let Some(client) = ClientId::from_str(&message.client) { - counts.0[client as usize] = - counts.0[client as usize].saturating_add(message.message_count.max(0)); - } + counts.0[message.client as usize] = + counts.0[message.client as usize].saturating_add(message.message_count.max(0)); } counts } @@ -70,6 +68,33 @@ fn input_cache_for_test_home(home: &Path) -> message_cache::InputMessageCache { )) } +fn test_decoder(decoder_id: message_cache::DecoderId) -> crate::adapters::DecoderSpec { + crate::adapters::DecoderSpec::plain(decoder_id, 0) +} + +fn plain_test_input( + decoder_id: message_cache::DecoderId, + path: PathBuf, +) -> crate::adapters::InputUnit { + crate::adapters::InputUnit::plain_file(path, test_decoder(decoder_id)) +} + +fn sqlite_test_input( + decoder_id: message_cache::DecoderId, + path: PathBuf, +) -> crate::adapters::InputUnit { + crate::adapters::InputUnit::sqlite_with_wal(path, test_decoder(decoder_id)) +} + +fn opencode_test_input(path: PathBuf) -> crate::adapters::InputUnit { + crate::adapters::InputUnit::sqlite_with_wal( + path, + crate::adapters::DecoderSpec::opencode_sqlite( + crate::adapters::OPENCODE_CURRENT_SQLITE_REVISION, + ), + ) +} + fn parse_all_messages_with_pricing_in_cache( home_dir: &str, clients: &[String], @@ -128,7 +153,7 @@ impl Drop for HomeEnvGuard { } fn make_workspace_message( - client: &str, + client: ClientId, model_id: &str, provider_id: &str, session_id: &str, @@ -160,7 +185,7 @@ fn make_workspace_message( #[allow(clippy::too_many_arguments)] fn make_message_with_tokens( - client: &str, + client: ClientId, model_id: &str, provider_id: &str, session_id: &str, @@ -336,7 +361,7 @@ fn test_streaming_tui_usage_matches_reference_aggregation() { #[test] #[serial_test::serial] -fn prepared_tui_bundle_data_size_matches_real_client_inventory_sum() { +fn prepared_tui_bundle_footprint_sums_the_two_client_fixture() { let home = tempfile::TempDir::new().unwrap(); let _home_guard = HomeEnvGuard::set(home.path()); let _pricing_guard = TestEnvGuard::set("TOKSCALE_PRICING_CACHE_ONLY", "1"); @@ -350,15 +375,46 @@ fn prepared_tui_bundle_data_size_matches_real_client_inventory_sum() { .block_on(super::load_prepared_tui_bundle_with_diagnostics(prepared)) .unwrap(); - assert!(result.client_space["opencode"] > 0); - assert!(result.client_space["codex"] > 0); - let client_space_total = result - .client_space - .values() - .copied() - .try_fold(0_u64, u64::checked_add) + let opencode_bytes = result.input_footprint.bytes_for(ClientId::OpenCode); + let codex_bytes = result.input_footprint.bytes_for(ClientId::Codex); + assert!(opencode_bytes > 0); + assert!(codex_bytes > 0); + assert_eq!( + result.input_footprint.total_bytes().unwrap(), + opencode_bytes.checked_add(codex_bytes).unwrap() + ); +} + +#[test] +#[serial_test::serial] +fn public_usage_report_carries_confirmed_input_footprint() { + let home = tempfile::TempDir::new().unwrap(); + let _home_guard = HomeEnvGuard::set(home.path()); + let _pricing_guard = TestEnvGuard::set("TOKSCALE_PRICING_CACHE_ONLY", "1"); + write_streaming_fold_fixture(home.path()); + + let report = tokio::runtime::Runtime::new() + .unwrap() + .block_on(super::get_usage_report(streaming_report_options( + home.path(), + vec!["opencode", "codex"], + ))) .unwrap(); - assert_eq!(result.health.input_data_bytes(), client_space_total); + + assert_eq!(report.health, report.data.health); + assert_ne!( + report.metadata.input_inventory_signature.as_bytes(), + &[0_u8; 32] + ); + let footprint = &report.metadata.input_footprint; + let opencode_bytes = footprint.bytes_for(ClientId::OpenCode); + let codex_bytes = footprint.bytes_for(ClientId::Codex); + assert!(opencode_bytes > 0); + assert!(codex_bytes > 0); + assert_eq!( + footprint.total_bytes().unwrap(), + opencode_bytes.checked_add(codex_bytes).unwrap() + ); } #[test] @@ -1080,7 +1136,7 @@ fn test_workspace_model_grouping_merges_same_workspace_and_model() { let entries = aggregate_model_usage_entries( vec![ make_workspace_message( - "claude", + ClientId::Claude, "claude-sonnet-4.5", "anthropic", "session-1", @@ -1089,7 +1145,7 @@ fn test_workspace_model_grouping_merges_same_workspace_and_model() { Some("repo-a"), ), make_workspace_message( - "qwen", + ClientId::Qwen, "claude-sonnet-4.5", "anthropic", "session-2", @@ -1114,7 +1170,7 @@ fn test_model_grouping_cleans_fast_variant() { let entries = aggregate_finalized_model_usage_entries( vec![ make_workspace_message( - "opencode", + ClientId::OpenCode, "gpt-5.5-fast", "openai", "session-1", @@ -1122,7 +1178,15 @@ fn test_model_grouping_cleans_fast_variant() { None, None, ), - make_workspace_message("codex", "gpt-5.5", "openai", "session-2", 2.0, None, None), + make_workspace_message( + ClientId::Codex, + "gpt-5.5", + "openai", + "session-2", + 2.0, + None, + None, + ), ], &GroupBy::Model, ); @@ -1137,7 +1201,7 @@ fn test_model_grouping_cleans_hyphenated_date_snapshot() { let entries = aggregate_finalized_model_usage_entries( vec![ make_workspace_message( - "qwen", + ClientId::Qwen, "qwen3.7-max-2026-05-20", "qwen", "session-1", @@ -1145,7 +1209,15 @@ fn test_model_grouping_cleans_hyphenated_date_snapshot() { None, None, ), - make_workspace_message("qwen", "qwen3.7-max", "qwen", "session-2", 2.75, None, None), + make_workspace_message( + ClientId::Qwen, + "qwen3.7-max", + "qwen", + "session-2", + 2.75, + None, + None, + ), ], &GroupBy::ClientModel, ); @@ -1160,7 +1232,7 @@ fn test_model_grouping_cleans_anthropic_prefixed_claude_variant() { let entries = aggregate_finalized_model_usage_entries( vec![ make_workspace_message( - "claude", + ClientId::Claude, "anthropic/claude-4-6-sonnet", "anthropic", "session-1", @@ -1169,7 +1241,7 @@ fn test_model_grouping_cleans_anthropic_prefixed_claude_variant() { Some("repo-a"), ), make_workspace_message( - "claude", + ClientId::Claude, "claude-sonnet-4.6", "anthropic", "session-2", @@ -1191,7 +1263,7 @@ fn test_model_grouping_uses_finalized_provider_ids() { let entries = aggregate_finalized_model_usage_entries( vec![ make_workspace_message( - "opencode", + ClientId::OpenCode, "xiaomi/mimo-v2.5-pro", "xiaomi", "session-1", @@ -1200,7 +1272,7 @@ fn test_model_grouping_uses_finalized_provider_ids() { None, ), make_workspace_message( - "opencode", + ClientId::OpenCode, "xiaomi/mimo-v2.5-pro", "xiaomi", "session-2", @@ -1223,7 +1295,7 @@ fn test_client_provider_model_grouping_uses_finalized_provider_ids() { let entries = aggregate_finalized_model_usage_entries( vec![ make_workspace_message( - "opencode", + ClientId::OpenCode, "xiaomi/mimo-v2.5-pro", "xiaomi", "session-1", @@ -1232,7 +1304,7 @@ fn test_client_provider_model_grouping_uses_finalized_provider_ids() { None, ), make_workspace_message( - "opencode", + ClientId::OpenCode, "xiaomi/mimo-v2.5-pro", "xiaomi", "session-2", @@ -1256,7 +1328,7 @@ fn test_model_grouping_orders_merged_clients_by_total_tokens() { let entries = aggregate_model_usage_entries( vec![ make_message_with_tokens( - "opencode", + ClientId::OpenCode, "gpt-5.5", "openai", "session-opencode", @@ -1267,7 +1339,7 @@ fn test_model_grouping_orders_merged_clients_by_total_tokens() { 0, ), make_message_with_tokens( - "codex", + ClientId::Codex, "gpt-5.5", "openai", "session-codex", @@ -1277,7 +1349,17 @@ fn test_model_grouping_orders_merged_clients_by_total_tokens() { 0, 0, ), - make_message_with_tokens("pi", "gpt-5.5", "openai", "session-pi", 100, 0, 0, 0, 0), + make_message_with_tokens( + ClientId::Pi, + "gpt-5.5", + "openai", + "session-pi", + 100, + 0, + 0, + 0, + 0, + ), ], &GroupBy::Model, ); @@ -1291,7 +1373,7 @@ fn test_model_grouping_ignores_negative_client_token_contribution() { let entries = aggregate_model_usage_entries( vec![ make_message_with_tokens( - "negative-client", + ClientId::Amp, "gpt-5.5", "openai", "session-negative", @@ -1302,7 +1384,7 @@ fn test_model_grouping_ignores_negative_client_token_contribution() { 0, ), make_message_with_tokens( - "positive-client", + ClientId::Codex, "gpt-5.5", "openai", "session-positive", @@ -1317,7 +1399,7 @@ fn test_model_grouping_ignores_negative_client_token_contribution() { ); assert_eq!(entries.len(), 1); - assert_eq!(entries[0].client, "positive-client, negative-client"); + assert_eq!(entries[0].client, "codex, amp"); } #[test] @@ -1325,7 +1407,7 @@ fn test_workspace_model_grouping_separates_different_workspaces() { let entries = aggregate_model_usage_entries( vec![ make_workspace_message( - "claude", + ClientId::Claude, "claude-sonnet-4-5-20250929", "anthropic", "session-1", @@ -1334,7 +1416,7 @@ fn test_workspace_model_grouping_separates_different_workspaces() { Some("repo-a"), ), make_workspace_message( - "claude", + ClientId::Claude, "claude-sonnet-4-5-20250929", "anthropic", "session-2", @@ -1359,7 +1441,7 @@ fn test_workspace_model_grouping_uses_unknown_bucket_without_workspace_metadata( let entries = aggregate_model_usage_entries( vec![ make_workspace_message( - "claude", + ClientId::Claude, "claude-sonnet-4-5-20250929", "anthropic", "session-1", @@ -1368,7 +1450,7 @@ fn test_workspace_model_grouping_uses_unknown_bucket_without_workspace_metadata( None, ), make_workspace_message( - "claude", + ClientId::Claude, "claude-sonnet-4-5-20250929", "anthropic", "session-2", @@ -1394,7 +1476,7 @@ fn test_workspace_model_grouping_keeps_real_unknown_workspace_separate() { let entries = aggregate_model_usage_entries( vec![ make_workspace_message( - "claude", + ClientId::Claude, "claude-sonnet-4-5-20250929", "anthropic", "session-1", @@ -1403,7 +1485,7 @@ fn test_workspace_model_grouping_keeps_real_unknown_workspace_separate() { Some("unknown-workspace"), ), make_workspace_message( - "claude", + ClientId::Claude, "claude-sonnet-4-5-20250929", "anthropic", "session-2", @@ -1433,7 +1515,7 @@ fn test_workspace_model_grouping_avoids_separator_key_collisions() { let entries = aggregate_model_usage_entries( vec![ make_workspace_message( - "claude", + ClientId::Claude, "c", "anthropic", "session-1", @@ -1442,7 +1524,7 @@ fn test_workspace_model_grouping_avoids_separator_key_collisions() { Some("workspace-ab"), ), make_workspace_message( - "claude", + ClientId::Claude, "b:c", "anthropic", "session-2", @@ -1467,23 +1549,6 @@ fn test_workspace_model_grouping_avoids_separator_key_collisions() { })); } -#[test] -fn test_retain_for_requested_clients_keeps_original_client_matches() { - let requested: HashSet<&str> = HashSet::from(["opencode"]); - assert!(retain_for_requested_clients( - "opencode", - "gpt-4o", - "anthropic", - &requested - )); - assert!(!retain_for_requested_clients( - "claude", - "gpt-4o", - "anthropic", - &requested - )); -} - fn write_kimi_code_usage_fixture(input_home: &std::path::Path) { let kimi_home = input_home.join(".kimi-code"); std::fs::create_dir_all(&kimi_home).unwrap(); @@ -1602,7 +1667,7 @@ fn kimi_unavailable_optional_config_preserves_current_wire_usage_in_production_p let wire_path = input_home .path() .join(".kimi-code/sessions/wd-project/session_1/agents/main/wire.jsonl"); - let parser_version = prepared.groups[0].units[0].parser_version; + let decoder_version = prepared.groups[0].units[0].decoder.version(); let mut cold_messages = Vec::new(); let cold_health = super::fold_prepared_local_inputs_with_pricing(prepared, None, &mut cold_messages) @@ -1611,7 +1676,7 @@ fn kimi_unavailable_optional_config_preserves_current_wire_usage_in_production_p assert_eq!(cold_messages.len(), 2); assert_eq!(cold_health.issue_count(), 0); assert!(input_cache_for_test_home(input_home.path()) - .get_meta(&wire_path, parser_version) + .get_meta(&wire_path, decoder_version) .unwrap() .is_some()); @@ -1659,7 +1724,7 @@ fn kimi_unavailable_optional_config_preserves_current_wire_usage_in_production_p assert_eq!(health.partial_inputs(), 1); assert_eq!(health.failed_inputs(), 0); assert!(input_cache_for_test_home(input_home.path()) - .get_meta(&wire_path, parser_version) + .get_meta(&wire_path, decoder_version) .unwrap() .is_none()); } @@ -1681,8 +1746,7 @@ fn test_input_cache_refreshes_stale_provider_on_cache_hit() { ); drop(conn); - let unit = crate::adapters::InputUnit::sqlite_with_wal(ClientId::OpenCode, path.clone()) - .with_meta(crate::adapters::InputUnitMeta::OpenCodeSqlite); + let unit = opencode_test_input(path.clone()); let fingerprint = unit.input_policy().fingerprint().unwrap(); // Provider deliberately wrong for the model: the cache-hit path // must re-run refresh_derived_fields (dates are derived from @@ -1706,7 +1770,7 @@ fn test_input_cache_refreshes_stale_provider_on_cache_hit() { let mut cache = input_cache_for_test_home(input_home.path()); cache.insert(message_cache::CachedInputEntry::new_with_version( &path, - unit.parser_version, + unit.decoder.version(), fingerprint, vec![stale_message], None, @@ -1738,8 +1802,20 @@ fn inventory_options(home: &Path, clients: &[&str]) -> LocalParseOptions { } } +#[test] +fn local_parse_request_resolves_client_strings_once_at_the_boundary() { + let home = tempfile::TempDir::new().unwrap(); + let (_, clients) = super::resolve_local_parse_request(&inventory_options( + home.path(), + &["opencode", "codex", "opencode"], + )) + .unwrap(); + + assert_eq!(clients, vec![ClientId::Codex, ClientId::OpenCode]); +} + fn signature_for_test_units( - requested_clients: &[String], + requested_clients: &[ClientId], client: ClientId, units: Vec, ) -> super::InputInventorySignature { @@ -1752,7 +1828,7 @@ fn prepared_test_group( units: Vec, ) -> crate::adapters::PreparedAdapterInputs { crate::adapters::PreparedAdapterInputs { - adapter: crate::adapters::adapter_for(client).unwrap(), + binding: crate::adapters::adapter_for(client).unwrap(), units: units .into_iter() .map(crate::adapters::InputUnit::prepare_snapshot) @@ -1762,26 +1838,28 @@ fn prepared_test_group( } #[test] -fn input_data_size_counts_related_inputs_once_by_file_identity() { +fn input_footprint_counts_related_inputs_once_by_file_identity() { let dir = tempfile::TempDir::new().unwrap(); let input = dir.path().join("input.jsonl"); let dependency = dir.path().join("dependency.json"); std::fs::write(&input, b"12345678").unwrap(); std::fs::write(&dependency, b"12345").unwrap(); - let with_dependency = crate::adapters::InputUnit::plain_file(ClientId::Amp, input.clone()) + let with_dependency = plain_test_input(message_cache::DecoderId::Amp, input.clone()) .with_dependency(dependency) .prepare_snapshot() .unwrap(); - let duplicate = crate::adapters::InputUnit::plain_file(ClientId::Amp, input) + let duplicate = plain_test_input(message_cache::DecoderId::Amp, input) .prepare_snapshot() .unwrap(); - assert_eq!(super::input_data_bytes([&with_dependency, &duplicate]), 13); + let group = prepared_test_group(ClientId::Amp, vec![with_dependency, duplicate]); + let footprint = super::prepared_input_footprint(&[ClientId::Amp], std::slice::from_ref(&group)); + assert_eq!(footprint.bytes_for(ClientId::Amp), 13); } #[test] -fn inventory_probe_refreshes_input_data_size_from_metadata() { +fn inventory_probe_refreshes_input_footprint_from_metadata() { let home = tempfile::TempDir::new().unwrap(); let amp_dir = home.path().join(".local/share/amp/threads"); std::fs::create_dir_all(&_dir).unwrap(); @@ -1790,11 +1868,11 @@ fn inventory_probe_refreshes_input_data_size_from_metadata() { let mut prepared = super::prepare_local_inputs(inventory_options(home.path(), &["amp"])).unwrap(); - assert_eq!(prepared.health.input_data_bytes(), 8); + assert_eq!(prepared.input_footprint().bytes_for(ClientId::Amp), 8); std::fs::write(&input, b"1234567890123").unwrap(); prepared.refresh_input_inventory_signature().unwrap(); - assert_eq!(prepared.health.input_data_bytes(), 13); + assert_eq!(prepared.input_footprint().bytes_for(ClientId::Amp), 13); } #[test] @@ -1978,7 +2056,7 @@ fn prepared_diagnostics_returns_signature_revalidated_after_pricing_boundary() { #[test] #[serial_test::serial] -fn prepared_tui_bundle_client_space_uses_confirmed_inventory() { +fn prepared_tui_bundle_input_footprint_uses_confirmed_inventory() { let home = tempfile::TempDir::new().unwrap(); let _home_guard = HomeEnvGuard::set(home.path()); let _pricing_guard = TestEnvGuard::set("TOKSCALE_PRICING_CACHE_ONLY", "1"); @@ -2014,8 +2092,10 @@ fn prepared_tui_bundle_client_space_uses_confirmed_inventory() { assert_ne!(stale_signature, result.input_inventory_signature); assert_eq!(confirmed_signature, result.input_inventory_signature); - assert_eq!(result.client_space.get("amp"), Some(&confirmed_bytes)); - assert_eq!(result.health.input_data_bytes(), confirmed_bytes); + assert_eq!( + result.input_footprint.bytes_for(ClientId::Amp), + confirmed_bytes + ); assert_eq!( result.accumulator.project(&GroupBy::Model).total_tokens, 113 @@ -2036,58 +2116,58 @@ fn inventory_signature_tracks_order_paths_related_stamps_and_unit_identity() { std::fs::write(&first, b"first").unwrap(); std::fs::write(&second, b"second").unwrap(); std::fs::write(&wal, b"wal-one").unwrap(); - let clients = vec!["zed".to_string(), "amp".to_string()]; + let clients = vec![ClientId::Zed, ClientId::Amp]; let ordered = signature_for_test_units( &clients, ClientId::Amp, vec![ - crate::adapters::InputUnit::plain_file(ClientId::Amp, first.clone()), - crate::adapters::InputUnit::plain_file(ClientId::Amp, second.clone()), + plain_test_input(message_cache::DecoderId::Amp, first.clone()), + plain_test_input(message_cache::DecoderId::Amp, second.clone()), ], ); let reordered = signature_for_test_units( &clients, ClientId::Amp, vec![ - crate::adapters::InputUnit::plain_file(ClientId::Amp, second), - crate::adapters::InputUnit::plain_file(ClientId::Amp, first.clone()), + plain_test_input(message_cache::DecoderId::Amp, second), + plain_test_input(message_cache::DecoderId::Amp, first.clone()), ], ); assert_ne!(ordered, reordered, "unit discovery order is significant"); let canonical_clients = signature_for_test_units( - &["amp".to_string(), "zed".to_string()], + &[ClientId::Amp, ClientId::Zed], ClientId::Amp, - vec![crate::adapters::InputUnit::plain_file( - ClientId::Amp, + vec![plain_test_input( + message_cache::DecoderId::Amp, first.clone(), )], ); let reversed_clients = signature_for_test_units( - &["zed".to_string(), "amp".to_string()], + &[ClientId::Zed, ClientId::Amp], ClientId::Amp, - vec![crate::adapters::InputUnit::plain_file( - ClientId::Amp, + vec![plain_test_input( + message_cache::DecoderId::Amp, first.clone(), )], ); assert_eq!(canonical_clients, reversed_clients); let sqlite_before = signature_for_test_units( - &["zed".to_string()], + &[ClientId::Zed], ClientId::Zed, - vec![crate::adapters::InputUnit::sqlite_with_wal( - ClientId::Zed, + vec![sqlite_test_input( + message_cache::DecoderId::Zed, first.clone(), )], ); std::fs::write(&wal, b"wal-two-and-longer").unwrap(); let sqlite_after = signature_for_test_units( - &["zed".to_string()], + &[ClientId::Zed], ClientId::Zed, - vec![crate::adapters::InputUnit::sqlite_with_wal( - ClientId::Zed, + vec![sqlite_test_input( + message_cache::DecoderId::Zed, first.clone(), )], ); @@ -2096,53 +2176,56 @@ fn inventory_signature_tracks_order_paths_related_stamps_and_unit_identity() { "related WAL stamp is significant" ); - let parser_changed = signature_for_test_units( - &["amp".to_string()], + let decoder_changed = signature_for_test_units( + &[ClientId::Amp], ClientId::Amp, - vec![ - crate::adapters::InputUnit::plain_file(ClientId::Amp, first).with_parser_version( - message_cache::ParserVersion::new(message_cache::ParserId::Amp, 999), - ), - ], + vec![crate::adapters::InputUnit::plain_file( + first, + crate::adapters::DecoderSpec::plain(message_cache::DecoderId::Amp, 999), + )], ); - assert_ne!(canonical_clients, parser_changed); + assert_ne!(canonical_clients, decoder_changed); let codebuddy_path = dir.path().join("codebuddy.jsonl"); std::fs::write(&codebuddy_path, b"codebuddy").unwrap(); - let jsonl_meta = signature_for_test_units( - &["codebuddy".to_string()], + let jsonl_decoder = signature_for_test_units( + &[ClientId::CodeBuddy], ClientId::CodeBuddy, - vec![ - crate::adapters::InputUnit::plain_file(ClientId::CodeBuddy, codebuddy_path.clone()) - .with_meta(crate::adapters::InputUnitMeta::CodeBuddyJsonl), - ], + vec![crate::adapters::InputUnit::plain_file( + codebuddy_path.clone(), + crate::adapters::DecoderSpec::codebuddy_jsonl(0), + )], ); - let extension_meta = signature_for_test_units( - &["codebuddy".to_string()], + let extension_decoder = signature_for_test_units( + &[ClientId::CodeBuddy], ClientId::CodeBuddy, - vec![ - crate::adapters::InputUnit::plain_file(ClientId::CodeBuddy, codebuddy_path.clone()) - .with_meta(crate::adapters::InputUnitMeta::CodeBuddyExtensionLog { - origin: crate::adapters::CodeBuddyLogOrigin::Extension, - }), - ], + vec![crate::adapters::InputUnit::plain_file( + codebuddy_path.clone(), + crate::adapters::DecoderSpec::codebuddy_extension_log( + 0, + crate::adapters::CodeBuddyLogOrigin::Extension, + ), + )], + ); + assert_ne!( + jsonl_decoder, extension_decoder, + "decoder route is significant" ); - assert_ne!(jsonl_meta, extension_meta, "unit subtype is significant"); let plain_policy = signature_for_test_units( - &["codebuddy".to_string()], + &[ClientId::CodeBuddy], ClientId::CodeBuddy, vec![crate::adapters::InputUnit::plain_file( - ClientId::CodeBuddy, codebuddy_path.clone(), + crate::adapters::DecoderSpec::codebuddy_jsonl(0), )], ); let no_cache_policy = signature_for_test_units( - &["codebuddy".to_string()], + &[ClientId::CodeBuddy], ClientId::CodeBuddy, vec![crate::adapters::InputUnit::no_message_cache( - ClientId::CodeBuddy, codebuddy_path, + crate::adapters::DecoderSpec::codebuddy_jsonl(0), )], ); assert_ne!(plain_policy, no_cache_policy, "input policy is significant"); @@ -2154,8 +2237,8 @@ fn inventory_signature_tracks_order_paths_related_stamps_and_unit_identity() { let amp_group = || { prepared_test_group( ClientId::Amp, - vec![crate::adapters::InputUnit::plain_file( - ClientId::Amp, + vec![plain_test_input( + message_cache::DecoderId::Amp, amp_group_path.clone(), )], ) @@ -2164,17 +2247,17 @@ fn inventory_signature_tracks_order_paths_related_stamps_and_unit_identity() { prepared_test_group( ClientId::CodeBuddy, vec![crate::adapters::InputUnit::plain_file( - ClientId::CodeBuddy, codebuddy_group_path.clone(), + crate::adapters::DecoderSpec::codebuddy_jsonl(0), )], ) }; let group_order = super::input_inventory_signature( - &["amp".to_string(), "codebuddy".to_string()], + &[ClientId::Amp, ClientId::CodeBuddy], &[amp_group(), codebuddy_group()], ); let reversed_group_order = super::input_inventory_signature( - &["amp".to_string(), "codebuddy".to_string()], + &[ClientId::Amp, ClientId::CodeBuddy], &[codebuddy_group(), amp_group()], ); assert_ne!( @@ -2187,7 +2270,7 @@ fn inventory_signature_tracks_order_paths_related_stamps_and_unit_identity() { fn inventory_preparation_rejects_unavailable_primary_snapshots() { let dir = tempfile::TempDir::new().unwrap(); let missing_a = dir.path().join("missing-a.json"); - let error = crate::adapters::InputUnit::plain_file(ClientId::Amp, missing_a.clone()) + let error = plain_test_input(message_cache::DecoderId::Amp, missing_a.clone()) .prepare_snapshot() .expect_err("missing primary inputs must fail before inventory hashing"); assert!(error.to_string().contains(missing_a.to_str().unwrap())); @@ -2210,17 +2293,14 @@ fn inventory_signature_hashes_native_non_utf8_paths() { std::fs::write(&second, b"same").unwrap(); let first_signature = signature_for_test_units( - &["amp".to_string()], + &[ClientId::Amp], ClientId::Amp, - vec![crate::adapters::InputUnit::plain_file(ClientId::Amp, first)], + vec![plain_test_input(message_cache::DecoderId::Amp, first)], ); let second_signature = signature_for_test_units( - &["amp".to_string()], + &[ClientId::Amp], ClientId::Amp, - vec![crate::adapters::InputUnit::plain_file( - ClientId::Amp, - second, - )], + vec![plain_test_input(message_cache::DecoderId::Amp, second)], ); assert_ne!(first_signature, second_signature); } @@ -2347,8 +2427,7 @@ fn test_warm_parse_taking_messages_keeps_outputs_and_cache_stable() { r#"{"id":"msg-1","sessionID":"session-1","role":"assistant","modelID":"accounts/fireworks/models/deepseek-v3-0324","providerID":"fireworks","cost":0,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#, ); drop(conn); - let unit = crate::adapters::InputUnit::sqlite_with_wal(ClientId::OpenCode, path.clone()) - .with_meta(crate::adapters::InputUnitMeta::OpenCodeSqlite); + let unit = opencode_test_input(path.clone()); let home = input_home.path().to_str().unwrap(); let clients = ["opencode".to_string()]; @@ -2370,7 +2449,7 @@ fn test_warm_parse_taking_messages_keeps_outputs_and_cache_stable() { cache .take_messages(&message_cache::CacheReadPlan::new( &path, - unit.parser_version, + unit.decoder.version(), fingerprint.clone(), )) .expect("saved warm cache shard must remain readable") @@ -2381,7 +2460,7 @@ fn test_warm_parse_taking_messages_keeps_outputs_and_cache_stable() { assert!(matches!( cache.take_messages(&message_cache::CacheReadPlan::new( std::path::Path::new("/nonexistent/opencode.db"), - unit.parser_version, + unit.decoder.version(), fingerprint, )), Err(message_cache::CacheReadFailure { @@ -2403,8 +2482,7 @@ fn test_opencode_database_open_errors_are_not_cached_as_empty_success() { { let path = input_home.path().join(".local/share/opencode/opencode.db"); std::fs::create_dir_all(&path).unwrap(); - let unit = crate::adapters::InputUnit::sqlite_with_wal(ClientId::OpenCode, path.clone()) - .with_meta(crate::adapters::InputUnitMeta::OpenCodeSqlite); + let unit = opencode_test_input(path.clone()); let scanner_settings = scanner::ScannerSettings { opencode_db_paths: vec![path.clone()], ..scanner::ScannerSettings::default() @@ -2439,7 +2517,7 @@ fn test_opencode_database_open_errors_are_not_cached_as_empty_success() { let cache = message_cache::InputMessageCache::load().unwrap(); assert!(cache - .get_meta(&path, unit.parser_version) + .get_meta(&path, unit.decoder.version()) .unwrap() .is_none()); @@ -2483,8 +2561,7 @@ fn test_clean_empty_opencode_scan_result_is_not_cached() { let conn = create_opencode_sqlite_db(&path); drop(conn); - let unit = crate::adapters::InputUnit::sqlite_with_wal(ClientId::OpenCode, path.clone()) - .with_meta(crate::adapters::InputUnitMeta::OpenCodeSqlite); + let unit = opencode_test_input(path.clone()); let first_messages = parse_all_messages_with_pricing( input_home.path().to_str().unwrap(), @@ -2496,7 +2573,7 @@ fn test_clean_empty_opencode_scan_result_is_not_cached() { let cache = message_cache::InputMessageCache::load().unwrap(); assert!(cache - .get_meta(&path, unit.parser_version) + .get_meta(&path, unit.decoder.version()) .unwrap() .is_none()); @@ -3237,8 +3314,8 @@ fn test_codex_cache_reparses_from_zero_when_incremental_prefix_is_stale() { assert!(input_cache_for_test_home(input_home.path()) .get_meta( &path, - message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION ) ) @@ -3334,8 +3411,8 @@ fn test_codex_untimestamped_token_row_is_partial_without_cache_shard() { .unwrap() .get_meta( &path, - message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION ) ) @@ -3402,8 +3479,8 @@ fn test_codex_malformed_json_suffix_keeps_prefix_without_cache_shard() { .unwrap() .get_meta( &path, - message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION ) ) @@ -3470,8 +3547,8 @@ fn test_codex_invalid_utf8_suffix_keeps_prefix_without_cache_shard() { assert!(cache .get_meta( &path, - message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION ) ) @@ -3535,8 +3612,8 @@ fn test_codex_unknown_model_prefix_is_partial_then_parses_when_completed() { assert!(input_cache_for_test_home(input_home.path()) .get_meta( &path, - message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION ) ) @@ -3578,8 +3655,8 @@ fn test_codex_unknown_model_prefix_is_partial_then_parses_when_completed() { assert!(input_cache_for_test_home(input_home.path()) .get_meta( &path, - message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION ) ) @@ -3622,8 +3699,8 @@ fn test_codex_cache_skips_non_newline_terminated_resume_prefix() { .unwrap() .get_meta( &path, - message_cache::ParserVersion::new( - message_cache::ParserId::Codex, + message_cache::DecoderVersion::new( + message_cache::DecoderId::Codex, crate::adapters::CODEX_EXEC_IDENTITY_REVISION ) ) @@ -3721,7 +3798,7 @@ fn test_input_cache_does_not_reuse_priced_cost_without_pricing_service() { #[test] fn test_apply_token_pricing_clears_existing_cost_without_pricing() { let mut msg = UnifiedMessage::new_with_agent( - "roocode", + ClientId::RooCode, "gpt-4o", "provider", "session-1", @@ -3833,7 +3910,7 @@ fn test_finalize_token_priced_messages_drops_rows_without_positive_tokens() { let mut messages = vec![ UnifiedMessage::new( - "gemini", + ClientId::Gemini, "gpt-4o", "openai", "zero", @@ -3842,7 +3919,7 @@ fn test_finalize_token_priced_messages_drops_rows_without_positive_tokens() { 0.42, ), UnifiedMessage::new( - "gemini", + ClientId::Gemini, "gpt-4o", "openai", "negative", @@ -3857,7 +3934,7 @@ fn test_finalize_token_priced_messages_drops_rows_without_positive_tokens() { 0.42, ), UnifiedMessage::new( - "gemini", + ClientId::Gemini, "gpt-4o", "openai", "mixed", @@ -3886,7 +3963,7 @@ fn test_finalize_token_priced_messages_drops_rows_without_positive_tokens() { fn test_finalize_token_priced_messages_canonicalizes_provider() { let mut messages = vec![ UnifiedMessage::new( - "pi", + ClientId::Pi, "gpt-5.5", "", "missing-provider", @@ -3901,7 +3978,7 @@ fn test_finalize_token_priced_messages_canonicalizes_provider() { 0.0, ), UnifiedMessage::new( - "mux", + ClientId::Mux, "some-model", "fireworks", "canonical-provider", @@ -3916,7 +3993,7 @@ fn test_finalize_token_priced_messages_canonicalizes_provider() { 0.0, ), UnifiedMessage::new( - "opencode", + ClientId::OpenCode, "grok-code-fast-1", "xai-oauth", "xai-oauth-provider", @@ -3931,7 +4008,7 @@ fn test_finalize_token_priced_messages_canonicalizes_provider() { 0.0, ), UnifiedMessage::new( - "opencode", + ClientId::OpenCode, "grok-code-fast-1", "grok-oauth", "grok-oauth-provider", @@ -3946,7 +4023,7 @@ fn test_finalize_token_priced_messages_canonicalizes_provider() { 0.0, ), UnifiedMessage::new( - "claude", + ClientId::Claude, "kimi-for-coding", "moonshotai", "moonshot-provider", @@ -3961,7 +4038,7 @@ fn test_finalize_token_priced_messages_canonicalizes_provider() { 0.0, ), UnifiedMessage::new( - "copilot", + ClientId::Copilot, "claude-sonnet-4.5", "github-copilot", "copilot-provider", @@ -3976,7 +4053,7 @@ fn test_finalize_token_priced_messages_canonicalizes_provider() { 0.0, ), UnifiedMessage::new( - "codex", + ClientId::Codex, "gpt-5.2", "azure", "azure-provider", @@ -3991,7 +4068,7 @@ fn test_finalize_token_priced_messages_canonicalizes_provider() { 0.0, ), UnifiedMessage::new( - "google-antigravity", + ClientId::Antigravity, "gemini-2.5-pro", "vertex", "vertex-provider", @@ -4006,7 +4083,7 @@ fn test_finalize_token_priced_messages_canonicalizes_provider() { 0.0, ), UnifiedMessage::new( - "opencode", + ClientId::OpenCode, "glm-5.1", "open.bigmodel.cn", "bigmodel-provider", @@ -4021,7 +4098,7 @@ fn test_finalize_token_priced_messages_canonicalizes_provider() { 0.0, ), UnifiedMessage::new( - "claude", + ClientId::Claude, "hy3-preview-agent", "", "hy3-missing-provider", @@ -4036,7 +4113,7 @@ fn test_finalize_token_priced_messages_canonicalizes_provider() { 0.0, ), UnifiedMessage::new( - "file", + ClientId::Gemini, "Jamba-1.5-Large", "unknown", "jamba-unknown-provider", @@ -4051,7 +4128,7 @@ fn test_finalize_token_priced_messages_canonicalizes_provider() { 0.0, ), UnifiedMessage::new( - "file", + ClientId::Gemini, "perplexity/llama-3", "", "route-prefix-must-not-drive-provider", @@ -4108,7 +4185,7 @@ fn test_finalize_token_priced_messages_preserves_custom_provider_literal_tag() { let pricing = pricing::PricingService::new(litellm, HashMap::new()); let mut messages = vec![UnifiedMessage::new( - "claude", + ClientId::Claude, "claude-sonnet-4.5", "venice", "custom-provider", @@ -4132,7 +4209,7 @@ fn test_finalize_token_priced_messages_preserves_custom_provider_literal_tag() { #[test] fn test_finalize_token_priced_messages_preserves_owl_provider_identity() { let mut messages = vec![UnifiedMessage::new( - "opencode", + ClientId::OpenCode, "gpt-5.2", "openai-owl", "owl-provider", @@ -4184,7 +4261,7 @@ fn test_token_breakdown_total_rejects_overflow() { fn test_tui_model_aggregation_uses_unsigned_token_capacity() { let message = || { UnifiedMessage::new( - "antigravity", + ClientId::Antigravity, "gemini-3-pro", "google", "overflow-session", @@ -4220,7 +4297,7 @@ fn test_apply_token_pricing_overrides_cost_when_pricing_exists() { let pricing = pricing::PricingService::new(litellm, HashMap::new()); let mut msg = UnifiedMessage::new( - "codex", + ClientId::Codex, "gpt-4o", "provider", "session-1", @@ -4254,7 +4331,7 @@ fn test_apply_token_pricing_resolves_canonical_longcat_model() { let pricing = pricing::PricingService::new(litellm, HashMap::new()); let mut msg = UnifiedMessage::new( - "claudecode", + ClientId::Claude, "longcat-flash-3b", "meituan", "session-1", @@ -4295,7 +4372,7 @@ fn test_apply_token_pricing_uses_same_price_for_zed_and_other_clients() { reasoning: 0, }; let mut zed_msg = UnifiedMessage::new( - "zed", + ClientId::Zed, "claude-sonnet-4-5", crate::sessions::zed::ZED_HOSTED_PROVIDER, "session-1", @@ -4304,7 +4381,7 @@ fn test_apply_token_pricing_uses_same_price_for_zed_and_other_clients() { 0.0, ); let mut claude_msg = UnifiedMessage::new( - "claudecode", + ClientId::Claude, "claude-sonnet-4-5", crate::sessions::zed::ZED_HOSTED_PROVIDER, "session-1", @@ -4338,7 +4415,7 @@ fn test_apply_token_pricing_custom_zed_price_is_final_price() { ); let mut msg = UnifiedMessage::new( - "zed", + ClientId::Zed, "claude-sonnet-4-5", crate::sessions::zed::ZED_HOSTED_PROVIDER, "session-1", @@ -4372,7 +4449,7 @@ fn test_apply_token_pricing_uses_upstream_provider_for_zed_byok() { let pricing = pricing::PricingService::new(litellm, HashMap::new()); let mut msg = UnifiedMessage::new( - "zed", + ClientId::Zed, "claude-sonnet-4-5", "anthropic", "session-1", @@ -4406,7 +4483,7 @@ fn test_apply_token_pricing_uses_reasoning_for_gemini() { let pricing = pricing::PricingService::new(litellm, HashMap::new()); let mut msg = UnifiedMessage::new( - "gemini", + ClientId::Gemini, "gemini-2.5-pro", "google", "session-1", @@ -4441,7 +4518,7 @@ fn test_apply_token_pricing_uses_cache_read_pricing_for_gemini() { let pricing = pricing::PricingService::new(litellm, HashMap::new()); let mut msg = UnifiedMessage::new( - "gemini", + ClientId::Gemini, "gemini-2.5-pro", "google", "session-1", @@ -4475,7 +4552,7 @@ fn test_finalize_token_pricing_cleans_free_variant_before_lookup() { let pricing = pricing::PricingService::new(litellm, HashMap::new()); let msg = UnifiedMessage::new( - "opencode", + ClientId::OpenCode, "nemotron-3-ultra-free", "nvidia", "session-1", @@ -4512,7 +4589,7 @@ fn test_finalize_token_pricing_cleans_date_variant_before_lookup() { let pricing = pricing::PricingService::new(litellm, HashMap::new()); let msg = UnifiedMessage::new( - "copilot", + ClientId::Copilot, "gpt-4o-mini-2024-07-18", "openai", "session-1", @@ -4550,7 +4627,7 @@ fn test_finalize_token_pricing_cleans_repeated_date_variant_before_lookup() { let mut messages = vec![ UnifiedMessage::new( - "copilot", + ClientId::Copilot, "gpt-4o-mini-2024-07-18", "openai", "session-1", @@ -4565,7 +4642,7 @@ fn test_finalize_token_pricing_cleans_repeated_date_variant_before_lookup() { 0.0, ), UnifiedMessage::new( - "copilot", + ClientId::Copilot, "gpt-4o-mini-2024-07-18", "openai", "session-2", @@ -4612,7 +4689,7 @@ fn test_apply_token_pricing_prefers_provider_aware_match() { let pricing = pricing::PricingService::new(litellm, HashMap::new()); let mut msg = UnifiedMessage::new( - "opencode", + ClientId::OpenCode, "grok-code", "azure", "session-1", @@ -4654,7 +4731,7 @@ fn test_apply_token_pricing_uses_nested_reseller_exact_match() { let pricing = pricing::PricingService::new(litellm, HashMap::new()); let mut msg = UnifiedMessage::new( - "opencode", + ClientId::OpenCode, "gpt-4", "azure", "session-1", @@ -4698,7 +4775,7 @@ fn test_apply_token_pricing_clears_cost_without_exact_pricing() { let pricing = pricing::PricingService::new(litellm, openrouter); let mut msg = UnifiedMessage::new( - "opencode", + ClientId::OpenCode, "accounts/fireworks/models/deepseek-v4-pro", "fireworks", "session-1", @@ -4745,7 +4822,7 @@ fn test_apply_token_pricing_prefers_provider_specific_exact_match_over_plain_exa let pricing = pricing::PricingService::new(litellm, openrouter); let mut msg = UnifiedMessage::new( - "opencode", + ClientId::OpenCode, "gemini-2.5-pro", "google", "session-1", @@ -4787,7 +4864,7 @@ fn test_apply_token_pricing_normalizes_openai_codex_provider() { let pricing = pricing::PricingService::new(litellm, HashMap::new()); let mut msg = UnifiedMessage::new( - "openclaw", + ClientId::OpenClaw, "gpt-5.2", "openai-codex", "session-1", @@ -4821,7 +4898,7 @@ fn test_apply_token_pricing_normalizes_openai_pro_provider() { let pricing = pricing::PricingService::new(litellm, HashMap::new()); let mut msg = UnifiedMessage::new( - "kimi", + ClientId::Kimi, "gpt-5.2", "openai-pro", "session-1", @@ -4855,7 +4932,7 @@ fn test_apply_token_pricing_honors_observed_owl_scope_for_gpt() { let pricing = pricing::PricingService::new(litellm, HashMap::new()); let mut msg = UnifiedMessage::new( - "opencode", + ClientId::OpenCode, "gpt-5.2", "owl", "session-1", @@ -4889,7 +4966,7 @@ fn test_apply_token_pricing_honors_observed_owl_scope_for_claude() { let pricing = pricing::PricingService::new(litellm, HashMap::new()); let mut msg = UnifiedMessage::new( - "opencode", + ClientId::OpenCode, "claude-sonnet-4-5", "owl", "session-1", @@ -4923,7 +5000,7 @@ fn test_apply_token_pricing_honors_observed_owl_scope_for_minimax() { let pricing = pricing::PricingService::new(litellm, HashMap::new()); let mut msg = UnifiedMessage::new( - "opencode", + ClientId::OpenCode, "MiniMax-M2.1", "owl", "session-1", @@ -4958,7 +5035,7 @@ fn test_apply_token_pricing_prices_claude_code_gpt_5_3_codex() { let pricing = pricing::PricingService::new(litellm, HashMap::new()); let mut msg = UnifiedMessage::new( - "claude", + ClientId::Claude, "gpt-5.3-codex", "openai", "session-1", @@ -4993,7 +5070,7 @@ fn test_apply_token_pricing_prices_claude_code_minimax_model() { let pricing = pricing::PricingService::new(litellm, HashMap::new()); let mut msg = UnifiedMessage::new( - "claude", + ClientId::Claude, "MiniMax-M2.1", "minimax", "session-1", @@ -5027,7 +5104,7 @@ fn test_apply_token_pricing_prices_canonical_kimi_k2_6() { let pricing = pricing::PricingService::new(HashMap::new(), openrouter); let mut msg = UnifiedMessage::new( - "kimi", + ClientId::Kimi, "kimi-k2.6", "kimi", "session-1", @@ -5065,7 +5142,7 @@ fn test_select_local_parse_pricing_prefers_fresh_service_for_new_models() { let selected = select_local_parse_pricing(Ok(Arc::clone(&fresh)), || Some(stale)).unwrap(); let mut msg = UnifiedMessage::new( - "opencode", + ClientId::OpenCode, "gpt-5.4", "openai", "session-1", @@ -5144,7 +5221,7 @@ fn test_parse_all_messages_with_pricing_keeps_gateway_message_under_real_client_ .unwrap(); assert_eq!(messages.len(), 1); - assert_eq!(messages[0].client.as_ref(), "opencode"); + assert_eq!(messages[0].client, ClientId::OpenCode); assert_eq!(messages[0].model_id.as_ref(), "deepseek-v3"); assert_eq!(messages[0].provider_id.as_ref(), "deepseek"); } @@ -5175,7 +5252,7 @@ fn test_local_message_loader_preserves_gateway_message_client_counts() { assert_eq!(parsed.counts.get(ClientId::OpenCode), 1); assert_eq!(parsed.messages.len(), 1); - assert_eq!(parsed.messages[0].client.as_ref(), "opencode"); + assert_eq!(parsed.messages[0].client, ClientId::OpenCode); assert_eq!(parsed.messages[0].model_id.as_ref(), "deepseek-v3"); assert_eq!(parsed.messages[0].provider_id.as_ref(), "fireworks"); } @@ -5256,7 +5333,7 @@ fn test_local_message_loader_honors_scanner_settings_opencode_db_paths() { "scanner.opencodeDbPaths must reach the local message loading path" ); assert_eq!(parsed_with_settings.messages.len(), 1); - assert_eq!(parsed_with_settings.messages[0].client.as_ref(), "opencode"); + assert_eq!(parsed_with_settings.messages[0].client, ClientId::OpenCode); assert_eq!( parsed_with_settings.messages[0].model_id.as_ref(), "claude-sonnet-4" @@ -5297,11 +5374,10 @@ fn test_missing_configured_opencode_database_is_an_explicit_error() { "failure must identify the failed input operation: {failure:?}" ); - let unit = crate::adapters::InputUnit::sqlite_with_wal(ClientId::OpenCode, missing_db.clone()) - .with_meta(crate::adapters::InputUnitMeta::OpenCodeSqlite); + let unit = opencode_test_input(missing_db.clone()); assert!(message_cache::InputMessageCache::load() .unwrap() - .get_meta(&missing_db, unit.parser_version) + .get_meta(&missing_db, unit.decoder.version()) .unwrap() .is_none()); } @@ -5365,6 +5441,11 @@ fn public_raw_message_report_preserves_input_health_and_metadata() { report.metadata.input_inventory_signature.as_bytes(), &[0_u8; 32] ); + assert!(report + .metadata + .input_footprint + .contains_client(ClientId::OpenCode)); + assert_eq!(report.metadata.input_footprint.total_bytes().unwrap(), 0); } #[test] @@ -5384,7 +5465,9 @@ fn test_opencode_auto_discovery_error_reaches_public_loader() { assert!(loaded.messages.is_empty()); assert_eq!(loaded.health.failed_inputs(), 1); - let failure = loaded.health.inputs()[0].status.failure().unwrap(); + let input_health = &loaded.health.inputs()[0]; + assert_eq!(input_health.client, ClientId::OpenCode); + let failure = input_health.status.failure().unwrap(); assert!( failure .message @@ -5445,7 +5528,7 @@ fn test_local_message_loader_honors_scanner_extra_scan_paths_for_hermes_profile_ assert_eq!(parsed_with_settings.counts.get(ClientId::Hermes), 2); assert_eq!(parsed_with_settings.messages.len(), 1); - assert_eq!(parsed_with_settings.messages[0].client.as_ref(), "hermes"); + assert_eq!(parsed_with_settings.messages[0].client, ClientId::Hermes); assert_eq!( parsed_with_settings.messages[0].agent.as_deref(), Some("Hermes Agent") @@ -5502,7 +5585,7 @@ fn test_local_message_loader_honors_scanner_extra_scan_paths_for_zed_threads_db( assert_eq!(parsed_with_settings.counts.get(ClientId::Zed), 1); assert_eq!(parsed_with_settings.messages.len(), 1); - assert_eq!(parsed_with_settings.messages[0].client.as_ref(), "zed"); + assert_eq!(parsed_with_settings.messages[0].client, ClientId::Zed); assert_eq!( parsed_with_settings.messages[0].session_id.as_ref(), "zed-extra-thread" @@ -5636,7 +5719,7 @@ fn test_driver_uses_zed_adapter_when_only_zed_requested() { "an explicit Zed-only request must not scan OpenCode SQLite" ); assert_eq!(parsed.messages.len(), 1); - assert_eq!(parsed.messages[0].client.as_ref(), "zed"); + assert_eq!(parsed.messages[0].client, ClientId::Zed); assert_eq!(parsed.messages[0].session_id.as_ref(), "zed-only-thread"); } @@ -5687,7 +5770,7 @@ fn test_driver_uses_simple_file_adapter_when_only_amp_requested() { "an explicit Amp-only request must not scan OpenCode SQLite" ); assert_eq!(parsed.messages.len(), 1); - assert_eq!(parsed.messages[0].client.as_ref(), "amp"); + assert_eq!(parsed.messages[0].client, ClientId::Amp); assert_eq!(parsed.messages[0].model_id.as_ref(), "claude-opus-4.7"); } @@ -5736,7 +5819,7 @@ fn test_driver_uses_custom_file_adapter_when_only_codebuff_requested() { "an explicit Codebuff-only request must not scan OpenCode SQLite" ); assert_eq!(parsed.messages.len(), 1); - assert_eq!(parsed.messages[0].client.as_ref(), "codebuff"); + assert_eq!(parsed.messages[0].client, ClientId::Codebuff); assert_eq!(parsed.messages[0].model_id.as_ref(), "claude-sonnet-4"); } @@ -5767,13 +5850,11 @@ fn test_driver_uses_pi_and_omp_adapters_when_requested() { assert_eq!(parsed.counts.get(ClientId::Pi), 1); assert_eq!(parsed.counts.get(ClientId::Omp), 2); - assert!(parsed - .messages - .iter() - .any(|message| message.client.as_ref() == "pi" - && message.session_id.as_ref() == "pi_ses_001")); + assert!(parsed.messages.iter().any( + |message| message.client == ClientId::Pi && message.session_id.as_ref() == "pi_ses_001" + )); assert!(parsed.messages.iter().any(|message| { - message.client.as_ref() == "omp" + message.client == ClientId::Omp && message.session_id.as_ref() == "child-session" && message.agent.as_deref() == Some("OMP Reviewer") })); @@ -5809,7 +5890,7 @@ fn test_driver_all_clients_includes_each_adapter_without_duplicate() { parsed .messages .iter() - .filter(|message| message.client.as_ref() == "zed") + .filter(|message| message.client == ClientId::Zed) .count(), 1 ); @@ -5958,12 +6039,12 @@ fn test_local_message_loader_claude_filter_ignores_scanner_settings_opencode_db_ "Claude message must still be counted" ); assert_eq!(parsed.messages.len(), 1); - assert_eq!(parsed.messages[0].client.as_ref(), "claude"); + assert_eq!(parsed.messages[0].client, ClientId::Claude); assert!( parsed .messages .iter() - .all(|m| m.client.as_ref() != "opencode"), + .all(|m| m.client != ClientId::OpenCode), "no OpenCode messages may leak into a Claude-only result, got {:?}", parsed.messages ); @@ -6003,7 +6084,7 @@ fn test_local_message_loader_claude_transcripts_count_only_usage_metadata() { assert_eq!(parsed.counts.get(ClientId::Claude), 1); assert_eq!(parsed.messages.len(), 1); - assert_eq!(parsed.messages[0].client.as_ref(), "claude"); + assert_eq!(parsed.messages[0].client, ClientId::Claude); assert_eq!( parsed.messages[0].session_id.as_ref(), "ses_123456789012345678901234567" @@ -6054,7 +6135,7 @@ fn test_local_message_loader_amp_reads_current_thread_files() { assert_eq!(parsed.counts.get(ClientId::Amp), 1); assert_eq!(parsed.messages.len(), 1); - assert_eq!(parsed.messages[0].client.as_ref(), "amp"); + assert_eq!(parsed.messages[0].client, ClientId::Amp); assert_eq!(parsed.messages[0].model_id.as_ref(), "claude-opus-4.7"); assert_eq!(parsed.messages[0].provider_id.as_ref(), "anthropic"); assert_eq!(parsed.messages[0].tokens.input, 10); diff --git a/crates/tokscale-core/src/message_cache.rs b/crates/tokscale-core/src/message_cache.rs index e6c523400..dc70db24d 100644 --- a/crates/tokscale-core/src/message_cache.rs +++ b/crates/tokscale-core/src/message_cache.rs @@ -15,10 +15,10 @@ use std::time::UNIX_EPOCH; #[cfg(not(any(unix, windows)))] compile_error!("input-message cache requires stable Unix or Windows file identity"); -// Input-message cache shards split serialization layout from parser/input -// semantics. Bump this only when the shard bincode layout changes; parser-only -// fixes should bump the relevant InputUnit parser revision instead. -const CACHE_FORMAT_VERSION: u32 = 11; +// Input-message cache shards split serialization layout from decoder/input +// semantics. Bump this only when the shard bincode layout changes; decoder-only +// fixes should bump the relevant InputUnit decoder revision instead. +const CACHE_FORMAT_VERSION: u32 = 12; #[cfg(test)] const UNSUPPORTED_CACHE_FORMAT_VERSION: u32 = CACHE_FORMAT_VERSION - 1; const SHARD_MAGIC: [u8; 8] = *b"TOKSHRD\0"; @@ -257,101 +257,112 @@ impl InputCachePruneError { } } -pub(crate) type ParserRevision = u32; +pub(crate) type DecoderRevision = u32; -// Persisted in input-cache shard headers. Shard filenames use the stable names -// below, but reordering or removing variants still changes header bincode and -// must bump CACHE_FORMAT_VERSION. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] -pub(crate) enum ParserId { - OpenCodeSqlite, - Claude, - Codex, - Gemini, - Amp, - Droid, - OpenClaw, - Pi, - Omp, - Kimi, - Qwen, - RooCode, - Mux, - Kilo, - Hermes, - Copilot, - Goose, - Codebuff, - AntigravityCliSqlite, - Zed, - Kiro, - KiroFile, - KiroSqlite, - KiroGlobalStorage, - Junie, - Cline, - CommandCode, - Grok, - Zcode, - Warp, - CodeBuddy, - OmpParentHealth, +macro_rules! define_decoder_ids { + ($($variant:ident => ($stable_name:literal, $plain:literal)),+ $(,)?) => { + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] + pub(crate) enum DecoderId { + $($variant),+ + } + + impl DecoderId { + pub(crate) const fn stable_name(self) -> &'static str { + match self { + $(Self::$variant => $stable_name),+ + } + } + + pub(crate) fn from_stable_name(stable_name: &str) -> Option { + match stable_name { + $($stable_name => Some(Self::$variant),)+ + _ => None, + } + } + + pub(crate) const fn supports_plain_route(self) -> bool { + match self { + $(Self::$variant => $plain),+ + } + } + } + + impl Serialize for DecoderId { + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + serializer.serialize_str(self.stable_name()) + } + } + + impl<'de> Deserialize<'de> for DecoderId { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + let stable_name = String::deserialize(deserializer)?; + Self::from_stable_name(&stable_name).ok_or_else(|| { + serde::de::Error::unknown_variant( + &stable_name, + &[$($stable_name),+], + ) + }) + } + } + }; +} + +define_decoder_ids! { + OpenCodeSqlite => ("opencode-sqlite", false), + Claude => ("claude", true), + Codex => ("codex", false), + Gemini => ("gemini", true), + Amp => ("amp", true), + Droid => ("droid", true), + OpenClaw => ("openclaw", true), + Pi => ("pi", true), + Omp => ("omp", true), + Kimi => ("kimi", true), + Qwen => ("qwen", true), + RooCode => ("roo-code", true), + Mux => ("mux", true), + Kilo => ("kilo", true), + Hermes => ("hermes", true), + Copilot => ("copilot", true), + Goose => ("goose", true), + Codebuff => ("codebuff", true), + AntigravityCliSqlite => ("antigravity-cli-sqlite", false), + Zed => ("zed", true), + Kiro => ("kiro", true), + KiroFile => ("kiro-file", false), + KiroSqlite => ("kiro-sqlite", false), + KiroGlobalStorage => ("kiro-global-storage", false), + Junie => ("junie", true), + Cline => ("cline", true), + CommandCode => ("command-code", true), + Grok => ("grok", true), + Zcode => ("zcode", true), + Warp => ("warp", true), + CodeBuddy => ("codebuddy", false), + OmpParentHealth => ("omp-parent-health", true), } #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] -pub(crate) struct ParserVersion { - pub parser_id: ParserId, - pub revision: ParserRevision, +pub(crate) struct DecoderVersion { + pub decoder_id: DecoderId, + pub revision: DecoderRevision, } -impl ParserVersion { - pub(crate) const fn new(parser_id: ParserId, revision: ParserRevision) -> Self { +impl DecoderVersion { + pub(crate) const fn new(decoder_id: DecoderId, revision: DecoderRevision) -> Self { Self { - parser_id, + decoder_id, revision, } } } -impl ParserId { - pub(crate) const fn stable_name(self) -> &'static str { - match self { - Self::OpenCodeSqlite => "opencode-sqlite", - Self::Claude => "claude", - Self::Codex => "codex", - Self::Gemini => "gemini", - Self::Amp => "amp", - Self::Droid => "droid", - Self::OpenClaw => "openclaw", - Self::Pi => "pi", - Self::Omp => "omp", - Self::OmpParentHealth => "omp-parent-health", - Self::Kimi => "kimi", - Self::Qwen => "qwen", - Self::RooCode => "roo-code", - Self::Mux => "mux", - Self::Kilo => "kilo", - Self::Hermes => "hermes", - Self::Copilot => "copilot", - Self::Goose => "goose", - Self::Codebuff => "codebuff", - Self::AntigravityCliSqlite => "antigravity-cli-sqlite", - Self::Zed => "zed", - Self::Kiro => "kiro", - Self::KiroFile => "kiro-file", - Self::KiroSqlite => "kiro-sqlite", - Self::KiroGlobalStorage => "kiro-global-storage", - Self::Junie => "junie", - Self::Cline => "cline", - Self::CommandCode => "command-code", - Self::Grok => "grok", - Self::Zcode => "zcode", - Self::Warp => "warp", - Self::CodeBuddy => "codebuddy", - } - } -} - fn cache_dir() -> Result { crate::paths::try_get_cache_dir() } @@ -1064,14 +1075,14 @@ pub(crate) struct CodexIncrementalCache { #[derive(Debug, Clone, PartialEq, Eq, Hash)] struct CachedInputKey { path: CachedPath, - parser_version: ParserVersion, + decoder_version: DecoderVersion, } impl CachedInputKey { - fn new(path: &Path, parser_version: ParserVersion) -> Self { + fn new(path: &Path, decoder_version: DecoderVersion) -> Self { Self { path: CachedPath::from_path(path), - parser_version, + decoder_version, } } @@ -1124,8 +1135,8 @@ pub(crate) enum CacheReadFailureReason { }, #[error("shard input path no longer matches the read plan")] InputPathMismatch, - #[error("shard parser version no longer matches the read plan")] - ParserVersionMismatch, + #[error("shard decoder version no longer matches the read plan")] + DecoderVersionMismatch, #[error("shard fingerprint no longer matches the read plan")] ShardFingerprintMismatch, #[error("failed to decode shard body: {source}")] @@ -1150,7 +1161,7 @@ impl CacheReadFailureReason { | Self::InvalidHeaderLength { .. } | Self::HeaderDecode { .. } | Self::InputPathMismatch - | Self::ParserVersionMismatch + | Self::DecoderVersionMismatch ) } } @@ -1158,7 +1169,7 @@ impl CacheReadFailureReason { #[derive(Debug)] pub(crate) struct CacheReadFailure { pub(crate) input_path: PathBuf, - pub(crate) parser_version: ParserVersion, + pub(crate) decoder_version: DecoderVersion, pub(crate) shard_path: Option, pub(crate) reason: CacheReadFailureReason, } @@ -1178,7 +1189,7 @@ impl CacheReadFailure { | CacheReadFailureReason::InvalidHeaderLength { .. } | CacheReadFailureReason::HeaderDecode { .. } | CacheReadFailureReason::InputPathMismatch - | CacheReadFailureReason::ParserVersionMismatch + | CacheReadFailureReason::DecoderVersionMismatch | CacheReadFailureReason::ShardFingerprintMismatch | CacheReadFailureReason::BodyDecode { .. } | CacheReadFailureReason::MessageCountMismatch { .. } => true, @@ -1206,7 +1217,7 @@ impl CacheReadFailure { | CacheReadFailureReason::InvalidHeaderLength { .. } | CacheReadFailureReason::HeaderDecode { .. } | CacheReadFailureReason::InputPathMismatch - | CacheReadFailureReason::ParserVersionMismatch + | CacheReadFailureReason::DecoderVersionMismatch | CacheReadFailureReason::ShardFingerprintMismatch => false, } } @@ -1215,7 +1226,7 @@ impl CacheReadFailure { #[derive(Debug)] pub(crate) struct CacheLookupFailure { pub(crate) input_path: PathBuf, - pub(crate) parser_version: ParserVersion, + pub(crate) decoder_version: DecoderVersion, pub(crate) shard_path: PathBuf, pub(crate) reason: CacheReadFailureReason, } @@ -1224,10 +1235,10 @@ impl std::fmt::Display for CacheLookupFailure { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!( formatter, - "input cache v{} header read failed for `{}` with parser {:?} at `{}`: {}", + "input cache v{} header read failed for `{}` with decoder {:?} at `{}`: {}", CACHE_FORMAT_VERSION, self.input_path.display(), - self.parser_version, + self.decoder_version, self.shard_path.display(), self.reason ) @@ -1244,9 +1255,9 @@ impl std::fmt::Display for CacheReadFailure { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!( formatter, - "input cache body read failed for `{}` with parser {:?}", + "input cache body read failed for `{}` with decoder {:?}", self.input_path.display(), - self.parser_version + self.decoder_version )?; if let Some(shard_path) = &self.shard_path { write!(formatter, " at `{}`", shard_path.display())?; @@ -1269,7 +1280,7 @@ impl CacheReadFailure { ) -> Self { Self { input_path: plan.path(), - parser_version: plan.parser_version(), + decoder_version: plan.decoder_version(), shard_path, reason, } @@ -1279,11 +1290,11 @@ impl CacheReadFailure { impl CacheReadPlan { pub(crate) fn new( path: &Path, - parser_version: ParserVersion, + decoder_version: DecoderVersion, fingerprint: InputFingerprint, ) -> Self { Self { - key: CachedInputKey::new(path, parser_version), + key: CachedInputKey::new(path, decoder_version), fingerprint, } } @@ -1292,15 +1303,15 @@ impl CacheReadPlan { self.key.to_path_buf() } - pub(crate) fn parser_version(&self) -> ParserVersion { - self.key.parser_version + pub(crate) fn decoder_version(&self) -> DecoderVersion { + self.key.decoder_version } } #[derive(Debug, Serialize, Deserialize)] pub(crate) struct CachedInputEntry { pub path: CachedPath, - pub parser_version: ParserVersion, + pub decoder_version: DecoderVersion, pub fingerprint: InputFingerprint, pub messages: Vec, pub codex_incremental: Option, @@ -1321,14 +1332,14 @@ impl CachedInputEntry { #[cfg(test)] pub(crate) fn new_with_revision( path: &Path, - parser_revision: ParserRevision, + decoder_revision: DecoderRevision, fingerprint: InputFingerprint, messages: Vec, codex_incremental: Option, ) -> Self { Self::new_with_version( path, - ParserVersion::new(ParserId::Amp, parser_revision), + DecoderVersion::new(DecoderId::Amp, decoder_revision), fingerprint, messages, codex_incremental, @@ -1338,14 +1349,14 @@ impl CachedInputEntry { #[cfg(test)] pub(crate) fn new_with_version( path: &Path, - parser_version: ParserVersion, + decoder_version: DecoderVersion, fingerprint: InputFingerprint, messages: Vec, codex_incremental: Option, ) -> Self { Self { path: CachedPath::from_path(path), - parser_version, + decoder_version, fingerprint, messages, codex_incremental, @@ -1356,7 +1367,7 @@ impl CachedInputEntry { fn plan(&self) -> CacheWritePlan { CacheWritePlan { path: self.path.clone(), - parser_version: self.parser_version, + decoder_version: self.decoder_version, fingerprint: self.fingerprint.clone(), codex_incremental: self.codex_incremental.clone(), rejections: self.rejections.clone(), @@ -1367,7 +1378,7 @@ impl CachedInputEntry { fn key(&self) -> CachedInputKey { CachedInputKey { path: self.path.clone(), - parser_version: self.parser_version, + decoder_version: self.decoder_version, } } } @@ -1375,7 +1386,7 @@ impl CachedInputEntry { #[derive(Debug, Clone)] pub(crate) struct CacheWritePlan { path: CachedPath, - parser_version: ParserVersion, + decoder_version: DecoderVersion, fingerprint: InputFingerprint, codex_incremental: Option, rejections: crate::input_health::RejectionSummary, @@ -1384,13 +1395,13 @@ pub(crate) struct CacheWritePlan { impl CacheWritePlan { pub(crate) fn new( path: &Path, - parser_version: ParserVersion, + decoder_version: DecoderVersion, fingerprint: InputFingerprint, codex_incremental: Option, ) -> Self { Self { path: CachedPath::from_path(path), - parser_version, + decoder_version, fingerprint, codex_incremental, rejections: Default::default(), @@ -1410,14 +1421,14 @@ impl CacheWritePlan { fn key(&self) -> CachedInputKey { CachedInputKey { path: self.path.clone(), - parser_version: self.parser_version, + decoder_version: self.decoder_version, } } } #[derive(Debug, Clone, Serialize, Deserialize)] struct CachedShardHeader { - parser_version: ParserVersion, + decoder_version: DecoderVersion, path: CachedPath, fingerprint: InputFingerprint, codex_incremental: Option, @@ -1513,9 +1524,9 @@ impl InputMessageCache { pub(crate) fn get_meta( &self, path: &Path, - parser_version: ParserVersion, + decoder_version: DecoderVersion, ) -> Result, CacheLookupFailure> { - let key = CachedInputKey::new(path, parser_version); + let key = CachedInputKey::new(path, decoder_version); if self.deleted_paths.contains(&key) || self.taken_paths.contains(&key) { return Ok(None); } @@ -1533,21 +1544,21 @@ impl InputMessageCache { Err(reason) => { return Err(CacheLookupFailure { input_path: key.to_path_buf(), - parser_version: key.parser_version, + decoder_version: key.decoder_version, shard_path, reason, }); } }; - if header.path != key.path || header.parser_version != key.parser_version { + if header.path != key.path || header.decoder_version != key.decoder_version { let reason = if header.path != key.path { CacheReadFailureReason::InputPathMismatch } else { - CacheReadFailureReason::ParserVersionMismatch + CacheReadFailureReason::DecoderVersionMismatch }; return Err(CacheLookupFailure { input_path: key.to_path_buf(), - parser_version: key.parser_version, + decoder_version: key.decoder_version, shard_path, reason, }); @@ -1632,8 +1643,8 @@ impl InputMessageCache { Ok(entry.messages) } - pub(crate) fn remove(&mut self, path: &Path, parser_version: ParserVersion) { - let key = CachedInputKey::new(path, parser_version); + pub(crate) fn remove(&mut self, path: &Path, decoder_version: DecoderVersion) { + let key = CachedInputKey::new(path, decoder_version); if self.is_protected(&key) { return; } @@ -1644,8 +1655,8 @@ impl InputMessageCache { self.dirty = true; } - pub(crate) fn invalidate_read(&mut self, path: &Path, parser_version: ParserVersion) { - let key = CachedInputKey::new(path, parser_version); + pub(crate) fn invalidate_read(&mut self, path: &Path, decoder_version: DecoderVersion) { + let key = CachedInputKey::new(path, decoder_version); self.invalidated_read_paths.insert(key.clone()); self.taken_paths.insert(key); } @@ -1740,7 +1751,7 @@ pub fn prune_input_message_cache() -> Result = HashMap::new(); + let mut latest_revisions: HashMap<(CachedPath, DecoderId), DecoderRevision> = HashMap::new(); let mut input_existence: HashMap = HashMap::new(); for shard_path in shard_paths { @@ -1755,7 +1766,7 @@ pub fn prune_input_message_cache() -> Result Result Result [u8; 32] { key.path.update_shard_key(&mut hasher); hash_inventory_bytes( &mut hasher, - key.parser_version.parser_id.stable_name().as_bytes(), + key.decoder_version.decoder_id.stable_name().as_bytes(), ); - hasher.update(key.parser_version.revision.to_le_bytes()); + hasher.update(key.decoder_version.revision.to_le_bytes()); hasher.finalize().into() } #[cfg(test)] fn shard_path( path: &Path, - parser_version: ParserVersion, + decoder_version: DecoderVersion, ) -> Result { let dir = cache_dir()?; Ok(shard_path_for_input_key( &dir, - &CachedInputKey::new(path, parser_version), + &CachedInputKey::new(path, decoder_version), )) } @@ -1866,18 +1877,18 @@ fn shard_path_for_input_key(cache_dir: &Path, key: &CachedInputKey) -> PathBuf { pub(crate) fn shard_path_for_test( cache_dir: &Path, input_path: &Path, - parser_version: ParserVersion, + decoder_version: DecoderVersion, ) -> PathBuf { - shard_path_for_input_key(cache_dir, &CachedInputKey::new(input_path, parser_version)) + shard_path_for_input_key(cache_dir, &CachedInputKey::new(input_path, decoder_version)) } #[cfg(test)] pub(crate) fn mark_current_key_shard_as_unsupported_format_for_test( cache_dir: &Path, input_path: &Path, - parser_version: ParserVersion, + decoder_version: DecoderVersion, ) -> PathBuf { - let shard_path = shard_path_for_test(cache_dir, input_path, parser_version); + let shard_path = shard_path_for_test(cache_dir, input_path, decoder_version); let mut file = std::fs::OpenOptions::new() .write(true) .open(&shard_path) @@ -1894,9 +1905,9 @@ pub(crate) fn mark_current_key_shard_as_unsupported_format_for_test( pub(crate) fn mark_current_key_shard_as_future_format_for_test( cache_dir: &Path, input_path: &Path, - parser_version: ParserVersion, + decoder_version: DecoderVersion, ) -> PathBuf { - let shard_path = shard_path_for_test(cache_dir, input_path, parser_version); + let shard_path = shard_path_for_test(cache_dir, input_path, decoder_version); let mut file = std::fs::OpenOptions::new() .write(true) .open(&shard_path) @@ -1913,9 +1924,9 @@ pub(crate) fn mark_current_key_shard_as_future_format_for_test( pub(crate) fn truncate_shard_after_header_for_test( cache_dir: &Path, input_path: &Path, - parser_version: ParserVersion, + decoder_version: DecoderVersion, ) -> PathBuf { - let shard_path = shard_path_for_test(cache_dir, input_path, parser_version); + let shard_path = shard_path_for_test(cache_dir, input_path, decoder_version); let mut file = std::fs::OpenOptions::new() .read(true) .write(true) @@ -1939,10 +1950,10 @@ pub(crate) fn truncate_shard_after_header_for_test( pub(crate) fn replace_shard_message_count_for_test( cache_dir: &Path, input_path: &Path, - parser_version: ParserVersion, + decoder_version: DecoderVersion, message_count: usize, ) -> PathBuf { - let shard_path = shard_path_for_test(cache_dir, input_path, parser_version); + let shard_path = shard_path_for_test(cache_dir, input_path, decoder_version); let mut file = std::fs::OpenOptions::new() .read(true) .write(true) @@ -1978,7 +1989,7 @@ fn hex_sha256(bytes: &[u8; 32]) -> String { fn header_from_plan(plan: &CacheWritePlan, message_count: usize) -> CachedShardHeader { CachedShardHeader { - parser_version: plan.parser_version, + decoder_version: plan.decoder_version, path: plan.path.clone(), fingerprint: plan.fingerprint.clone(), codex_incremental: plan.codex_incremental.clone(), @@ -2027,8 +2038,8 @@ fn read_shard_entry_with_plan( if header.path != plan.key.path { return Err(CacheReadFailureReason::InputPathMismatch); } - if header.parser_version != plan.key.parser_version { - return Err(CacheReadFailureReason::ParserVersionMismatch); + if header.decoder_version != plan.key.decoder_version { + return Err(CacheReadFailureReason::DecoderVersionMismatch); } if header.fingerprint != plan.fingerprint { return Err(CacheReadFailureReason::ShardFingerprintMismatch); @@ -2046,7 +2057,7 @@ fn read_shard_entry_with_plan( Ok(CachedInputEntry { path: header.path, - parser_version: header.parser_version, + decoder_version: header.decoder_version, fingerprint: header.fingerprint, messages: body.messages, rejections: header.rejections, @@ -2446,35 +2457,51 @@ mod tests { restore_env_var("TOKSCALE_CONFIG_DIR", prev.3); } - fn test_parser_version(revision: ParserRevision) -> ParserVersion { - ParserVersion::new(ParserId::Amp, revision) + fn test_decoder_version(revision: DecoderRevision) -> DecoderVersion { + DecoderVersion::new(DecoderId::Amp, revision) + } + + #[test] + fn decoder_id_bincode_uses_its_stable_name() { + let encoded_id = bincode::options().serialize(&DecoderId::Amp).unwrap(); + let encoded_name = bincode::options() + .serialize(DecoderId::Amp.stable_name()) + .unwrap(); + + assert_eq!(encoded_id, encoded_name); + assert_eq!( + bincode::options() + .deserialize::(&encoded_id) + .unwrap(), + DecoderId::Amp + ); } fn test_cache_read_failure(reason: CacheReadFailureReason) -> CacheReadFailure { CacheReadFailure { input_path: PathBuf::from("/test/input"), - parser_version: test_parser_version(1), + decoder_version: test_decoder_version(1), shard_path: Some(PathBuf::from("/test/shard")), reason, } } #[test] - fn current_shard_key_uses_stable_path_parser_tag_and_revision_fields() { + fn current_shard_key_uses_stable_path_decoder_tag_and_revision_fields() { let path = Path::new("/test/input"); - let amp_v1 = CachedInputKey::new(path, ParserVersion::new(ParserId::Amp, 1)); - let amp_v2 = CachedInputKey::new(path, ParserVersion::new(ParserId::Amp, 2)); - let claude_v1 = CachedInputKey::new(path, ParserVersion::new(ParserId::Claude, 1)); + let amp_v1 = CachedInputKey::new(path, DecoderVersion::new(DecoderId::Amp, 1)); + let amp_v2 = CachedInputKey::new(path, DecoderVersion::new(DecoderId::Amp, 2)); + let claude_v1 = CachedInputKey::new(path, DecoderVersion::new(DecoderId::Claude, 1)); let other_path = CachedInputKey::new( Path::new("/test/other-input"), - ParserVersion::new(ParserId::Amp, 1), + DecoderVersion::new(DecoderId::Amp, 1), ); assert_eq!( shard_key_for_input_key(&_v1), shard_key_for_input_key(&CachedInputKey::new( path, - ParserVersion::new(ParserId::Amp, 1), + DecoderVersion::new(DecoderId::Amp, 1), )) ); assert_ne!( @@ -2526,7 +2553,7 @@ mod tests { )), }, CacheReadFailureReason::InputPathMismatch, - CacheReadFailureReason::ParserVersionMismatch, + CacheReadFailureReason::DecoderVersionMismatch, CacheReadFailureReason::FingerprintMismatch, CacheReadFailureReason::ShardFingerprintMismatch, ] { @@ -3203,7 +3230,7 @@ mod tests { cache.insert(entry); cache.save_if_dirty().unwrap(); - let shard = shard_path(file.path(), test_parser_version(1)).unwrap(); + let shard = shard_path(file.path(), test_decoder_version(1)).unwrap(); assert!(shard.exists()); let mut envelope = [0_u8; 12]; File::open(&shard) @@ -3218,7 +3245,7 @@ mod tests { let mut loaded = InputMessageCache::load().unwrap(); let meta = loaded - .get_meta(file.path(), test_parser_version(1)) + .get_meta(file.path(), test_decoder_version(1)) .unwrap() .unwrap(); assert_eq!(meta.fingerprint, expected_fingerprint); @@ -3228,7 +3255,7 @@ mod tests { let messages = loaded .take_messages(&CacheReadPlan::new( file.path(), - test_parser_version(1), + test_decoder_version(1), expected_fingerprint, )) .unwrap(); @@ -3255,7 +3282,7 @@ mod tests { let fingerprint = InputFingerprint::from_path(file.path()).unwrap(); let plan = CacheWritePlan::new( file.path(), - test_parser_version(3), + test_decoder_version(3), fingerprint.clone(), None, ); @@ -3279,19 +3306,19 @@ mod tests { assert!(!cache.dirty); assert!(cache.dirty_entries.is_empty()); - let shard = shard_path(file.path(), test_parser_version(3)).unwrap(); + let shard = shard_path(file.path(), test_decoder_version(3)).unwrap(); assert!(shard.exists()); let mut loaded = InputMessageCache::load().unwrap(); let meta = loaded - .get_meta(file.path(), test_parser_version(3)) + .get_meta(file.path(), test_decoder_version(3)) .unwrap() .unwrap(); assert_eq!(meta.fingerprint, fingerprint); let restored = loaded .take_messages(&CacheReadPlan::new( file.path(), - test_parser_version(3), + test_decoder_version(3), fingerprint, )) .unwrap(); @@ -3302,7 +3329,7 @@ mod tests { #[test] #[serial_test::serial] - fn test_explicit_prune_removes_orphans_and_old_parser_revisions() { + fn test_explicit_prune_removes_orphans_and_old_decoder_revisions() { let temp_home = TempDir::new().unwrap(); let prev_env = sandbox_cache_env(temp_home.path()); @@ -3337,9 +3364,10 @@ mod tests { )); } cache.save_if_dirty().unwrap(); - let stale_revision_shard = shard_path(live_input.path(), test_parser_version(1)).unwrap(); - let current_revision_shard = shard_path(live_input.path(), test_parser_version(3)).unwrap(); - let orphan_shard = shard_path(&orphan_path, test_parser_version(2)).unwrap(); + let stale_revision_shard = shard_path(live_input.path(), test_decoder_version(1)).unwrap(); + let current_revision_shard = + shard_path(live_input.path(), test_decoder_version(3)).unwrap(); + let orphan_shard = shard_path(&orphan_path, test_decoder_version(2)).unwrap(); assert!(stale_revision_shard.exists()); assert!(current_revision_shard.exists()); assert!(orphan_shard.exists()); @@ -3378,7 +3406,7 @@ mod tests { None, )); cache.save_if_dirty().unwrap(); - let orphan_shard = shard_path(&orphan_path, test_parser_version(1)).unwrap(); + let orphan_shard = shard_path(&orphan_path, test_decoder_version(1)).unwrap(); drop(orphan_input); let invalid_shard = cache_dir() @@ -3483,17 +3511,17 @@ mod tests { fn prune_classifier_accepts_current_envelope_and_rejects_unsupported_version() { let cache_home = TempDir::new().unwrap(); let input = write_temp_file(b"primary"); - let parser_version = test_parser_version(1); + let decoder_version = test_decoder_version(1); let mut cache = InputMessageCache::with_cache_dir(cache_home.path()); cache.insert(CachedInputEntry::new_with_version( input.path(), - parser_version, + decoder_version, InputFingerprint::from_path(input.path()).unwrap(), Vec::new(), None, )); cache.save_if_dirty().unwrap(); - let current_shard = shard_path_for_test(cache_home.path(), input.path(), parser_version); + let current_shard = shard_path_for_test(cache_home.path(), input.path(), decoder_version); assert!(read_shard_header_for_prune(¤t_shard) .unwrap() .is_some()); @@ -3544,7 +3572,7 @@ mod tests { None, )); cache.save_if_dirty().unwrap(); - let shard = shard_path(&path, test_parser_version(1)).unwrap(); + let shard = shard_path(&path, test_decoder_version(1)).unwrap(); assert!(shard.exists()); drop(input); @@ -3592,11 +3620,11 @@ mod tests { fn save_reports_invalidated_shard_removal_failure_with_path() { let cache_home = TempDir::new().unwrap(); let input = write_temp_file(b"primary"); - let parser_version = test_parser_version(31); - let shard_path = shard_path_for_test(cache_home.path(), input.path(), parser_version); + let decoder_version = test_decoder_version(31); + let shard_path = shard_path_for_test(cache_home.path(), input.path(), decoder_version); let mut cache = InputMessageCache::with_cache_dir(cache_home.path()); ensure_cache_dir(&shard_path).unwrap(); - cache.remove(input.path(), parser_version); + cache.remove(input.path(), decoder_version); let error = cache .save_if_dirty() @@ -3627,7 +3655,7 @@ mod tests { None, )); seed.save_if_dirty().unwrap(); - let shard = shard_path(input.path(), test_parser_version(1)).unwrap(); + let shard = shard_path(input.path(), test_decoder_version(1)).unwrap(); let file = std::fs::OpenOptions::new() .write(true) .open(&shard) @@ -3636,7 +3664,7 @@ mod tests { let loaded = InputMessageCache::load().unwrap(); let failure = loaded - .get_meta(input.path(), test_parser_version(1)) + .get_meta(input.path(), test_decoder_version(1)) .expect_err("oversized shard lookup must fail explicitly"); assert_eq!(failure.input_path, input.path()); assert_eq!(failure.shard_path, shard); @@ -3657,10 +3685,10 @@ mod tests { let input = write_temp_file(b"input\n"); let _initialized = InputMessageCache::load().unwrap(); - let shard = shard_path(input.path(), test_parser_version(1)).unwrap(); + let shard = shard_path(input.path(), test_decoder_version(1)).unwrap(); ensure_cache_dir(shard.parent().unwrap()).unwrap(); let header = CachedShardHeader { - parser_version: test_parser_version(1), + decoder_version: test_decoder_version(1), path: CachedPath::from_path(input.path()), fingerprint: InputFingerprint::from_path(input.path()).unwrap(), codex_incremental: None, @@ -3679,7 +3707,7 @@ mod tests { let loaded = InputMessageCache::load().unwrap(); assert!(loaded - .get_meta(input.path(), test_parser_version(1)) + .get_meta(input.path(), test_decoder_version(1)) .is_err()); assert!(shard.exists()); @@ -3693,12 +3721,12 @@ mod tests { let prev_env = sandbox_cache_env(temp_home.path()); let input = write_temp_file(b"input\n"); - let parser_version = test_parser_version(1); + let decoder_version = test_decoder_version(1); let fingerprint = InputFingerprint::from_path(input.path()).unwrap(); let mut seed = InputMessageCache::load().unwrap(); seed.insert(CachedInputEntry::new_with_version( input.path(), - parser_version, + decoder_version, fingerprint.clone(), vec![ParsedMessage::new( "gpt-5", @@ -3711,7 +3739,7 @@ mod tests { None, )); seed.save_if_dirty().unwrap(); - let shard = shard_path(input.path(), parser_version).unwrap(); + let shard = shard_path(input.path(), decoder_version).unwrap(); let mut file = std::fs::OpenOptions::new() .write(true) .open(&shard) @@ -3724,7 +3752,7 @@ mod tests { let original_bytes = std::fs::read(&shard).unwrap(); let mut loaded = InputMessageCache::load().unwrap(); - assert!(loaded.get_meta(input.path(), parser_version).is_err()); + assert!(loaded.get_meta(input.path(), decoder_version).is_err()); assert_eq!( std::fs::read(&shard).unwrap(), original_bytes, @@ -3741,7 +3769,7 @@ mod tests { )]; assert!(loaded .write_messages( - CacheWritePlan::new(input.path(), parser_version, fingerprint.clone(), None,), + CacheWritePlan::new(input.path(), decoder_version, fingerprint.clone(), None,), &replacement, ) .is_ok()); @@ -3752,13 +3780,13 @@ mod tests { ); let mut warm = InputMessageCache::load().unwrap(); let meta = warm - .get_meta(input.path(), parser_version) + .get_meta(input.path(), decoder_version) .expect("successful atomic replacement must read without error") .expect("successful atomic replacement must produce a current-format hit"); let messages = warm .take_messages(&CacheReadPlan::new( input.path(), - parser_version, + decoder_version, meta.fingerprint, )) .unwrap(); @@ -3775,10 +3803,10 @@ mod tests { let input = write_temp_file(b"input\n"); let _initialized = InputMessageCache::load().unwrap(); - for (parser_version, bytes) in [ - (test_parser_version(11), b"raw-blob".to_vec()), + for (decoder_version, bytes) in [ + (test_decoder_version(11), b"raw-blob".to_vec()), ( - test_parser_version(12), + test_decoder_version(12), [ SHARD_MAGIC.as_slice(), CACHE_FORMAT_VERSION.to_le_bytes().as_slice(), @@ -3788,11 +3816,11 @@ mod tests { .concat(), ), ] { - let shard = shard_path(input.path(), parser_version).unwrap(); + let shard = shard_path(input.path(), decoder_version).unwrap(); ensure_cache_dir(shard.parent().unwrap()).unwrap(); std::fs::write(&shard, &bytes).unwrap(); let cache = InputMessageCache::load().unwrap(); - assert!(cache.get_meta(input.path(), parser_version).is_err()); + assert!(cache.get_meta(input.path(), decoder_version).is_err()); assert_eq!(std::fs::read(&shard).unwrap(), bytes); } @@ -3805,21 +3833,21 @@ mod tests { let temp_home = TempDir::new().unwrap(); let prev_env = sandbox_cache_env(temp_home.path()); let input = write_temp_file(b"input\n"); - let parser_version = test_parser_version(13); + let decoder_version = test_decoder_version(13); let _initialized = InputMessageCache::load().unwrap(); - let shard = shard_path(input.path(), parser_version).unwrap(); + let shard = shard_path(input.path(), decoder_version).unwrap(); ensure_cache_dir(shard.parent().unwrap()).unwrap(); let unknown_bytes = b"unknown!"; std::fs::write(&shard, unknown_bytes).unwrap(); let fingerprint = InputFingerprint::from_path(input.path()).unwrap(); let mut cache = InputMessageCache::load().unwrap(); - assert!(cache.get_meta(input.path(), parser_version).is_err()); + assert!(cache.get_meta(input.path(), decoder_version).is_err()); let real_cache_dir = cache.cache_dir.clone(); cache.cache_dir = PathBuf::from(OsString::from("invalid\0cache-dir")); let error = cache .write_messages( - CacheWritePlan::new(input.path(), parser_version, fingerprint, None), + CacheWritePlan::new(input.path(), decoder_version, fingerprint, None), &[ParsedMessage::new( "gpt-5", "provider", @@ -3847,7 +3875,7 @@ mod tests { fn cache_lookup_error_retains_path_version_and_decode_root_cause() { let failure = CacheLookupFailure { input_path: PathBuf::from("/test/input"), - parser_version: test_parser_version(7), + decoder_version: test_decoder_version(7), shard_path: PathBuf::from("/test/shard"), reason: CacheReadFailureReason::HeaderDecode { source: Box::new(bincode::ErrorKind::Custom("bad header".to_string())), @@ -3870,7 +3898,7 @@ mod tests { #[test] #[serial_test::serial] - fn test_get_meta_ignores_stale_parser_revision() { + fn test_get_meta_ignores_stale_decoder_revision() { let temp_home = TempDir::new().unwrap(); let prev_env = sandbox_cache_env(temp_home.path()); @@ -3901,11 +3929,11 @@ mod tests { let loaded = InputMessageCache::load().unwrap(); assert!(loaded - .get_meta(input.path(), test_parser_version(7)) + .get_meta(input.path(), test_decoder_version(7)) .unwrap() .is_some()); assert!(loaded - .get_meta(input.path(), test_parser_version(8)) + .get_meta(input.path(), test_decoder_version(8)) .unwrap() .is_none()); @@ -3914,7 +3942,7 @@ mod tests { #[test] #[serial_test::serial] - fn test_get_meta_ignores_stale_parser_id() { + fn test_get_meta_ignores_stale_decoder_id() { let temp_home = TempDir::new().unwrap(); let prev_env = sandbox_cache_env(temp_home.path()); @@ -3923,7 +3951,7 @@ mod tests { let mut cache = InputMessageCache::load().unwrap(); cache.insert(CachedInputEntry::new_with_version( input.path(), - ParserVersion::new(ParserId::Copilot, 1), + DecoderVersion::new(DecoderId::Copilot, 1), fingerprint, vec![ParsedMessage::new( "gpt-5", @@ -3945,11 +3973,11 @@ mod tests { let loaded = InputMessageCache::load().unwrap(); assert!(loaded - .get_meta(input.path(), ParserVersion::new(ParserId::Copilot, 1)) + .get_meta(input.path(), DecoderVersion::new(DecoderId::Copilot, 1)) .unwrap() .is_some()); assert!(loaded - .get_meta(input.path(), ParserVersion::new(ParserId::Gemini, 1)) + .get_meta(input.path(), DecoderVersion::new(DecoderId::Gemini, 1)) .unwrap() .is_none()); @@ -4013,17 +4041,17 @@ mod tests { let loaded = InputMessageCache::load().unwrap(); assert!(loaded - .get_meta(file_one.path(), test_parser_version(1)) + .get_meta(file_one.path(), test_decoder_version(1)) .unwrap() .is_some()); assert!(loaded - .get_meta(file_two.path(), test_parser_version(1)) + .get_meta(file_two.path(), test_decoder_version(1)) .unwrap() .is_some()); - assert!(shard_path(file_one.path(), test_parser_version(1)) + assert!(shard_path(file_one.path(), test_decoder_version(1)) .unwrap() .exists()); - assert!(shard_path(file_two.path(), test_parser_version(1)) + assert!(shard_path(file_two.path(), test_decoder_version(1)) .unwrap() .exists()); } @@ -4033,14 +4061,14 @@ mod tests { #[test] #[serial_test::serial] - fn test_same_path_different_parser_versions_use_distinct_shards() { + fn test_same_path_different_decoder_versions_use_distinct_shards() { let temp_home = TempDir::new().unwrap(); let prev_env = sandbox_cache_env(temp_home.path()); let input = write_temp_file(b"input\n"); let fingerprint = InputFingerprint::from_path(input.path()).unwrap(); - let copilot_version = ParserVersion::new(ParserId::Copilot, 1); - let gemini_version = ParserVersion::new(ParserId::Gemini, 1); + let copilot_version = DecoderVersion::new(DecoderId::Copilot, 1); + let gemini_version = DecoderVersion::new(DecoderId::Gemini, 1); let mut cache = InputMessageCache::load().unwrap(); cache.insert(CachedInputEntry::new_with_version( input.path(), @@ -4126,12 +4154,12 @@ mod tests { let prev_env = sandbox_cache_env(temp_home.path()); let input = write_temp_file(b"input-one\n"); - let parser_version = ParserVersion::new(ParserId::Copilot, 1); + let decoder_version = DecoderVersion::new(DecoderId::Copilot, 1); let initial_fingerprint = InputFingerprint::from_path(input.path()).unwrap(); let mut seed = InputMessageCache::load().unwrap(); seed.insert(CachedInputEntry::new_with_version( input.path(), - parser_version, + decoder_version, initial_fingerprint.clone(), vec![ParsedMessage::new( "gpt-5", @@ -4153,17 +4181,17 @@ mod tests { let mut reader = InputMessageCache::load().unwrap(); let meta = reader - .get_meta(input.path(), parser_version) + .get_meta(input.path(), decoder_version) .unwrap() .unwrap(); - let read_plan = CacheReadPlan::new(input.path(), parser_version, meta.fingerprint); + let read_plan = CacheReadPlan::new(input.path(), decoder_version, meta.fingerprint); std::fs::write(input.path(), b"input-two\n").unwrap(); let replacement_fingerprint = InputFingerprint::from_path(input.path()).unwrap(); let mut writer = InputMessageCache::load().unwrap(); writer.insert(CachedInputEntry::new_with_version( input.path(), - parser_version, + decoder_version, replacement_fingerprint, vec![ParsedMessage::new( "gpt-5", @@ -4196,7 +4224,7 @@ mod tests { let replacement_messages = reader .take_messages(&CacheReadPlan::new( input.path(), - parser_version, + decoder_version, InputFingerprint::from_path(input.path()).unwrap(), )) .expect("failed stale read plan must not poison the input key"); diff --git a/crates/tokscale-core/src/sessions/mod.rs b/crates/tokscale-core/src/sessions/mod.rs index 3b9d1ba28..b50f5e58e 100644 --- a/crates/tokscale-core/src/sessions/mod.rs +++ b/crates/tokscale-core/src/sessions/mod.rs @@ -34,60 +34,80 @@ pub mod warp; pub mod zcode; pub mod zed; +use std::ops::{Deref, DerefMut}; + use crate::{clients::ClientId, TokenBreakdown}; -macro_rules! define_message { - ($(#[$meta:meta])* $name:ident, $($source_field:tt)*) => { - $(#[$meta])* - #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] - pub struct $name { - $($source_field)* - #[serde(deserialize_with = "intern::de_intern")] - pub model_id: std::sync::Arc, - #[serde(deserialize_with = "intern::de_intern")] - pub provider_id: std::sync::Arc, - #[serde(deserialize_with = "intern::de_intern")] - pub session_id: std::sync::Arc, - /// Whether the input record directly identifies this usage as belonging to a - /// top-level session. `session_id` remains the legacy Total grouping key. - #[serde(default = "default_true")] - pub is_main_session: bool, - #[serde(default, deserialize_with = "intern::de_intern_opt")] - pub workspace_key: Option>, - #[serde(default, deserialize_with = "intern::de_intern_opt")] - pub workspace_label: Option>, - pub timestamp: i64, - pub tokens: TokenBreakdown, - pub cost: f64, - #[serde(default = "default_message_count")] - pub message_count: i32, - #[serde(default, deserialize_with = "intern::de_intern_opt")] - pub agent: Option>, - #[serde(default, deserialize_with = "intern::de_intern_opt")] - pub agent_instance: Option>, - pub dedup_key: Option, - /// True if this message is the first assistant response after a user turn. - /// Used to count user interaction turns (as opposed to API message count). - #[serde(default)] - pub is_turn_start: bool, - } - }; +/// Source-neutral usage payload produced by session parsers and persisted in +/// the input-message cache. +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct UsageRecord { + #[serde(deserialize_with = "intern::de_intern")] + pub model_id: std::sync::Arc, + #[serde(deserialize_with = "intern::de_intern")] + pub provider_id: std::sync::Arc, + #[serde(deserialize_with = "intern::de_intern")] + pub session_id: std::sync::Arc, + /// Whether the input record directly identifies this usage as belonging to a + /// top-level session. `session_id` remains the legacy Total grouping key. + #[serde(default = "default_true")] + pub is_main_session: bool, + #[serde(default, deserialize_with = "intern::de_intern_opt")] + pub workspace_key: Option>, + #[serde(default, deserialize_with = "intern::de_intern_opt")] + pub workspace_label: Option>, + pub timestamp: i64, + pub tokens: TokenBreakdown, + pub cost: f64, + #[serde(default = "default_message_count")] + pub message_count: i32, + #[serde(default, deserialize_with = "intern::de_intern_opt")] + pub agent: Option>, + #[serde(default, deserialize_with = "intern::de_intern_opt")] + pub agent_instance: Option>, + pub dedup_key: Option, + /// True if this message is the first assistant response after a user turn. + /// Used to count user interaction turns (as opposed to API message count). + #[serde(default)] + pub is_turn_start: bool, } -define_message!( - /// Source-neutral usage produced by a session parser. - /// - /// The local-input adapter owns source identity and attributes it only after - /// parsing, cache persistence, and adapter-specific filtering have completed. - ParsedMessage, -); - -define_message!( - /// Usage attributed to a concrete client and ready for public aggregation. - UnifiedMessage, - #[serde(deserialize_with = "intern::de_intern")] - pub client: std::sync::Arc, -); +/// Parser/cache name for the source-neutral usage payload. +pub type ParsedMessage = UsageRecord; + +/// Usage attributed to a concrete local client and ready for aggregation. +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct UnifiedMessage { + pub client: ClientId, + #[serde(flatten)] + pub(crate) usage: UsageRecord, +} + +impl Deref for UnifiedMessage { + type Target = UsageRecord; + + fn deref(&self) -> &Self::Target { + &self.usage + } +} + +impl DerefMut for UnifiedMessage { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.usage + } +} + +impl AsMut for UsageRecord { + fn as_mut(&mut self) -> &mut UsageRecord { + self + } +} + +impl AsMut for UnifiedMessage { + fn as_mut(&mut self) -> &mut UsageRecord { + &mut self.usage + } +} const fn default_message_count() -> i32 { 1 @@ -260,7 +280,7 @@ fn titlecase_agent(name: &str) -> String { .join(" ") } -impl ParsedMessage { +impl UsageRecord { pub fn new( model_id: impl AsRef, provider_id: impl AsRef, @@ -355,49 +375,16 @@ impl ParsedMessage { } pub(crate) fn attribute(self, client: ClientId) -> UnifiedMessage { - self.with_client(client.as_str()) - } - - fn with_client(self, client: impl AsRef) -> UnifiedMessage { - let Self { - model_id, - provider_id, - session_id, - is_main_session, - workspace_key, - workspace_label, - timestamp, - tokens, - cost, - message_count, - agent, - agent_instance, - dedup_key, - is_turn_start, - } = self; UnifiedMessage { - client: intern::intern(client.as_ref()), - model_id, - provider_id, - session_id, - is_main_session, - workspace_key, - workspace_label, - timestamp, - tokens, - cost, - message_count, - agent, - agent_instance, - dedup_key, - is_turn_start, + client, + usage: self, } } } impl UnifiedMessage { pub fn new( - client: impl AsRef, + client: ClientId, model_id: impl AsRef, provider_id: impl AsRef, session_id: impl AsRef, @@ -405,13 +392,13 @@ impl UnifiedMessage { tokens: TokenBreakdown, cost: f64, ) -> Self { - ParsedMessage::new(model_id, provider_id, session_id, timestamp, tokens, cost) - .with_client(client) + UsageRecord::new(model_id, provider_id, session_id, timestamp, tokens, cost) + .attribute(client) } #[allow(clippy::too_many_arguments)] pub fn new_with_agent( - client: impl AsRef, + client: ClientId, model_id: impl AsRef, provider_id: impl AsRef, session_id: impl AsRef, @@ -420,7 +407,7 @@ impl UnifiedMessage { cost: f64, agent: Option, ) -> Self { - ParsedMessage::new_with_agent( + UsageRecord::new_with_agent( model_id, provider_id, session_id, @@ -429,12 +416,12 @@ impl UnifiedMessage { cost, agent, ) - .with_client(client) + .attribute(client) } #[allow(clippy::too_many_arguments)] pub fn new_with_dedup( - client: impl AsRef, + client: ClientId, model_id: impl AsRef, provider_id: impl AsRef, session_id: impl AsRef, @@ -443,7 +430,7 @@ impl UnifiedMessage { cost: f64, dedup_key: Option, ) -> Self { - ParsedMessage::new_with_dedup( + UsageRecord::new_with_dedup( model_id, provider_id, session_id, @@ -452,61 +439,50 @@ impl UnifiedMessage { cost, dedup_key, ) - .with_client(client) + .attribute(client) } } -macro_rules! impl_message_methods { - ($name:ident) => { - impl $name { - /// Local calendar date derived from `timestamp`; `None` when the - /// timestamp is outside the representable range. - pub fn local_date(&self) -> Option { - use chrono::TimeZone; - match chrono::Local.timestamp_millis_opt(self.timestamp) { - chrono::LocalResult::Single(dt) => Some(dt.date_naive()), - _ => None, - } - } +impl UsageRecord { + /// Local calendar date derived from `timestamp`; `None` when the + /// timestamp is outside the representable range. + pub fn local_date(&self) -> Option { + use chrono::TimeZone; + match chrono::Local.timestamp_millis_opt(self.timestamp) { + chrono::LocalResult::Single(dt) => Some(dt.date_naive()), + _ => None, + } + } - /// Local `YYYY-MM-DD` string derived from `timestamp` (empty when out of - /// range). Allocates; prefer [`Self::local_date`] in hot paths. - pub fn date_string(&self) -> String { - timestamp_to_date(self.timestamp) - } + /// Local `YYYY-MM-DD` string derived from `timestamp` (empty when out of + /// range). Allocates; prefer [`Self::local_date`] in hot paths. + pub fn date_string(&self) -> String { + timestamp_to_date(self.timestamp) + } - pub fn set_workspace( - &mut self, - workspace_key: Option, - workspace_label: Option, - ) { - self.workspace_key = workspace_key.as_deref().map(intern::intern); - self.workspace_label = workspace_label.as_deref().map(intern::intern); - } + pub fn set_workspace( + &mut self, + workspace_key: Option, + workspace_label: Option, + ) { + self.workspace_key = workspace_key.as_deref().map(intern::intern); + self.workspace_label = workspace_label.as_deref().map(intern::intern); + } - pub fn set_agent_instance(&mut self, agent_instance: Option) { - self.agent_instance = agent_instance.as_deref().map(intern::intern); - } + pub fn set_agent_instance(&mut self, agent_instance: Option) { + self.agent_instance = agent_instance.as_deref().map(intern::intern); + } - #[cfg_attr(not(test), allow(dead_code))] - pub(crate) fn refresh_derived_fields(&mut self) { - if let Some(provider) = - crate::provider_identity::provider_override_from_model_and_provider( - &self.model_id, - &self.provider_id, - ) - { - self.provider_id = intern::intern(provider); - } - } + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn refresh_derived_fields(&mut self) { + if let Some(provider) = crate::provider_identity::provider_override_from_model_and_provider( + &self.model_id, + &self.provider_id, + ) { + self.provider_id = intern::intern(provider); } - }; -} - -impl_message_methods!(ParsedMessage); -impl_message_methods!(UnifiedMessage); + } -impl ParsedMessage { pub(crate) fn set_timestamp(&mut self, timestamp: i64) { self.timestamp = timestamp; self.refresh_derived_fields(); @@ -629,7 +605,7 @@ mod tests { }; let msg = UnifiedMessage::new( - "opencode", + ClientId::OpenCode, "claude-sonnet-4.6", "anthropic", "test-session-id", @@ -638,7 +614,7 @@ mod tests { 0.05, ); - assert_eq!(msg.client.as_ref(), "opencode"); + assert_eq!(msg.client, ClientId::OpenCode); assert_eq!(msg.model_id.as_ref(), "claude-sonnet-4.6"); assert_eq!(msg.session_id.as_ref(), "test-session-id"); assert_eq!(msg.date_string(), timestamp_to_date(1733011200000)); @@ -646,6 +622,15 @@ mod tests { assert_eq!(msg.agent, None); assert_eq!(msg.workspace_key, None); assert_eq!(msg.workspace_label, None); + + let serialized = serde_json::to_value(&msg).unwrap(); + assert_eq!(serialized["client"], "opencode"); + assert_eq!(serialized["model_id"], "claude-sonnet-4.6"); + assert!(serialized.get("usage").is_none()); + assert_eq!( + serde_json::from_value::(serialized).unwrap(), + msg + ); } #[test] @@ -667,7 +652,7 @@ mod tests { assert!(serialized.get("client").is_none()); let attributed = parsed.attribute(ClientId::Codex); - assert_eq!(attributed.client.as_ref(), ClientId::Codex.as_str()); + assert_eq!(attributed.client, ClientId::Codex); assert_eq!(attributed.model_id.as_ref(), "gpt-5.5"); } @@ -675,7 +660,7 @@ mod tests { fn test_refresh_derived_fields_overrides_historical_model_provider() { let tokens = TokenBreakdown::default(); let mut msg = UnifiedMessage::new( - "pi", + ClientId::Pi, "model1", "pandora-deepseek", "test-session-id", @@ -689,7 +674,7 @@ mod tests { assert_eq!(msg.provider_id.as_ref(), "deepseek"); let mut non_alias = UnifiedMessage::new( - "pi", + ClientId::Pi, "model10", "pandora-deepseek", "test-session-id", @@ -703,7 +688,7 @@ mod tests { assert_eq!(non_alias.provider_id.as_ref(), "pandora-deepseek"); let mut non_claude = UnifiedMessage::new( - "droid", + ClientId::Droid, "glm-5.1", "anthropic", "test-session-id", diff --git a/docs/adr/0007-client-identity-catalog.md b/docs/adr/0007-client-identity-catalog.md index 854a715cc..ed3d5a33b 100644 --- a/docs/adr/0007-client-identity-catalog.md +++ b/docs/adr/0007-client-identity-catalog.md @@ -40,26 +40,33 @@ model names, providers, or display labels. ### Local-input authority -Every catalog entry has exactly one registered local-input adapter and one scan +Every catalog entry has exactly one registry `AdapterBinding` and one scan definition. These three sets must have exact parity and no duplicate entries. -The scan definition and the adapter's discovery implementation are the sole -authority for fixed default roots beneath the selected home, filename -selection, companion files, database sidecars, custom-root support, and source -identity. An adapter stamps its generated `InputUnit` with its catalog -`ClientId`; that value is the only source identity carried through acquisition. -Additional roots come only from `scanner.extraScanPaths`; OpenCode database -files come only from `scanner.opencodeDbPaths`. +The adapter registry stores an `AdapterBinding` containing one typed `ClientId` +and one adapter. This binding is the sole authority associating discovered +inputs and emitted usage with a client. The scan definition and adapter +discovery implementation remain authoritative for fixed default roots beneath +the selected home, filename selection, companion files, database sidecars, and +custom-root support. Additional roots come only from +`scanner.extraScanPaths`; OpenCode database files come only from +`scanner.opencodeDbPaths`. -The adapter's session schema/parser is the sole authority for accepted +The adapter's session decoder and schema are the sole authority for accepted envelopes, database schemas, required fields, record semantics, deduplication, -and token interpretation, but not source identity. Parsers produce -source-neutral `ParsedMessage` values. Derived message-cache shards also store -only source-neutral messages. The sequential adapter fold attributes each -message from `InputUnit.client` when it emits the final `UnifiedMessage`; parser -modules do not repeat catalog IDs as message data. Before any adapter executes, -the pipeline verifies that every discovered unit carries that adapter's -`ClientId` and rejects a mismatched batch as a contract error. +and token interpretation, but not source identity. Each source-neutral +`InputUnit` contains a `DecoderSpec` that atomically binds decoder ID, semantic +revision, and decode route. There is no client-to-decoder default mapping that +can drift from the selected route. + +Decoders produce one source-neutral `UsageRecord` representation. Derived +message-cache shards store that same type. The runner constructs a +`BoundMessageSink` from the registry binding; the sink accepts only +`UsageRecord` and attaches the binding's typed `ClientId` to construct a +`UnifiedMessage`. `UnifiedMessage` composes client attribution with a +`UsageRecord` instead of duplicating the usage fields. Neither an input, +decoder, adapter fold, nor cache shard can supply a competing client identity, +so mismatched source attribution is not representable. `docs/clients.md` is the user-facing discovery map generated from that contract. It does not create a second path or schema authority. @@ -74,7 +81,7 @@ An absent automatically discovered root means that the client has no input at that location. Once a root or configured input exists, discovery, open, query, snapshot, and parse failures remain visible through Data Health. Record-level schema failures reject the affected records and preserve valid records from the -same input when the parser can continue. A failure for one client does not abort +same input when the decoder can continue. A failure for one client does not abort unrelated clients. Accepted usage records preserve their observed token values and canonical model @@ -102,22 +109,24 @@ Adding a client requires one atomic contract change containing: 1. one catalog identity; 2. one scan definition; -3. one registered local-input adapter; -4. one current source-neutral session schema/parser; -5. focused discovery, parser, and health tests; +3. one registered `AdapterBinding`; +4. one current source-neutral decoder specification and session schema; +5. focused discovery, decoder, and health tests; 6. catalog/scan/adapter parity checks; and 7. a current discovery row in `docs/clients.md`. Changing a root, filename rule, companion dependency, database schema, or record envelope requires an adapter/schema change with focused tests and a matching -documentation update. A parser behavior change that affects cached output also -requires a parser revision change. +documentation update. A decoder behavior change that affects cached output also +requires a decoder revision change. ## Consequences Each accepted local integration has one public identity and one executable input contract. Reports, filters, scanner configuration, caches, and TUI views -therefore share the same client namespace. A parser cannot silently create a -second source namespace because its output and persisted cache representation -have no client field; path and format evolution stays owned by the adapter and -schema that can validate it. +therefore share the same client namespace. The production pipeline gives only +the registry binding authority to attach `ClientId`; the source-neutral input, +decoder output, persisted shard, and diagnostic errors carry no competing +client field. Decoder identity, revision, and route cannot disagree because +they are one `DecoderSpec`; path and format evolution stays owned by the adapter +and schema that can validate it. diff --git a/docs/adr/0008-single-copy-memory-pipeline.md b/docs/adr/0008-single-copy-memory-pipeline.md index 73b0815af..6c1355109 100644 --- a/docs/adr/0008-single-copy-memory-pipeline.md +++ b/docs/adr/0008-single-copy-memory-pipeline.md @@ -19,8 +19,9 @@ current formats, and ADR 0028 owns the TUI lifecycle built on this pipeline. ### Acquisition authority -Selected client adapters own input discovery, input identity, parsing, and -error attribution. Public report paths do not use a central `ScanResult`, +Selected registry `AdapterBinding` values own client attribution, while their +adapters own input discovery, input identity, parsing, and error attribution. +Public report paths do not use a central `ScanResult`, `scan_all_clients*`, generic scanner error, or dead per-client database slots. `ScannerSettings` and focused scanner primitives remain adapter and test seams; they do not define another product command or discovery authority. @@ -29,17 +30,20 @@ Discovery produces one consumptive `PreparedLocalInputs` inventory. It records: - requested clients in canonical order; - selected-adapter and per-adapter unit order; -- parser and unit identity; -- one `InputPolicy` per unit, including its primary and parser-relevant related +- decoder and unit identity; +- one `InputPolicy` per unit, including its primary and decoder-relevant related inputs; and - compact pre-execution metadata snapshots. -Each `InputUnit.client` is adapter-issued provenance from ADR 0007. Session -parsers produce `ParsedMessage`, which intentionally has no source identity, and -message-cache bodies persist that same source-neutral type. After cache writes, -adapter-specific enrichment, filtering, and deduplication, the sequential fold -combines each parsed message with the unit's `ClientId` and emits a -`UnifiedMessage`. No parser or cache shard can override that attribution. +`InputUnit` is source-neutral. Its `DecoderSpec` atomically binds decoder ID, +semantic revision, and decode route; decoder selection is never reconstructed +from a client default. Session decoders produce `UsageRecord`, and message-cache +bodies persist that same source-neutral type. Each adapter's fold owns its cache +and enrichment order, then applies filtering and deduplication before sending +accepted records through the binding's `BoundMessageSink`. The sink alone +combines its typed `ClientId` with the record and emits a `UnifiedMessage`. No +production pipeline path lets an input, decoder, adapter fold, or cache shard +override that attribution. Freshness probing and execution consume the same inventory and never rediscover inputs. A file added after preparation belongs to the next inventory. Related @@ -71,7 +75,7 @@ matches the prepared snapshot. Each inventory has two freshness keys: - a versioned SHA-256 `InputInventorySignature` over canonical clients, adapter - and unit order, parser/unit identity, and every declared input's native path, + and unit order, decoder/unit identity, and every declared input's native path, label, presence, size, mtime, and native identity; and - a process-local `u64` digest over those stable signature bytes, which is never persisted. @@ -83,7 +87,7 @@ fresh TUI generation establishes the initial process digest from its persisted signature without startup discovery; stale or missing generations prepare and execute in the background. -Codex computes its full digest during the parser's read. Append handling verifies +Codex computes its full digest during the decoder's read. Append handling verifies the previous full digest as the expected prefix, then continues the same hasher across the tail. Exact hits require stamp and cached-digest consistency without another input-byte pass. @@ -119,12 +123,12 @@ hit bodies before that index is built, and then folds hits and misses in bounded order. Codex retains its dedicated exact-hit, stale, append, and recovery path because its incremental parse state is not the generic adapter contract. -Codex cold parses, append merges, and race reparses own one raw -`ParsedMessage` vector. +Codex cold parses, append merges, and race reparses own one raw `UsageRecord` +vector. When cacheable, the fold serializes a borrowed raw slice before applying timestamp completion, token normalization/filtering, canonical identity, pricing, and exec-session attribution in place. Client attribution occurs only -after those operations when the adapter emits the messages. +after those operations when the bound sink emits the messages. OpenCode borrows potentially large message TEXT. It stream-validates the full JSON document and required role envelope, then fully decodes assistant payloads @@ -133,14 +137,16 @@ instead of becoming empty usage. ### Message representation and aggregation -`ParsedMessage` is the source-neutral parser and shard representation. -`UnifiedMessage` adds the adapter-attributed client for public aggregation and -otherwise stores no redundant derivable value: +`UsageRecord` is the single source-neutral decoder, shard, and enrichment +representation. `UnifiedMessage` composes a typed adapter-attributed `ClientId` +with one `UsageRecord` for public aggregation; it does not repeat the record's +fields or accept an untyped client string. The record otherwise stores no +redundant derivable value: - date is derived from timestamp; - `dedup_key` is a 64-bit hash rather than a formatted string; and -- repeated client, model, provider, session, workspace, and agent identities - use interned `Arc`. +- repeated model, provider, session, workspace, and agent identities use + interned `Arc`. The process interner indexes `Weak`, confirms hash matches with full string equality, and is swept after transient messages and Arc-backed accumulators are @@ -154,16 +160,15 @@ buckets. Composite identities never use delimiter-concatenated public keys. Persisted map keys are versioned, variant-tagged, byte-length-prefixed values with a distinct unknown-workspace tag. -Session client-space accounting is the byte size of the input snapshots -confirmed at the final cache-decision/fold boundary, deduplicated within each -client inventory. Overview Data Size is exactly the checked sum of that -per-client map; it does not apply a second cross-client physical-file -deduplication rule. The TUI generation schema changes with this accounting -contract so an older generation cannot preserve the previous total alongside -current per-client values. Current-schema cache writes and reads both reject a -Data Size that differs from the checked sum of `clientSpace`. Usage, Sessions, -Data Health, input space, and the inventory signature all derive from those same -confirmed snapshots. +`InputFootprint` is the sole input-space fact. It maps typed `ClientId` values to +the byte size of snapshots confirmed at the final cache-decision/fold boundary, +deduplicated within each binding's inventory. Overview derives Data Size as its +checked sum; there is no separately stored total. Data Health does not store +input bytes. The TUI generation persists the footprint map (`clientSpace`) once, +and Sessions and Overview project their values from it. Headless local reports +expose the same confirmed map as `metadata.inputFootprint`. Usage, Sessions, +Data Health, input footprint, and the inventory signature all derive from the +same confirmed snapshots. ### Shard contract and recovery @@ -172,7 +177,7 @@ and body. Header discovery does not materialize the body. A planned hit succeeds only after the body is opened, its identity is checked, it decodes, and its message count matches the header. -Body failures retain the input, parser revision, shard path, and root cause. +Body failures retain the input, decoder revision, shard path, and root cause. The CLI emits an explicit diagnostic and reparses the authoritative current input through its registered adapter. A successful cacheable reparse atomically replaces the shard. @@ -184,9 +189,9 @@ bypassed, but a potentially valid replacement shard remains unless the reparse independently proves it invalid or non-cacheable. Partial and unavailable input never publishes a shard. -Cache writes serialize borrowed source-neutral `ParsedMessage` slices and do +Cache writes serialize borrowed source-neutral `UsageRecord` slices and do not clone messages merely to construct a cache representation. -Parser-semantic changes bump the owning parser revision; +Decoder-semantic changes bump the owning decoder revision; serialization-layout changes bump the shard format. The message-shard envelope is the `TOKSHRD\0` magic, a little-endian format @@ -201,16 +206,17 @@ truncated envelope, malformed header, undecodable header, oversized shard, or filesystem inspection failure aborts the operation before deletion begins. After successful classification, pruning removes a shard only when its authoritative input is absent, its path is not the canonical path -derived from the input and parser key, or a higher parser revision exists for -the same live input and parser. Once deletion begins, an unlink failure is +derived from the input and decoder key, or a higher decoder revision exists for +the same live input and decoder. Once deletion begins, an unlink failure is reported explicitly; already completed removals are not rolled back. ### Atomic TUI generation storage -One local fold produces health data, client-space accounting, sessions, +One local fold produces health data, one `InputFootprint`, sessions, client-aware canonical accumulator, one group-agnostic Common projection, and four full-universe Grouped projections. The current TUI cache schema stores them -in one atomic JSON bundle. +in one atomic JSON bundle. It stores only the per-client footprint map for input +space; Overview derives its total and Health carries no duplicate byte count. Common stores Agents, daily and hourly totals with Client membership, the contribution graph, report totals, and streaks exactly once. Each Grouped @@ -256,5 +262,5 @@ indexes, and the consumer's required aggregate rather than an adapter-wide message collection. Exact warm hits read no authoritative input bytes, while native file identity prevents same-size/same-mtime path replacement from reusing stale data. Cache faults remain visible and recover from authoritative -input, and explicit pruning deletes shards outside the current input and parser +input, and explicit pruning deletes shards outside the current input and decoder inventory. diff --git a/docs/adr/0010-period-views-derive-from-daily.md b/docs/adr/0010-period-views-derive-from-daily.md index 583b07a04..eb0cf46ac 100644 --- a/docs/adr/0010-period-views-derive-from-daily.md +++ b/docs/adr/0010-period-views-derive-from-daily.md @@ -62,7 +62,7 @@ reinterpret runtime labels. Agent aggregation uses the structured `(client, agent)` identity and every public Agent entry carries exactly one Client, so equal labels from different Clients never merge. Identity-semantic changes invalidate affected message shards and the persisted TUI generation -through the appropriate parser revision and schema/version change. +through the appropriate decoder revision and schema/version change. ### Group By projections diff --git a/docs/cli.md b/docs/cli.md index 6d865f55d..9d149ce8f 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -110,15 +110,21 @@ All local Models JSON uses this top-level envelope: "rejectedRecords": 0, "partialInputs": 0, "failedInputs": 0, - "inputDataBytes": 0, "issues": [] }, "metadata": { + "inputFootprint": { + "codex": 0 + }, "processingTimeMs": 0 } } ``` +`metadata.inputFootprint` is the confirmed byte count keyed by canonical Client +ID. Its checked sum is the report's Data Size; the report does not persist a +second total. + Stdout contains only the table or JSON document. Progress, `--benchmark` timing, Data Health summaries, warnings, and errors go to stderr. A degraded report still exits `0` when its payload was produced; inspect `health` when @@ -200,7 +206,7 @@ tokscale cache prune Models never writes that aggregate cache. Scan-input message shards remain an internal derived cache and are written while parsing. -`cache prune` removes orphaned Inputs and superseded parser revisions. +`cache prune` removes orphaned Inputs and superseded decoder revisions. Unreadable or unclassifiable shards make the explicit maintenance command fail instead of reporting partial success. diff --git a/docs/configuration.md b/docs/configuration.md index e7219abae..c802659a0 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -92,11 +92,11 @@ want a fresh local rebuild: - `images/` Scan-input message cache writes use the current shard envelope and stable -explicit parser keys. Ordinary reports and `tokscale cache prune` accept only +explicit decoder keys. Ordinary reports and `tokscale cache prune` accept only shards in the format supported by the running binary. Pruning explicitly traverses the shard directory and removes current shards whose authoritative -input is absent, whose path is not canonical for the input and parser key, or -whose parser revision has been superseded. +input is absent, whose path is not canonical for the input and decoder key, or +whose decoder revision has been superseded. Traversal and classification complete before deletion; an unknown, future, truncated, malformed, undecodable, or oversized shard aborts pruning without deleting anything.