diff --git a/crates/tokscale-core/src/message_cache.rs b/crates/tokscale-core/src/message_cache.rs index 1d17d09cc..8764d764f 100644 --- a/crates/tokscale-core/src/message_cache.rs +++ b/crates/tokscale-core/src/message_cache.rs @@ -10,7 +10,7 @@ use std::io::{BufReader, BufWriter, Read, Seek, SeekFrom, Write}; use std::path::{Path, PathBuf}; use std::time::UNIX_EPOCH; -const CACHE_SCHEMA_VERSION: u32 = 11; +const CACHE_SCHEMA_VERSION: u32 = 12; const CACHE_FILENAME: &str = "source-message-cache.bin"; const CACHE_LOCK_FILENAME: &str = "source-message-cache.lock"; const MAX_CACHE_FILE_BYTES: u64 = 256 * 1024 * 1024; diff --git a/crates/tokscale-core/src/sessions/codex.rs b/crates/tokscale-core/src/sessions/codex.rs index 0d933dcf7..20cf3b458 100644 --- a/crates/tokscale-core/src/sessions/codex.rs +++ b/crates/tokscale-core/src/sessions/codex.rs @@ -61,6 +61,7 @@ pub struct CodexTokenUsage { pub cached_input_tokens: Option, pub cache_read_input_tokens: Option, pub reasoning_output_tokens: Option, + pub total_tokens: Option, } #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)] @@ -160,6 +161,9 @@ pub(crate) struct CodexParseState { pub session_agent: Option, pub session_workspace_key: Option, pub session_workspace_label: Option, + pub forked_child_waiting_for_turn_context: bool, + pub forked_child_inherited_baseline: Option, + pub forked_child_inherited_reported_total: Option, } #[derive(Debug, Clone)] @@ -261,6 +265,21 @@ fn parse_codex_reader( }; let event_model = payload_model.clone().or(info_model.clone()); + if state.forked_child_waiting_for_turn_context { + if entry.entry_type == "turn_context" { + state.forked_child_waiting_for_turn_context = false; + state.current_model = payload_model.clone(); + handled = true; + } else { + if is_token_count { + if let Some(info) = payload.info.as_ref() { + remember_forked_child_inherited_baseline(&mut state, info); + } + } + continue; + } + } + if !pending_model_messages.is_empty() && event_model.is_none() && !is_token_count @@ -283,6 +302,9 @@ fn parse_codex_reader( } if let Some(ref forked_from_id) = payload.forked_from_id { state.session_forked_from_id = Some(forked_from_id.clone()); + state.forked_child_waiting_for_turn_context = true; + state.forked_child_inherited_baseline = None; + state.forked_child_inherited_reported_total = None; } if let Some(ref provider) = payload.model_provider { state.session_provider = Some(provider.clone()); @@ -337,6 +359,21 @@ fn parse_codex_reader( let total_usage = info.total_token_usage.as_ref().map(CodexTotals::from_usage); let last_usage = info.last_token_usage.as_ref().map(CodexTotals::from_usage); + if forked_child_matches_inherited_baseline( + &state, + info.total_token_usage.as_ref(), + total_usage, + ) { + if let Some(total) = total_usage { + state.previous_totals = Some(total); + } + state.forked_child_inherited_baseline = None; + state.forked_child_inherited_reported_total = None; + continue; + } + state.forked_child_inherited_baseline = None; + state.forked_child_inherited_reported_total = None; + let (tokens, next_totals) = match (total_usage, last_usage, state.previous_totals) { // Both present with previous baseline (standard path) @@ -446,6 +483,13 @@ fn parse_codex_reader( continue; } + if state.forked_child_waiting_for_turn_context { + let mut json_probe = trimmed.as_bytes().to_vec(); + if simd_json::from_slice::(&mut json_probe).is_ok() { + continue; + } + } + let headless_message = parse_codex_headless_line( trimmed, session_id, @@ -590,6 +634,41 @@ pub fn parse_codex_file(path: &Path) -> Vec { parsed.messages } +fn reported_total_tokens(usage: &CodexTokenUsage) -> Option { + usage.total_tokens.filter(|total| *total >= 0) +} + +fn remember_forked_child_inherited_baseline(state: &mut CodexParseState, info: &CodexInfo) { + let Some(total_usage) = info.total_token_usage.as_ref() else { + return; + }; + + let totals = CodexTotals::from_usage(total_usage); + state.previous_totals = Some(totals); + state.forked_child_inherited_baseline = Some(totals); + state.forked_child_inherited_reported_total = reported_total_tokens(total_usage); +} + +fn forked_child_matches_inherited_baseline( + state: &CodexParseState, + total_usage: Option<&CodexTokenUsage>, + totals: Option, +) -> bool { + if let (Some(usage), Some(baseline)) = + (total_usage, state.forked_child_inherited_reported_total) + { + if reported_total_tokens(usage) == Some(baseline) { + return true; + } + } + + if let (Some(totals), Some(baseline)) = (totals, state.forked_child_inherited_baseline) { + return totals == baseline; + } + + false +} + pub(crate) fn parse_codex_file_incremental( path: &Path, start_offset: u64, @@ -1392,6 +1471,80 @@ mod tests { assert!(messages[0].dedup_key.is_some()); } + #[test] + fn test_forked_child_ignores_inherited_records_before_turn_context() { + let file = create_test_file(concat!( + r#"{"timestamp":"2026-05-05T21:51:57.991Z","type":"session_meta","payload":{"id":"child-session","forked_from_id":"parent-session","source":{"subagent":{"thread_spawn":{"parent_thread_id":"parent-session","depth":1}}},"model_provider":"openai","agent_nickname":"worker","cwd":"/repo-child"}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:51:57.992Z","type":"session_meta","payload":{"id":"parent-session","source":"interactive","model_provider":"azure","agent_nickname":"parent","cwd":"/repo-parent"}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:51:57.993Z","type":"event_msg","payload":{"type":"user_message","message":"parent prompt copied into child log"}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:51:57.994Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":116000,"cached_input_tokens":114000,"output_tokens":1000,"total_tokens":117000},"last_token_usage":{"input_tokens":73000,"cached_input_tokens":72000,"output_tokens":500,"total_tokens":73500}}}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:51:58.947Z","type":"turn_context","payload":{"model":"gpt-5.5","cwd":"/repo-child"}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:51:58.948Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":116000,"cached_input_tokens":114000,"output_tokens":1000,"total_tokens":117000},"last_token_usage":{"input_tokens":73000,"cached_input_tokens":72000,"output_tokens":500,"total_tokens":73500}}}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:51:59.253Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":117500,"cached_input_tokens":115000,"output_tokens":1200,"reasoning_output_tokens":50,"total_tokens":118700},"last_token_usage":{"input_tokens":1500,"cached_input_tokens":1000,"output_tokens":200,"reasoning_output_tokens":50,"total_tokens":1700}}}}"#, + "\n" + )); + + let messages = parse_codex_file(file.path()); + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].model_id, "gpt-5.5"); + assert_eq!(messages[0].provider_id, "openai"); + assert_eq!(messages[0].agent.as_deref(), Some("worker")); + assert_eq!(messages[0].workspace_key.as_deref(), Some("/repo-child")); + assert_eq!(messages[0].tokens.input, 500); + assert_eq!(messages[0].tokens.cache_read, 1000); + assert_eq!(messages[0].tokens.output, 200); + assert_eq!(messages[0].tokens.reasoning, 50); + } + + #[test] + fn test_forked_child_incremental_state_skips_inherited_prefix() { + let file = create_test_file(concat!( + r#"{"timestamp":"2026-05-05T21:51:57.991Z","type":"session_meta","payload":{"id":"child-session","forked_from_id":"parent-session","source":{"subagent":{"thread_spawn":{"parent_thread_id":"parent-session","depth":1}}},"model_provider":"openai","agent_nickname":"worker","cwd":"/repo-child"}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:51:57.992Z","type":"session_meta","payload":{"id":"parent-session","source":"interactive","model_provider":"azure","agent_nickname":"parent","cwd":"/repo-parent"}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:51:57.994Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":116000,"cached_input_tokens":114000,"output_tokens":1000,"total_tokens":117000},"last_token_usage":{"input_tokens":73000,"cached_input_tokens":72000,"output_tokens":500,"total_tokens":73500}}}}"#, + "\n" + )); + let prefix_size = file.as_file().metadata().unwrap().len(); + let prefix = parse_codex_file_incremental(file.path(), 0, CodexParseState::default()); + + assert!(prefix.parse_succeeded); + assert!(!prefix.unresolved_model_events); + assert!(prefix.messages.is_empty()); + + let appended = concat!( + r#"{"timestamp":"2026-05-05T21:51:58.947Z","type":"turn_context","payload":{"model":"gpt-5.5","cwd":"/repo-child"}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:51:58.948Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":116000,"cached_input_tokens":114000,"output_tokens":1000,"total_tokens":117000},"last_token_usage":{"input_tokens":73000,"cached_input_tokens":72000,"output_tokens":500,"total_tokens":73500}}}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:51:59.253Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":117500,"cached_input_tokens":115000,"output_tokens":1200,"reasoning_output_tokens":50,"total_tokens":118700},"last_token_usage":{"input_tokens":1500,"cached_input_tokens":1000,"output_tokens":200,"reasoning_output_tokens":50,"total_tokens":1700}}}}"#, + "\n" + ); + let mut reopened = file.reopen().unwrap(); + reopened.seek(SeekFrom::End(0)).unwrap(); + reopened.write_all(appended.as_bytes()).unwrap(); + reopened.flush().unwrap(); + + let incremental = + parse_codex_file_incremental(file.path(), prefix_size, prefix.state.clone()); + let full = parse_codex_file(file.path()); + + assert_eq!(incremental.messages, full); + assert_eq!(incremental.messages.len(), 1); + assert_eq!(incremental.messages[0].tokens.input, 500); + assert_eq!(incremental.messages[0].tokens.cache_read, 1000); + assert_eq!(incremental.messages[0].tokens.output, 200); + assert_eq!(incremental.messages[0].tokens.reasoning, 50); + } + #[test] fn test_session_meta_cwd_sets_workspace_metadata() { let line1 = r#"{"timestamp":"2026-01-01T00:00:00Z","type":"session_meta","payload":{"source":"interactive","cwd":"/Users/alice/demo-repo"}}"#; @@ -1490,6 +1643,7 @@ mod tests { cached_input_tokens: Some(10), cache_read_input_tokens: Some(20), reasoning_output_tokens: Some(5), + total_tokens: None, }; let totals = CodexTotals::from_usage(&usage); assert_eq!(totals.cached, 20);