diff --git a/codex-rs/rollout-trace/src/protocol_event.rs b/codex-rs/rollout-trace/src/protocol_event.rs index a911b2b5f3..58481aacb1 100644 --- a/codex-rs/rollout-trace/src/protocol_event.rs +++ b/codex-rs/rollout-trace/src/protocol_event.rs @@ -11,6 +11,7 @@ //! variants a compile-time prompt to decide whether the trace should capture //! them. +use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::ExecCommandBeginEvent; use codex_protocol::protocol::ExecCommandEndEvent; @@ -41,6 +42,7 @@ pub(crate) fn codex_turn_trace_event( thread_id: AgentThreadId, default_turn_id: &str, event: &EventMsg, + store_agent_message: impl FnOnce(&ResponseItem) -> Option, ) -> Option { match event { EventMsg::TurnStarted(event) => { @@ -84,6 +86,10 @@ pub(crate) fn codex_turn_trace_event( context_turn_id: default_turn_id.to_string(), payload: RawTraceEventPayload::HistoryItemAccepted { item_id: event.item.id()?.to_string(), + agent_message_payload: match &event.item { + item @ ResponseItem::AgentMessage { .. } => Some(store_agent_message(item)?), + _ => None, + }, }, }), _ => None, diff --git a/codex-rs/rollout-trace/src/raw_event.rs b/codex-rs/rollout-trace/src/raw_event.rs index ba634b63a1..044a8a122f 100644 --- a/codex-rs/rollout-trace/src/raw_event.rs +++ b/codex-rs/rollout-trace/src/raw_event.rs @@ -102,6 +102,9 @@ pub enum RawTraceEventPayload { /// evidence of its rendered content, which may be truncated before sampling. HistoryItemAccepted { item_id: String, + /// A delivery may never appear in a subsequent inference or compaction snapshot. + #[serde(default, skip_serializing_if = "Option::is_none")] + agent_message_payload: Option, }, InferenceStarted { inference_call_id: InferenceCallId, @@ -260,7 +263,6 @@ impl RawTraceEventPayload { | RawTraceEventPayload::ThreadResumed { .. } | RawTraceEventPayload::CodexTurnStarted { .. } | RawTraceEventPayload::CodexTurnEnded { .. } - | RawTraceEventPayload::HistoryItemAccepted { .. } | RawTraceEventPayload::CompactionRequestFailed { .. } | RawTraceEventPayload::CodeCellStarted { .. } | RawTraceEventPayload::McpToolCallCorrelationAssigned { .. } @@ -269,6 +271,10 @@ impl RawTraceEventPayload { carried_payload: None, .. } => Vec::new(), + RawTraceEventPayload::HistoryItemAccepted { + agent_message_payload, + .. + } => agent_message_payload.iter().collect(), RawTraceEventPayload::ThreadStarted { metadata_payload, .. } => metadata_payload.iter().collect(), diff --git a/codex-rs/rollout-trace/src/reducer/acceptance_tests.rs b/codex-rs/rollout-trace/src/reducer/acceptance_tests.rs index 7651bac6be..5c1113ebc3 100644 --- a/codex-rs/rollout-trace/src/reducer/acceptance_tests.rs +++ b/codex-rs/rollout-trace/src/reducer/acceptance_tests.rs @@ -41,6 +41,7 @@ fn identical_text_from_a_discarded_attempt_does_not_gain_acceptance() -> anyhow: trace_context("turn-1"), RawTraceEventPayload::HistoryItemAccepted { item_id: "accepted".into(), + agent_message_payload: None, }, )?; let response = writer.write_json_payload( @@ -220,6 +221,7 @@ fn inherited_output_is_child_prefill_without_copying_parent_sampling_provenance( ); let accepted = RawTraceEventPayload::HistoryItemAccepted { item_id: "msg-parent".into(), + agent_message_payload: None, }; let context = trace_context("turn-parent"); let first = writer.append_with_context(context.clone(), accepted.clone())?; diff --git a/codex-rs/rollout-trace/src/reducer/conversation.rs b/codex-rs/rollout-trace/src/reducer/conversation.rs index c08753cb7d..6f00a5b53d 100644 --- a/codex-rs/rollout-trace/src/reducer/conversation.rs +++ b/codex-rs/rollout-trace/src/reducer/conversation.rs @@ -27,6 +27,40 @@ use crate::payload::RawPayloadRef; mod normalize; impl TraceReducer { + pub(super) fn record_agent_receipt( + &mut self, + thread_id: &str, + turn_id: Option, + at_unix_ms: i64, + wire_id: &str, + payload: &RawPayloadRef, + ) -> Result<()> { + // A request can be recorded before its history notification. Its item already owns the edge. + if self.rollout.threads[thread_id] + .conversation_item_ids + .iter() + .rev() + .any(|id| { + self.rollout.conversation_items[id] + .responses_item_id + .as_deref() + == Some(wire_id) + }) + { + return Ok(()); + } + let item = normalize::normalize_model_item(&self.read_payload_json(payload)?, payload)?; + if item.agent_message.is_none() || item.responses_item_id.as_deref() != Some(wire_id) { + bail!("accepted agent message does not match receipt {wire_id}"); + } + let id = self.create_conversation_item(thread_id, turn_id, at_unix_ms, item, Vec::new()); + self.append_thread_conversation_items(thread_id, std::slice::from_ref(&id))?; + self.resolve_pending_agent_edges_for_item(&id)?; + self.receipt_only_items + .insert((thread_id.to_owned(), wire_id.to_owned()), id); + Ok(()) + } + /// Reduces an inference request input snapshot into model-visible conversation items. /// /// Full snapshots reuse matching items even if the prompt builder reorders them. @@ -434,6 +468,16 @@ impl TraceReducer { item: NormalizedConversationItem, produced_by: Vec, ) -> String { + // The first request or checkpoint that sees a receipt supplies the model-visible + // representation (which can be trimmed). Keep the ID already used by the delivery edge. + if item.agent_message.is_some() + && let Some(wire_id) = &item.responses_item_id + && let Some(id) = self + .receipt_only_items + .remove(&(thread_id.to_owned(), wire_id.clone())) + { + return id; + } let item_id = self.next_conversation_item_id(); self.rollout.conversation_items.insert( item_id.clone(), @@ -474,6 +518,10 @@ impl TraceReducer { if item.kind == ConversationItemKind::Reasoning { merge_reasoning_body(&mut item.body, &normalized.body)?; } + // A request or checkpoint may contain a trimmed version of an accepted delivery. + if normalized.agent_message.is_some() && item.body != normalized.body { + item.body.clone_from(&normalized.body); + } for producer in produced_by { if !item.produced_by.contains(producer) { item.produced_by.push(producer.clone()); diff --git a/codex-rs/rollout-trace/src/reducer/conversation/normalize.rs b/codex-rs/rollout-trace/src/reducer/conversation/normalize.rs index 6171dd87e2..92484d8a89 100644 --- a/codex-rs/rollout-trace/src/reducer/conversation/normalize.rs +++ b/codex-rs/rollout-trace/src/reducer/conversation/normalize.rs @@ -63,7 +63,7 @@ pub(super) fn token_usage_from_value(value: &Value) -> Option { }) } -fn normalize_model_item( +pub(super) fn normalize_model_item( item: &Value, raw_payload: &RawPayloadRef, ) -> Result { diff --git a/codex-rs/rollout-trace/src/reducer/conversation_tests.rs b/codex-rs/rollout-trace/src/reducer/conversation_tests.rs index ed9b8c4471..4dc95a766c 100644 --- a/codex-rs/rollout-trace/src/reducer/conversation_tests.rs +++ b/codex-rs/rollout-trace/src/reducer/conversation_tests.rs @@ -223,6 +223,7 @@ fn response_outputs_require_harness_acceptance() -> anyhow::Result<()> { trace_context("turn-1"), RawTraceEventPayload::HistoryItemAccepted { item_id: "msg-accepted".into(), + agent_message_payload: None, }, )?; diff --git a/codex-rs/rollout-trace/src/reducer/inference_tests.rs b/codex-rs/rollout-trace/src/reducer/inference_tests.rs index d940a9a5b5..930c85b88e 100644 --- a/codex-rs/rollout-trace/src/reducer/inference_tests.rs +++ b/codex-rs/rollout-trace/src/reducer/inference_tests.rs @@ -31,6 +31,7 @@ fn cancelled_inference_retains_accepted_partial_output() -> anyhow::Result<()> { trace_context("turn-1"), RawTraceEventPayload::HistoryItemAccepted { item_id: "msg-partial".into(), + agent_message_payload: None, }, )?; let partial_response = writer.write_json_payload( diff --git a/codex-rs/rollout-trace/src/reducer/mod.rs b/codex-rs/rollout-trace/src/reducer/mod.rs index 5f5f163502..cfd0dac15a 100644 --- a/codex-rs/rollout-trace/src/reducer/mod.rs +++ b/codex-rs/rollout-trace/src/reducer/mod.rs @@ -83,6 +83,7 @@ pub fn replay_bundle(bundle_dir: impl AsRef) -> Result { pending_code_cell_lifecycle_events: BTreeMap::new(), pending_agent_interaction_edges: Vec::new(), first_acceptances: BTreeMap::new(), + receipt_only_items: BTreeMap::new(), }; let event_log_path = bundle_dir.join(RAW_EVENT_LOG_FILE_NAME); @@ -212,6 +213,8 @@ struct TraceReducer { bundle_dir: PathBuf, /// First (wall time, sequence) by receiving thread and Responses item ID. first_acceptances: BTreeMap>, + /// Delivered items awaiting their first model-visible sighting, keyed by thread and wire ID. + receipt_only_items: BTreeMap<(String, String), String>, next_conversation_item_ordinal: u64, next_terminal_operation_ordinal: u64, /// Latest session for each thread-local runtime handle; each exec replaces it. @@ -334,11 +337,23 @@ impl TraceReducer { status, )?; } - RawTraceEventPayload::HistoryItemAccepted { item_id } => { + RawTraceEventPayload::HistoryItemAccepted { + item_id, + agent_message_payload, + } => { let thread_id = event .thread_id .context("history acceptance has no thread")?; self.thread_mut(&thread_id)?; + if let Some(payload) = agent_message_payload { + self.record_agent_receipt( + &thread_id, + event.codex_turn_id, + event.wall_time_unix_ms, + &item_id, + &payload, + )?; + } self.first_acceptances .entry(thread_id) .or_default() diff --git a/codex-rs/rollout-trace/src/reducer/tool/agents_tests.rs b/codex-rs/rollout-trace/src/reducer/tool/agents_tests.rs index f7fec30c44..4733233ce1 100644 --- a/codex-rs/rollout-trace/src/reducer/tool/agents_tests.rs +++ b/codex-rs/rollout-trace/src/reducer/tool/agents_tests.rs @@ -246,27 +246,68 @@ fn sub_agent_started_activity_creates_spawn_edge() -> anyhow::Result<()> { )?; start_thread(&writer, child_thread_id, "/root/reviewer")?; start_turn_for_thread(&writer, child_thread_id, "turn-child-1")?; + let mut delivery = json!({ + "type": "agent_message", + "id": "amsg-delivery", + "author": "/root", + "recipient": "/root/reviewer", + "content": [ + { + "type": "input_text", + "text": "Message Type: NEW_TASK\nTask name: /root/reviewer\nSender: /root\nPayload:\n" + }, + {"type": "encrypted_content", "encrypted_content": "review this"} + ] + }); + let receipt = RawTraceEventPayload::HistoryItemAccepted { + item_id: "amsg-delivery".into(), + agent_message_payload: Some( + writer.write_json_payload(RawPayloadKind::ProtocolEvent, &delivery)?, + ), + }; + let context = trace_context_for_thread(child_thread_id, "turn-child-1"); + writer.append_with_context(context.clone(), receipt.clone())?; + + // Interrupt or reset here: the receipt already has its delivery edge without any child request. + let received = replay_bundle(temp.path())?; + let edge_id = format!("edge:spawn:019d0000-0000-7000-8000-000000000001:{child_thread_id}"); + let receipt_id = target_conversation_item_id(&received.interaction_edges[&edge_id].target); + let delivered_item_ids = std::slice::from_ref(receipt_id); + assert_eq!( + received.threads[child_thread_id].conversation_item_ids, + delivered_item_ids + ); + assert!( + received.conversation_items[receipt_id] + .accepted_seq + .is_some() + ); + assert!(received.inference_calls.is_empty()); + + // The request trims the encrypted part. A repeated receipt must retain the model's view. + delivery["content"] = json!([delivery["content"][0]]); append_inference_request( &writer, child_thread_id, "turn-child-1", "inference-child-1", - vec![json!({ - "type": "agent_message", - "author": "/root", - "recipient": "/root/reviewer", - "content": [ - { - "type": "input_text", - "text": "Message Type: NEW_TASK\nTask name: /root/reviewer\nSender: /root\nPayload:\n" - }, - {"type": "encrypted_content", "encrypted_content": "review this"} - ] - })], + vec![delivery], )?; + writer.append_with_context(context, receipt)?; let replayed = replay_bundle(temp.path())?; - let edge_id = format!("edge:spawn:019d0000-0000-7000-8000-000000000001:{child_thread_id}"); + assert_eq!( + replayed.inference_calls["inference-child-1"].request_item_ids, + delivered_item_ids + ); + assert_eq!( + replayed.threads[child_thread_id].conversation_item_ids, + delivered_item_ids + ); + assert_eq!( + replayed.conversation_items[receipt_id].body.parts, + received.conversation_items[receipt_id].body.parts[..1] + ); let edge = &replayed.interaction_edges[&edge_id]; assert_eq!(edge.kind, InteractionEdgeKind::SpawnAgent); let target_item_id = target_conversation_item_id(&edge.target); diff --git a/codex-rs/rollout-trace/src/thread.rs b/codex-rs/rollout-trace/src/thread.rs index c96cc30176..6c911c161b 100644 --- a/codex-rs/rollout-trace/src/thread.rs +++ b/codex-rs/rollout-trace/src/thread.rs @@ -226,7 +226,9 @@ impl ThreadTraceContext { return; }; let Some(trace_event) = - codex_turn_trace_event(context.thread_id.clone(), default_turn_id, event) + codex_turn_trace_event(context.thread_id.clone(), default_turn_id, event, |item| { + context.write_json_payload_best_effort(RawPayloadKind::ProtocolEvent, item) + }) else { return; };