Persist response token usage in rollout history (#41912)

## Why

Resumed threads need to continue per-turn and per-thread token totals without scanning arbitrarily far beyond the latest compaction checkpoint.

## What changed

- Add durable `TokenUsageRecord` rollout items with response, turn, thread, session, and root-turn attribution.
- Restore accumulated usage on resume and snapshot the latest record plus the compaction response ID in compaction checkpoints.
- Preserve root-turn lineage in persisted turn context, while ensuring forked child threads start with their own usage totals.

## Testing

- Cover usage accumulation across multiple responses and resumed turns.
- Cover local and remote compaction checkpoints, invalid remote compaction output, rollout reconstruction, and fork isolation.

GitOrigin-RevId: ef9e0c4a9102a08a2c382be4cdac68c84353c90a
This commit is contained in:
Owen Lin
2026-08-31 21:15:38 +00:00
committed by copyberry
parent a7913390f7
commit 5f79a92e39
67 changed files with 932 additions and 162 deletions
@@ -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(),
@@ -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(_)
@@ -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)]),
@@ -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,
@@ -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,
@@ -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::<serde_json::Value>(&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::<serde_json::Value>(&requests[7])?;
assert_ne!(
continuation_request["client_metadata"]["turn_id"].as_str(),
+1 -2
View File
@@ -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
+5 -1
View File
@@ -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
+129
View File
@@ -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::<RolloutLine>(line).expect("parse rollout line"))
.collect::<Vec<_>>();
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 {
+7 -12
View File
@@ -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<String>,
+12 -13
View File
@@ -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<String>,
}
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<String> {
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),
+1
View File
@@ -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;
+16 -20
View File
@@ -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<TokenUsage>,
usage_metadata: Option<ResponseUsageMetadata>,
}
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<RemoteCompactionV2Output> {
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!(
+2 -14
View File
@@ -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<ResponseItem>,
pub(super) prompt_input_metadata: Vec<Option<CodexHarnessMetadata>>,
pub(super) compaction_output: ResponseItem,
pub(super) compaction_response_id: String,
pub(super) token_usage: Option<TokenUsage>,
/// Keeps a session created for standalone compaction alive through lifecycle completion.
pub(super) owned_client_session: Option<ModelClientSession>,
@@ -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,
})
@@ -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")
+2 -4
View File
@@ -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)?;
+1 -12
View File
@@ -769,21 +769,10 @@ pub(crate) async fn drain_async_hook_results(
.collect::<Vec<_>>();
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;
}
+2 -13
View File
@@ -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<ResponseItem>,
source_turn_id: Option<&str>,
) -> Result<(), Vec<ResponseItem>> {
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::<Vec<_>>();
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(),
+10 -17
View File
@@ -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,
+59
View File
@@ -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<TokenUsageRecord> {
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<PreviousTurnSettings> {
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,
@@ -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")
@@ -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.
+3 -1
View File
@@ -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(_) => {}
}
}
+141 -4
View File
@@ -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");
+4 -7
View File
@@ -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
@@ -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(),
+7 -3
View File
@@ -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(
+43
View File
@@ -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<BaseInstructionsProvenance>,
pub(crate) history: ContextManager,
pub(crate) latest_rate_limits: Option<RateLimitSnapshot>,
pub(crate) latest_token_usage_record: Option<TokenUsageRecord>,
pub(crate) server_reasoning_included: bool,
pub(crate) mcp_dependency_prompted: HashSet<String>,
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<TurnContextItem>) {
self.history.set_reference_context_item(item);
}
+46
View File
@@ -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() {
+5 -8
View File
@@ -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;
+1 -2
View File
@@ -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(),
+1 -10
View File
@@ -129,7 +129,6 @@ pub(crate) struct TurnMetadataState {
turn_started_at_unix_ms: RwLock<Option<i64>>,
responses_api_metadata: RwLock<BTreeMap<String, String>>,
responsesapi_client_metadata: RwLock<BTreeMap<String, String>>,
root_turn_ambiguous: AtomicBool,
user_input_requested_during_turn: AtomicBool,
enrichment_task: Mutex<Option<JoinHandle<()>>>,
git_enrichment_complete: watch::Sender<bool>,
@@ -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<String> {
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 {
+34 -2
View File
@@ -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::<RolloutLine>(line).ok())
.map(|line| line.item)
.collect::<Vec<_>>();
let records = rollout_items
.iter()
.filter_map(|item| match item {
RolloutItem::TokenUsageRecord(record) => Some(record),
_ => None,
})
.collect::<Vec<_>>();
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)]
@@ -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::<RolloutLine>(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(())
}
+5 -1
View File
@@ -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
+1
View File
@@ -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;
+2 -2
View File
@@ -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| {
@@ -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,
@@ -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<TokenUsageRecord> {
std::fs::read_to_string(path)
.expect("read rollout")
.lines()
.filter_map(|line| serde_json::from_str::<RolloutLine>(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<_>>(),
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(())
}
+2 -2
View File
@@ -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);
+1 -4
View File
@@ -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::<TurnStartOptions>();
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<SemaphorePermit<'_>, String> {
@@ -219,6 +219,7 @@ fn source_model_items(items: &[RolloutItem]) -> Option<Vec<SourceModelItem<'_>>>
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<Vec<&ResponseItem>> {
match item {
RolloutItem::SessionMeta(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::TokenUsageRecord(_)
| RolloutItem::RealtimeItem(_)
| RolloutItem::SecurityRiskScore(_) => {}
RolloutItem::ResponseItem(response_item) => model_items.push(&response_item.item),
+10
View File
@@ -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<String>,
pub previous_window_id: Option<String>,
pub window_id: Option<String>,
/// Responses API ID for the model-backed compaction request, when one exists.
pub compaction_response_id: Option<String>,
/// 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<TokenUsageRecord>,
}
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(_)
+18
View File
@@ -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<RolloutItemWire<'_>> 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<String>")]
window_id: Option<WindowIdWire<'a>>,
#[serde(default)]
compaction_response_id: Option<Cow<'a, str>>,
#[serde(default)]
latest_token_usage_record: Option<Cow<'a, TokenUsageRecord>>,
}
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<CompactedItemWire<'_>> 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),
})
}
}
+58 -5
View File
@@ -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(())
+1
View File
@@ -417,6 +417,7 @@ mod job {
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::TokenUsageRecord(_)
| RolloutItem::RealtimeItem(_)
| RolloutItem::WorldState(_)
| RolloutItem::SecurityRiskScore(_)
+19
View File
@@ -2210,6 +2210,19 @@ pub struct TokenUsage {
pub codex_rollout_budget_units: Option<serde_json::Number>,
}
/// 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<String>,
/// 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<String>,
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,
+5
View File
@@ -1174,6 +1174,9 @@ async fn read_head_summary(path: &Path, head_limit: usize) -> io::Result<HeadTai
RolloutItem::TurnContext(_) => {
// 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<Vec<serde_json::Va
RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::TokenUsageRecord(_)
| RolloutItem::WorldState(_)
| RolloutItem::RealtimeItem(_)
| RolloutItem::SecurityRiskScore(_)
@@ -1303,6 +1307,7 @@ pub async fn read_session_meta_line(path: &Path) -> io::Result<SessionMetaLine>
RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_)
| RolloutItem::TurnContext(_)
| RolloutItem::TokenUsageRecord(_)
| RolloutItem::WorldState(_)
| RolloutItem::RealtimeItem(_)
| RolloutItem::SecurityRiskScore(_)
+2
View File
@@ -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,
}),
+2
View File
@@ -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");
+2 -1
View File
@@ -160,7 +160,8 @@ impl ModelContextScan {
| RolloutItem::SessionMeta(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::RealtimeItem(_)
| RolloutItem::SecurityRiskScore(_) => {}
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::TokenUsageRecord(_) => {}
}
if self.has_bounded_cutoff() {
@@ -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 {
+1
View File
@@ -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,
+1
View File
@@ -2092,6 +2092,7 @@ async fn resume_candidate_matches_cwd(
| RolloutItem::Compacted(_)
| RolloutItem::WorldState(_)
| RolloutItem::RealtimeItem(_)
| RolloutItem::TokenUsageRecord(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::EventMsg(_) => None,
})
+1
View File
@@ -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,
+1
View File
@@ -294,6 +294,7 @@ fn conversation_text_from_item(item: &RolloutItem) -> Option<String> {
| RolloutItem::Compacted(_)
| RolloutItem::RealtimeItem(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::TokenUsageRecord(_)
| RolloutItem::WorldState(_) => None,
}
}
+6
View File
@@ -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")
+1
View File
@@ -1323,6 +1323,7 @@ pub(super) fn extract_memory_mode(items: &[RolloutItem]) -> Option<String> {
| RolloutItem::WorldState(_)
| RolloutItem::RealtimeItem(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::TokenUsageRecord(_)
| RolloutItem::EventMsg(_) => None,
})
}
+1
View File
@@ -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,
@@ -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<Vec<ResponseItem>>) -> R
first_window_id: None,
previous_window_id: None,
window_id: None,
compaction_response_id: None,
latest_token_usage_record: None,
})
}
@@ -292,6 +292,7 @@ impl LegacyRolloutCanonicalizer {
}
item @ (RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::TurnContext(_)
| RolloutItem::TokenUsageRecord(_)
| RolloutItem::RealtimeItem(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::WorldState(_)) => {
@@ -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,
}
@@ -76,6 +76,7 @@ impl ModelReplayPlanner {
| RolloutItem::EventMsg(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::RealtimeItem(_)
| RolloutItem::TokenUsageRecord(_)
| RolloutItem::SecurityRiskScore(_)
| RolloutItem::WorldState(_) => return,
};
@@ -235,6 +235,8 @@ fn compacted(replacement_history: Vec<ResponseItem>) -> 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,
@@ -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
@@ -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));