diff --git a/src/daemon/token_usage_worker.rs b/src/daemon/token_usage_worker.rs index f2a527515c..84673bb80a 100644 --- a/src/daemon/token_usage_worker.rs +++ b/src/daemon/token_usage_worker.rs @@ -55,7 +55,7 @@ use crate::error::GitAiError; use crate::metrics::{EventAttributes, MetricEvent, PosEncoded, TokenUsageValues}; use crate::streams::db::{StreamRecord, StreamsDatabase}; use crate::token_usage::db::{BatchCommit, TokenUsageDatabase, TrackedFile}; -use crate::token_usage::extractor_for_tool; +use crate::token_usage::{Speed, extractor_for_tool}; /// One raw JSONL line read as bytes: unlike UTF-8-strict `read_line`, a /// single invalid byte cannot wedge the cursor forever (the line decodes @@ -729,6 +729,38 @@ fn reconcile_flagged_sessions( Ok(()) } +/// Speed fallback from the rollout's codex-home `config.toml` (recorded +/// service tiers win; this covers unmarked entries). Cached by config mtime: +/// a pass over many rollouts stats the file once per rollout but reads and +/// parses it only when it changed — never per line. +fn codex_config_fallback_speed(stream_path: &Path) -> Option { + use std::sync::{Mutex, OnceLock}; + use std::time::SystemTime; + type Cache = HashMap, Option)>; + static CACHE: OnceLock> = OnceLock::new(); + + let config_path = + crate::streams::model_extraction::codex_home_from_transcript_path(stream_path)? + .join("config.toml"); + let modified = std::fs::metadata(&config_path) + .ok() + .and_then(|meta| meta.modified().ok()); + let cache = CACHE.get_or_init(Default::default); + if let Ok(cache) = cache.lock() + && let Some((cached_modified, speed)) = cache.get(&config_path) + && *cached_modified == modified + { + return *speed; + } + let speed = crate::streams::model_extraction::read_codex_config(&config_path) + .as_deref() + .and_then(crate::token_usage::codex::config_fallback_speed); + if let Ok(mut cache) = cache.lock() { + cache.insert(config_path, (modified, speed)); + } + speed +} + /// Incrementally read one transcript file, persist deduplicated entries, and /// emit changed buckets through `sink`. Split out (with injectable repo /// resolution and sink) for direct testing without a daemon. @@ -789,6 +821,12 @@ fn process_file( return Ok(()); } + // Only lines fed below consume the fallback speed, so quiet-skipped + // passes never pay the config lookup. + if identity.tool == "codex" { + extractor.set_fallback_speed(codex_config_fallback_speed(Path::new(stream_path))); + } + let file = std::fs::File::open(stream_path)?; let mut reader = BufReader::with_capacity(128 * 1024, file); reader.seek(SeekFrom::Start(offset))?; @@ -1476,6 +1514,63 @@ mod tests { ); } + #[test] + fn codex_config_service_tier_is_the_fallback_for_unmarked_entries() { + // A rollout under a codex home whose config.toml requests the fast + // tier: unmarked entries resolve to fast (inferred), while a recorded + // tier still wins. + let (dir, db, _) = setup(); + let codex_home = dir.path().join(".codex"); + fs::create_dir_all(codex_home.join("sessions")).unwrap(); + fs::write( + codex_home.join("config.toml"), + "model = \"gpt-5.1\"\nservice_tier = \"fast\" # priority lane\n", + ) + .unwrap(); + let transcript = codex_home.join("sessions").join("rollout-test.jsonl"); + let usage = |minute: u32, totals: (u64, u64, u64, u64, u64)| { + format!( + r#"{{"timestamp":"{}","type":"event_msg","payload":{{"type":"token_count","info":{{"total_token_usage":{{"input_tokens":{},"cached_input_tokens":{},"output_tokens":{},"reasoning_output_tokens":{},"total_tokens":{}}}}}}}}}"#, + recent_ts(minute, 0), + totals.0, + totals.1, + totals.2, + totals.3, + totals.4 + ) + }; + fs::write( + &transcript, + format!( + "{}\n{}\n{}\n", + usage(1, (100, 40, 50, 10, 150)), + r#"{"timestamp":"2026-01-01T00:02:00Z","type":"event_msg","payload":{"type":"thread_settings_applied","thread_settings":{"service_tier":"standard"}}}"#, + usage(3, (300, 140, 90, 30, 390)), + ), + ) + .unwrap(); + let identity = SessionIdentity { + tool: "codex".to_string(), + ..identity() + }; + run_as(&db, &identity, &transcript).unwrap(); + + let conn = + crate::sqlite::open_with_memory_limits(dir.path().join("token-usage-db")).unwrap(); + // Both turns land in the same 5-minute bucket, so order by insertion + // (file) order rather than the tying bucket_ts. + let rows: Vec<(u64, i64, bool)> = conn + .prepare("SELECT input_tokens, speed, speed_inferred FROM usage_entries ORDER BY rowid") + .unwrap() + .query_map([], |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?))) + .unwrap() + .collect::>() + .unwrap(); + // First turn: unmarked -> config fast, inferred. Second: recorded + // standard wins, not inferred. + assert_eq!(rows, vec![(60, 1, true), (100, 0, false)]); + } + #[test] fn cross_session_replacement_reconciles_the_previous_owner_without_its_file() { let (dir, db, transcript_a) = setup(); diff --git a/src/streams/model_extraction.rs b/src/streams/model_extraction.rs index edce141a4d..e3be0aa662 100644 --- a/src/streams/model_extraction.rs +++ b/src/streams/model_extraction.rs @@ -205,15 +205,7 @@ fn extract_model_from_codex_jsonl_line(line: &str) -> Option { fn extract_model_from_codex_config(path: &Path) -> Option { let codex_home = codex_home_from_transcript_path(path)?; - let config_path = codex_home.join("config.toml"); - let file = File::open(config_path).ok()?; - let mut content = String::new(); - file.take(MAX_CODEX_CONFIG_BYTES + 1) - .read_to_string(&mut content) - .ok()?; - if content.len() as u64 > MAX_CODEX_CONFIG_BYTES { - return None; - } + let content = read_codex_config(&codex_home.join("config.toml"))?; let config: toml::Value = toml::from_str(&content).ok()?; config @@ -224,7 +216,21 @@ fn extract_model_from_codex_config(path: &Path) -> Option { .or_else(|| toml_string_candidate(config.get("model"))) } -fn codex_home_from_transcript_path(path: &Path) -> Option { +/// Read a codex `config.toml`, size-capped so an absurd file can't balloon +/// memory. `None` when missing, unreadable, or over the cap. +pub(crate) fn read_codex_config(config_path: &Path) -> Option { + let file = File::open(config_path).ok()?; + let mut content = String::new(); + file.take(MAX_CODEX_CONFIG_BYTES + 1) + .read_to_string(&mut content) + .ok()?; + if content.len() as u64 > MAX_CODEX_CONFIG_BYTES { + return None; + } + Some(content) +} + +pub(crate) fn codex_home_from_transcript_path(path: &Path) -> Option { let configured_home = crate::mdm::utils::codex_home_dir(); if path.starts_with(&configured_home) { return Some(configured_home); diff --git a/src/token_usage/claude.rs b/src/token_usage/claude.rs index f31b1b3323..00d88e2bc6 100644 --- a/src/token_usage/claude.rs +++ b/src/token_usage/claude.rs @@ -8,17 +8,15 @@ //! - Deduplication and the replacement policy run in the token-usage database //! (entries stream in across runs), via [`should_replace_entry`] and the //! `message_id` fallback; ccusage dedups in memory over whole files. -//! - Fast-speed entries keep their base model name (no "-fast" suffix or fast -//! pricing multiplier); `usage.speed` remains the replacement tie-breaker. +//! - Fast-speed entries keep their base model name (no "-fast" suffix); +//! `usage.speed` remains the replacement tie-breaker. //! - Entries whose model is missing or `` are attributed to //! [`UNKNOWN_MODEL`] instead of carrying no model, so tokens aren't lost. -//! - Long-context tiered pricing is not applied (git-ai's models.dev catalog -//! has flat rates); the 1h-ephemeral cache multiplier is (see `cost.rs`). use serde::Deserialize; use super::extractor::UsageExtractor; -use super::types::{TokenCounts, UsageEntry, parse_rfc3339_secs}; +use super::types::{PricingShape, Speed, TokenCounts, UsageEntry, parse_rfc3339_secs}; /// Model recorded for entries with no usable model name. pub const UNKNOWN_MODEL: &str = "unknown"; @@ -52,7 +50,7 @@ impl From<&UsageEntry> for ReplacementCandidate { Self { is_sidechain: entry.is_sidechain, token_total: entry.dedupe_token_total(), - has_speed: entry.has_speed, + has_speed: entry.speed.is_some(), } } } @@ -103,21 +101,14 @@ struct RawUsage { cache_creation_input_tokens: u64, #[serde(default)] cache_read_input_tokens: u64, + /// ccusage `Speed`: unknown values fail the line's parse, matching + /// ccusage's strict schema. #[serde(default)] speed: Option, #[serde(default)] cache_creation: Option, } -/// ccusage `Speed`: unknown values fail the line's parse, matching ccusage's -/// strict schema. -#[derive(Deserialize, Clone, Copy)] -#[serde(rename_all = "lowercase")] -enum Speed { - Standard, - Fast, -} - #[derive(Deserialize, Default, Clone, Copy)] struct RawCacheCreation { #[serde(default)] @@ -230,7 +221,12 @@ fn usage_entry( .filter(|cost| cost.is_finite() && *cost >= 0.0) .map(super::cost::micro_usd), is_sidechain, - has_speed: usage.speed.is_some(), + speed: usage.speed, + // The wire flag means "tier not recorded in the transcript" — an + // unmarked Claude entry lands in the standard bucket by default, + // exactly like an unmarked Codex entry, and must report the same. + speed_inferred: usage.speed.is_none(), + pricing_shape: PricingShape::Claude, } } @@ -447,7 +443,9 @@ mod tests { assert_eq!(e.cache_write_1h, 0); assert_eq!(e.transcript_cost_micro_usd, None); assert!(!e.is_sidechain); - assert!(!e.has_speed); + assert_eq!(e.speed, None); + assert!(e.speed_inferred, "no usage.speed marker was recorded"); + assert_eq!(e.pricing_shape, PricingShape::Claude); } #[test] @@ -566,7 +564,7 @@ mod tests { // Tie: speed marker wins. let mut with_speed = base.clone(); - with_speed.has_speed = true; + with_speed.speed = Some(Speed::Standard); assert!(should_replace_entry(&with_speed, base)); assert!(!should_replace_entry(base, &with_speed)); assert!(!should_replace_entry(&base.clone(), base)); @@ -582,9 +580,13 @@ mod tests { fn speed_marker_is_extracted() { let line = r#"{"timestamp":"2026-01-01T00:00:00Z","message":{"id":"m","model":"x","usage":{"input_tokens":1,"output_tokens":1,"speed":"fast"}},"requestId":"r"}"#; let e = &extract(line)[0]; - assert!(e.has_speed); + assert_eq!(e.speed, Some(Speed::Fast)); + assert!(!e.speed_inferred); // Model name keeps its base form (deviation from ccusage's "-fast"). assert_eq!(e.model, "x"); + + let standard = r#"{"timestamp":"2026-01-01T00:00:00Z","message":{"id":"m","model":"x","usage":{"input_tokens":1,"output_tokens":1,"speed":"standard"}},"requestId":"r"}"#; + assert_eq!(extract(standard)[0].speed, Some(Speed::Standard)); } #[test] diff --git a/src/token_usage/codex.rs b/src/token_usage/codex.rs index 3fd6ebdcd1..9db10e4b83 100644 --- a/src/token_usage/codex.rs +++ b/src/token_usage/codex.rs @@ -9,8 +9,13 @@ //! - Session rollout format only (`event_msg`/`token_count`, `turn_context`, //! `session_meta`); the headless `codex exec` log format is not tracked by //! git-ai's streams and is not parsed. -//! - No service-tier / fast pricing multipliers and no `codex-auto-review` -//! release-date fallback table; model ids price through git-ai's catalog. +//! - The service tier is resolved per entry at extraction time (recorded +//! `thread_settings_applied` tier, else the injected config fallback, else +//! standard) and stored on the entry; ccusage's auto mode instead applies +//! the config in force at *report* time to unmarked usage, retroactively +//! repricing history when `~/.codex/config.toml` changes. +//! - No `codex-auto-review` release-date fallback table; model ids price +//! through git-ai's catalog. //! - Fork replay: ccusage matches a forked session's leading usage against //! the parent log's usage prefix, which requires reading other files. That //! is not possible incrementally, so forks always take ccusage's fallback @@ -28,7 +33,7 @@ use serde::{Deserialize, Serialize}; use super::extractor::UsageExtractor; -use super::types::{TokenCounts, UsageEntry}; +use super::types::{PricingShape, Speed, TokenCounts, UsageEntry}; /// ccusage `CODEX_REWRITTEN_BURST_PAUSE_MS`: the longest pause tolerated /// inside a burst of replayed usage. Codex rewrites replayed history to the @@ -42,6 +47,10 @@ const FALLBACK_MODEL: &str = "gpt-5"; #[derive(Default)] pub struct CodexUsageExtractor { state: CodexState, + /// Speed for entries whose transcript records no service tier, resolved + /// from `~/.codex/config.toml` and injected per pass by the worker (not + /// persisted — the config in force when an entry is extracted decides). + fallback_speed: Option, } /// Parser state persisted between incremental runs. @@ -59,6 +68,10 @@ struct CodexState { /// subtraction. #[serde(default)] prev_totals: Option, + /// Sticky service tier recorded by the last `thread_settings_applied` + /// event that carried one (ccusage `current_service_tier`). + #[serde(default)] + service_tier: Option, #[serde(default)] replay: ReplayState, } @@ -147,6 +160,12 @@ struct PendingEvent { ts_ms: i64, model: String, delta: CodexTotals, + /// Resolved speed at event time (defaults keep pre-speed persisted state + /// readable). + #[serde(default)] + speed: Speed, + #[serde(default)] + speed_inferred: bool, } impl UsageExtractor for CodexUsageExtractor { @@ -154,6 +173,7 @@ impl UsageExtractor for CodexUsageExtractor { line.contains("token_count") || line.contains("turn_context") || line.contains("session_meta") + || line.contains("thread_settings_applied") } fn extract_line(&mut self, line: &str) -> Vec { @@ -201,6 +221,10 @@ impl UsageExtractor for CodexUsageExtractor { } } + fn set_fallback_speed(&mut self, speed: Option) { + self.fallback_speed = speed; + } + fn has_pending(&self) -> bool { matches!( self.state.replay, @@ -250,6 +274,22 @@ impl CodexUsageExtractor { let Some(payload) = raw.payload.as_ref() else { return Vec::new(); }; + if payload.payload_type.as_deref() == Some("thread_settings_applied") { + // A settings event that carries no `service_tier` at all says + // nothing about the tier, so the previous one stands (Codex emits + // such events for auto-review threads). A tier that is present + // but unrecognized is different: it means the tier changed to + // something unknown, so the stale value must not be inherited + // (ccusage `visit_codex_session_entry`). + if let Some(recorded) = payload + .thread_settings + .as_ref() + .and_then(|settings| settings.service_tier.as_deref()) + { + self.state.service_tier = service_tier_speed(recorded); + } + return Vec::new(); + } if payload.payload_type.as_deref() != Some("token_count") { return Vec::new(); } @@ -295,10 +335,15 @@ impl CodexUsageExtractor { .model .clone() .unwrap_or_else(|| FALLBACK_MODEL.to_string()); + // Recorded tier wins, else the config fallback, else standard + // (ccusage's auto speed policy, resolved at event time). + let recorded = self.state.service_tier; PendingEvent { ts_ms, model, delta, + speed: recorded.or(self.fallback_speed).unwrap_or_default(), + speed_inferred: recorded.is_none(), } }); // The replay filter sees every usage-carrying event, including @@ -360,6 +405,8 @@ fn make_entry(event: PendingEvent) -> UsageEntry { ts_ms, model, delta, + speed, + speed_inferred, } = event; // ccusage clamps cached to input; normalized input excludes cache. let cached = delta.cached_input_tokens.min(delta.input_tokens); @@ -379,10 +426,45 @@ fn make_entry(event: PendingEvent) -> UsageEntry { cache_write_1h: 0, transcript_cost_micro_usd: None, is_sidechain: false, - has_speed: false, + speed: Some(speed), + speed_inferred, + pricing_shape: PricingShape::Codex, } } +/// The speed a recorded or configured service-tier value maps to (ccusage +/// `codex_service_tier`): exact strings, no case-folding. Unknown values map +/// to `None`. +fn service_tier_speed(value: &str) -> Option { + match value { + // Both spellings mean non-priority pricing and occur in the same + // Codex version on the same day; which one is written depends on the + // client (Codex Desktop writes "standard"), not on the CLI version. + "default" | "standard" => Some(Speed::Standard), + "fast" | "priority" => Some(Speed::Fast), + _ => None, + } +} + +/// Speed fallback implied by a Codex `config.toml`: the `service_tier` of +/// the active profile (`profiles..service_tier`, matching Codex's +/// own config resolution and `streams::model_extraction`'s model lookup), +/// falling back to the top-level key. Only `fast`/`priority` produce a +/// fallback — a configured `standard` resolves like an absent key, since +/// the auto policy already defaults to standard. Deviation from ccusage, +/// whose table-blind line scan would mark the whole home fast for a +/// `service_tier = "fast"` inside an *inactive* profile table. +pub fn config_fallback_speed(config_toml: &str) -> Option { + let config: toml::Value = toml::from_str(config_toml).ok()?; + let tier = config + .get("profile") + .and_then(toml::Value::as_str) + .and_then(|profile| config.get("profiles")?.get(profile)?.get("service_tier")) + .or_else(|| config.get("service_tier")) + .and_then(toml::Value::as_str)?; + service_tier_speed(tier).filter(|speed| *speed == Speed::Fast) +} + /// Content-derived dedup key over the event's full identity (timestamp, /// model, all counts), matching ccusage's codex dedup key. A fork that /// replays the parent's events verbatim maps them to the same keys @@ -421,12 +503,19 @@ struct RawPayload { #[serde(alias = "modelId")] model_id: Option, metadata: Option, + // thread_settings_applied fields: + thread_settings: Option, // session_meta fields: id: Option, forked_from_id: Option, source: Option, } +#[derive(Deserialize)] +struct RawThreadSettings { + service_tier: Option, +} + #[derive(Deserialize)] struct RawInfo { total_token_usage: Option, @@ -547,12 +636,206 @@ mod tests { ) } + fn thread_settings_line(service_tier: Option<&str>) -> String { + let settings = match service_tier { + Some(tier) => format!(r#"{{"service_tier":"{tier}"}}"#), + None => "{}".to_string(), + }; + format!( + r#"{{"timestamp":"2026-01-01T00:00:00Z","type":"event_msg","payload":{{"type":"thread_settings_applied","thread_settings":{settings}}}}}"# + ) + } + + #[test] + fn service_tier_is_sticky_and_maps_all_spellings() { + // "default"/"standard" and "fast"/"priority" are spelling pairs; the + // recorded tier applies to every following usage event until changed. + for (tier, speed) in [ + ("default", Speed::Standard), + ("standard", Speed::Standard), + ("fast", Speed::Fast), + ("priority", Speed::Fast), + ] { + let mut e = CodexUsageExtractor::default(); + e.extract_line(&thread_settings_line(Some(tier))); + let entries = e.extract_line(&token_count_line( + "2026-01-01T00:00:10Z", + (100, 40, 50, 10, 150), + )); + assert_eq!(entries[0].speed, Some(speed), "tier {tier}"); + assert!(!entries[0].speed_inferred, "tier {tier} is recorded"); + } + + // Sticky across turns, and a later settings event switches it. + let mut e = CodexUsageExtractor::default(); + e.extract_line(&thread_settings_line(Some("fast"))); + let first = e.extract_line(&token_count_line( + "2026-01-01T00:00:10Z", + (100, 40, 50, 10, 150), + )); + assert_eq!(first[0].speed, Some(Speed::Fast)); + e.extract_line(&thread_settings_line(Some("standard"))); + let second = e.extract_line(&token_count_line( + "2026-01-01T00:01:10Z", + (300, 140, 90, 30, 390), + )); + assert_eq!(second[0].speed, Some(Speed::Standard)); + assert!(!second[0].speed_inferred); + } + + #[test] + fn unknown_tier_resets_but_an_absent_tier_says_nothing() { + let mut e = CodexUsageExtractor::default(); + e.extract_line(&thread_settings_line(Some("fast"))); + // No `service_tier` at all (Codex writes such events for auto-review + // threads): the previous tier stands. + e.extract_line(&thread_settings_line(None)); + let kept = e.extract_line(&token_count_line( + "2026-01-01T00:00:10Z", + (100, 40, 50, 10, 150), + )); + assert_eq!(kept[0].speed, Some(Speed::Fast)); + assert!(!kept[0].speed_inferred); + + // A present but unrecognized tier means it changed to something + // unknown; the stale value must not be inherited. + e.extract_line(&thread_settings_line(Some("turbo"))); + let reset = e.extract_line(&token_count_line( + "2026-01-01T00:01:10Z", + (300, 140, 90, 30, 390), + )); + assert_eq!(reset[0].speed, Some(Speed::Standard)); + assert!(reset[0].speed_inferred, "falls back to the default"); + } + + #[test] + fn fallback_speed_covers_unmarked_entries_only() { + let mut e = CodexUsageExtractor::default(); + e.set_fallback_speed(Some(Speed::Fast)); + let unmarked = e.extract_line(&token_count_line( + "2026-01-01T00:00:10Z", + (100, 40, 50, 10, 150), + )); + assert_eq!(unmarked[0].speed, Some(Speed::Fast)); + assert!(unmarked[0].speed_inferred); + + // A recorded tier wins over the config fallback. + e.extract_line(&thread_settings_line(Some("standard"))); + let recorded = e.extract_line(&token_count_line( + "2026-01-01T00:01:10Z", + (300, 140, 90, 30, 390), + )); + assert_eq!(recorded[0].speed, Some(Speed::Standard)); + assert!(!recorded[0].speed_inferred); + + // No fallback and no recorded tier: standard, inferred. + let mut bare = CodexUsageExtractor::default(); + let entries = bare.extract_line(&token_count_line( + "2026-01-01T00:00:10Z", + (100, 40, 50, 10, 150), + )); + assert_eq!(entries[0].speed, Some(Speed::Standard)); + assert!(entries[0].speed_inferred); + } + + #[test] + fn config_fallback_detects_explicit_fast_service_tier_values() { + // ccusage `detects_explicit_fast_service_tier_values`. + assert_eq!( + config_fallback_speed(r#"service_tier = "fast""#), + Some(Speed::Fast) + ); + assert_eq!( + config_fallback_speed(r#"service_tier = 'priority' # use higher tier"#), + Some(Speed::Fast) + ); + } + + #[test] + fn config_fallback_ignores_unrelated_or_substring_service_tier_values() { + // ccusage `ignores_unrelated_or_substring_service_tier_values`; a + // configured "standard" resolves like an absent key. + assert_eq!( + config_fallback_speed(r#"service_tier_override = "fast""#), + None + ); + assert_eq!(config_fallback_speed(r#"service_tier = "breakfast""#), None); + assert_eq!(config_fallback_speed(r#"service_tier = "standard""#), None); + assert_eq!(config_fallback_speed(""), None); + assert_eq!(config_fallback_speed("not [valid toml"), None); + } + + #[test] + fn config_fallback_respects_profile_scoping() { + // A fast tier inside an INACTIVE profile table must not mark the + // whole codex home fast (deviation from ccusage's line scan), while + // the active profile's tier wins over the top-level key. + assert_eq!( + config_fallback_speed( + "profile = \"work\"\n[profiles.turbo]\nservice_tier = \"fast\"\n" + ), + None + ); + assert_eq!( + config_fallback_speed( + "profile = \"turbo\"\n[profiles.turbo]\nservice_tier = \"fast\"\n" + ), + Some(Speed::Fast) + ); + assert_eq!( + config_fallback_speed( + "profile = \"calm\"\nservice_tier = \"fast\"\n[profiles.calm]\nservice_tier = \"standard\"\n" + ), + None, + "the active profile's standard tier wins over the top-level fast" + ); + assert_eq!( + config_fallback_speed("service_tier = \"priority\"\n[profiles.idle]\nmodel = \"x\"\n"), + Some(Speed::Fast), + "top-level tier applies when no profile is selected" + ); + } + + #[test] + fn service_tier_roundtrips_through_persisted_state() { + let mut e = CodexUsageExtractor::default(); + e.extract_line(&thread_settings_line(Some("fast"))); + let state = e.state_json().unwrap(); + + let mut restored = CodexUsageExtractor::default(); + assert!(restored.restore_state(&state)); + let entries = restored.extract_line(&token_count_line( + "2026-01-01T00:00:10Z", + (100, 40, 50, 10, 150), + )); + assert_eq!(entries[0].speed, Some(Speed::Fast)); + assert!(!entries[0].speed_inferred); + } + + #[test] + fn pre_speed_persisted_state_still_restores() { + // State written before the service-tier fields existed must restore + // (serde defaults), not reset the cursor. + let mut e = CodexUsageExtractor::default(); + assert!(e.restore_state( + r#"{"model":"gpt-5.1","prev_totals":{"input_tokens":100,"cached_input_tokens":40,"output_tokens":50,"reasoning_output_tokens":10,"total_tokens":150},"replay":{"kind":"done"}}"# + )); + let entries = e.extract_line(&token_count_line( + "2026-01-01T00:01:10Z", + (300, 140, 90, 30, 390), + )); + assert_eq!(entries[0].tokens.input, 100); + assert_eq!(entries[0].speed, Some(Speed::Standard)); + assert!(entries[0].speed_inferred); + } + #[test] fn prefilter_matches_relevant_lines() { let e = CodexUsageExtractor::default(); assert!(e.wants_line(&token_count_line("2026-01-01T00:00:00Z", (1, 0, 1, 0, 2)))); assert!(e.wants_line(&turn_context_line("gpt-5.1"))); assert!(e.wants_line(r#"{"type":"session_meta","payload":{"id":"x"}}"#)); + assert!(e.wants_line(&thread_settings_line(Some("fast")))); assert!(!e.wants_line(r#"{"type":"response_item","payload":{"type":"message"}}"#)); } diff --git a/src/token_usage/cost.rs b/src/token_usage/cost.rs index e7b0277389..4ef04df7ec 100644 --- a/src/token_usage/cost.rs +++ b/src/token_usage/cost.rs @@ -76,7 +76,9 @@ mod tests { cache_write_1h: 0, transcript_cost_micro_usd: None, is_sidechain: false, - has_speed: false, + speed: None, + speed_inferred: false, + pricing_shape: crate::token_usage::PricingShape::Claude, } } diff --git a/src/token_usage/db.rs b/src/token_usage/db.rs index e76e35f966..679a55f9bd 100644 --- a/src/token_usage/db.rs +++ b/src/token_usage/db.rs @@ -41,7 +41,7 @@ use rusqlite::{Connection, OptionalExtension, Transaction, params}; use super::claude::{ReplacementCandidate, should_replace}; use super::cost::entry_cost_micro_usd; -use super::types::{UsageEntry, bucket_ts}; +use super::types::{Speed, UsageEntry, bucket_ts}; use crate::error::GitAiError; /// Schema migrations - each entry is SQL to apply for that version. @@ -130,6 +130,80 @@ const MIGRATIONS: &[&str] = &[ INSERT INTO schema_version (version) VALUES (2); "#, + // Version 3: per-entry pricing dimensions (speed/service tier, tier + // inference, 1h cache-write split, transcript-vs-catalog cost provenance, + // long-context tier decision, pricing catalog id). Pre-release reset: v2 + // rows carry neither the new extraction facts nor tier-aware costs, so + // the tables are rebuilt and the dropped cursors make the next pass + // re-extract every transcript under the new pricing rules (entry-level + // dedup keeps that idempotent; emission reconciles by fingerprint). + r#" + DROP TABLE IF EXISTS tracked_files; + DROP TABLE IF EXISTS usage_entries; + DROP TABLE IF EXISTS bucket_state; + + CREATE TABLE tracked_files ( + session_id TEXT NOT NULL, + stream_path TEXT NOT NULL, + tool TEXT NOT NULL, + byte_offset INTEGER NOT NULL DEFAULT 0, + state_json TEXT, + last_known_size INTEGER NOT NULL DEFAULT 0, + last_modified INTEGER, + processing_errors INTEGER NOT NULL DEFAULT 0, + last_error TEXT, + pending_flush INTEGER NOT NULL DEFAULT 0, + last_error_at INTEGER, + external_session_id TEXT NOT NULL DEFAULT '', + needs_reconcile INTEGER NOT NULL DEFAULT 0, + repo_url TEXT, + PRIMARY KEY (session_id, stream_path) + ); + + CREATE TABLE usage_entries ( + session_id TEXT NOT NULL, + entry_key TEXT NOT NULL, + message_id TEXT, + model TEXT NOT NULL, + bucket_ts INTEGER NOT NULL, + input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_write_tokens INTEGER NOT NULL DEFAULT 0, + cache_write_1h_tokens INTEGER NOT NULL DEFAULT 0, + reasoning_output_tokens INTEGER, + total_tokens INTEGER NOT NULL DEFAULT 0, + cost_micro_usd INTEGER, + transcript_cost_micro_usd INTEGER, + is_sidechain INTEGER NOT NULL DEFAULT 0, + speed INTEGER, + speed_inferred INTEGER NOT NULL DEFAULT 0, + long_context INTEGER NOT NULL DEFAULT 0, + pricing_catalog TEXT, + PRIMARY KEY (session_id, entry_key) + ); + + CREATE INDEX idx_usage_entries_bucket + ON usage_entries(session_id, model, bucket_ts); + CREATE INDEX idx_usage_entries_message + ON usage_entries(session_id, message_id) WHERE message_id IS NOT NULL; + CREATE UNIQUE INDEX idx_usage_entries_key + ON usage_entries(entry_key); + CREATE INDEX idx_usage_entries_message_global + ON usage_entries(message_id) WHERE message_id IS NOT NULL; + + CREATE TABLE bucket_state ( + session_id TEXT NOT NULL, + model TEXT NOT NULL, + bucket_ts INTEGER NOT NULL, + emitted_fingerprint TEXT NOT NULL, + last_emitted_at INTEGER NOT NULL, + emit_seq INTEGER NOT NULL DEFAULT 0, + PRIMARY KEY (session_id, model, bucket_ts) + ); + + INSERT INTO schema_version (version) VALUES (3); + "#, ]; const TRACKED_FILE_COLUMNS: &str = "session_id, stream_path, tool, byte_offset, state_json, \ @@ -323,6 +397,17 @@ impl TokenUsageDatabase { } else { 0 }; + // A database from a NEWER binary must fail closed here: migrations + // are destructive from v3 on (columns renamed/dropped), so a + // downgraded daemon would otherwise open fine and then fail every + // pass at runtime on missing columns, churning error backoff. + if current_version > MIGRATIONS.len() as u32 { + return Err(GitAiError::Generic(format!( + "token-usage database schema v{current_version} is newer than this binary \ + (supports up to v{}); upgrade git-ai or delete the database", + MIGRATIONS.len() + ))); + } for (version, migration_sql) in MIGRATIONS.iter().enumerate() { if current_version < (version + 1) as u32 { // Each migration commits atomically: a crash between the @@ -749,9 +834,11 @@ fn upsert_entry( session_id = ?1, entry_key = ?2, message_id = ?3, model = ?4, bucket_ts = ?5, input_tokens = ?6, output_tokens = ?7, cache_read_tokens = ?8, cache_write_tokens = ?9, - reasoning_output_tokens = ?10, total_tokens = ?11, - cost_micro_usd = ?12, is_sidechain = ?13, has_speed = ?14 - WHERE rowid = ?15", + cache_write_1h_tokens = ?10, reasoning_output_tokens = ?11, + total_tokens = ?12, cost_micro_usd = ?13, + transcript_cost_micro_usd = ?14, is_sidechain = ?15, + speed = ?16, speed_inferred = ?17 + WHERE rowid = ?18", params![ session_id, entry.entry_key, @@ -762,11 +849,14 @@ fn upsert_entry( to_db_i64(entry.tokens.output), to_db_i64(entry.tokens.cache_read), to_db_i64(entry.tokens.cache_write), + to_db_i64(entry.cache_write_1h), entry.tokens.reasoning_output.map(to_db_i64), to_db_i64(entry.tokens.total), entry_cost_micro_usd(entry).map(to_db_i64), + entry.transcript_cost_micro_usd.map(to_db_i64), entry.is_sidechain, - entry.has_speed, + entry.speed.map(speed_to_db), + entry.speed_inferred, row.rowid, ], )?; @@ -781,9 +871,11 @@ fn upsert_entry( "INSERT INTO usage_entries ( session_id, entry_key, message_id, model, bucket_ts, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, - reasoning_output_tokens, total_tokens, cost_micro_usd, - is_sidechain, has_speed - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)", + cache_write_1h_tokens, reasoning_output_tokens, total_tokens, + cost_micro_usd, transcript_cost_micro_usd, is_sidechain, + speed, speed_inferred + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, + ?15, ?16, ?17)", params![ session_id, entry.entry_key, @@ -794,11 +886,14 @@ fn upsert_entry( to_db_i64(entry.tokens.output), to_db_i64(entry.tokens.cache_read), to_db_i64(entry.tokens.cache_write), + to_db_i64(entry.cache_write_1h), entry.tokens.reasoning_output.map(to_db_i64), to_db_i64(entry.tokens.total), entry_cost_micro_usd(entry).map(to_db_i64), + entry.transcript_cost_micro_usd.map(to_db_i64), entry.is_sidechain, - entry.has_speed, + entry.speed.map(speed_to_db), + entry.speed_inferred, ], )?; Ok(None) @@ -806,6 +901,14 @@ fn upsert_entry( } } +/// Column encoding of a recorded speed (NULL = the transcript carried none). +fn speed_to_db(speed: Speed) -> i64 { + match speed { + Speed::Standard => 0, + Speed::Fast => 1, + } +} + fn find_dedupe_target( tx: &Transaction<'_>, entry: &UsageEntry, @@ -826,7 +929,8 @@ fn find_dedupe_target( }) }; const ROW_COLUMNS: &str = "rowid, session_id, input_tokens, output_tokens, \ - cache_read_tokens, cache_write_tokens, is_sidechain, has_speed"; + cache_read_tokens, cache_write_tokens, is_sidechain, \ + speed IS NOT NULL"; let exact = tx .query_row( @@ -880,7 +984,9 @@ mod tests { cache_write_1h: 0, transcript_cost_micro_usd: Some(1_000), is_sidechain: false, - has_speed: false, + speed: None, + speed_inferred: false, + pricing_shape: crate::token_usage::PricingShape::Claude, } } @@ -918,10 +1024,32 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("token-usage-db"); let db = TokenUsageDatabase::open(&path).unwrap(); - assert_eq!(db.schema_version().unwrap(), 2); + assert_eq!(db.schema_version().unwrap(), 3); drop(db); let db = TokenUsageDatabase::open(&path).unwrap(); - assert_eq!(db.schema_version().unwrap(), 2); + assert_eq!(db.schema_version().unwrap(), 3); + } + + #[test] + fn newer_schema_versions_fail_closed_on_open() { + // Migrations are destructive from v3 on, so a downgraded binary must + // refuse a database written by a newer one instead of failing every + // pass at runtime on missing columns. + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("token-usage-db"); + drop(TokenUsageDatabase::open(&path).unwrap()); + crate::sqlite::open_with_memory_limits(&path) + .unwrap() + .execute("INSERT INTO schema_version (version) VALUES (99)", []) + .unwrap(); + + let Err(err) = TokenUsageDatabase::open(&path) else { + panic!("opening a newer-schema database must fail closed"); + }; + assert!( + err.to_string().contains("newer than this binary"), + "unexpected error: {err}" + ); } #[test] @@ -1189,10 +1317,11 @@ mod tests { #[test] fn v1_databases_with_legacy_cross_session_duplicates_upgrade_cleanly() { // A v1 database (session-scoped dedup) can hold the same entry_key - // under several sessions. The v2 migration keeps the first-seen row - // and enforces global uniqueness, so a later cross-session - // replacement can never collide on the (session_id, entry_key) - // primary key and poison the transcript. + // under several sessions; the intermediate v2 migration must still + // purge those so its unique index can be created, and the v3 + // pre-release rebuild then drops everything: cursors reset so the + // next pass re-extracts every transcript under the pricing-aware + // schema. let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("token-usage-db"); { @@ -1216,31 +1345,17 @@ mod tests { } let db = TokenUsageDatabase::open(&path).unwrap(); - assert_eq!(db.schema_version().unwrap(), 2); - // Sessions whose rows were purged are flagged so their (now lower) - // aggregates re-emit; the surviving session is not. - let flagged: Vec = db - .sessions_needing_reconcile() - .unwrap() - .into_iter() - .map(|session| session.session_id) - .collect(); - assert_eq!(flagged, vec!["s2".to_string(), "s3".to_string()]); - // First-seen row (s1) survives; the duplicates are gone. + assert_eq!(db.schema_version().unwrap(), 3); + // Rebuilt empty: no legacy rows, cursors, or reconcile flags survive. + assert!(db.all_files().unwrap().is_empty()); + assert!(db.sessions_needing_reconcile().unwrap().is_empty()); assert_eq!( db.aggregate_bucket("s1", "claude-sonnet-4-20250514", 600) .unwrap() .message_count, - 1 - ); - assert_eq!( - db.aggregate_bucket("s2", "claude-sonnet-4-20250514", 600) - .unwrap() - .message_count, 0 ); - // A cross-session replacement over the legacy key works (this used - // to violate the primary key when a duplicate existed). + // The rebuilt schema accepts fresh commits, including a legacy key. commit_for( &db, "s2", diff --git a/src/token_usage/extractor.rs b/src/token_usage/extractor.rs index 3e25914359..67a7ace6fc 100644 --- a/src/token_usage/extractor.rs +++ b/src/token_usage/extractor.rs @@ -1,6 +1,6 @@ //! Per-agent usage extraction trait. -use super::types::UsageEntry; +use super::types::{Speed, UsageEntry}; /// Incremental, line-oriented extractor of token-usage entries from an agent /// transcript. Fed complete JSONL lines in file order; may keep per-session @@ -13,6 +13,12 @@ pub trait UsageExtractor: Send { /// Consume one raw JSONL line, returning any usage entries it completes. fn extract_line(&mut self, line: &str) -> Vec; + /// Inject the configuration-derived speed for entries whose transcript + /// records no service tier (Codex `~/.codex/config.toml`). Resolved by + /// the worker once per pass, before any line is fed. No-op for agents + /// without the concept. + fn set_fallback_speed(&mut self, _speed: Option) {} + /// Serialized parser state to persist between incremental runs. `None` /// for stateless extractors. fn state_json(&self) -> Option { diff --git a/src/token_usage/mod.rs b/src/token_usage/mod.rs index 6f00cbbdc5..a2d17b5a21 100644 --- a/src/token_usage/mod.rs +++ b/src/token_usage/mod.rs @@ -17,4 +17,4 @@ pub mod extractor; pub mod types; pub use extractor::{UsageExtractor, extractor_for_tool}; -pub use types::{TokenCounts, UsageEntry, bucket_ts}; +pub use types::{PricingShape, Speed, TokenCounts, UsageEntry, bucket_ts}; diff --git a/src/token_usage/types.rs b/src/token_usage/types.rs index e22b57391c..35a7fd549e 100644 --- a/src/token_usage/types.rs +++ b/src/token_usage/types.rs @@ -25,6 +25,26 @@ pub struct TokenCounts { pub total: u64, } +/// Request speed / service tier (ccusage `Speed`): fast/priority requests +/// bill at the model's fast multiplier over the whole request cost. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum Speed { + #[default] + Standard, + Fast, +} + +/// Which of ccusage's two cost formulas prices an entry — the producing +/// extractor's shape, not the model's (the paths differ in how the +/// long-context tier is selected and how unpublished cache-read rates +/// default). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PricingShape { + Claude, + Codex, +} + /// One usage entry extracted from a transcript, before deduplication. #[derive(Debug, Clone, PartialEq)] pub struct UsageEntry { @@ -46,9 +66,14 @@ pub struct UsageEntry { pub transcript_cost_micro_usd: Option, /// Claude sidechain (subagent) entry; loses to non-sidechain duplicates. pub is_sidechain: bool, - /// Entry carried an explicit `usage.speed` marker; wins ties on - /// replacement. - pub has_speed: bool, + /// Request speed. `None` when the transcript carried no marker (Claude + /// without `usage.speed`); a present marker wins ties on replacement. + pub speed: Option, + /// The speed was not recorded in the transcript but resolved from + /// configuration or the standard default (unmarked entries of any tool). + pub speed_inferred: bool, + /// Cost formula for this entry (see [`PricingShape`]). + pub pricing_shape: PricingShape, } impl UsageEntry {