diff --git a/CHANGELOG.md b/CHANGELOG.md index c054768f2c..2dca9c43a0 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,8 @@ menu-bar layout. ### Fixed - Claude: when Hide Personal Info is enabled, keep saved account rows distinguishable with stable localized `Account N` labels and matching redacted tooltips. +- Codex: a fork of a fork whose parent has no token snapshot before the child forked now inherits the whole chain's cumulative counter origin instead of staying unresolved; ancestor chains deeper than 64 stay unresolved, and the Codex cost cache is rebuilt once (cache schema 5). +- Codex: a fork larger than the per-refresh scan budget continues where the last refresh stopped instead of restarting from the top and never finishing, and a finished fork that grows is no longer billed for usage it inherited from its parent. --- diff --git a/rust/src/core/jsonl_scanner.rs b/rust/src/core/jsonl_scanner.rs index 13e78bcfe9..a5dad9ec04 100755 --- a/rust/src/core/jsonl_scanner.rs +++ b/rust/src/core/jsonl_scanner.rs @@ -396,12 +396,35 @@ pub struct CodexForkAccountingState { pub history_base_thread_id: Option, pub fork_timestamp: Option, pub inherited_totals: Option, + /// Timestamp of the first token event in this fork's log. A descendant + /// forked before it uses this fork's inherited counter origin. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub first_token_timestamp: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub remaining_inherited_totals: Option, /// True when the child log itself supplied enough copied-prefix history to /// establish the inherited baseline without consulting a parent cache row. #[serde(default, skip_serializing_if = "std::ops::Not::not")] pub locally_resolved: bool, + /// Parser state of an unfinished bounded parse, so the next pass continues + /// at `parsed_bytes` instead of rereading the fork from byte zero. Absent + /// once the parse reaches its target, and in caches written before it. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub resume: Option, +} + +/// Fork parser state that the per-file cache fields do not already carry. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CodexForkResumeState { + /// Parent baseline the interrupted parse started from. A validated parent + /// that now supplies a different baseline restarts the parse. + pub parse_baseline: CodexTotals, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub totals_watermark: Option, + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub saw_interleaved_totals: bool, + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub paginated_baseline_checked: bool, } /// Snapshot of the last validated cost report, persisted so spend surfaces keep @@ -449,6 +472,10 @@ pub struct CodexParseResult { pub token_timestamps_monotonic: Option, /// Last token timestamp observed by the parser. pub last_token_timestamp: Option, + /// First token timestamp of the parsed history. A resumed fork parse + /// carries its saved value forward; resumed standard parses see only their + /// suffix, so fork accounting never reads it from them. + pub first_token_timestamp: Option, /// Number of timestamp comparisons performed while validating this parse. pub token_timestamp_comparisons: u64, /// Newly consumed bytes in this parse pass. @@ -463,6 +490,9 @@ pub struct CodexParseResult { /// Remaining inherited counters used when a fork emits last-only rows. pub remaining_inherited_totals: Option, pub fork_baseline_locally_resolved: bool, + /// State to continue an unfinished parent-baseline fork parse. `None` once + /// the parse is complete and for every other parse mode. + pub fork_resume_state: Option, } /// A billable Codex token-count delta. @@ -514,6 +544,7 @@ impl CostUsageDayRange { /// JSONL Scanner for cost/usage logs pub struct JsonlScanner; pub(crate) mod codex; +pub(crate) use codex::CodexForkParseResume; pub(crate) use codex::source_rows::{ read_source_rows, recover_rows, row_cache, row_cache_matches, row_cache_needs_recovery, }; diff --git a/rust/src/core/jsonl_scanner/codex.rs b/rust/src/core/jsonl_scanner/codex.rs index 108d4284d7..d6059fa4a8 100644 --- a/rust/src/core/jsonl_scanner/codex.rs +++ b/rust/src/core/jsonl_scanner/codex.rs @@ -10,12 +10,32 @@ use helpers::{ }; use parser::{CodexParseMode, CodexParserState}; +/// Saved cursor and parser state of an unfinished parent-baseline fork parse. +/// Restoring all of it accounts the remaining bytes exactly as one +/// uninterrupted parse would. +#[derive(Debug, Clone)] +pub(crate) struct CodexForkParseResume { + pub start_offset: i64, + pub paginated_continuation: bool, + /// Effective inherited baseline, after any paginated-continuation raise. + pub inherited_totals: CodexTotals, + pub remaining_inherited_totals: Option, + pub last_model: Option, + pub last_totals: CodexTotals, + pub last_token_timestamp: Option, + pub first_token_timestamp: Option, + pub token_timestamps_monotonic: Option, + pub state: CodexForkResumeState, +} + /// Persisted Codex cache schema version. Version 0 predates 64-bit totals; /// version 1 can retain a terminal pause after treating a paginated v2 /// subagent's independent counters as an inherited fork. Version 3 adds /// persisted paginated-fork accounting state. Version 4 reparses copied-prefix -/// subagents with locally inferred component baselines. Rebuild older artifacts. -pub(crate) const CODEX_CACHE_SCHEMA_VERSION: u32 = 4; +/// subagents with locally inferred component baselines. Version 5 records a +/// fork's first token timestamp so direct-fork chains resolve through parents +/// with no token snapshot at the descendant's fork time. Rebuild older artifacts. +pub(crate) const CODEX_CACHE_SCHEMA_VERSION: u32 = 5; /// Whether a persisted Codex cache artifact matches the current schema. /// A mismatched artifact (e.g. a pre-64-bit cache from an older release) is @@ -293,6 +313,18 @@ impl JsonlScanner { } } + /// Strict counterpart of [`Self::codex_timestamp_at_or_before`]: `earlier` + /// is before `later`. Malformed timestamps fail closed. + pub(crate) fn codex_timestamp_before(earlier: &str, later: &str) -> bool { + match ( + parse_rfc3339_timestamp(earlier), + parse_rfc3339_timestamp(later), + ) { + (Some(earlier), Some(later)) => earlier < later, + _ => false, + } + } + /// Parse a Codex JSONL file pub fn parse_codex_file( file_path: &Path, @@ -514,6 +546,26 @@ impl JsonlScanner { ) } + /// Continue an unfinished fork parse at its saved cursor. Upstream 0.67.0 + /// resumes a resolved fork the same way instead of rereading its prefix. + pub(crate) fn parse_codex_fork_resume( + file_path: &Path, + range: &CostUsageDayRange, + resume: CodexForkParseResume, + cancel: Option<&AtomicBool>, + scan_target_size: Option, + max_bytes_to_read: Option, + ) -> std::io::Result { + Self::parse_codex_file_with_state_bounded_internal( + file_path, + range, + cancel, + scan_target_size, + max_bytes_to_read, + CodexParseMode::ResumeParentBaseline(resume), + ) + } + fn parse_codex_file_with_state_bounded_internal( file_path: &Path, range: &CostUsageDayRange, @@ -631,6 +683,11 @@ impl JsonlScanner { let is_complete = !cancelled && !budget_exhausted && parsed_bytes >= effective_target_size; let bytes_read = parsed_bytes.saturating_sub(safe_start_offset).max(0); let fork_baseline_locally_resolved = parser.fork_baseline_locally_resolved(); + let fork_resume_state = if is_complete { + None + } else { + parser.fork_resume_state() + }; Ok(CodexParseResult { records: parser.records, parsed_bytes, @@ -643,6 +700,7 @@ impl JsonlScanner { last_totals: parser.previous_totals, token_timestamps_monotonic: parser.token_timestamps_monotonic, last_token_timestamp: parser.previous_token_timestamp, + first_token_timestamp: parser.first_token_timestamp, token_timestamp_comparisons: parser.token_timestamp_comparisons, bytes_read, is_complete, @@ -650,6 +708,7 @@ impl JsonlScanner { fork_baseline: parser.fork_baseline, remaining_inherited_totals: parser.remaining_inherited_totals, fork_baseline_locally_resolved, + fork_resume_state, }) } diff --git a/rust/src/core/jsonl_scanner/codex/parser.rs b/rust/src/core/jsonl_scanner/codex/parser.rs index bdce9c68d8..407bf01a36 100644 --- a/rust/src/core/jsonl_scanner/codex/parser.rs +++ b/rust/src/core/jsonl_scanner/codex/parser.rs @@ -1,5 +1,7 @@ use super::helpers::*; -use super::{CodexTotals, CodexUsageRecord, CostUsageDayRange}; +use super::{ + CodexForkParseResume, CodexForkResumeState, CodexTotals, CodexUsageRecord, CostUsageDayRange, +}; use crate::core::CostUsagePricing; use chrono::{DateTime, Utc}; use serde_json::Value; @@ -14,6 +16,10 @@ pub(super) struct CodexParserState { saw_interleaved_totals: bool, pub(super) records: Vec<(CodexUsageRecord, i64)>, pub(super) previous_token_timestamp: Option, + /// First token timestamp seen by this parse. A fork uses it to prove that + /// its inherited origin is still the counter state at a descendant's + /// cutoff. Resumed non-fork parses do not observe the whole prefix. + pub(super) first_token_timestamp: Option, previous_token_timestamp_parsed: Option>, pub(super) token_timestamps_monotonic: Option, pub(super) token_timestamp_comparisons: u64, @@ -23,6 +29,8 @@ pub(super) struct CodexParserState { paginated_baseline_checked: bool, pub(super) fork_baseline_ambiguous: bool, fork_baseline_inference: Option, + /// Parent baseline a parent-baseline fork parse started from. + fork_parse_baseline: Option, } pub(super) enum CodexParseMode { @@ -38,6 +46,8 @@ pub(super) enum CodexParseMode { paginated_continuation: bool, remaining_inherited_totals: Option, }, + /// Continue an unfinished `ParentBaseline` parse at its saved cursor. + ResumeParentBaseline(CodexForkParseResume), InferSubagent { start_ordinal: Option, }, @@ -47,6 +57,7 @@ impl CodexParseMode { pub(super) fn start_offset(&self) -> i64 { match self { Self::Standard { start_offset, .. } => *start_offset, + Self::ResumeParentBaseline(resume) => resume.start_offset, Self::ParentBaseline { .. } | Self::InferSubagent { .. } => 0, } } @@ -213,6 +224,10 @@ impl CodexParserState { } pub(super) fn from_mode(mode: CodexParseMode) -> Self { + let fork_parse_baseline = match &mode { + CodexParseMode::ParentBaseline { baseline, .. } => Some(baseline.clone()), + _ => None, + }; let ( initial_model, initial_totals, @@ -257,6 +272,7 @@ impl CodexParserState { None, ) } + CodexParseMode::ResumeParentBaseline(resume) => return Self::resume_fork(resume), CodexParseMode::InferSubagent { start_ordinal } => ( None, None, @@ -278,6 +294,7 @@ impl CodexParserState { saw_interleaved_totals: false, records: Vec::new(), previous_token_timestamp, + first_token_timestamp: None, previous_token_timestamp_parsed, // A parser always validates a fresh prefix. `None` is only an // input marker for the legacy-cache path, not an output state. @@ -289,9 +306,54 @@ impl CodexParserState { paginated_baseline_checked: false, fork_baseline_ambiguous: false, fork_baseline_inference, + fork_parse_baseline, } } + /// Restore an unfinished fork parse at its saved cursor. A resumable parse + /// never latched an ambiguous baseline: that result is cached unresolved. + fn resume_fork(resume: CodexForkParseResume) -> Self { + let CodexForkParseResume { + start_offset, + paginated_continuation, + inherited_totals, + remaining_inherited_totals, + last_model, + last_totals, + last_token_timestamp, + first_token_timestamp, + token_timestamps_monotonic, + state, + } = resume; + let mut parser = Self::from_mode(CodexParseMode::Standard { + start_offset, + initial_model: last_model, + initial_totals: Some(last_totals), + previous_token_timestamp: last_token_timestamp, + token_timestamps_monotonic, + }); + parser.totals_watermark = state.totals_watermark; + parser.saw_interleaved_totals = state.saw_interleaved_totals; + parser.first_token_timestamp = first_token_timestamp; + parser.fork_baseline = Some(inherited_totals); + // Kept verbatim: `None` means the inherited counters are used up. + parser.remaining_inherited_totals = remaining_inherited_totals; + parser.paginated_continuation = paginated_continuation; + parser.paginated_baseline_checked = state.paginated_baseline_checked; + parser.fork_parse_baseline = Some(state.parse_baseline); + parser + } + + /// State a later pass needs to continue this parent-baseline fork parse. + pub(super) fn fork_resume_state(&self) -> Option { + Some(CodexForkResumeState { + parse_baseline: self.fork_parse_baseline.clone()?, + totals_watermark: self.totals_watermark.clone(), + saw_interleaved_totals: self.saw_interleaved_totals, + paginated_baseline_checked: self.paginated_baseline_checked, + }) + } + pub(super) fn fork_baseline_locally_resolved(&self) -> bool { self.fork_baseline_inference .as_ref() @@ -846,6 +908,9 @@ impl CodexParserState { self.token_timestamps_monotonic = Some(false); } } + if self.first_token_timestamp.is_none() { + self.first_token_timestamp = Some(timestamp.to_string()); + } self.previous_token_timestamp = Some(timestamp.to_string()); self.previous_token_timestamp_parsed = current_parsed; } diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index 6c0ebc5efa..673450dcfe 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -2,11 +2,13 @@ use super::*; use crate::core::{CodexForkAccountingState, CodexSessionLineage, CodexSessionMetadata}; mod cache_days; +mod fork_resume; mod logical_target; mod pending_range; mod reconciliation; mod scan; use cache_days::rebuild_cache_days; +use fork_resume::CodexForkResume; use logical_target::*; use pending_range::{ CodexPendingScanContext, CodexPendingScanDisposition, codex_cache_has_validated_state, @@ -19,7 +21,6 @@ enum CodexAccountingMode { Baseline { baseline: crate::core::CodexTotals, paginated_continuation: bool, - remaining_inherited_totals: Option, provenance: CodexBaselineProvenance, }, InferSubagent { @@ -639,6 +640,19 @@ impl CostScanner { } } + let fork_resume = cached + .as_ref() + .filter(|entry| { + is_fork + && cache_covers_range + && !cached_identity_changed + && identity_matches_cached(entry) + }) + .zip(matching_cached_fork_state) + .and_then(|(entry, state)| { + CodexForkResume::for_entry(path, entry, state, &accounting_mode, size, mtime_ms) + }); + let resumed = fork_resume.is_some(); let parse_target_size = (!accounting_mode.requires_cached_reparse()) .then(|| { cached @@ -661,18 +675,32 @@ impl CostScanner { CodexAccountingMode::Baseline { baseline, paginated_continuation, - remaining_inherited_totals, .. - } => JsonlScanner::parse_codex_file_with_state_bounded_fork_target_with_accounting( - path, - range, - baseline.clone(), - *paginated_continuation, - remaining_inherited_totals.clone(), - cancel, - parse_target_size, - max_bytes_to_read, - ), + } => match fork_resume { + Some(resume) => JsonlScanner::parse_codex_fork_resume( + path, + range, + resume.parse, + cancel, + Some(resume.target_size), + max_bytes_to_read, + ), + None => { + JsonlScanner::parse_codex_file_with_state_bounded_fork_target_with_accounting( + path, + range, + baseline.clone(), + *paginated_continuation, + // A parse from byte zero replays the inherited counters + // itself. Only a resumed parse carries their used-up + // remainder, as upstream restores fork state only then. + None, + cancel, + parse_target_size, + max_bytes_to_read, + ) + } + }, CodexAccountingMode::InferSubagent { start_ordinal } => { JsonlScanner::parse_codex_file_with_inferred_fork_baseline( path, @@ -724,10 +752,18 @@ impl CostScanner { is_complete: false, }; } - let mut days = HashMap::new(); + // A resumed parse returns only the suffix; its prefix is the cached day + // map, which is billed as the non-fork resume path bills it. + let mut days = match (resumed, cached.as_ref()) { + (true, Some(entry)) => entry.days.clone(), + _ => HashMap::new(), + }; merge_codex_records_into_days(&mut days, &parse_result.records); - let (session_cost, has_tokens) = - add_codex_records_to_summary(summary, &parse_result.records, range); + let (session_cost, has_tokens) = if resumed { + add_codex_days_map_to_summary(summary, &days, range) + } else { + add_codex_records_to_summary(summary, &parse_result.records, range) + }; if has_tokens { summary.total_cost_usd += session_cost; summary.sessions_count += 1; @@ -746,8 +782,10 @@ impl CostScanner { history_base_thread_id: history_base_thread_id.clone(), fork_timestamp: codex_fork_timestamp.clone(), inherited_totals: parse_result.fork_baseline.clone(), + first_token_timestamp: parse_result.first_token_timestamp.clone(), remaining_inherited_totals: parse_result.remaining_inherited_totals.clone(), locally_resolved, + resume: parse_result.fork_resume_state.clone(), }) } else { None @@ -773,7 +811,11 @@ impl CostScanner { codex_unresolved_fork_parent: false, }, ); - stats.files_parsed = stats.files_parsed.saturating_add(1); + if resumed { + stats.files_resumed = stats.files_resumed.saturating_add(1); + } else { + stats.files_parsed = stats.files_parsed.saturating_add(1); + } outcome } } diff --git a/rust/src/cost_scanner/codex/fork_resume.rs b/rust/src/cost_scanner/codex/fork_resume.rs new file mode 100644 index 0000000000..81f3c68e69 --- /dev/null +++ b/rust/src/cost_scanner/codex/fork_resume.rs @@ -0,0 +1,75 @@ +//! Bounded fork parses continue at their saved cursor. +//! +//! Upstream 0.67.0 resumes a resolved fork at its cached offset while its +//! parent dependency is unchanged (`canResumeCodexForkAccounting`). A fork that +//! restarted at byte zero on every pass would never finish once it is larger +//! than its byte allowance. + +use super::*; +use crate::core::CodexForkParseResume; + +/// An unfinished fork parse that can continue at its saved cursor. +pub(super) struct CodexForkResume { + pub(super) parse: CodexForkParseResume, + pub(super) target_size: i64, +} + +impl CodexForkResume { + /// Resume point for a cached fork, or `None` when the fork must be parsed + /// from byte zero: no saved parser state, a validated parent that now + /// supplies a different baseline, a rewritten file, or a cursor that is + /// not on a line boundary. + pub(super) fn for_entry( + path: &Path, + entry: &CostUsageFileUsage, + state: &CodexForkAccountingState, + accounting_mode: &CodexAccountingMode, + size: i64, + mtime_ms: i64, + ) -> Option { + let CodexAccountingMode::Baseline { + baseline, + paginated_continuation, + provenance, + .. + } = accounting_mode + else { + return None; + }; + let saved = state.resume.as_ref()?; + // A cached-parent baseline is the saved state itself. + let same_parent_baseline = match provenance { + CodexBaselineProvenance::ValidatedParent { .. } => &saved.parse_baseline == baseline, + CodexBaselineProvenance::CachedValidatedParent => true, + }; + let start_offset = entry.parsed_bytes.unwrap_or(0); + let target_size = codex_resumable_scan_target_size(size, entry)?; + let same_partial = size == entry.size && mtime_ms == entry.mtime_unix_ms; + let growing = size > entry.size; + let resumable = same_parent_baseline + && !state.locally_resolved + && !entry.codex_unresolved_fork_parent + && (same_partial || growing) + && start_offset > 0 + && entry.codex_token_timestamps_monotonic.is_some() + && JsonlScanner::is_line_boundary_offset(path, start_offset); + if !resumable { + return None; + } + Some(Self { + parse: CodexForkParseResume { + start_offset, + paginated_continuation: *paginated_continuation, + inherited_totals: state.inherited_totals.clone()?, + remaining_inherited_totals: state.remaining_inherited_totals.clone(), + last_model: entry.last_model.clone(), + last_totals: entry.last_totals.clone()?, + last_token_timestamp: entry.codex_last_token_timestamp.clone(), + first_token_timestamp: state.first_token_timestamp.clone(), + token_timestamps_monotonic: entry.codex_token_timestamps_monotonic, + state: saved.clone(), + }, + target_size, + }) + } +} diff --git a/rust/src/cost_scanner/codex/logical_target.rs b/rust/src/cost_scanner/codex/logical_target.rs index dd42b96a98..c0a8ad5b3e 100644 --- a/rust/src/cost_scanner/codex/logical_target.rs +++ b/rust/src/cost_scanner/codex/logical_target.rs @@ -21,6 +21,9 @@ pub(super) enum CodexLineageDecision { Unsafe, } +/// Longest ancestor chain a parent baseline is resolved through. +const CODEX_MAX_LINEAGE_DEPTH: usize = 64; + struct CodexLineageNode { path: String, session_id: Option, @@ -446,6 +449,22 @@ impl CodexLineagePlanner { node_index: usize, child_fork_timestamp: Option<&str>, ) -> Option { + self.parent_owner_baseline_at_depth(cache, node_index, child_fork_timestamp, 0) + } + + /// Cumulative counters of one parent at a descendant's fork time. Ancestors + /// are revalidated on the way up, and a chain deeper than + /// [`CODEX_MAX_LINEAGE_DEPTH`] stays unresolved. + fn parent_owner_baseline_at_depth( + &self, + cache: &CostUsageCache, + node_index: usize, + child_fork_timestamp: Option<&str>, + depth: usize, + ) -> Option { + if depth >= CODEX_MAX_LINEAGE_DEPTH { + return None; + } let graph = self.graph()?; let node = graph.nodes.get(node_index)?; if graph.gates[node_index] == CodexLineageGate::Unsafe || !node.may_author_parent { @@ -459,22 +478,23 @@ impl CodexLineagePlanner { return None; } - if super::codex_usage_uses_parent(usage) { - let inherited = usage - .codex_fork_accounting_state - .as_ref()? - .inherited_totals - .as_ref()?; + let fork_origin = if super::codex_usage_uses_parent(usage) { + let state = usage.codex_fork_accounting_state.as_ref()?; + let inherited = state.inherited_totals.as_ref()?; let parent_index = graph.parent_indices[node_index]?; - let baseline = self.parent_owner_baseline( + let baseline = self.parent_owner_baseline_at_depth( cache, parent_index, usage.codex_fork_timestamp.as_deref(), + depth + 1, )?; if &baseline != inherited { return None; } - } + Some((state, inherited)) + } else { + None + }; let metadata = fs::metadata(&node.path).ok()?; let expected_identity = usage.codex_file_identity.as_ref()?; @@ -490,11 +510,28 @@ impl CodexLineagePlanner { { return None; } - let last_totals = usage.last_totals.clone()?; - let last_token_timestamp = usage.codex_last_token_timestamp.as_deref()?; + let child_fork_timestamp = child_fork_timestamp?; - JsonlScanner::codex_timestamp_at_or_before(last_token_timestamp, child_fork_timestamp) - .then_some(last_totals) + if let Some(last_token_timestamp) = usage.codex_last_token_timestamp.as_deref() + && JsonlScanner::codex_timestamp_at_or_before( + last_token_timestamp, + child_fork_timestamp, + ) + { + return usage.last_totals.clone(); + } + + // No token snapshot of this fork exists at the descendant's cutoff, so + // its counters are still the origin it inherited when it forked. + let (state, inherited) = fork_origin?; + let forked_before_cutoff = usage.codex_fork_timestamp.as_deref().is_some_and(|forked| { + JsonlScanner::codex_timestamp_at_or_before(forked, child_fork_timestamp) + }); + let first_own_token_after_cutoff = usage.codex_last_token_timestamp.is_none() + || state.first_token_timestamp.as_deref().is_some_and(|first| { + JsonlScanner::codex_timestamp_before(child_fork_timestamp, first) + }); + (forked_before_cutoff && first_own_token_after_cutoff).then(|| inherited.clone()) } } @@ -512,14 +549,9 @@ impl CodexLineageDecision { let replaces_cached_state = matching_cached_state.is_some_and(|state| { state.locally_resolved || state.inherited_totals.as_ref() != Some(baseline) }); - let cached_parent_state = matching_cached_state.filter(|state| { - !state.locally_resolved && state.inherited_totals.as_ref() == Some(baseline) - }); CodexAccountingMode::Baseline { baseline: baseline.clone(), paginated_continuation, - remaining_inherited_totals: cached_parent_state - .and_then(|state| state.remaining_inherited_totals.clone()), provenance: CodexBaselineProvenance::ValidatedParent { replaces_cached_state, }, @@ -532,7 +564,6 @@ impl CodexLineageDecision { return CodexAccountingMode::Baseline { baseline, paginated_continuation, - remaining_inherited_totals: state.remaining_inherited_totals.clone(), provenance: CodexBaselineProvenance::CachedValidatedParent, }; } diff --git a/rust/src/cost_scanner/tests.rs b/rust/src/cost_scanner/tests.rs index 69f0c8dffa..a4dce082b8 100644 --- a/rust/src/cost_scanner/tests.rs +++ b/rust/src/cost_scanner/tests.rs @@ -3345,6 +3345,12 @@ fn incomplete_or_buffered_empty_codex_fragment_is_not_marked_complete() { #[path = "tests/copied_prefix.rs"] mod copied_prefix; #[cfg(test)] +#[path = "tests/direct_fork.rs"] +mod direct_fork; +#[cfg(test)] +#[path = "tests/fork_resume.rs"] +mod fork_resume; +#[cfg(test)] #[path = "tests/lineage_cache.rs"] mod lineage_cache; #[cfg(test)] diff --git a/rust/src/cost_scanner/tests/direct_fork.rs b/rust/src/cost_scanner/tests/direct_fork.rs new file mode 100644 index 0000000000..209f90397e --- /dev/null +++ b/rust/src/cost_scanner/tests/direct_fork.rs @@ -0,0 +1,357 @@ +//! Direct fork chains: `root -> parent (fork) -> child (fork)`. +//! +//! Port of upstream `CodexDirectForkBaselineTests`. The child of a fork must +//! inherit the cumulative-counter origin of the whole chain, including when +//! the intermediate parent has no token snapshot before the child forks. +use super::*; +use serde_json::{Value, json}; + +const CACHE_SCHEMA_BEFORE_DIRECT_FORK_BASELINES: u64 = 4; + +pub(super) fn timestamp(at: DateTime) -> String { + at.to_rfc3339_opts(chrono::SecondsFormat::Millis, true) +} + +/// Upstream fixture shape: the session (and fork) time lives in the payload. +pub(super) fn session_meta(id: &str, parent: Option<&str>, at: DateTime) -> Value { + let mut payload = json!({ + "id": id, + "source": "vscode", + "thread_source": "user", + "timestamp": timestamp(at), + }); + if let Some(parent) = parent { + payload["forked_from_id"] = json!(parent); + } + json!({"type": "session_meta", "payload": payload}) +} + +/// Upstream fixture shape: a full `event_msg` token_count with a model and +/// both cumulative and last-turn counters. +pub(super) fn token_count(at: DateTime, total_input: i64, last_input: i64) -> Value { + json!({ + "type": "event_msg", + "timestamp": timestamp(at), + "payload": {"type": "token_count", "info": { + "model": "gpt-5.4", + "total_token_usage": {"input_tokens": total_input, "output_tokens": 0}, + "last_token_usage": {"input_tokens": last_input, "output_tokens": 0}, + }}, + }) +} + +pub(super) fn write_rows(dir: &Path, name: &str, rows: &[Value]) -> PathBuf { + std::fs::create_dir_all(dir).unwrap(); + let path = dir.join(name); + let body = rows + .iter() + .map(|row| format!("{row}\n")) + .collect::(); + std::fs::write(&path, body).unwrap(); + path +} + +struct Chain { + _root: tempfile::TempDir, + sessions: PathBuf, + cache_root: PathBuf, + root_file: PathBuf, + parent_file: PathBuf, + child_file: PathBuf, + base: DateTime, +} + +impl Chain { + fn at(&self, seconds: i64) -> DateTime { + self.base + Duration::seconds(seconds) + } + + fn root_rows(&self, total_input: i64) -> Vec { + vec![ + session_meta("root", None, self.at(0)), + token_count(self.at(1), total_input, total_input), + ] + } +} + +/// Upstream timeline: `root` (1000 at t=1) -> `parent` forked at t=2, with an +/// optional own token event at `parent_event_time` (total 1040 / last 40) -> +/// `child` forked at t=4. Upstream covers no parent event and events at t=3 +/// and t=8; t=2 additionally puts the parent's event at its own fork instant. +fn write_chain(parent_event_time: Option) -> Chain { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let day = base.with_timezone(&Local).date_naive(); + let day_dir = sessions + .join(day.format("%Y").to_string()) + .join(day.format("%m").to_string()) + .join(day.format("%d").to_string()); + let mut chain = Chain { + sessions, + cache_root, + root_file: PathBuf::new(), + parent_file: PathBuf::new(), + child_file: PathBuf::new(), + base, + _root: root, + }; + + chain.root_file = write_rows(&day_dir, "root.jsonl", &chain.root_rows(1_000)); + let mut parent_rows = vec![session_meta("parent", Some("root"), chain.at(2))]; + if let Some(parent_event_time) = parent_event_time { + parent_rows.push(token_count(chain.at(parent_event_time), 1_040, 40)); + } + chain.parent_file = write_rows(&day_dir, "parent.jsonl", &parent_rows); + + let (inherited, copied_last) = match parent_event_time { + Some(2 | 3) => (1_040, 40), + None | Some(8) => (1_000, 1_000), + Some(_) => unreachable!("the fixture uses parent events at t=2, 3, or 8"), + }; + chain.child_file = write_rows( + &day_dir, + "child.jsonl", + &[ + session_meta("child", Some("parent"), chain.at(4)), + token_count(chain.at(5), inherited, copied_last), + token_count(chain.at(6), inherited + 20, 20), + token_count(chain.at(7), inherited + 20, 20), + ], + ); + chain +} + +fn scanner(sessions: &Path, cache_root: &Path, bounded: bool) -> CostScanner { + let mut options = CostScanOptions::app_driven(); + if bounded { + // Upstream bounds only the per-refresh byte budget. + options.codex_max_scan_bytes_per_refresh = 512; + } + CostScanner::new(7) + .with_options(options) + .with_cache_root(cache_root) + .with_sessions_dirs(vec![sessions.to_path_buf()]) +} + +fn scan_to_completion(scanner: &CostScanner) -> (CostSummary, CostUsageCache) { + let (summary, cache, _) = scan_to_completion_with_passes(scanner); + (summary, cache) +} + +/// Passes and Codex history bytes read until the scan stops deferring work. +pub(super) struct Progress { + pub(super) passes: usize, + pub(super) bytes_read: u64, +} + +pub(super) fn scan_to_completion_with_passes( + scanner: &CostScanner, +) -> (CostSummary, CostUsageCache, Progress) { + let mut bytes_read = 0; + for passes in 1..=80 { + let (summary, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + bytes_read += stats.codex_bytes_read; + if !cache.codex_scan_incomplete { + return (summary, cache, Progress { passes, bytes_read }); + } + } + panic!("bounded Codex scan never completed"); +} + +pub(super) fn file_len(path: &Path) -> u64 { + std::fs::metadata(path).unwrap().len() +} + +pub(super) fn billed_input(cache: &CostUsageCache, path: &Path) -> i64 { + let usage = &cache.files[&path.to_string_lossy().to_string()]; + assert!( + !usage.codex_unresolved_fork_parent, + "{path:?} is unresolved" + ); + usage + .days + .values() + .flat_map(HashMap::values) + .map(|row| row[0]) + .sum() +} + +fn downgrade_cache_schema(cache_root: &Path) { + let path = cache_root.join("cost-usage").join("codex-v1.json"); + let mut cache: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + cache["codex_cache_schema_version"] = json!(CACHE_SCHEMA_BEFORE_DIRECT_FORK_BASELINES); + std::fs::write(&path, serde_json::to_vec(&cache).unwrap()).unwrap(); +} + +fn assert_chain_billing( + cache: &CostUsageCache, + summary: &CostSummary, + chain: &Chain, + parent_event_time: Option, +) { + assert_eq!(billed_input(cache, &chain.child_file), 20); + let expected = 1_020 + if parent_event_time.is_some() { 40 } else { 0 }; + assert_eq!(summary.input_tokens, expected); +} + +#[test] +fn direct_fork_chain_preserves_cumulative_inheritance() { + for parent_event_time in [None, Some(2_i64), Some(3), Some(8)] { + for bounded in [false, true] { + let chain = write_chain(parent_event_time); + let scanner = scanner(&chain.sessions, &chain.cache_root, bounded); + + // Cold scan, then an unchanged warm scan from the persisted cache. + let (summary, cache, progress) = scan_to_completion_with_passes(&scanner); + if bounded { + assert!( + progress.passes > 1, + "512-byte budget should defer part of the chain" + ); + // Deferred forks continue where they stopped: every history + // byte is read once, as in one unbounded pass. + let chain_bytes: u64 = [&chain.root_file, &chain.parent_file, &chain.child_file] + .into_iter() + .map(|path| file_len(path)) + .sum(); + assert_eq!(progress.bytes_read, chain_bytes); + } + assert_chain_billing(&cache, &summary, &chain, parent_event_time); + let (summary, cache) = scan_to_completion(&scanner); + assert_chain_billing(&cache, &summary, &chain, parent_event_time); + + // A cache written before this correction is rebuilt, not trusted. + downgrade_cache_schema(&chain.cache_root); + let (summary, cache) = scan_to_completion(&scanner); + assert_chain_billing(&cache, &summary, &chain, parent_event_time); + + // A forced rescan is a cold scan against an empty cache. + let forced_cache_root = chain.cache_root.with_file_name("cache-forced"); + let forced = self::scanner(&chain.sessions, &forced_cache_root, bounded); + let (summary, cache) = scan_to_completion(&forced); + assert_chain_billing(&cache, &summary, &chain, parent_event_time); + } + } +} + +#[test] +fn changed_root_revalidates_empty_parent_descendants() { + let chain = write_chain(None); + let scanner = scanner(&chain.sessions, &chain.cache_root, false); + let (summary, cache) = scan_to_completion(&scanner); + assert_chain_billing(&cache, &summary, &chain, None); + let root_dir = chain.root_file.parent().unwrap(); + + // Root rewritten with the same counters: the grandchild is revalidated + // through the empty parent and keeps its own 20 tokens. + write_rows(root_dir, "root.jsonl", &chain.root_rows(1_000)); + let (summary, cache) = scan_to_completion(&scanner); + assert_chain_billing(&cache, &summary, &chain, None); + + // Root history rewritten to a counter above what the child replays. The + // child must not keep the baseline it inherited through the empty + // parent; it fails closed instead of billing against a stale origin. + write_rows(root_dir, "root.jsonl", &chain.root_rows(1_010)); + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + let child = &cache.files[&chain.child_file.to_string_lossy().to_string()]; + assert!(child.codex_unresolved_fork_parent); + assert!(child.days.is_empty()); + assert_eq!(summary.input_tokens, 1_010); +} + +#[test] +fn cyclic_empty_fork_ancestry_stays_unresolved() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let day = base.with_timezone(&Local).date_naive(); + let day_dir = sessions + .join(day.format("%Y").to_string()) + .join(day.format("%m").to_string()) + .join(day.format("%d").to_string()); + let files = [("a", "b"), ("b", "a")].map(|(id, parent)| { + write_rows( + &day_dir, + &format!("{id}.jsonl"), + &[session_meta(id, Some(parent), base)], + ) + }); + let scanner = scanner(&sessions, &cache_root, false); + + for _ in 0..2 { + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(summary.input_tokens, 0); + for file in &files { + let usage = &cache.files[&file.to_string_lossy().to_string()]; + assert!(usage.codex_unresolved_fork_parent); + assert!(usage.days.is_empty()); + } + } +} + +/// `root(1000)` -> `depth` empty forks -> `leaf` (20 own tokens). Returns the +/// scan result for the leaf. +fn scan_empty_fork_ladder(depth: usize) -> (CostSummary, CostUsageCache, PathBuf) { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let base = Utc::now() - Duration::hours(1); + let day = base.with_timezone(&Local).date_naive(); + let day_dir = sessions + .join(day.format("%Y").to_string()) + .join(day.format("%m").to_string()) + .join(day.format("%d").to_string()); + write_rows( + &day_dir, + "s000.jsonl", + &[ + session_meta("s000", None, base), + token_count(base + Duration::seconds(1), 1_000, 1_000), + ], + ); + for level in 1..=depth { + write_rows( + &day_dir, + &format!("s{level:03}.jsonl"), + &[session_meta( + &format!("s{level:03}"), + Some(&format!("s{:03}", level - 1)), + base + Duration::seconds(2), + )], + ); + } + let leaf = write_rows( + &day_dir, + "leaf.jsonl", + &[ + session_meta( + "leaf", + Some(&format!("s{depth:03}")), + base + Duration::seconds(3), + ), + token_count(base + Duration::seconds(4), 1_000, 1_000), + token_count(base + Duration::seconds(5), 1_020, 20), + ], + ); + let scanner = scanner(&sessions, &root.path().join("cache"), false); + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + (summary, cache, leaf) +} + +#[test] +fn empty_fork_ladder_resolves_within_depth_guard() { + let (summary, cache, leaf) = scan_empty_fork_ladder(8); + assert_eq!(billed_input(&cache, &leaf), 20); + assert_eq!(summary.input_tokens, 1_020); +} + +#[test] +fn empty_fork_ladder_beyond_depth_guard_stays_unresolved() { + let (summary, cache, leaf) = scan_empty_fork_ladder(70); + let usage = &cache.files[&leaf.to_string_lossy().to_string()]; + assert!(usage.codex_unresolved_fork_parent); + assert!(usage.days.is_empty()); + assert_eq!(summary.input_tokens, 1_000); +} diff --git a/rust/src/cost_scanner/tests/fork_resume.rs b/rust/src/cost_scanner/tests/fork_resume.rs new file mode 100644 index 0000000000..d463f84cd1 --- /dev/null +++ b/rust/src/cost_scanner/tests/fork_resume.rs @@ -0,0 +1,174 @@ +//! Bounded fork parses continue at their saved cursor. +//! +//! Upstream resumes an unfinished fork at its cached offset while its parent +//! dependency is unchanged (`canResumeCodexForkAccounting`), and parses every +//! other fork from byte zero with fresh accounting state. Either way a fork +//! bills exactly what one uninterrupted cold parse bills. +use super::direct_fork::{ + billed_input, file_len, scan_to_completion_with_passes, session_meta, timestamp, token_count, + write_rows, +}; +use super::*; +use serde_json::{Value, json}; + +/// A token row without cumulative counters. A fork's first rows replay usage +/// it inherited, which the parent baseline absorbs before the fork bills. +fn last_only(at: DateTime, input: i64) -> Value { + json!({ + "type": "event_msg", + "timestamp": timestamp(at), + "payload": {"type": "token_count", "info": { + "model": "gpt-5.4", + "last_token_usage": {"input_tokens": input, "output_tokens": 0}, + }}, + }) +} + +struct Sessions { + root: tempfile::TempDir, + sessions: PathBuf, + day_dir: PathBuf, + base: DateTime, +} + +impl Sessions { + fn new() -> Self { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let base = Utc::now() - Duration::hours(1); + let day = base.with_timezone(&Local).date_naive(); + let day_dir = sessions + .join(day.format("%Y").to_string()) + .join(day.format("%m").to_string()) + .join(day.format("%d").to_string()); + Self { + root, + sessions, + day_dir, + base, + } + } + + fn at(&self, seconds: i64) -> DateTime { + self.base + Duration::seconds(seconds) + } + + fn write(&self, name: &str, rows: &[Value]) -> PathBuf { + write_rows(&self.day_dir, name, rows) + } + + /// Root session holding one cumulative snapshot of `total_input` at t=1. + fn write_root(&self, total_input: i64) -> PathBuf { + self.write( + "root.jsonl", + &[ + session_meta("root", None, self.at(0)), + token_count(self.at(1), total_input, total_input), + ], + ) + } + + /// Scanner with its own cache directory. A one-byte refresh budget reads + /// one JSONL line per refresh. + fn scanner(&self, cache: &str, bytes_per_refresh: Option) -> CostScanner { + let mut options = CostScanOptions::app_driven(); + options.prefer_newest_codex_sessions_first = false; + if let Some(limit) = bytes_per_refresh { + options.codex_max_scan_bytes_per_refresh = limit; + } + CostScanner::new(7) + .with_options(options) + .with_cache_root(self.root.path().join(cache)) + .with_sessions_dirs(vec![self.sessions.clone()]) + } +} + +#[test] +fn bounded_fork_keeps_used_up_inherited_counters() { + let sessions = Sessions::new(); + let root = sessions.write_root(1_000); + let child = sessions.write( + "child.jsonl", + &[ + session_meta("child", Some("root"), sessions.at(2)), + last_only(sessions.at(3), 1_000), + last_only(sessions.at(4), 20), + last_only(sessions.at(5), 20), + ], + ); + + for (cache, bytes_per_refresh) in [("cache", None), ("cache-bounded", Some(1))] { + let scanner = sessions.scanner(cache, bytes_per_refresh); + let (summary, cache, progress) = scan_to_completion_with_passes(&scanner); + // The first row replays the inherited 1000; a resumed pass must not + // refill the used-up counters and absorb the child's own 20 + 20. + assert_eq!(billed_input(&cache, &child), 40); + assert_eq!(summary.input_tokens, 1_040); + assert_eq!(progress.bytes_read, file_len(&root) + file_len(&child)); + } +} + +#[test] +fn partial_fork_restarts_when_its_parent_baseline_changes() { + let sessions = Sessions::new(); + sessions.write_root(1_000); + let child = sessions.write( + "child.jsonl", + &[ + session_meta("child", Some("root"), sessions.at(2)), + token_count(sessions.at(3), 1_000, 1_000), + token_count(sessions.at(4), 1_020, 20), + token_count(sessions.at(5), 1_040, 20), + ], + ); + let child_key = child.to_string_lossy().to_string(); + + // Read one line per refresh until the child holds a saved cursor. + let bounded = sessions.scanner("cache", Some(1)); + let saved_cursor = (0..20).any(|_| { + let (_, _, cache) = bounded.scan_codex_detailed_with_cache(None); + cache + .files + .get(&child_key) + .and_then(|usage| usage.codex_fork_accounting_state.as_ref()) + .is_some_and(|state| state.resume.is_some()) + }); + assert!(saved_cursor, "the child never paused mid-parse"); + + // The root now supplies another baseline. The saved cursor was accounted + // against the old one, so the child is parsed again from byte zero. + sessions.write_root(990); + let (summary, cache, _) = scan_to_completion_with_passes(&sessions.scanner("cache", None)); + let (cold_summary, cold_cache, _) = + scan_to_completion_with_passes(&sessions.scanner("cache-cold", None)); + assert_eq!(billed_input(&cold_cache, &child), 50); + assert_eq!(billed_input(&cache, &child), 50); + assert_eq!(summary.input_tokens, cold_summary.input_tokens); +} + +#[test] +fn grown_fork_reparse_starts_from_the_full_inherited_baseline() { + let sessions = Sessions::new(); + sessions.write_root(1_000); + let mut rows = vec![ + session_meta("child", Some("root"), sessions.at(2)), + last_only(sessions.at(3), 600), + last_only(sessions.at(4), 20), + ]; + let child = sessions.write("child.jsonl", &rows); + let scanner = sessions.scanner("cache", None); + let (_, cache, _) = scan_to_completion_with_passes(&scanner); + // Both rows still replay the 1000 the child inherited. + assert_eq!(billed_input(&cache, &child), 0); + + // A finished fork that grows is parsed again from byte zero, so it must + // start from the whole inherited baseline, not the part left over. + rows.push(last_only(sessions.at(5), 500)); + sessions.write("child.jsonl", &rows); + let (summary, cache, _) = scan_to_completion_with_passes(&scanner); + let (cold_summary, cold_cache, _) = + scan_to_completion_with_passes(&sessions.scanner("cache-cold", None)); + assert_eq!(billed_input(&cold_cache, &child), 120); + assert_eq!(billed_input(&cache, &child), 120); + assert_eq!(summary.input_tokens, cold_summary.input_tokens); +} diff --git a/rust/src/cost_scanner/tests/paginated.rs b/rust/src/cost_scanner/tests/paginated.rs index 592d484823..c02b2088eb 100644 --- a/rust/src/cost_scanner/tests/paginated.rs +++ b/rust/src/cost_scanner/tests/paginated.rs @@ -196,6 +196,55 @@ fn paginated_continuation_raises_inherited_baseline_from_total_last() { ); } +#[test] +fn bounded_paginated_continuation_raises_its_baseline_once() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let base = Utc::now() - Duration::hours(1); + write_codex_fork_session_fixture( + &sessions, + "ancestor.jsonl", + "original-ancestor", + None, + base, + base, + &[1_539_046], + ); + let continuation = write_codex_paginated_continuation_fixture( + &sessions, + "continuation.jsonl", + "thread-session", + "original-ancestor", + "thread-session", + base + Duration::seconds(10), + ); + let mut options = CostScanOptions::app_driven(); + options.prefer_newest_codex_sessions_first = false; + // One JSONL line per refresh, so each page is read by a resumed parse. + options.codex_max_scan_bytes_per_refresh = 1; + let scanner = CostScanner::new(7) + .with_options(options) + .with_cache_root(root.path().join("cache")) + .with_sessions_dirs(vec![sessions]); + + let (summary, cache) = (0..20) + .find_map(|_| { + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + (!cache.codex_scan_incomplete).then_some((summary, cache)) + }) + .expect("bounded Codex scan never completed"); + // The first page raises the inherited baseline; the second page bills its + // cumulative growth, exactly as one unbounded parse does. + assert_eq!(summary.input_tokens, 19_533_671); + assert_eq!(summary.cached_tokens, 17_620_864); + assert_eq!(summary.output_tokens, 53_217); + let usage = &cache.files[&continuation.to_string_lossy().to_string()]; + let state = usage.codex_fork_accounting_state.as_ref().unwrap(); + assert_eq!(state.inherited_totals.as_ref().unwrap().input, 739_823_760); + assert!(state.resume.is_none(), "a finished parse keeps no cursor"); + assert_eq!(cached_input_total(usage), 17_994_625); +} + #[test] fn paginated_history_base_equal_parent_keeps_true_fork_subtraction() { let root = tempfile::tempdir().unwrap();