diff --git a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst index 7893421a83..c43dd8756d 100644 Binary files a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst and b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst differ diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history.rs b/codex-rs/app-server-protocol/src/protocol/thread_history.rs index 7a1d70c85f..40d3d9fc7e 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -396,6 +396,7 @@ impl ThreadHistoryBuilder { RolloutItem::InterAgentCommunication(_) | RolloutItem::InterAgentCommunicationMetadata { .. } | RolloutItem::TurnContext(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::WorldState(_) | RolloutItem::RealtimeItem(_) | RolloutItem::SecurityRiskScore(_) @@ -4060,6 +4061,8 @@ mod tests { first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent { turn_id: "turn-compact".into(), diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history_projection.rs b/codex-rs/app-server-protocol/src/protocol/thread_history_projection.rs index 6bb1138dc0..7fc8648b32 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history_projection.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history_projection.rs @@ -82,6 +82,7 @@ pub fn project_rollout_line(line: &RolloutLine) -> ThreadHistoryChangeSet { | RolloutItem::InterAgentCommunicationMetadata { .. } | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::WorldState(_) | RolloutItem::RealtimeItem(_) | RolloutItem::SecurityRiskScore(_) diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history_projection_tests.rs b/codex-rs/app-server-protocol/src/protocol/thread_history_projection_tests.rs index 1c1cda3129..4eb2c02267 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history_projection_tests.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history_projection_tests.rs @@ -205,6 +205,8 @@ fn ignores_legacy_abort_without_turn_id_and_context_only_records() { first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, })); let security_risk = project(RolloutItem::SecurityRiskScore(SecurityRiskScore { scores: BTreeMap::from([("action_risk".to_string(), 0.92)]), diff --git a/codex-rs/app-server/src/request_processors/feedback_processor.rs b/codex-rs/app-server/src/request_processors/feedback_processor.rs index d2948c4ade..c68aa1d23e 100644 --- a/codex-rs/app-server/src/request_processors/feedback_processor.rs +++ b/codex-rs/app-server/src/request_processors/feedback_processor.rs @@ -704,6 +704,7 @@ mod tests { ordinal: None, item: RolloutItem::TurnContext(TurnContextItem { turn_id: Some((*turn_id).to_string()), + root_turn_id: None, cwd: AbsolutePathBuf::from_absolute_path(tempdir.path()) .expect("absolute feedback rollout directory"), workspace_roots: None, diff --git a/codex-rs/app-server/src/request_processors/persisted_resume_settings_tests.rs b/codex-rs/app-server/src/request_processors/persisted_resume_settings_tests.rs index 167e770af5..792fc344a1 100644 --- a/codex-rs/app-server/src/request_processors/persisted_resume_settings_tests.rs +++ b/codex-rs/app-server/src/request_processors/persisted_resume_settings_tests.rs @@ -62,6 +62,7 @@ fn turn_context_item( ) -> RolloutItem { RolloutItem::TurnContext(TurnContextItem { turn_id: Some(turn_id.to_string()), + root_turn_id: None, cwd: cwd(), workspace_roots: Some(vec![cwd()]), current_date: None, diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume.rs b/codex-rs/app-server/tests/suite/v2/thread_resume.rs index 6338c57b45..485336edd7 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -201,6 +201,8 @@ async fn thread_resume_paginated_model_context_preserves_original_metadata() -> first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), ) .await?; @@ -3238,7 +3240,10 @@ async fn thread_goal_keeps_original_root_until_external_objective_edit() -> Resu reopened_original_request["client_metadata"]["turn_id"].as_str(), Some(original_turn.turn.id.as_str()) ); - responses::assert_root_turn(&reopened_original_request, /*expected*/ None)?; + responses::assert_root_turn( + &reopened_original_request, + Some(original_turn.turn.id.as_str()), + )?; let intervening_request = serde_json::from_slice::(&requests[3])?; let intervening_turn_id = intervening_request["client_metadata"]["turn_id"] .as_str() @@ -3265,7 +3270,7 @@ async fn thread_goal_keeps_original_root_until_external_objective_edit() -> Resu reopened_request["client_metadata"]["turn_id"].as_str(), Some(edited_turn_id) ); - responses::assert_root_turn(&reopened_request, /*expected*/ None)?; + responses::assert_root_turn(&reopened_request, Some(original_turn.turn.id.as_str()))?; let continuation_request = serde_json::from_slice::(&requests[7])?; assert_ne!( continuation_request["client_metadata"]["turn_id"].as_str(), diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 8dbd4d329a..141aaddc6e 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -116,8 +116,7 @@ pub(crate) struct ListedAgent { /// which keeps the registry scoped to that root thread rather than the entire `ThreadManager`. #[derive(Clone)] pub(crate) struct AgentControl { - /// ID shared by the whole agent control session. This means every sub-agents from a common - /// root share the same session ID. + /// session_id is equal to the root thread's ID. session_id: SessionId, /// Weak handle back to the global thread registry/state. /// This is `Weak` to avoid reference cycles and shadow persistence of the form diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index 2f31539798..b0dbd8a162 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -96,6 +96,8 @@ fn keep_forked_rollout_item(item: &RolloutItem, preserve_reference_context_item: // from the parent's durable baseline. Truncated forks drop part of that prompt, // so they must rebuild context on their first child turn. RolloutItem::TurnContext(_) | RolloutItem::WorldState(_) => preserve_reference_context_item, + // Child threads inherit model context, not the parent's cumulative usage state. + RolloutItem::TokenUsageRecord(_) => false, RolloutItem::Compacted(_) | RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) => true, } } @@ -992,6 +994,8 @@ impl AgentControl { retain_forked_item(response_item, &mut replaced_parent_developer_instructions) } RolloutItem::Compacted(compacted) => { + // This checkpoint belongs to the inherited parent prefix. + compacted.latest_token_usage_record = None; if let Some(replacement_history) = compacted.replacement_history.as_mut() { // Matches before this checkpoint cannot survive its replacement history. replaced_parent_developer_instructions = false; @@ -1016,7 +1020,7 @@ impl AgentControl { | RolloutItem::TurnContext(_) | RolloutItem::InterAgentCommunication(_) | RolloutItem::InterAgentCommunicationMetadata { .. } => true, - RolloutItem::SecurityRiskScore(_) => false, + RolloutItem::TokenUsageRecord(_) | RolloutItem::SecurityRiskScore(_) => false, } }); // Full forks reuse the parent's reference context instead of rebuilding it. If that diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 99200ac47a..2b2d81eda5 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -56,6 +56,8 @@ use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadSettingsAppliedEvent; use codex_protocol::protocol::ThreadSettingsSnapshot; +use codex_protocol::protocol::TokenUsage; +use codex_protocol::protocol::TokenUsageRecord; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnCompleteEvent; @@ -276,6 +278,7 @@ async fn persisted_originator(thread: &CodexThread) -> String { | RolloutItem::WorldState(_) | RolloutItem::RealtimeItem(_) | RolloutItem::SecurityRiskScore(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::TurnContext(_) => None, }) .expect("session metadata should be persisted") @@ -1382,6 +1385,124 @@ async fn spawn_agent_without_fork_from_paginated_parent_stays_fresh_and_paginate .expect("parent shutdown should submit"); } +#[tokio::test] +async fn spawn_agent_fork_drops_inherited_token_usage_state() { + let harness = AgentControlHarness::new().await; + let (parent_thread_id, parent_thread) = harness.start_paginated_thread().await; + let parent_usage = TokenUsage { + total_tokens: 120, + ..TokenUsage::default() + }; + let parent_record = TokenUsageRecord { + thread_id: parent_thread_id, + turn_id: "parent-turn".to_string(), + session_id: parent_thread.session.session_id(), + root_turn_id: "parent-turn".to_string(), + response_id: "parent-response".to_string(), + usage: parent_usage.clone(), + turn_token_usage: parent_usage.clone(), + thread_token_usage: parent_usage, + }; + let parent_spawn_call_id = "spawn-call-token-usage".to_string(); + parent_thread + .session + .persist_rollout_items(&[ + RolloutItem::Compacted(CompactedItem { + message: String::new(), + replacement_history: Some(vec![user_message("compacted parent context").into()]), + mcp_resource_origins: None, + window_number: None, + first_window_id: None, + previous_window_id: None, + window_id: None, + compaction_response_id: None, + latest_token_usage_record: Some(parent_record.clone()), + }), + RolloutItem::TokenUsageRecord(parent_record), + rollout_response_item(spawn_agent_call(&parent_spawn_call_id)), + ]) + .await; + + let child_thread_id = harness + .spawn_anonymous_child( + parent_thread_id, + SpawnAgentOptions { + fork_parent_spawn_call_id: Some(parent_spawn_call_id), + fork_mode: Some(SpawnAgentForkMode::FullHistory), + ..Default::default() + }, + ) + .await; + let child_thread = harness + .manager + .get_thread(child_thread_id) + .await + .expect("child thread should be registered"); + + let child_usage = TokenUsage { + total_tokens: 80, + ..TokenUsage::default() + }; + let turn_context = child_thread.session.new_default_turn().await; + child_thread + .session + .record_observed_response_completed( + turn_context.as_ref(), + "child-response", + Some(&child_usage), + /*usage_metadata*/ None, + ) + .await; + child_thread + .flush_rollout() + .await + .expect("child rollout should flush"); + let rollout_path = child_thread + .rollout_path() + .expect("child rollout should exist"); + let lines = std::fs::read_to_string(&rollout_path) + .expect("read child rollout") + .lines() + .map(|line| serde_json::from_str::(line).expect("parse rollout line")) + .collect::>(); + assert!( + !lines.iter().any(|line| { + matches!( + &line.item, + RolloutItem::TokenUsageRecord(record) if record.thread_id == parent_thread_id + ) + }), + "child rollout should not inherit parent token usage records" + ); + assert!( + lines.iter().all(|line| { + !matches!( + &line.item, + RolloutItem::Compacted(compacted) + if compacted.latest_token_usage_record.is_some() + ) + }), + "child rollout should not inherit parent token usage checkpoints" + ); + let child_record = lines.iter().rev().find_map(|line| match &line.item { + RolloutItem::TokenUsageRecord(record) => Some(record), + _ => None, + }); + assert_eq!( + child_record, + Some(&TokenUsageRecord { + thread_id: child_thread_id, + turn_id: turn_context.sub_id.clone(), + session_id: child_thread.session.session_id(), + root_turn_id: turn_context.sub_id.clone(), + response_id: "child-response".to_string(), + usage: child_usage.clone(), + turn_token_usage: child_usage.clone(), + thread_token_usage: child_usage, + }) + ); +} + #[tokio::test] async fn spawn_agent_numeric_fork_from_compacted_paginated_parent_clamps_to_provable_turns() { let harness = AgentControlHarness::new().await; @@ -1409,6 +1530,8 @@ async fn spawn_agent_numeric_fork_from_compacted_paginated_parent_clamps_to_prov first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), rollout_response_item(ResponseItem::Message { id: None, @@ -1905,6 +2028,8 @@ async fn spawn_agent_fork_strips_parent_usage_hints_from_compacted_history() { first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), RolloutItem::TurnContext(turn_context.to_turn_context_item()), rollout_response_item(spawn_agent_call(&parent_spawn_call_id)), @@ -2087,6 +2212,8 @@ async fn spawn_agent_full_fork_restores_instructions_after_compaction_discards_p first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), RolloutItem::TurnContext(turn_context.to_turn_context_item()), rollout_response_item(spawn_agent_call(&parent_spawn_call_id)), @@ -2240,6 +2367,8 @@ async fn spawn_agent_full_fork_legacy_compaction_rebuilds_child_instructions_onc first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), ]; if let Some(instructions) = parent_developer_instructions { diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 563d73c7ec..3ef4c219f9 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -390,8 +390,15 @@ impl CodexThread { trace, cyber_access_program, } = request; + let root_turn_id = self + .session + .reference_context_item() + .await + .filter(|context| context.turn_id.as_deref() == Some(turn_id.as_str())) + .and_then(|context| context.root_turn_id); let start_options = TurnStartOptions { cyber_access_program, + root_turn_id, ..Default::default() }; match self @@ -514,18 +521,6 @@ impl CodexThread { task.turn_context.turn_metadata_state.root_turn_id() } - /// Invalidates the trusted root when the expected turn is currently active. - pub async fn invalidate_turn_lineage(&self, expected_turn_id: &str) { - let active = self.session.active_turn.lock().await; - if let Some(task) = active.as_ref().and_then(|turn| turn.task.as_ref()) - && task.turn_context.sub_id == expected_turn_id - { - task.turn_context - .turn_metadata_state - .mark_root_turn_ambiguous(); - } - } - pub async fn set_app_server_client_info( &self, app_server_client_name: Option, diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index 1651ab1b9e..49306c357c 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -46,7 +46,6 @@ use codex_protocol::models::InternalChatMessageMetadataPassthrough; use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; -use codex_protocol::protocol::RawResponseCompletedEvent; use codex_protocol::protocol::TurnStartedEvent; use codex_protocol::protocol::WarningEvent; use codex_protocol::user_input::UserInput; @@ -86,6 +85,7 @@ pub(crate) struct CompactedHistoryMetadata { pub(crate) message: String, pub(crate) window_number: u64, pub(crate) window_ids: AutoCompactWindowIds, + pub(crate) compaction_response_id: Option, } pub(crate) async fn build_compaction_initial_context( @@ -273,7 +273,7 @@ async fn run_compact_task_inner_impl( ) .await; - loop { + let compaction_response_id = loop { // Clone is required because of the loop let turn_input = history .clone() @@ -294,8 +294,8 @@ async fn run_compact_task_inner_impl( .await; match attempt_result { - Ok(()) => { - break; + Ok(response_id) => { + break response_id; } Err(err) if matches!( @@ -347,7 +347,7 @@ async fn run_compact_task_inner_impl( } } } - } + }; let history_snapshot = sess.clone_history().await; let history_items = history_snapshot.annotated_items(); @@ -384,6 +384,7 @@ async fn run_compact_task_inner_impl( message: summary_text, window_number, window_ids, + compaction_response_id: Some(compaction_response_id), }, ) .await; @@ -738,7 +739,7 @@ async fn drain_to_completed( client_session: &mut ModelClientSession, responses_metadata: &CodexResponsesMetadata, prompt: &Prompt, -) -> CodexResult<()> { +) -> CodexResult { let mut stream = client_session .stream( prompt, @@ -777,18 +778,16 @@ async fn drain_to_completed( usage_metadata, .. }) => { - sess.send_event( + sess.record_observed_response_completed( turn_context, - EventMsg::RawResponseCompleted(RawResponseCompletedEvent { - response_id, - token_usage: token_usage.clone(), - usage_metadata, - }), + &response_id, + token_usage.as_ref(), + usage_metadata.as_ref(), ) .await; sess.update_token_usage_info(turn_context, token_usage.as_ref()) .await?; - return Ok(()); + return Ok(response_id); } Ok(_) => continue, Err(e) => return Err(e), diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index ea1d5a372d..3d16dafd92 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -300,6 +300,7 @@ async fn run_remote_compact_task_inner_impl( message: String::new(), window_number: new_window_number, window_ids: new_window_ids, + compaction_response_id: None, }, ) .await; diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index e36d11387f..def91ff2f4 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -38,7 +38,6 @@ use codex_context_fragments::to_annotated_content; use codex_features::Feature; use codex_history::CodexHarnessMetadata; use codex_history::ResponseItemEnvelope; -use codex_protocol::ResponseUsageMetadata; use codex_protocol::error::CodexErr; use codex_protocol::error::CodexErrorDetails; use codex_protocol::error::Result as CodexResult; @@ -298,6 +297,7 @@ async fn run_remote_compact_task_inner_impl( prompt_input, prompt_input_metadata, compaction_output, + compaction_response_id, token_usage, owned_client_session: _owned_client_session, } = attempt; @@ -350,6 +350,7 @@ async fn run_remote_compact_task_inner_impl( message: String::new(), window_number: new_window_number, window_ids: new_window_ids, + compaction_response_id: Some(compaction_response_id), }, ) .await; @@ -364,7 +365,6 @@ struct RemoteCompactionV2Output { compaction_output: ResponseItem, response_id: String, token_usage: Option, - usage_metadata: Option, } async fn run_remote_compaction_request_v2( @@ -395,7 +395,7 @@ async fn run_remote_compaction_request_v2( ) .await { - Ok(stream) => collect_compaction_output(stream).await, + Ok(stream) => collect_compaction_output(sess, turn_context, stream).await, Err(err) => Err(err), }; @@ -419,15 +419,15 @@ async fn run_remote_compaction_request_v2( } async fn collect_compaction_output( + sess: &Session, + turn_context: &TurnContext, mut stream: ResponseStream, ) -> CodexResult { let mut output_item_count = 0usize; let mut compaction_count = 0usize; let mut compaction_output = None; - let mut saw_completed = false; let mut completed_response_id = None; let mut completed_token_usage = None; - let mut completed_usage_metadata = None; while let Some(event) = stream.next().await { match event? { ResponseEvent::OutputItemDone(item) => { @@ -445,21 +445,26 @@ async fn collect_compaction_output( usage_metadata, .. } => { - saw_completed = true; + sess.record_observed_response_completed( + turn_context, + &response_id, + token_usage.as_ref(), + usage_metadata.as_ref(), + ) + .await; completed_response_id = Some(response_id); completed_token_usage = token_usage; - completed_usage_metadata = usage_metadata; break; } _ => {} } } - if !saw_completed { + let Some(response_id) = completed_response_id else { return Err(CodexErr::Stream( "remote compaction v2 stream closed before response.completed".to_string(), )); - } + }; if compaction_count != 1 { return Err(CodexErr::Fatal(format!( @@ -470,14 +475,10 @@ async fn collect_compaction_output( let Some(compaction_output) = compaction_output else { unreachable!("compaction output must exist when count is exactly one"); }; - let Some(response_id) = completed_response_id else { - unreachable!("response id must exist after response.completed"); - }; Ok(RemoteCompactionV2Output { compaction_output, response_id, token_usage: completed_token_usage, - usage_metadata: completed_usage_metadata, }) } @@ -1179,16 +1180,11 @@ mod tests { }), ]); - let output = collect_compaction_output(stream) + let (sess, turn_context) = crate::session::tests::make_session_and_context().await; + let output = collect_compaction_output(&sess, &turn_context, stream) .await .expect("compaction should be collected"); - assert_eq!( - output.usage_metadata, - Some(codex_protocol::ResponseUsageMetadata { - amount: Some("0.125".to_string()), - }), - ); assert_eq!(output.compaction_output, compaction); assert_eq!(output.response_id, "resp-compact"); assert_eq!( diff --git a/codex-rs/core/src/compact_remote_v2_attempt.rs b/codex-rs/core/src/compact_remote_v2_attempt.rs index 1c58c7e4c4..8827ba1d21 100644 --- a/codex-rs/core/src/compact_remote_v2_attempt.rs +++ b/codex-rs/core/src/compact_remote_v2_attempt.rs @@ -13,8 +13,6 @@ use crate::session::step_context::StepContext; use codex_history::CodexHarnessMetadata; use codex_protocol::error::Result as CodexResult; use codex_protocol::models::ResponseItem; -use codex_protocol::protocol::EventMsg; -use codex_protocol::protocol::RawResponseCompletedEvent; use codex_protocol::protocol::TokenUsage; use codex_rollout_trace::CompactionTraceContext; use tracing::info; @@ -24,6 +22,7 @@ pub(super) struct RemoteCompactV2Attempt { pub(super) prompt_input: Vec, pub(super) prompt_input_metadata: Vec>, pub(super) compaction_output: ResponseItem, + pub(super) compaction_response_id: String, pub(super) token_usage: Option, /// Keeps a session created for standalone compaction alive through lifecycle completion. pub(super) owned_client_session: Option, @@ -119,19 +118,7 @@ pub(super) async fn run_remote_compact_v2_attempt( compaction_output, response_id, token_usage, - usage_metadata, } = compaction_output_result?; - // TODO: Emit this before compaction output validation so malformed completed - // responses still surface their raw upstream usage. - sess.send_event( - turn_context, - EventMsg::RawResponseCompleted(RawResponseCompletedEvent { - response_id, - token_usage: token_usage.clone(), - usage_metadata, - }), - ) - .await; let mut prompt_input = prompt.input; prompt_input.pop(); Ok(RemoteCompactV2Attempt { @@ -139,6 +126,7 @@ pub(super) async fn run_remote_compact_v2_attempt( prompt_input, prompt_input_metadata, compaction_output, + compaction_response_id: response_id, token_usage, owned_client_session, }) diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index 2c661588f6..e7c6491db5 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -314,6 +314,7 @@ fn developer_msg_with_fragments(texts: &[&str]) -> ResponseItem { fn reference_context_item() -> TurnContextItem { TurnContextItem { turn_id: Some("reference-turn".to_string()), + root_turn_id: None, cwd: AbsolutePathBuf::try_from( std::env::current_dir() .expect("current directory") diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index 6ee7fd4bad..fe043d3bab 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/codex-rs/core/src/guardian/tests.rs @@ -3544,8 +3544,6 @@ async fn guardian_ephemeral_retry_preserves_parallel_trunk_and_fork_history() -> ) .await; - // A conflicting input removes the known root while the trunk review is in flight. - turn.turn_metadata_state.mark_root_turn_ambiguous(); let third_decision = review_approval_request( &session, &turn, @@ -3570,8 +3568,8 @@ async fn guardian_ephemeral_retry_preserves_parallel_trunk_and_fork_history() -> for (body, expected_root) in [ (&first_request_body, Some("causal-root-turn")), (&second_request_body, Some("causal-root-turn")), - (&failed_ephemeral_request_body, None), - (&retried_ephemeral_request_body, None), + (&failed_ephemeral_request_body, Some("causal-root-turn")), + (&retried_ephemeral_request_body, Some("causal-root-turn")), ] { assert_parent_turn(body, Some(turn.sub_id.as_str()))?; assert_root_turn(body, expected_root)?; diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index f3a0ec11c9..83711b97f3 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -769,21 +769,10 @@ pub(crate) async fn drain_async_hook_results( .collect::>(); if before_user_prompt { - // A fresh user turn owns its root; automatic turns only inherit one. - let current_turn_id = Some(turn_context.sub_id.as_str()); - if !additional_contexts.is_empty() - && result.turn_id.as_deref() != current_turn_id - && turn_context.turn_metadata_state.root_turn_id().as_deref() != current_turn_id - { - turn_context.turn_metadata_state.mark_root_turn_ambiguous(); - } record_additional_contexts(sess, turn_context, additional_contexts).await; } else if !additional_contexts.is_empty() { let _ = sess - .inject_hook_context_if_running( - additional_context_messages(additional_contexts), - result.turn_id.as_deref(), - ) + .inject_hook_context_if_running(additional_context_messages(additional_contexts)) .await; } diff --git a/codex-rs/core/src/session/inject.rs b/codex-rs/core/src/session/inject.rs index 3c57e7033f..6c103f266c 100644 --- a/codex-rs/core/src/session/inject.rs +++ b/codex-rs/core/src/session/inject.rs @@ -35,7 +35,7 @@ impl Session { } } - /// Injects hook context while classifying its actual receiving turn atomically. + /// Injects hook context into the running turn atomically. #[expect( clippy::await_holding_invalid_type, reason = "active turn provenance and turn state updates must remain atomic" @@ -43,19 +43,13 @@ impl Session { pub(crate) async fn inject_hook_context_if_running( &self, input: Vec, - source_turn_id: Option<&str>, ) -> Result<(), Vec> { let mut active = self.active_turn.lock().await; let Some(active_turn) = active.as_mut() else { return Err(input); }; - let Some(task) = active_turn.task.as_ref() else { + if active_turn.task.is_none() { return Err(input); - }; - if source_turn_id != Some(task.turn_context.sub_id.as_str()) { - task.turn_context - .turn_metadata_state - .mark_root_turn_ambiguous(); } self.input_queue .extend_pending_input_and_accept_mailbox_delivery_for_turn_state( @@ -86,11 +80,6 @@ impl Session { .collect::>(); let mut active = self.active_turn.lock().await; if let Some(active_turn) = active.as_mut() { - if let Some(task) = active_turn.task.as_ref() { - task.turn_context - .turn_metadata_state - .mark_root_turn_ambiguous(); - } self.input_queue .extend_pending_input_and_accept_mailbox_delivery_for_turn_state( active_turn.turn_state.as_ref(), diff --git a/codex-rs/core/src/session/input_queue.rs b/codex-rs/core/src/session/input_queue.rs index 6e71857aa3..e6bd3ba388 100644 --- a/codex-rs/core/src/session/input_queue.rs +++ b/codex-rs/core/src/session/input_queue.rs @@ -169,16 +169,16 @@ impl InputQueue { .and_then(|id| id.filter(|id| !id.trim().is_empty()).map(str::to_string)); start_options.root_turn_id = pending_mails .iter() - .filter(|mail| mail.communication.trigger_turn) - .map(|mail| { + .find(|mail| mail.communication.trigger_turn) + .and_then(|mail| { mail.start_options .parent_turn_id .as_deref() .filter(|id| !id.trim().is_empty()) .and(mail.start_options.root_turn_id.as_deref()) + .filter(|id| !id.trim().is_empty()) }) - .reduce(|expected, candidate| expected.filter(|id| candidate == Some(*id))) - .and_then(|id| id.filter(|id| !id.trim().is_empty()).map(str::to_string)); + .map(str::to_string); let items = pending_mails .into_iter() .map(|mail| TurnInput::InterAgentCommunication(mail.communication)) @@ -320,17 +320,10 @@ impl InputQueue { } let (mailbox_items, start_options) = self.drain_mailbox_input_items().await; if let Some(active_turn_metadata) = active_turn_metadata - && mailbox_items.iter().any(|item| { - matches!( - item, - TurnInput::InterAgentCommunication(communication) - if communication.trigger_turn - ) - }) - && (start_options.root_turn_id.is_none() - || active_turn_metadata.root_turn_id() != start_options.root_turn_id) + && active_turn_metadata.root_turn_id().is_none() + && let Some(root_turn_id) = start_options.root_turn_id.as_ref() { - active_turn_metadata.mark_root_turn_ambiguous(); + active_turn_metadata.set_root_turn_id(root_turn_id.clone()); } if pending_input.is_empty() { (mailbox_items, start_options) @@ -561,7 +554,7 @@ mod tests { } #[tokio::test] - async fn input_queue_requires_one_unambiguous_trigger_parent() { + async fn input_queue_uses_unambiguous_trigger_parent_and_first_root() { let (parent, peer, root, root2) = (Some("a"), Some("b"), Some("r"), Some("s")); for (pending_mails, expected_parent_turn_id, expected_root_turn_id) in [ (Vec::new(), None, None), @@ -572,8 +565,8 @@ mod tests { (vec![(true, parent, None)], parent, None), (vec![(true, parent, Some(""))], parent, None), (vec![(true, parent, root), (true, peer, root)], None, root), - (vec![(true, parent, root), (true, peer, root2)], None, None), - (vec![(true, parent, root), (true, None, root)], None, None), + (vec![(true, parent, root), (true, peer, root2)], None, root), + (vec![(true, parent, root), (true, None, root)], None, root), ( vec![(true, parent, root), (true, parent, root)], parent, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index c99599b7d6..f7ab1271a0 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -90,6 +90,7 @@ use codex_network_proxy::normalize_host; use codex_otel::current_span_trace_id; use codex_otel::current_span_w3c_trace_context; use codex_otel::set_parent_from_w3c_trace_context; +use codex_protocol::ResponseUsageMetadata; use codex_protocol::SessionId; use codex_protocol::ThreadId; use codex_protocol::approvals::ElicitationRequest; @@ -349,6 +350,7 @@ use codex_protocol::protocol::ModelVerificationEvent; use codex_protocol::protocol::NetworkApprovalContext; use codex_protocol::protocol::Op; use codex_protocol::protocol::RateLimitSnapshot; +use codex_protocol::protocol::RawResponseCompletedEvent; use codex_protocol::protocol::RequestUserInputEvent; use codex_protocol::protocol::ReviewDecision; use codex_protocol::protocol::SandboxPolicy; @@ -360,6 +362,7 @@ use codex_protocol::protocol::ThreadMemoryMode; use codex_protocol::protocol::TokenCountEvent; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; +use codex_protocol::protocol::TokenUsageRecord; use codex_protocol::protocol::TurnModerationMetadataEvent; use codex_protocol::protocol::WarningEvent; use codex_protocol::turn_input::TurnInputMode; @@ -1440,6 +1443,8 @@ impl Session { let mut state = self.state.lock().await; state.set_token_info(Some(info)); } + self.state.lock().await.latest_token_usage_record = + Self::last_token_usage_record_from_rollout(&rollout_items); // Defer seeding the session's initial context until the first turn starts so // turn/start overrides can be merged before we write to the rollout. @@ -1460,6 +1465,8 @@ impl Session { let mut state = self.state.lock().await; state.set_token_info(Some(info)); } + self.state.lock().await.latest_token_usage_record = + Self::last_token_usage_record_from_rollout(&rollout_items); let thread_settings_applied = RolloutItem::EventMsg(thread_settings::applied_event(self).await); @@ -1624,6 +1631,21 @@ impl Session { }) } + fn last_token_usage_record_from_rollout( + rollout_items: &[RolloutItem], + ) -> Option { + for item in rollout_items.iter().rev() { + match item { + RolloutItem::TokenUsageRecord(record) => return Some(record.clone()), + RolloutItem::Compacted(compacted) => { + return compacted.latest_token_usage_record.clone(); + } + _ => {} + } + } + None + } + async fn previous_turn_settings(&self) -> Option { let state = self.state.lock().await; state.previous_turn_settings() @@ -3668,6 +3690,8 @@ impl Session { .previous_window_id .map(|id| id.to_string()), window_id: Some(metadata.window_ids.window_id.to_string()), + compaction_response_id: metadata.compaction_response_id, + latest_token_usage_record: self.state.lock().await.latest_token_usage_record.clone(), }; // Wait for accepted updates to finish persisting, then keep later updates from // overtaking the current settings snapshot while its checkpoint is written. @@ -4118,6 +4142,7 @@ impl Session { message: String::new(), window_number, window_ids, + compaction_response_id: None, }, ) .await; @@ -4232,6 +4257,40 @@ impl Session { result } + pub(crate) async fn record_observed_response_completed( + &self, + turn_context: &TurnContext, + response_id: &str, + usage: Option<&TokenUsage>, + usage_metadata: Option<&ResponseUsageMetadata>, + ) { + self.send_event( + turn_context, + EventMsg::RawResponseCompleted(RawResponseCompletedEvent { + response_id: response_id.to_string(), + token_usage: usage.cloned(), + usage_metadata: usage_metadata.cloned(), + }), + ) + .await; + let Some(usage) = usage else { + return; + }; + let record = self.state.lock().await.record_token_usage( + self.thread_id, + &turn_context.sub_id, + self.session_id(), + turn_context + .turn_metadata_state + .root_turn_id() + .unwrap_or_else(|| turn_context.sub_id.clone()), + response_id.to_string(), + usage, + ); + self.persist_rollout_items(&[RolloutItem::TokenUsageRecord(record)]) + .await; + } + pub(crate) async fn record_token_usage_info( &self, turn_context: &TurnContext, diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index e087240ccd..46e94eb6e4 100644 --- a/codex-rs/core/src/session/rollout_reconstruction.rs +++ b/codex-rs/core/src/session/rollout_reconstruction.rs @@ -294,6 +294,7 @@ impl Session { | RolloutItem::SessionMeta(_) | RolloutItem::RealtimeItem(_) | RolloutItem::SecurityRiskScore(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::InterAgentCommunicationMetadata { .. } => {} } @@ -385,6 +386,7 @@ impl Session { | RolloutItem::RealtimeItem(_) | RolloutItem::WorldState(_) | RolloutItem::SecurityRiskScore(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::SessionMeta(_) => {} } } @@ -424,6 +426,7 @@ impl Session { | RolloutItem::InterAgentCommunicationMetadata { .. } | RolloutItem::TurnContext(_) | RolloutItem::RealtimeItem(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::SecurityRiskScore(_) | RolloutItem::EventMsg(_) => { unreachable!("only world-state replay items are collected") diff --git a/codex-rs/core/src/session/rollout_reconstruction_tests.rs b/codex-rs/core/src/session/rollout_reconstruction_tests.rs index 3f0b410b9e..6d7f2863e3 100644 --- a/codex-rs/core/src/session/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/session/rollout_reconstruction_tests.rs @@ -265,6 +265,7 @@ async fn record_initial_history_resumed_bare_turn_context_does_not_hydrate_previ let previous_model = "previous-rollout-model"; let previous_context_item = TurnContextItem { turn_id: Some(turn_context.sub_id.clone()), + root_turn_id: None, #[allow(deprecated)] cwd: turn_context.cwd.clone(), workspace_roots: None, @@ -314,6 +315,7 @@ async fn record_initial_history_resumed_hydrates_previous_turn_settings_from_lif let previous_model = "previous-rollout-model"; let mut previous_context_item = TurnContextItem { turn_id: Some(turn_context.sub_id.clone()), + root_turn_id: None, #[allow(deprecated)] cwd: turn_context.cwd.clone(), workspace_roots: None, @@ -1064,6 +1066,8 @@ async fn record_initial_history_resumed_rollback_drops_incomplete_user_turn_comp first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), RolloutItem::EventMsg(EventMsg::ThreadRolledBack( codex_protocol::protocol::ThreadRolledBackEvent { num_turns: 1 }, @@ -1124,6 +1128,8 @@ async fn record_initial_history_requires_surviving_full_snapshot_without_user_tu first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), ], }; @@ -1156,6 +1162,8 @@ async fn record_initial_history_resumed_does_not_seed_reference_context_item_aft first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), ]; @@ -1226,6 +1234,8 @@ async fn reconstruct_history_prefers_compacted_window_over_session_meta() { first_window_id: Some(compacted_first_window_id.to_string()), previous_window_id: Some(compacted_previous_window_id.to_string()), window_id: Some(compacted_window_id.to_string()), + compaction_response_id: None, + latest_token_usage_record: None, }), ]; @@ -1262,6 +1272,8 @@ async fn reconstruct_history_replays_world_state_from_latest_compaction_window() first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), RolloutItem::WorldState(WorldStateItem::full(object!({ "environment": {"status": "starting", "cwd": "/workspace"} @@ -1324,6 +1336,8 @@ async fn bounded_replay_matches_full_replay_after_empty_turn_compactions() { first_window_id: Some(window_ids[0].to_string()), previous_window_id: Some(window_ids[window_number - 1].to_string()), window_id: Some(window_ids[window_number].to_string()), + compaction_response_id: None, + latest_token_usage_record: None, }), RolloutItem::WorldState(WorldStateItem::full(object!({ "environment": {"window": window_number, "status": "starting"} @@ -1420,6 +1434,8 @@ async fn reconstruct_history_preserves_legacy_compaction_count_with_session_meta first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), ]; @@ -1448,6 +1464,8 @@ async fn reconstruct_history_legacy_compaction_without_replacement_history_does_ first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), ]; @@ -1484,6 +1502,8 @@ async fn reconstruct_history_legacy_compaction_without_replacement_history_clear first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), RolloutItem::EventMsg(EventMsg::TurnStarted( codex_protocol::protocol::TurnStartedEvent { @@ -1532,6 +1552,7 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis let previous_model = "previous-rollout-model"; let previous_context_item = TurnContextItem { turn_id: Some(turn_context.sub_id.clone()), + root_turn_id: Some("root-turn".to_string()), #[allow(deprecated)] cwd: turn_context.cwd.clone(), workspace_roots: None, @@ -1588,6 +1609,8 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), RolloutItem::TurnContext(previous_context_item), RolloutItem::EventMsg(EventMsg::TurnComplete( @@ -1624,6 +1647,7 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis .expect("serialize seeded reference context item"), serde_json::to_value(Some(TurnContextItem { turn_id: Some(turn_context.sub_id.clone()), + root_turn_id: Some("root-turn".to_string()), #[allow(deprecated)] cwd: turn_context.cwd.clone(), workspace_roots: None, @@ -1658,6 +1682,7 @@ async fn record_initial_history_resumed_aborted_turn_without_id_clears_active_tu let previous_model = "previous-rollout-model"; let previous_context_item = TurnContextItem { turn_id: Some(turn_context.sub_id.clone()), + root_turn_id: None, #[allow(deprecated)] cwd: turn_context.cwd.clone(), workspace_roots: None, @@ -1755,6 +1780,8 @@ async fn record_initial_history_resumed_aborted_turn_without_id_clears_active_tu first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), ]; @@ -1791,6 +1818,7 @@ async fn record_initial_history_resumed_unmatched_abort_preserves_active_turn_fo let unmatched_abort_turn_id = "other-turn".to_string(); let current_context_item = TurnContextItem { turn_id: Some(current_turn_id.clone()), + root_turn_id: None, #[allow(deprecated)] cwd: turn_context.cwd.clone(), workspace_roots: None, @@ -1920,6 +1948,7 @@ async fn record_initial_history_resumed_trailing_incomplete_turn_compaction_clea let previous_model = "previous-rollout-model"; let previous_context_item = TurnContextItem { turn_id: Some(turn_context.sub_id.clone()), + root_turn_id: None, #[allow(deprecated)] cwd: turn_context.cwd.clone(), workspace_roots: None, @@ -2008,6 +2037,8 @@ async fn record_initial_history_resumed_trailing_incomplete_turn_compaction_clea first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), ]; @@ -2093,6 +2124,7 @@ async fn record_initial_history_resumed_replaced_incomplete_compacted_turn_clear let previous_model = "previous-rollout-model"; let previous_context_item = TurnContextItem { turn_id: Some(turn_context.sub_id.clone()), + root_turn_id: None, #[allow(deprecated)] cwd: turn_context.cwd.clone(), workspace_roots: None, @@ -2182,6 +2214,8 @@ async fn record_initial_history_resumed_replaced_incomplete_compacted_turn_clear first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), // A newer TurnStarted replaces the incomplete compacted turn without a matching // completion/abort for the old one. diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 8dbfad88d3..3b3da7440b 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -770,11 +770,12 @@ impl Session { } InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_) => None, }; - // Legacy subagent rollouts synthesize session_id from their own thread id. + // Legacy subagent rollouts synthesized session_id from their own thread ID. let resumed_session_id = resumed_session_id.filter(|session_id| { !session_configuration.session_source.is_non_root_agent() || *session_id != SessionId::from(thread_id) }); + // session_id is equal to the root thread's ID. let session_id = resumed_session_id.unwrap_or_else(|| { if session_configuration.session_source.is_non_root_agent() { agent_control.session_id() @@ -1357,6 +1358,7 @@ impl Session { | RolloutItem::TurnContext(_) | RolloutItem::WorldState(_) | RolloutItem::RealtimeItem(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::SecurityRiskScore(_) => {} } } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index e962382eb7..b8d97c87e5 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -164,6 +164,7 @@ use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::protocol::TokenCountEvent; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; +use codex_protocol::protocol::TokenUsageRecord; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnCompleteEvent; use codex_protocol::protocol::TurnStartedEvent; @@ -2227,6 +2228,8 @@ async fn reconstruct_history_uses_replacement_history_verbatim() { first_window_id: Some(first_window_id.to_string()), previous_window_id: Some(previous_window_id.to_string()), window_id: Some(window_id.to_string()), + compaction_response_id: None, + latest_token_usage_record: None, })]; let reconstructed = session @@ -2883,6 +2886,49 @@ async fn record_initial_history_seeds_token_info_from_rollout() { assert_eq!(actual, Some(info2)); } +#[test] +fn latest_token_usage_record_stops_at_compaction_checkpoint() { + let thread_id = ThreadId::new(); + let checkpoint_record = TokenUsageRecord { + thread_id, + turn_id: "turn-1".to_string(), + session_id: SessionId::from(thread_id), + root_turn_id: "turn-1".to_string(), + response_id: "response-1".to_string(), + usage: TokenUsage::default(), + turn_token_usage: TokenUsage::default(), + thread_token_usage: TokenUsage::default(), + }; + let checkpoint = |latest_token_usage_record| { + RolloutItem::Compacted(CompactedItem { + message: String::new(), + replacement_history: None, + mcp_resource_origins: None, + window_number: None, + first_window_id: None, + previous_window_id: None, + window_id: None, + compaction_response_id: None, + latest_token_usage_record, + }) + }; + + assert_eq!( + Session::last_token_usage_record_from_rollout(&[ + RolloutItem::TokenUsageRecord(checkpoint_record.clone()), + checkpoint(Some(checkpoint_record.clone())), + ]), + Some(checkpoint_record.clone()) + ); + assert_eq!( + Session::last_token_usage_record_from_rollout(&[ + RolloutItem::TokenUsageRecord(checkpoint_record), + checkpoint(None), + ]), + None + ); +} + #[tokio::test] async fn recompute_token_usage_uses_session_base_instructions() { let (session, turn_context) = make_session_and_context().await; @@ -3412,7 +3458,7 @@ async fn record_initial_history_reconstructs_forked_transcript() { } #[tokio::test] -async fn start_new_context_window_assigns_and_persists_item_ids() { +async fn start_new_context_window_persists_checkpoint_state() { let (mut session, turn_context, _rx) = make_session_and_context_with_auth_and_config_and_rx( CodexAuth::from_api_key("Test API Key"), Vec::new(), @@ -3421,6 +3467,18 @@ async fn start_new_context_window_assigns_and_persists_item_ids() { .await; let rollout_path = attach_thread_persistence(Arc::get_mut(&mut session).expect("unique session")).await; + let thread_id = ThreadId::new(); + let token_usage_record = TokenUsageRecord { + thread_id, + turn_id: "turn-1".to_string(), + session_id: SessionId::from(thread_id), + root_turn_id: "turn-1".to_string(), + response_id: "response-1".to_string(), + usage: TokenUsage::default(), + turn_token_usage: TokenUsage::default(), + thread_token_usage: TokenUsage::default(), + }; + session.state.lock().await.latest_token_usage_record = Some(token_usage_record.clone()); let step_context = session .capture_step_context(Arc::clone(&turn_context), &CancellationToken::new()) .await @@ -3447,8 +3505,8 @@ async fn start_new_context_window_assigns_and_persists_item_ids() { else { panic!("expected resumed rollout history"); }; - let persisted_replacement_history = resumed.history.iter().rev().find_map(|item| match item { - RolloutItem::Compacted(compacted) => compacted.replacement_history.as_ref(), + let persisted_compacted = resumed.history.iter().rev().find_map(|item| match item { + RolloutItem::Compacted(compacted) => Some(compacted), RolloutItem::SessionMeta(_) | RolloutItem::ResponseItem(_) | RolloutItem::InterAgentCommunication(_) @@ -3456,13 +3514,23 @@ async fn start_new_context_window_assigns_and_persists_item_ids() { | RolloutItem::TurnContext(_) | RolloutItem::WorldState(_) | RolloutItem::SecurityRiskScore(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::RealtimeItem(_) | RolloutItem::EventMsg(_) => None, }); assert_eq!( - persisted_replacement_history.cloned(), + persisted_compacted.and_then(|compacted| compacted.replacement_history.clone()), Some(live_history.annotated_items().to_vec()) ); + assert_eq!( + persisted_compacted.map(|compacted| { + ( + compacted.compaction_response_id.as_deref(), + compacted.latest_token_usage_record.as_ref(), + ) + }), + Some((None, Some(&token_usage_record))) + ); } #[tokio::test] @@ -3535,6 +3603,7 @@ async fn record_initial_history_assigns_and_persists_id_for_forked_response_item | RolloutItem::TurnContext(_) | RolloutItem::WorldState(_) | RolloutItem::SecurityRiskScore(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::RealtimeItem(_) | RolloutItem::EventMsg(_) => None, }); @@ -3691,6 +3760,7 @@ async fn record_initial_history_forked_hydrates_previous_turn_settings() { let previous_model = "forked-rollout-model"; let previous_context_item = TurnContextItem { turn_id: Some(turn_context.sub_id.clone()), + root_turn_id: None, #[allow(deprecated)] cwd: turn_context.cwd.clone(), workspace_roots: None, @@ -4102,6 +4172,8 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio first_window_id: Some(first_window_id.to_string()), previous_window_id: Some(previous_window_id.to_string()), window_id: Some(compacted_window_id.to_string()), + compaction_response_id: None, + latest_token_usage_record: None, }), RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent { turn_id: compact_turn_id, @@ -5864,6 +5936,7 @@ async fn compaction_checkpoint_waits_for_accepted_settings_persistence() { message: "summary".to_string(), window_number, window_ids, + compaction_response_id: None, }, ), )); @@ -11907,6 +11980,66 @@ async fn trigger_turn_mailbox_mail_waits_for_next_turn_after_answer_boundary() { assert!(sess.input_queue.has_trigger_turn_mailbox_items().await); } +#[tokio::test] +async fn active_turn_keeps_first_root_when_mail_coalesces() { + let (sess, tc, _rx) = make_session_and_context_with_rx().await; + tc.turn_metadata_state + .set_root_turn_id("root-a".to_string()); + sess.spawn_task( + Arc::clone(&tc), + Vec::new(), + NeverEndingTask { + kind: TaskKind::Regular, + listen_to_cancellation_token: true, + }, + ) + .await; + let first = InterAgentCommunication::new( + AgentPath::try_from("/root/worker_a").expect("worker path should parse"), + AgentPath::root(), + Vec::new(), + "first".to_string(), + /*trigger_turn*/ true, + ); + let second = InterAgentCommunication::new( + AgentPath::try_from("/root/worker_b").expect("worker path should parse"), + AgentPath::root(), + Vec::new(), + "second".to_string(), + /*trigger_turn*/ true, + ); + for (communication, parent_turn_id, root_turn_id) in [ + (first.clone(), "parent-a", "root-a"), + (second.clone(), "parent-b", "root-b"), + ] { + sess.input_queue + .enqueue_mailbox_communication( + communication, + codex_protocol::turn_input::TurnStartOptions { + parent_turn_id: Some(parent_turn_id.to_string()), + root_turn_id: Some(root_turn_id.to_string()), + ..Default::default() + }, + ) + .await; + } + + assert_eq!( + (sess.input_queue.get_pending_input(&sess.active_turn).await).0, + vec![ + TurnInput::InterAgentCommunication(first), + TurnInput::InterAgentCommunication(second), + ] + ); + assert_eq!( + tc.turn_metadata_state.root_turn_id().as_deref(), + Some("root-a") + ); + assert!(!sess.input_queue.has_pending_mailbox_items().await); + + sess.abort_all_tasks(TurnAbortReason::Replaced).await; +} + #[tokio::test] async fn steered_input_reopens_mailbox_delivery_for_current_turn() { let (sess, tc, _rx) = make_session_and_context_with_rx().await; @@ -12271,6 +12404,8 @@ async fn sample_rollout( first_window_id: Some(window_ids.first_window_id.to_string()), previous_window_id: window_ids.previous_window_id.map(|id| id.to_string()), window_id: Some(window_ids.window_id.to_string()), + compaction_response_id: None, + latest_token_usage_record: None, })); let user2 = user_message("second user"); @@ -12301,6 +12436,8 @@ async fn sample_rollout( first_window_id: Some(window_ids.first_window_id.to_string()), previous_window_id: window_ids.previous_window_id.map(|id| id.to_string()), window_id: Some(window_ids.window_id.to_string()), + compaction_response_id: None, + latest_token_usage_record: None, })); let user3 = user_message("third user"); diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 9ab25ca151..f43ad3cc35 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -100,7 +100,6 @@ use codex_protocol::protocol::ErrorEvent; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::InternalSessionSource; use codex_protocol::protocol::PlanDeltaEvent; -use codex_protocol::protocol::RawResponseCompletedEvent; use codex_protocol::protocol::ReasoningContentDeltaEvent; use codex_protocol::protocol::ReasoningRawContentDeltaEvent; use codex_protocol::protocol::SafetyBufferingEvent; @@ -2590,13 +2589,11 @@ async fn try_run_sampling_request( &mut assistant_message_stream_parsers, ) .await; - sess.send_event( + sess.record_observed_response_completed( &turn_context, - EventMsg::RawResponseCompleted(RawResponseCompletedEvent { - response_id, - token_usage: token_usage.clone(), - usage_metadata, - }), + &response_id, + token_usage.as_ref(), + usage_metadata.as_ref(), ) .await; let budget_result = sess diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index cc3681d420..44c62cac67 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -578,6 +578,7 @@ impl TurnContext { let cwd = self.cwd.clone(); TurnContextItem { turn_id: Some(self.sub_id.clone()), + root_turn_id: self.turn_metadata_state.root_turn_id(), cwd, workspace_roots: (!workspace_roots.is_empty()).then_some(workspace_roots), current_date: self.current_date.clone(), diff --git a/codex-rs/core/src/session/turn_input.rs b/codex-rs/core/src/session/turn_input.rs index b83b175953..2882931bca 100644 --- a/codex-rs/core/src/session/turn_input.rs +++ b/codex-rs/core/src/session/turn_input.rs @@ -624,13 +624,17 @@ impl Session { input => pending_turn_input(input.clone()), }; pending_input.push(input); - if let Some(incoming_root_turn_id) = incoming_root_turn_id - && active_task.turn_context.turn_metadata_state.root_turn_id() != incoming_root_turn_id + if active_task + .turn_context + .turn_metadata_state + .root_turn_id() + .is_none() + && let Some(Some(incoming_root_turn_id)) = incoming_root_turn_id { active_task .turn_context .turn_metadata_state - .mark_root_turn_ambiguous(); + .set_root_turn_id(incoming_root_turn_id); } self.input_queue .extend_pending_input_and_accept_mailbox_delivery_for_turn_state( diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index c65eafadbd..768e8d9fe3 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -19,9 +19,12 @@ use crate::session::session::SessionConfiguration; use crate::session::time_reminder::CurrentTimeReminderState; use crate::session_startup_prewarm::SessionStartupPrewarmHandle; use codex_history::ResponseItemEnvelope; +use codex_protocol::SessionId; +use codex_protocol::ThreadId; use codex_protocol::protocol::RateLimitSnapshot; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; +use codex_protocol::protocol::TokenUsageRecord; use codex_protocol::protocol::TurnContextItem; use codex_utils_output_truncation::TruncationPolicy; @@ -32,6 +35,7 @@ pub(crate) struct SessionState { pub(crate) base_instructions_provenance: Option, pub(crate) history: ContextManager, pub(crate) latest_rate_limits: Option, + pub(crate) latest_token_usage_record: Option, pub(crate) server_reasoning_included: bool, pub(crate) mcp_dependency_prompted: HashSet, pub(crate) additional_context: AdditionalContextStore, @@ -70,6 +74,7 @@ impl SessionState { base_instructions_provenance: None, history, latest_rate_limits: None, + latest_token_usage_record: None, server_reasoning_included: false, mcp_dependency_prompted: HashSet::new(), additional_context: AdditionalContextStore::default(), @@ -148,6 +153,44 @@ impl SessionState { self.history.set_token_info(info); } + pub(crate) fn record_token_usage( + &mut self, + thread_id: ThreadId, + turn_id: &str, + session_id: SessionId, + root_turn_id: String, + response_id: String, + usage: &TokenUsage, + ) -> TokenUsageRecord { + let mut turn_token_usage = self + .latest_token_usage_record + .as_ref() + .filter(|record| record.turn_id == turn_id) + .map_or_else(TokenUsage::default, |record| { + record.turn_token_usage.clone() + }); + turn_token_usage.add_assign(usage); + let mut thread_token_usage = self + .latest_token_usage_record + .as_ref() + .map_or_else(TokenUsage::default, |record| { + record.thread_token_usage.clone() + }); + thread_token_usage.add_assign(usage); + let record = TokenUsageRecord { + thread_id, + turn_id: turn_id.to_string(), + session_id, + root_turn_id, + response_id, + usage: usage.clone(), + turn_token_usage, + thread_token_usage, + }; + self.latest_token_usage_record = Some(record.clone()); + record + } + pub(crate) fn set_reference_context_item(&mut self, item: Option) { self.history.set_reference_context_item(item); } diff --git a/codex-rs/core/src/state/session_tests.rs b/codex-rs/core/src/state/session_tests.rs index 44011b7d82..533bd2df26 100644 --- a/codex-rs/core/src/state/session_tests.rs +++ b/codex-rs/core/src/state/session_tests.rs @@ -1,11 +1,57 @@ use super::*; use crate::session::tests::make_session_configuration_for_tests; use crate::state::AutoCompactWindowSnapshot; +use codex_protocol::SessionId; +use codex_protocol::ThreadId; use codex_protocol::protocol::CreditsSnapshot; use codex_protocol::protocol::RateLimitWindow; use codex_protocol::protocol::SpendControlLimitSnapshot; +use codex_protocol::protocol::TokenUsage; +use codex_protocol::protocol::TokenUsageRecord; use pretty_assertions::assert_eq; +#[tokio::test] +async fn record_token_usage_continues_restored_totals() { + let thread_id = ThreadId::new(); + let session_id = SessionId::from(ThreadId::new()); + let usage = |total_tokens| TokenUsage { + total_tokens, + ..TokenUsage::default() + }; + let mut restored = SessionState::new(make_session_configuration_for_tests().await); + restored.latest_token_usage_record = Some(TokenUsageRecord { + thread_id, + turn_id: "turn-b".to_string(), + session_id, + root_turn_id: "root-turn".to_string(), + response_id: "response-c".to_string(), + usage: usage(30), + turn_token_usage: usage(30), + thread_token_usage: usage(230), + }); + let after_resume = restored.record_token_usage( + thread_id, + "turn-b", + session_id, + "root-turn".to_string(), + "response-d".to_string(), + &usage(20), + ); + assert_eq!( + after_resume, + TokenUsageRecord { + thread_id, + turn_id: "turn-b".to_string(), + session_id, + root_turn_id: "root-turn".to_string(), + response_id: "response-d".to_string(), + usage: usage(20), + turn_token_usage: usage(50), + thread_token_usage: usage(250), + } + ); +} + #[tokio::test] // Verifies connector merging deduplicates repeated IDs. async fn merge_connector_selection_deduplicates_entries() { diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index c51db874ee..5bf41019d5 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -307,16 +307,13 @@ impl Session { .await .clear_turn(&turn_context.sub_id); - // Reserved turn input already has its context; only newly arriving mail can change lineage. let (pending_items, start_options) = self.input_queue.drain_mailbox_input_items().await; - if pending_items.iter().any(|item| { - matches!( - item, - TurnInput::InterAgentCommunication(communication) if communication.trigger_turn - ) - }) && turn_context.turn_metadata_state.root_turn_id() != start_options.root_turn_id + if turn_context.turn_metadata_state.root_turn_id().is_none() + && let Some(root_turn_id) = start_options.root_turn_id { - turn_context.turn_metadata_state.mark_root_turn_ambiguous(); + turn_context + .turn_metadata_state + .set_root_turn_id(root_turn_id); } let turn_state = { let mut active = self.active_turn.lock().await; diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index 65c19f6ae3..9808432f86 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/codex-rs/core/src/tools/registry_tests.rs @@ -728,7 +728,6 @@ async fn dispatch_uses_canonical_tool_names_for_lifecycle_contributors() -> anyh /*terminal_outcome_reached*/ None, ) .await?; - turn.turn_metadata_state.mark_root_turn_ambiguous(); let err = match registry .dispatch_any_with_terminal_outcome( test_invocation( @@ -760,7 +759,7 @@ async fn dispatch_uses_canonical_tool_names_for_lifecycle_contributors() -> anyh RecordedToolLifecycle::Start { call_id: "failing-call".to_string(), tool_name: failing_tool.clone(), - root_turn_id: None, + root_turn_id: Some("root-turn".to_string()), }, RecordedToolLifecycle::Finish { call_id: "failing-call".to_string(), diff --git a/codex-rs/core/src/turn_metadata.rs b/codex-rs/core/src/turn_metadata.rs index afae92f22f..2f7b540d06 100644 --- a/codex-rs/core/src/turn_metadata.rs +++ b/codex-rs/core/src/turn_metadata.rs @@ -129,7 +129,6 @@ pub(crate) struct TurnMetadataState { turn_started_at_unix_ms: RwLock>, responses_api_metadata: RwLock>, responsesapi_client_metadata: RwLock>, - root_turn_ambiguous: AtomicBool, user_input_requested_during_turn: AtomicBool, enrichment_task: Mutex>>, git_enrichment_complete: watch::Sender, @@ -199,7 +198,6 @@ impl TurnMetadataState { turn_started_at_unix_ms: RwLock::new(None), responses_api_metadata: RwLock::new(BTreeMap::new()), responsesapi_client_metadata: RwLock::new(BTreeMap::new()), - root_turn_ambiguous: AtomicBool::new(false), user_input_requested_during_turn: AtomicBool::new(false), enrichment_task: Mutex::new(None), git_enrichment_complete: watch::channel(/*init*/ true).0, @@ -311,14 +309,7 @@ impl TurnMetadataState { } pub(crate) fn root_turn_id(&self) -> Option { - self.root_turn_id - .get() - .filter(|_| !self.root_turn_ambiguous.load(Ordering::Relaxed)) - .cloned() - } - - pub(crate) fn mark_root_turn_ambiguous(&self) { - self.root_turn_ambiguous.store(true, Ordering::Relaxed); + self.root_turn_id.get().cloned() } pub(crate) fn can_start_root_turn(&self, session_source: &SessionSource) -> bool { diff --git a/codex-rs/core/tests/suite/compact.rs b/codex-rs/core/tests/suite/compact.rs index 0a2c462d51..3473e39628 100644 --- a/codex-rs/core/tests/suite/compact.rs +++ b/codex-rs/core/tests/suite/compact.rs @@ -989,7 +989,7 @@ async fn manual_compact_uses_custom_prompt() { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn manual_compact_emits_api_and_local_token_usage_events() { +async fn manual_compact_records_durable_and_local_token_usage() { skip_if_no_network!(); let server = start_mock_server().await; @@ -1008,7 +1008,9 @@ async fn manual_compact_emits_api_and_local_token_usage_events() { config.model_provider = model_provider; set_test_compact_prompt(config); }); - let codex = builder.build(&server).await.unwrap().codex; + let test = builder.build(&server).await.unwrap(); + let rollout_path = test.codex.rollout_path().expect("rollout path"); + let codex = test.codex; // Trigger manual compact and collect TokenCount events for the compact turn. codex.submit(Op::Compact).await.unwrap(); @@ -1044,6 +1046,36 @@ async fn manual_compact_emits_api_and_local_token_usage_events() { last > 0, "second TokenCount should reflect a non-zero estimated context size after compaction" ); + let rollout_items = fs::read_to_string(rollout_path) + .expect("read rollout") + .lines() + .filter_map(|line| serde_json::from_str::(line).ok()) + .map(|line| line.item) + .collect::>(); + let records = rollout_items + .iter() + .filter_map(|item| match item { + RolloutItem::TokenUsageRecord(record) => Some(record), + _ => None, + }) + .collect::>(); + assert_eq!(records.len(), 1); + assert_eq!(records[0].response_id, "r1"); + assert_eq!(records[0].usage.total_tokens, 0); + let compacted = rollout_items + .iter() + .find_map(|item| match item { + RolloutItem::Compacted(compacted) => Some(compacted), + _ => None, + }) + .expect("compaction checkpoint"); + assert_eq!( + ( + compacted.compaction_response_id.as_deref(), + compacted.latest_token_usage_record.as_ref(), + ), + (Some("r1"), records.first().copied()) + ); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index 053ccecb59..dd905d4f67 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -685,6 +685,67 @@ async fn remote_compact_v2_retains_only_client_developer_messages_when_enabled( }) .count(); assert_eq!(retained_client_developers, usize::from(enabled)); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn remote_compact_v2_records_usage_before_output_validation() -> Result<()> { + skip_if_no_network!(Ok(())); + + let harness = TestCodexHarness::with_auto_env_builder( + test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(|config| { + config + .features + .enable(Feature::RemoteCompactionV2) + .expect("remote compaction v2 should be configurable"); + }), + ) + .await?; + let codex = &harness.test().codex; + let rollout_path = codex.rollout_path().context("rollout path")?; + responses::mount_sse_sequence( + harness.server(), + vec![ + sse(vec![responses::ev_completed("before-compact")]), + sse(vec![ + json!({ + "type": "response.output_item.done", + "item": { + "type": "compaction", + "encrypted_content": "FIRST_SUMMARY", + }, + }), + json!({ + "type": "response.output_item.done", + "item": { + "type": "compaction", + "encrypted_content": "SECOND_SUMMARY", + }, + }), + responses::ev_completed_with_tokens("invalid-compact", /*total_tokens*/ 8_200), + ]), + ], + ) + .await; + + harness.test().submit_turn("before compact").await?; + codex.submit(Op::Compact).await?; + wait_for_event(codex, |event| matches!(event, EventMsg::Error(_))).await; + wait_for_event(codex, |event| matches!(event, EventMsg::TurnComplete(_))).await; + codex.shutdown_and_wait().await?; + + let record = fs::read_to_string(&rollout_path)? + .lines() + .filter_map(|line| serde_json::from_str::(line).ok()) + .filter_map(|line| match line.item { + RolloutItem::TokenUsageRecord(record) => Some(record), + _ => None, + }) + .find(|record| record.response_id == "invalid-compact") + .context("remote compaction usage record")?; + assert_eq!(record.usage.total_tokens, 8_200); Ok(()) } diff --git a/codex-rs/core/tests/suite/hooks.rs b/codex-rs/core/tests/suite/hooks.rs index e4874b489c..90f6832481 100644 --- a/codex-rs/core/tests/suite/hooks.rs +++ b/codex-rs/core/tests/suite/hooks.rs @@ -1843,7 +1843,11 @@ async fn async_hook_finishing_while_idle_waits_for_the_next_turn( assert_ne!(first_turn_id, second_turn_id); responses::assert_root_turn( &requests[1].body_json(), - (!automatic_continuation).then_some(second_turn_id.as_str()), + Some(if automatic_continuation { + first_turn_id.as_str() + } else { + second_turn_id.as_str() + }), )?; assert_eq!( warning_event diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 5fc867a5cb..4d434aed55 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -166,6 +166,7 @@ mod stream_no_completed; mod subagent_notifications; mod subagent_service_tier; mod token_budget; +mod token_usage_rollout; mod tool_harness; mod tool_lifecycle; mod tool_parallelism; diff --git a/codex-rs/core/tests/suite/pending_input.rs b/codex-rs/core/tests/suite/pending_input.rs index 372fab5ce2..a636b2c248 100644 --- a/codex-rs/core/tests/suite/pending_input.rs +++ b/codex-rs/core/tests/suite/pending_input.rs @@ -1102,8 +1102,8 @@ async fn injected_response_item_reopens_turn_after_final_answer() { responses::assert_root_turn(&first, Some(first_turn_id)) .expect("initial root should be trusted"); let second: Value = from_slice(&requests[1]).expect("parse second request"); - responses::assert_root_turn(&second, /*expected*/ None) - .expect("external injection should invalidate the active turn root"); + responses::assert_root_turn(&second, Some(first_turn_id)) + .expect("external injection should preserve the active turn root"); let relevant_user_input = message_input_texts(&second, "user") .into_iter() .filter(|text| { diff --git a/codex-rs/core/tests/suite/resume_warning.rs b/codex-rs/core/tests/suite/resume_warning.rs index 1be835b6fe..e5dd6f61bc 100644 --- a/codex-rs/core/tests/suite/resume_warning.rs +++ b/codex-rs/core/tests/suite/resume_warning.rs @@ -30,6 +30,7 @@ fn resume_history( let turn_id = "resume-warning-seed-turn".to_string(); let turn_ctx = TurnContextItem { turn_id: Some(turn_id.clone()), + root_turn_id: None, cwd: config.cwd.clone(), workspace_roots: None, current_date: None, diff --git a/codex-rs/core/tests/suite/token_usage_rollout.rs b/codex-rs/core/tests/suite/token_usage_rollout.rs new file mode 100644 index 0000000000..db9e6c50e1 --- /dev/null +++ b/codex-rs/core/tests/suite/token_usage_rollout.rs @@ -0,0 +1,117 @@ +//! Verifies observed Responses API usage is durably recorded in rollout history. + +use anyhow::Result; +use codex_history::RolloutItem; +use codex_history::RolloutLine; +use codex_protocol::SessionId; +use codex_protocol::protocol::TokenUsageRecord; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed_with_tokens; +use core_test_support::responses::ev_function_call; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_sse_sequence; +use core_test_support::responses::sse; +use core_test_support::responses::start_mock_server; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex::test_codex; +use pretty_assertions::assert_eq; +use serde_json::json; + +fn token_usage_records(path: &std::path::Path) -> Vec { + std::fs::read_to_string(path) + .expect("read rollout") + .lines() + .filter_map(|line| serde_json::from_str::(line).ok()) + .filter_map(|line| match line.item { + RolloutItem::TokenUsageRecord(record) => Some(record), + _ => None, + }) + .collect() +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn observed_response_usage_accumulates_per_turn_and_thread() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let plan_args = json!({ + "plan": [{ + "step": "keep sampling", + "status": "in_progress" + }] + }) + .to_string(); + mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("response-a"), + ev_function_call("call-a", "update_plan", &plan_args), + ev_completed_with_tokens("response-a", /*total_tokens*/ 120), + ]), + sse(vec![ + ev_response_created("response-b"), + ev_assistant_message("message-b", "done"), + ev_completed_with_tokens("response-b", /*total_tokens*/ 80), + ]), + sse(vec![ + ev_response_created("response-c"), + ev_assistant_message("message-c", "next"), + ev_completed_with_tokens("response-c", /*total_tokens*/ 30), + ]), + sse(vec![ + ev_response_created("response-without-usage"), + ev_assistant_message("message-d", "no usage"), + json!({ + "type": "response.completed", + "response": { + "id": "response-without-usage" + } + }), + ]), + ], + ) + .await; + let test = test_codex().build_with_auto_env(&server).await?; + let rollout_path = test.codex.rollout_path().expect("rollout path"); + let home = test.home.clone(); + + test.submit_turn("first").await?; + test.codex.shutdown_and_wait().await?; + + let resumed = test_codex() + .resume(&server, home, rollout_path.clone()) + .await?; + for prompt in ["second", "third"] { + resumed.submit_turn(prompt).await?; + } + resumed.codex.shutdown_and_wait().await?; + + let records = token_usage_records(&rollout_path); + assert_eq!(records.len(), 3); + assert_eq!( + records + .iter() + .map(|record| { + ( + record.response_id.as_str(), + record.turn_token_usage.total_tokens, + record.thread_token_usage.total_tokens, + ) + }) + .collect::>(), + vec![ + ("response-a", 120, 120), + ("response-b", 200, 200), + ("response-c", 30, 230), + ] + ); + assert_eq!(records[0].turn_id, records[1].turn_id); + assert_ne!(records[1].turn_id, records[2].turn_id); + assert!(records.iter().all(|record| { + record.session_id == SessionId::from(record.thread_id) + && record.root_turn_id == record.turn_id + })); + + Ok(()) +} diff --git a/codex-rs/ext/goal/src/api.rs b/codex-rs/ext/goal/src/api.rs index a643b90ff2..b414e9ccfb 100644 --- a/codex-rs/ext/goal/src/api.rs +++ b/codex-rs/ext/goal/src/api.rs @@ -276,7 +276,7 @@ impl GoalService { }; if let Some(runtime) = runtime.as_ref() { - runtime.invalidate_turn_lineage().await; + runtime.clear_pending_turn_start_options().await; } if objective.is_some() { @@ -321,7 +321,7 @@ impl GoalService { })?; let cleared = cleared_goal.is_some(); if cleared && let Some(runtime) = runtime.as_ref() { - runtime.invalidate_turn_lineage().await; + runtime.clear_pending_turn_start_options().await; } drop(goal_state_permit); drop(runtime); diff --git a/codex-rs/ext/goal/src/runtime.rs b/codex-rs/ext/goal/src/runtime.rs index 9a3594b760..84fa4204ba 100644 --- a/codex-rs/ext/goal/src/runtime.rs +++ b/codex-rs/ext/goal/src/runtime.rs @@ -135,7 +135,7 @@ impl GoalRuntimeHandle { self.inner.root_accounting_state.clone() } - pub(crate) async fn invalidate_turn_lineage(&self) { + pub(crate) async fn clear_pending_turn_start_options(&self) { let Some(thread_manager) = self.inner.thread_manager.upgrade() else { return; }; @@ -143,9 +143,6 @@ impl GoalRuntimeHandle { return; }; thread.thread_extension_data().remove::(); - if let Some(turn_id) = self.inner.accounting_state.current_turn_id() { - thread.invalidate_turn_lineage(turn_id.as_str()).await; - } } pub(crate) async fn goal_state_permit(&self) -> Result, String> { diff --git a/codex-rs/external-agent-migration/src/sessions/append.rs b/codex-rs/external-agent-migration/src/sessions/append.rs index d76a54b7ce..04dd151279 100644 --- a/codex-rs/external-agent-migration/src/sessions/append.rs +++ b/codex-rs/external-agent-migration/src/sessions/append.rs @@ -219,6 +219,7 @@ fn source_model_items(items: &[RolloutItem]) -> Option>> match item { RolloutItem::SessionMeta(_) | RolloutItem::InterAgentCommunicationMetadata { .. } + | RolloutItem::TokenUsageRecord(_) | RolloutItem::RealtimeItem(_) => {} RolloutItem::ResponseItem(response_item) => { model_items.push(SourceModelItem { @@ -257,6 +258,7 @@ fn history_model_items(items: &[RolloutItem]) -> Option> { match item { RolloutItem::SessionMeta(_) | RolloutItem::InterAgentCommunicationMetadata { .. } + | RolloutItem::TokenUsageRecord(_) | RolloutItem::RealtimeItem(_) | RolloutItem::SecurityRiskScore(_) => {} RolloutItem::ResponseItem(response_item) => model_items.push(&response_item.item), diff --git a/codex-rs/history/src/lib.rs b/codex-rs/history/src/lib.rs index 0241394fcb..a87c83a279 100644 --- a/codex-rs/history/src/lib.rs +++ b/codex-rs/history/src/lib.rs @@ -21,6 +21,7 @@ use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadSource; +use codex_protocol::protocol::TokenUsageRecord; use codex_protocol::protocol::TurnContextItem; use codex_protocol::protocol::WorldStateItem; use codex_protocol::realtime::RealtimeItem; @@ -107,6 +108,7 @@ pub enum RolloutItem { }, Compacted(CompactedItem), TurnContext(TurnContextItem), + TokenUsageRecord(TokenUsageRecord), WorldState(WorldStateItem), SecurityRiskScore(SecurityRiskScore), EventMsg(EventMsg), @@ -157,6 +159,13 @@ pub struct CompactedItem { pub first_window_id: Option, pub previous_window_id: Option, pub window_id: Option, + /// Responses API ID for the model-backed compaction request, when one exists. + pub compaction_response_id: Option, + /// Snapshot of the latest reachable token usage record when this compaction was written. + /// + /// `thread/resume` can restore token usage totals from this field without scanning arbitrarily + /// far past the compaction. + pub latest_token_usage_record: Option, } impl Serialize for CompactedItem { @@ -423,6 +432,7 @@ fn multi_agent_version_from_items( | RolloutItem::InterAgentCommunication(_) | RolloutItem::InterAgentCommunicationMetadata { .. } | RolloutItem::Compacted(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::WorldState(_) | RolloutItem::SecurityRiskScore(_) | RolloutItem::RealtimeItem(_) diff --git a/codex-rs/history/src/rollout_payload.rs b/codex-rs/history/src/rollout_payload.rs index 8f65619363..d68f51c3e1 100644 --- a/codex-rs/history/src/rollout_payload.rs +++ b/codex-rs/history/src/rollout_payload.rs @@ -11,6 +11,7 @@ use super::ResponseItemEnvelope; use super::RolloutItem; use super::SecurityRiskScore; use super::SessionMetaLine; +use super::TokenUsageRecord; use super::TurnContextItem; use super::WorldStateItem; use schemars::JsonSchema; @@ -41,6 +42,9 @@ pub(super) enum RolloutItemWire<'a> { TurnContext { payload: Cow<'a, TurnContextItem>, }, + TokenUsageRecord { + payload: Cow<'a, TokenUsageRecord>, + }, WorldState { payload: Cow<'a, WorldStateItem>, }, @@ -81,6 +85,9 @@ impl<'a> From<&'a RolloutItem> for RolloutItemWire<'a> { RolloutItem::TurnContext(payload) => Self::TurnContext { payload: Cow::Borrowed(payload), }, + RolloutItem::TokenUsageRecord(payload) => Self::TokenUsageRecord { + payload: Cow::Borrowed(payload), + }, RolloutItem::WorldState(payload) => Self::WorldState { payload: Cow::Borrowed(payload), }, @@ -117,6 +124,9 @@ impl From> for RolloutItem { } RolloutItemWire::Compacted { payload } => Self::Compacted(payload.into_owned()), RolloutItemWire::TurnContext { payload } => Self::TurnContext(payload.into_owned()), + RolloutItemWire::TokenUsageRecord { payload } => { + Self::TokenUsageRecord(payload.into_owned()) + } RolloutItemWire::WorldState { payload } => Self::WorldState(payload.into_owned()), RolloutItemWire::SecurityRiskScore { payload } => { Self::SecurityRiskScore(payload.into_owned()) @@ -150,6 +160,10 @@ pub(super) struct CompactedItemWire<'a> { #[serde(default, skip_serializing_if = "Option::is_none")] #[schemars(with = "Option")] window_id: Option>, + #[serde(default)] + compaction_response_id: Option>, + #[serde(default)] + latest_token_usage_record: Option>, } impl<'a> From<&'a CompactedItem> for CompactedItemWire<'a> { @@ -185,6 +199,8 @@ impl<'a> From<&'a CompactedItem> for CompactedItemWire<'a> { .window_id .as_deref() .map(|window_id| WindowIdWire::Id(Cow::Borrowed(window_id))), + compaction_response_id: item.compaction_response_id.as_deref().map(Cow::Borrowed), + latest_token_usage_record: item.latest_token_usage_record.as_ref().map(Cow::Borrowed), } } } @@ -246,6 +262,8 @@ impl TryFrom> for CompactedItem { first_window_id: item.first_window_id.map(Cow::into_owned), previous_window_id: item.previous_window_id.map(Cow::into_owned), window_id, + compaction_response_id: item.compaction_response_id.map(Cow::into_owned), + latest_token_usage_record: item.latest_token_usage_record.map(Cow::into_owned), }) } } diff --git a/codex-rs/history/src/tests.rs b/codex-rs/history/src/tests.rs index cca1050cd7..de015b014a 100644 --- a/codex-rs/history/src/tests.rs +++ b/codex-rs/history/src/tests.rs @@ -142,8 +142,8 @@ fn response_item_envelope_ignores_unknown_harness_metadata_fields() -> Result<() } #[test] -/// Keeps legacy compacted replacement histories readable and shape-compatible. -fn response_item_replacement_history_preserves_shape() -> Result<()> { +/// Keeps legacy compacted replacement histories readable while adding nullable checkpoints. +fn response_item_replacement_history_accepts_legacy_shape() -> Result<()> { let legacy_item = json!({ "message": "summary", "replacement_history": [{ @@ -166,7 +166,10 @@ fn response_item_replacement_history_preserves_shape() -> Result<()> { ResponseItem::Message { .. } )); assert_eq!(replacement_history[0].metadata, None); - assert_eq!(serde_json::to_value(item)?, legacy_item); + let mut expected_item = legacy_item; + expected_item["compaction_response_id"] = json!(null); + expected_item["latest_token_usage_record"] = json!(null); + assert_eq!(serde_json::to_value(item)?, expected_item); Ok(()) } @@ -196,6 +199,8 @@ fn compacted_replacement_history_stores_metadata_in_an_aligned_sidecar() -> Resu first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }; let serialized = serde_json::to_value(item)?; @@ -208,6 +213,8 @@ fn compacted_replacement_history_stores_metadata_in_an_aligned_sidecar() -> Resu { "client_authored": true }, { "client_authored": false }, ], + "compaction_response_id": null, + "latest_token_usage_record": null, }) ); @@ -300,6 +307,8 @@ fn compacted_metadata_remains_compatible_with_legacy_response_item_readers() -> first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }))?; let LegacyRolloutItem::Compacted(legacy) = @@ -343,7 +352,11 @@ fn rollout_item_variants_preserve_existing_payload_shapes() -> Result<()> { }), json!({ "type": "compacted", - "payload": { "message": "summary" }, + "payload": { + "message": "summary", + "compaction_response_id": null, + "latest_token_usage_record": null, + }, }), json!({ "type": "turn_context", @@ -355,6 +368,40 @@ fn rollout_item_variants_preserve_existing_payload_shapes() -> Result<()> { "summary": "auto", }, }), + json!({ + "type": "token_usage_record", + "payload": { + "thread_id": "0195cda5-433d-7f9a-9d7b-a9f15b60c2e2", + "turn_id": "turn-1", + "session_id": "0195cda5-433d-7f9a-9d7b-a9f15b60c2e2", + "root_turn_id": "turn-1", + "response_id": "response-1", + "usage": { + "input_tokens": 10, + "cached_input_tokens": 2, + "cache_write_input_tokens": 0, + "output_tokens": 3, + "reasoning_output_tokens": 1, + "total_tokens": 13, + }, + "turn_token_usage": { + "input_tokens": 10, + "cached_input_tokens": 2, + "cache_write_input_tokens": 0, + "output_tokens": 3, + "reasoning_output_tokens": 1, + "total_tokens": 13, + }, + "thread_token_usage": { + "input_tokens": 10, + "cached_input_tokens": 2, + "cache_write_input_tokens": 0, + "output_tokens": 3, + "reasoning_output_tokens": 1, + "total_tokens": 13, + }, + }, + }), json!({ "type": "world_state", "payload": { "full": true, "state": { "cwd": "/tmp" } }, @@ -396,7 +443,7 @@ fn rollout_item_variants_preserve_existing_payload_shapes() -> Result<()> { fn rollout_item_schema_matches_tagged_payload_and_sibling_metadata() -> Result<()> { let schema = serde_json::to_value(schemars::schema_for!(RolloutItem))?; let variants = schema["oneOf"].as_array().expect("rollout variants"); - assert_eq!(variants.len(), 10); + assert_eq!(variants.len(), 11); for variant in variants { let required = variant["required"].as_array().expect("required fields"); @@ -452,6 +499,8 @@ fn compacted_item_serializes_window_number_and_id() -> Result<()> { first_window_id: Some("019b3f6e-0000-7000-8000-000000000001".to_string()), previous_window_id: Some("019b3f6e-0000-7000-8000-000000000002".to_string()), window_id: Some("019b3f6e-7a10-7cc3-8b6e-1d09e2f7a001".to_string()), + compaction_response_id: None, + latest_token_usage_record: None, }; assert_eq!( @@ -462,6 +511,8 @@ fn compacted_item_serializes_window_number_and_id() -> Result<()> { "first_window_id": "019b3f6e-0000-7000-8000-000000000001", "previous_window_id": "019b3f6e-0000-7000-8000-000000000002", "window_id": "019b3f6e-7a10-7cc3-8b6e-1d09e2f7a001", + "compaction_response_id": null, + "latest_token_usage_record": null, }) ); Ok(()) @@ -485,6 +536,8 @@ fn compacted_item_migrates_legacy_numeric_window_id() -> Result<()> { first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, } ); Ok(()) diff --git a/codex-rs/memories/write/src/phase1.rs b/codex-rs/memories/write/src/phase1.rs index 0646979939..f4724201ce 100644 --- a/codex-rs/memories/write/src/phase1.rs +++ b/codex-rs/memories/write/src/phase1.rs @@ -417,6 +417,7 @@ mod job { | RolloutItem::InterAgentCommunicationMetadata { .. } | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::RealtimeItem(_) | RolloutItem::WorldState(_) | RolloutItem::SecurityRiskScore(_) diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 3a8c087efa..db862f866a 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2210,6 +2210,19 @@ pub struct TokenUsage { pub codex_rollout_budget_units: Option, } +/// Best-effort Responses API usage observed for one completed response. +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)] +pub struct TokenUsageRecord { + pub thread_id: ThreadId, + pub turn_id: String, + pub session_id: SessionId, + pub root_turn_id: String, + pub response_id: String, + pub usage: TokenUsage, + pub turn_token_usage: TokenUsage, + pub thread_token_usage: TokenUsage, +} + #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)] pub struct TokenUsageInfo { pub total_token_usage: TokenUsage, @@ -2995,6 +3008,7 @@ pub struct HistoryPosition { /// and should be used when there is no config override. #[derive(Serialize, Deserialize, Clone, Debug, JsonSchema, TS)] pub struct SessionMeta { + /// session_id is equal to the root thread's ID. pub session_id: SessionId, pub id: ThreadId, #[serde(skip_serializing_if = "Option::is_none")] @@ -3158,6 +3172,10 @@ pub struct TurnContextNetworkItem { pub struct TurnContextItem { #[serde(default, skip_serializing_if = "Option::is_none")] pub turn_id: Option, + /// Root turn that owns this subagent turn's attribution. + /// Only set for subagent turns; persisted so resume keeps the scope frozen at turn start. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub root_turn_id: Option, pub cwd: AbsolutePathBuf, /// Effective workspace roots used to materialize symbolic /// `:workspace_roots` filesystem permissions in `permission_profile`. @@ -5957,6 +5975,7 @@ mod tests { fn turn_context_item_serializes_network_when_present() -> Result<()> { let item = TurnContextItem { turn_id: None, + root_turn_id: None, cwd: test_path_buf("/tmp").abs(), workspace_roots: None, current_date: None, diff --git a/codex-rs/rollout/src/list.rs b/codex-rs/rollout/src/list.rs index 3c775e3124..f01fee8200 100644 --- a/codex-rs/rollout/src/list.rs +++ b/codex-rs/rollout/src/list.rs @@ -1174,6 +1174,9 @@ async fn read_head_summary(path: &Path, head_limit: usize) -> io::Result { // Not included in `head`; skip. } + RolloutItem::TokenUsageRecord(_) => { + // Not included in `head`; skip. + } RolloutItem::WorldState(_) | RolloutItem::SecurityRiskScore(_) => { // Not included in `head`; skip. } @@ -1249,6 +1252,7 @@ pub async fn read_head_for_summary(path: &Path) -> io::Result io::Result RolloutItem::InterAgentCommunicationMetadata { .. } | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::WorldState(_) | RolloutItem::RealtimeItem(_) | RolloutItem::SecurityRiskScore(_) diff --git a/codex-rs/rollout/src/metadata.rs b/codex-rs/rollout/src/metadata.rs index 3ac5a744cb..f76652329d 100644 --- a/codex-rs/rollout/src/metadata.rs +++ b/codex-rs/rollout/src/metadata.rs @@ -77,6 +77,7 @@ pub fn builder_from_items( | RolloutItem::TurnContext(_) | RolloutItem::WorldState(_) | RolloutItem::RealtimeItem(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::SecurityRiskScore(_) | RolloutItem::EventMsg(_) => None, }) && let Some(builder) = builder_from_session_meta(session_meta, rollout_path) @@ -167,6 +168,7 @@ pub async fn extract_metadata_from_rollout( | RolloutItem::TurnContext(_) | RolloutItem::WorldState(_) | RolloutItem::RealtimeItem(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::SecurityRiskScore(_) | RolloutItem::EventMsg(_) => None, }), diff --git a/codex-rs/rollout/src/metadata_tests.rs b/codex-rs/rollout/src/metadata_tests.rs index 9ab474b13a..3d493c3327 100644 --- a/codex-rs/rollout/src/metadata_tests.rs +++ b/codex-rs/rollout/src/metadata_tests.rs @@ -296,6 +296,8 @@ fn builder_from_items_falls_back_to_filename() { first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, })]; let builder = builder_from_items(items.as_slice(), path.as_path()).expect("builder"); diff --git a/codex-rs/rollout/src/model_context.rs b/codex-rs/rollout/src/model_context.rs index f41021442e..f5328ca608 100644 --- a/codex-rs/rollout/src/model_context.rs +++ b/codex-rs/rollout/src/model_context.rs @@ -160,7 +160,8 @@ impl ModelContextScan { | RolloutItem::SessionMeta(_) | RolloutItem::InterAgentCommunicationMetadata { .. } | RolloutItem::RealtimeItem(_) - | RolloutItem::SecurityRiskScore(_) => {} + | RolloutItem::SecurityRiskScore(_) + | RolloutItem::TokenUsageRecord(_) => {} } if self.has_bounded_cutoff() { diff --git a/codex-rs/rollout/src/persistence_metrics.rs b/codex-rs/rollout/src/persistence_metrics.rs index 968578b60f..ca046e1f61 100644 --- a/codex-rs/rollout/src/persistence_metrics.rs +++ b/codex-rs/rollout/src/persistence_metrics.rs @@ -235,6 +235,7 @@ fn rollout_item_type(item: &RolloutItem) -> String { } RolloutItem::Compacted(_) => "compacted".to_string(), RolloutItem::TurnContext(_) => "turn_context".to_string(), + RolloutItem::TokenUsageRecord(_) => "token_usage_record".to_string(), RolloutItem::WorldState(_) => "world_state".to_string(), RolloutItem::SecurityRiskScore(_) => "security_risk_score".to_string(), RolloutItem::RealtimeItem(item) => match &item.content { diff --git a/codex-rs/rollout/src/policy.rs b/codex-rs/rollout/src/policy.rs index 56be88b738..cc54ac5554 100644 --- a/codex-rs/rollout/src/policy.rs +++ b/codex-rs/rollout/src/policy.rs @@ -17,6 +17,7 @@ pub fn is_persisted_rollout_item(item: &RolloutItem, history_mode: ThreadHistory // Persist Codex executive markers so we can analyze flows (e.g., compaction, API turns). RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::WorldState(_) | RolloutItem::SecurityRiskScore(_) | RolloutItem::SessionMeta(_) => true, diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 1ca019356e..dff3f52d3a 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -2092,6 +2092,7 @@ async fn resume_candidate_matches_cwd( | RolloutItem::Compacted(_) | RolloutItem::WorldState(_) | RolloutItem::RealtimeItem(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::SecurityRiskScore(_) | RolloutItem::EventMsg(_) => None, }) diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index 0056d0ef77..23afea6153 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -1718,6 +1718,7 @@ async fn resume_candidate_matches_cwd_reads_latest_turn_context() -> std::io::Re ordinal: None, item: RolloutItem::TurnContext(TurnContextItem { turn_id: Some("turn-1".to_string()), + root_turn_id: None, cwd: serde_json::from_value(serde_json::json!(&latest_cwd)) .expect("absolute latest cwd"), workspace_roots: None, diff --git a/codex-rs/rollout/src/search.rs b/codex-rs/rollout/src/search.rs index 946e2672fd..cbfd54de8c 100644 --- a/codex-rs/rollout/src/search.rs +++ b/codex-rs/rollout/src/search.rs @@ -294,6 +294,7 @@ fn conversation_text_from_item(item: &RolloutItem) -> Option { | RolloutItem::Compacted(_) | RolloutItem::RealtimeItem(_) | RolloutItem::SecurityRiskScore(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::WorldState(_) => None, } } diff --git a/codex-rs/state/src/extract.rs b/codex-rs/state/src/extract.rs index d2bd33bb3e..94370044b8 100644 --- a/codex-rs/state/src/extract.rs +++ b/codex-rs/state/src/extract.rs @@ -28,6 +28,7 @@ pub fn apply_rollout_item( RolloutItem::WorldState(_) => {} RolloutItem::SecurityRiskScore(_) => {} RolloutItem::RealtimeItem(_) => {} + RolloutItem::TokenUsageRecord(_) => {} } if metadata.model_provider.is_empty() { metadata.model_provider = default_provider.to_string(); @@ -56,6 +57,7 @@ pub fn rollout_item_affects_thread_metadata(item: &RolloutItem) -> bool { | RolloutItem::Compacted(_) | RolloutItem::RealtimeItem(_) | RolloutItem::SecurityRiskScore(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::WorldState(_) => false, } } @@ -424,6 +426,7 @@ mod tests { &mut metadata, &RolloutItem::TurnContext(TurnContextItem { turn_id: Some("turn-1".to_string()), + root_turn_id: None, cwd: serde_json::from_value(serde_json::json!( std::env::current_dir() .expect("current directory") @@ -472,6 +475,7 @@ mod tests { &mut metadata, &RolloutItem::TurnContext(TurnContextItem { turn_id: Some("turn-1".to_string()), + root_turn_id: None, cwd: serde_json::from_value(serde_json::json!( std::env::current_dir() .expect("current directory") @@ -520,6 +524,7 @@ mod tests { &mut metadata, &RolloutItem::TurnContext(TurnContextItem { turn_id: Some("turn-1".to_string()), + root_turn_id: None, cwd: serde_json::from_value(serde_json::json!(&fallback_cwd)) .expect("absolute fallback cwd"), workspace_roots: None, @@ -557,6 +562,7 @@ mod tests { &mut metadata, &RolloutItem::TurnContext(TurnContextItem { turn_id: Some("turn-1".to_string()), + root_turn_id: None, cwd: serde_json::from_value(serde_json::json!( std::env::current_dir() .expect("current directory") diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 3d2e0d923c..a18c051d8e 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -1323,6 +1323,7 @@ pub(super) fn extract_memory_mode(items: &[RolloutItem]) -> Option { | RolloutItem::WorldState(_) | RolloutItem::RealtimeItem(_) | RolloutItem::SecurityRiskScore(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::EventMsg(_) => None, }) } diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index de9f72e618..e83e7ccd46 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -933,6 +933,7 @@ mod tests { let turn_context = |model: &str, approval_policy| { RolloutItem::TurnContext(TurnContextItem { turn_id: Some("turn-1".to_string()), + root_turn_id: None, cwd: serde_json::from_value(serde_json::json!(cwd)).expect("absolute cwd"), workspace_roots: None, current_date: None, diff --git a/codex-rs/thread-store/src/local/model_context_tests.rs b/codex-rs/thread-store/src/local/model_context_tests.rs index a6ec02169d..73f8274118 100644 --- a/codex-rs/thread-store/src/local/model_context_tests.rs +++ b/codex-rs/thread-store/src/local/model_context_tests.rs @@ -729,6 +729,7 @@ fn agent_message(message: &str) -> RolloutItem { fn turn_context(root: &Path, turn_id: &str) -> RolloutItem { RolloutItem::TurnContext(TurnContextItem { turn_id: Some(turn_id.to_string()), + root_turn_id: None, cwd: serde_json::from_value(serde_json::json!(root)).expect("absolute cwd"), workspace_roots: None, current_date: None, @@ -763,5 +764,7 @@ fn compacted(message: &str, replacement_history: Option>) -> R first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }) } diff --git a/codex-rs/thread-store/src/local/rollout_migration/canonicalizer.rs b/codex-rs/thread-store/src/local/rollout_migration/canonicalizer.rs index 9eed043bd0..0607411c6e 100644 --- a/codex-rs/thread-store/src/local/rollout_migration/canonicalizer.rs +++ b/codex-rs/thread-store/src/local/rollout_migration/canonicalizer.rs @@ -292,6 +292,7 @@ impl LegacyRolloutCanonicalizer { } item @ (RolloutItem::InterAgentCommunicationMetadata { .. } | RolloutItem::TurnContext(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::RealtimeItem(_) | RolloutItem::SecurityRiskScore(_) | RolloutItem::WorldState(_)) => { diff --git a/codex-rs/thread-store/src/local/rollout_migration/rollback_plan.rs b/codex-rs/thread-store/src/local/rollout_migration/rollback_plan.rs index f8404b6177..f015caaa32 100644 --- a/codex-rs/thread-store/src/local/rollout_migration/rollback_plan.rs +++ b/codex-rs/thread-store/src/local/rollout_migration/rollback_plan.rs @@ -219,6 +219,9 @@ impl RollbackPlanner { self.pending_turn_records.push(index); } } + RolloutItem::TokenUsageRecord(record) => { + self.assign_targeted_record(index, Some(record.turn_id.as_str())); + } RolloutItem::WorldState(_) | RolloutItem::RealtimeItem(_) => {} RolloutItem::SecurityRiskScore(_) => self.record_boundaries[index] = None, } diff --git a/codex-rs/thread-store/src/local/rollout_migration/rollback_replay.rs b/codex-rs/thread-store/src/local/rollout_migration/rollback_replay.rs index b8d47f36bf..80dfdbcac7 100644 --- a/codex-rs/thread-store/src/local/rollout_migration/rollback_replay.rs +++ b/codex-rs/thread-store/src/local/rollout_migration/rollback_replay.rs @@ -76,6 +76,7 @@ impl ModelReplayPlanner { | RolloutItem::EventMsg(_) | RolloutItem::InterAgentCommunicationMetadata { .. } | RolloutItem::RealtimeItem(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::SecurityRiskScore(_) | RolloutItem::WorldState(_) => return, }; diff --git a/codex-rs/thread-store/src/local/rollout_migration_tests.rs b/codex-rs/thread-store/src/local/rollout_migration_tests.rs index 5f1263d6ed..b1b38ba060 100644 --- a/codex-rs/thread-store/src/local/rollout_migration_tests.rs +++ b/codex-rs/thread-store/src/local/rollout_migration_tests.rs @@ -235,6 +235,8 @@ fn compacted(replacement_history: Vec) -> RolloutItem { first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }) } @@ -1411,6 +1413,8 @@ async fn migration_compacts_subagent_prefix_and_does_not_project_it() { first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), RolloutItem::Compacted(CompactedItem { message: "latest checkpoint".to_string(), @@ -1431,10 +1435,13 @@ async fn migration_compacts_subagent_prefix_and_does_not_project_it() { first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), started("child-turn"), RolloutItem::TurnContext(TurnContextItem { turn_id: Some("child-turn".to_string()), + root_turn_id: None, cwd: serde_json::from_value(json!(home.path())).expect("absolute cwd"), workspace_roots: None, current_date: None, diff --git a/codex-rs/thread-store/src/thread_metadata_sync.rs b/codex-rs/thread-store/src/thread_metadata_sync.rs index e2557dbd52..dbfac21ef5 100644 --- a/codex-rs/thread-store/src/thread_metadata_sync.rs +++ b/codex-rs/thread-store/src/thread_metadata_sync.rs @@ -305,6 +305,7 @@ impl ThreadMetadataSync { | RolloutItem::InterAgentCommunicationMetadata { .. } | RolloutItem::Compacted(_) | RolloutItem::RealtimeItem(_) + | RolloutItem::TokenUsageRecord(_) | RolloutItem::SecurityRiskScore(_) | RolloutItem::WorldState(_) => {} } @@ -587,6 +588,8 @@ mod tests { first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }); let first = sync diff --git a/codex-rs/tui/src/resume_picker_transcript_preview_tests.rs b/codex-rs/tui/src/resume_picker_transcript_preview_tests.rs index efc9a4b043..d40ac1565c 100644 --- a/codex-rs/tui/src/resume_picker_transcript_preview_tests.rs +++ b/codex-rs/tui/src/resume_picker_transcript_preview_tests.rs @@ -78,6 +78,8 @@ fn legacy_transcript_preview_scans_tail_across_compaction() { first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), rollout_user_message("recent user"), rollout_agent_message("commentary one\ncommentary two"), @@ -209,6 +211,8 @@ fn legacy_transcript_preview_falls_back_for_oversized_hidden_record() { first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }), rollout_agent_message("recent assistant"), ], @@ -236,6 +240,8 @@ fn legacy_transcript_preview_falls_back_when_scan_budget_is_exhausted() { first_window_id: None, previous_window_id: None, window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, }); let mut items = vec![rollout_user_message("older user")]; items.extend(std::iter::repeat_n(compacted, 9));