diff --git a/crates/tokscale-core/src/message_cache.rs b/crates/tokscale-core/src/message_cache.rs index 3187e687b..ca772eee5 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 = 15; +const CACHE_SCHEMA_VERSION: u32 = 16; 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 2f865da01..94ba0598e 100644 --- a/crates/tokscale-core/src/sessions/codex.rs +++ b/crates/tokscale-core/src/sessions/codex.rs @@ -39,6 +39,11 @@ pub struct CodexPayload { pub model_provider: Option, /// Agent name from session_meta pub agent_nickname: Option, + /// Free-text body of an `event_msg` `user_message` payload. Used to detect + /// human turn boundaries: real human input is plain text, whereas + /// system-injected context (``, ``, + /// ``, …) begins with `<`. + pub message: Option, } #[derive(Debug, Deserialize)] @@ -166,6 +171,11 @@ pub(crate) struct CodexParseState { pub forked_child_waiting_for_turn_context: bool, pub forked_child_inherited_baseline: Option, pub forked_child_inherited_reported_total: Option, + /// Set when a human `user_message` event is seen; consumed by the next + /// token_count-derived message to mark it as a turn start. `#[serde(default)]` + /// keeps a pending turn alive across incremental re-parses of appended chunks. + #[serde(default)] + pub pending_turn_start: bool, } #[derive(Debug, Clone)] @@ -336,6 +346,26 @@ fn parse_codex_reader( handled = true; } + // A human `user_message` event starts a new turn. The event + // itself carries no tokens, so we defer the flag to the next + // token_count-derived message (the assistant's reply). This + // counts `codex exec` one-shots too: they are headless but still + // carry a real human prompt, so each is one turn. Only + // system-injected messages (leading `<`, e.g. + // , ) are excluded as + // non-human input. Forked-child replays of the parent prompt + // arrive before turn_context and are skipped by the + // `forked_child_waiting_for_turn_context` branch above, so they + // never reach here. + if entry.entry_type == "event_msg" + && payload.payload_type.as_deref() == Some("user_message") + { + if codex_message_is_human_turn(payload.message.as_deref()) { + state.pending_turn_start = true; + } + handled = true; + } + // Process token_count events if is_token_count { let info = match payload.info { @@ -456,6 +486,13 @@ fn parse_codex_reader( agent, ); message.duration_ms = duration_ms; + // Apply a deferred human-turn marker from a preceding + // user_message to this assistant reply — the first + // token-bearing message after the human input. + if state.pending_turn_start { + message.is_turn_start = true; + state.pending_turn_start = false; + } if parsed_timestamp.is_some() { if let Some(model) = model.as_deref() { set_codex_dedup_key(&mut message, model); @@ -863,12 +900,63 @@ fn extract_timestamp_from_value(value: &Value) -> Option { .and_then(parse_timestamp_value) } +/// Prefixes Codex prepends to context it injects as `user_message` events. +/// These are the bodies that must NOT be counted as human turns. +const CODEX_SYSTEM_INJECTED_PREFIXES: [&str; 3] = [ + "", + "", + "", +]; + +/// Returns true when a Codex `user_message` payload represents real human input +/// rather than system-injected context. Codex stores the body as a plain string +/// in `payload.message`; the harness injects context blocks that open with one of +/// the known tags in [`CODEX_SYSTEM_INJECTED_PREFIXES`] after trimming. Matching +/// those specific prefixes — rather than any leading `<` — avoids dropping +/// legitimate human prompts that happen to start with markup (asking about a +/// `
`, pasting an XML snippet, etc.). The `kind` field can't be used to +/// distinguish them: both human and injected bodies appear as `kind:"plain"` or +/// with no `kind` at all. +fn codex_message_is_human_turn(message: Option<&str>) -> bool { + match message { + Some(text) => { + let trimmed = text.trim_start(); + !CODEX_SYSTEM_INJECTED_PREFIXES + .iter() + .any(|prefix| trimmed.starts_with(prefix)) + } + None => false, + } +} + #[cfg(test)] mod tests { use super::*; use std::io::{BufRead, Cursor, Error, ErrorKind, Seek, SeekFrom, Write}; use tempfile::NamedTempFile; + #[test] + fn codex_human_turn_matches_only_known_system_tags() { + // Real human prompts that happen to start with markup must still count. + assert!(codex_message_is_human_turn(Some( + "how do I center a
?" + ))); + assert!(codex_message_is_human_turn(Some("
hi
"))); + assert!(codex_message_is_human_turn(Some(" plain question"))); + // Known system-injected context blocks are not human turns. + assert!(!codex_message_is_human_turn(Some( + "cwd=/tmp" + ))); + assert!(!codex_message_is_human_turn(Some( + " be concise" + ))); + assert!(!codex_message_is_human_turn(Some( + "do X" + ))); + // A missing body is never a human turn. + assert!(!codex_message_is_human_turn(None)); + } + fn create_test_file(content: &str) -> NamedTempFile { let mut file = NamedTempFile::new().unwrap(); file.write_all(content.as_bytes()).unwrap(); @@ -1768,4 +1856,118 @@ mod tests { assert_eq!(parsed.messages.len(), 1); assert_eq!(parsed.messages[0].model_id, "gpt-5.5"); } + + #[test] + fn test_user_message_marks_next_token_count_as_turn_start() { + let content = [ + r#"{"timestamp":"2026-01-01T00:00:01Z","type":"turn_context","payload":{"model":"gpt-5.2"}}"#, + r#"{"timestamp":"2026-01-01T00:00:02Z","type":"event_msg","payload":{"type":"user_message","message":"continue please"}}"#, + r#"{"timestamp":"2026-01-01T00:00:03Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, + r#"{"timestamp":"2026-01-01T00:00:04Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":20,"cached_input_tokens":4,"output_tokens":6},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, + ] + .join("\n"); + let file = create_test_file(&content); + + let messages = parse_codex_file(file.path()); + + assert_eq!(messages.len(), 2); + assert!( + messages[0].is_turn_start, + "first reply after a human user_message is a turn start" + ); + assert!( + !messages[1].is_turn_start, + "a later reply with no new user_message is not a turn start" + ); + } + + #[test] + fn test_xml_user_message_does_not_mark_turn_start() { + let content = [ + r#"{"timestamp":"2026-01-01T00:00:01Z","type":"turn_context","payload":{"model":"gpt-5.2"}}"#, + r#"{"timestamp":"2026-01-01T00:00:02Z","type":"event_msg","payload":{"type":"user_message","message":"\n\n /tmp\n"}}"#, + r#"{"timestamp":"2026-01-01T00:00:03Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, + ] + .join("\n"); + let file = create_test_file(&content); + + let messages = parse_codex_file(file.path()); + + assert_eq!(messages.len(), 1); + assert!( + !messages[0].is_turn_start, + "a system-injected <...> message is not a human turn" + ); + } + + #[test] + fn test_exec_user_message_still_marks_turn_start() { + // A `codex exec` one-shot is headless but still carries a real human + // prompt, so it counts as exactly one turn (verified against a real + // `codex exec` session: 1 user_message -> turn_count 1). + let content = [ + r#"{"timestamp":"2026-01-01T00:00:00Z","type":"session_meta","payload":{"source":"exec"}}"#, + r#"{"timestamp":"2026-01-01T00:00:01Z","type":"turn_context","payload":{"model":"gpt-5.2"}}"#, + r#"{"timestamp":"2026-01-01T00:00:02Z","type":"event_msg","payload":{"type":"user_message","message":"hello"}}"#, + // A real `codex exec` interleaves an agent_message between the user + // prompt and the token_count; the deferred turn flag must survive it. + r#"{"timestamp":"2026-01-01T00:00:02Z","type":"event_msg","payload":{"type":"agent_message","message":"hi"}}"#, + r#"{"timestamp":"2026-01-01T00:00:03Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, + ] + .join("\n"); + let file = create_test_file(&content); + + let messages = parse_codex_file(file.path()); + + assert_eq!(messages.len(), 1); + assert!( + messages[0].is_turn_start, + "an exec one-shot with a human prompt counts as one turn" + ); + assert_eq!(messages[0].agent.as_deref(), Some("headless")); + } + + #[test] + fn test_incremental_parse_preserves_pending_turn_start() { + let content = [ + r#"{"timestamp":"2026-01-01T00:00:01Z","type":"turn_context","payload":{"model":"gpt-5.2"}}"#, + r#"{"timestamp":"2026-01-01T00:00:02Z","type":"event_msg","payload":{"type":"user_message","message":"hello"}}"#, + "", + ] + .join("\n"); + let file = create_test_file(&content); + let initial_size = file.as_file().metadata().unwrap().len(); + + let initial = parse_codex_file_incremental(file.path(), 0, CodexParseState::default()); + assert!( + initial.messages.is_empty(), + "no token_count yet, so no message" + ); + assert!( + initial.state.pending_turn_start, + "a pending turn survives a chunk that ends before the token_count" + ); + + let appended = format!( + "{}\n", + r#"{"timestamp":"2026-01-01T00:00:03Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"# + ); + 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(), initial_size, initial.state.clone()); + + assert_eq!(incremental.messages.len(), 1); + assert!( + incremental.messages[0].is_turn_start, + "the deferred turn applies to the message parsed in the next chunk" + ); + assert!( + !incremental.state.pending_turn_start, + "the pending flag is consumed once applied" + ); + } }