From 70fb10676a2cc4799a98ec75243457949cec51da Mon Sep 17 00:00:00 2001 From: Albin Cassirer Date: Thu, 16 Apr 2026 10:06:41 -0700 Subject: [PATCH] Fix edge issues when sub agent threads are complete failures. The edges were assuming that there are assistant messages to attach to but if everything fails straight away then we now fall back to attaching the edges to the sub agent thread instead. --- codex-rs/rollout-trace/src/reducer/mod.rs | 4 + .../rollout-trace/src/reducer/tool/agents.rs | 74 +++- .../src/reducer/tool/agents_tests.rs | 317 +++++++++++++----- 3 files changed, 306 insertions(+), 89 deletions(-) diff --git a/codex-rs/rollout-trace/src/reducer/mod.rs b/codex-rs/rollout-trace/src/reducer/mod.rs index 3dd480cc87..e4f6b837f4 100644 --- a/codex-rs/rollout-trace/src/reducer/mod.rs +++ b/codex-rs/rollout-trace/src/reducer/mod.rs @@ -77,6 +77,10 @@ pub fn replay_bundle(bundle_dir: impl AsRef) -> Result { .with_context(|| format!("parse trace event line {}", line_index + 1))?; reducer.apply_event(event)?; } + // Spawn edges prefer the child task message as their target, but a child can + // fail before that message is ever reduced. Only after replaying the whole + // bundle do we know which spawn deliveries need the child-thread fallback. + reducer.resolve_pending_spawn_edge_fallbacks()?; Ok(reducer.rollout) } diff --git a/codex-rs/rollout-trace/src/reducer/tool/agents.rs b/codex-rs/rollout-trace/src/reducer/tool/agents.rs index 463a4e7f9c..a49b794d14 100644 --- a/codex-rs/rollout-trace/src/reducer/tool/agents.rs +++ b/codex-rs/rollout-trace/src/reducer/tool/agents.rs @@ -23,13 +23,15 @@ use crate::payload::RawPayloadRef; /// /// Multi-agent v2 records the sender tool before the target thread necessarily /// includes the delivered mailbox message in a model-visible request. The edge -/// stays pending until it can target that exact conversation item. +/// stays pending so it can target that exact conversation item when possible. pub(in crate::reducer) struct PendingAgentInteractionEdge { pub(in crate::reducer) edge_id: String, pub(in crate::reducer) kind: InteractionEdgeKind, pub(in crate::reducer) source: TraceAnchor, pub(in crate::reducer) target_thread_id: String, pub(in crate::reducer) message_content: String, + /// Spawn-only fallback for children that fail before their task message is model-visible. + pub(in crate::reducer) unresolved_spawn_thread_id: Option, pub(in crate::reducer) started_at_unix_ms: i64, pub(in crate::reducer) ended_at_unix_ms: Option, pub(in crate::reducer) carried_raw_payload_ids: Vec, @@ -225,8 +227,9 @@ impl TraceReducer { source: TraceAnchor::ToolCall { tool_call_id: tool_call_id.to_string(), }, - target_thread_id: child_thread_id, + target_thread_id: child_thread_id.clone(), message_content: payload.prompt.clone(), + unresolved_spawn_thread_id: Some(child_thread_id), started_at_unix_ms: tool_call.execution.started_at_unix_ms, ended_at_unix_ms: Some(wall_time_unix_ms), carried_raw_payload_ids: self.agent_tool_payload_ids(tool_call_id)?, @@ -266,6 +269,7 @@ impl TraceReducer { }, target_thread_id, message_content, + unresolved_spawn_thread_id: None, started_at_unix_ms: tool_call.execution.started_at_unix_ms, ended_at_unix_ms, carried_raw_payload_ids: self.agent_tool_payload_ids(tool_call_id)?, @@ -326,29 +330,36 @@ impl TraceReducer { }) } - /// Queues or resolves the edge from a child agent's final message to its parent notification. + /// Queues or resolves the edge from a child completion to its parent notification. pub(in crate::reducer) fn queue_agent_result_interaction_edge( &mut self, observed: ObservedAgentResultEdge, ) -> Result<()> { - let Some(source_item_id) = self.latest_assistant_message_item_for_turn( + let source = if let Some(source_item_id) = self.latest_assistant_message_item_for_turn( &observed.child_thread_id, &observed.child_codex_turn_id, - ) else { - bail!( - "agent result edge {} could not find a child result message", - observed.edge_id - ); + ) { + TraceAnchor::ConversationItem { + item_id: source_item_id, + } + } else { + // Child completion is delivered from AgentStatus, not from transcript + // content. Failed or cancelled children can therefore notify the parent + // without producing a final assistant message. Anchor those edges to + // the child thread so the trace keeps the valid delivery instead of + // inventing a missing conversation item. + TraceAnchor::Thread { + thread_id: observed.child_thread_id, + } }; self.queue_or_resolve_agent_interaction_edge(PendingAgentInteractionEdge { edge_id: observed.edge_id, kind: InteractionEdgeKind::AgentResult, - source: TraceAnchor::ConversationItem { - item_id: source_item_id, - }, + source, target_thread_id: observed.parent_thread_id, message_content: observed.message, + unresolved_spawn_thread_id: None, started_at_unix_ms: observed.wall_time_unix_ms, ended_at_unix_ms: Some(observed.wall_time_unix_ms), carried_raw_payload_ids: observed @@ -399,6 +410,7 @@ impl TraceReducer { || existing.source != pending.source || existing.target_thread_id != pending.target_thread_id || existing.message_content != pending.message_content + || existing.unresolved_spawn_thread_id != pending.unresolved_spawn_thread_id { bail!( "pending interaction edge {} was observed with conflicting delivery data", @@ -425,6 +437,44 @@ impl TraceReducer { Ok(()) } + /// Materializes unresolved spawn edges that have a valid child-thread fallback target. + pub(in crate::reducer) fn resolve_pending_spawn_edge_fallbacks(&mut self) -> Result<()> { + let pending_edges = std::mem::take(&mut self.pending_agent_interaction_edges); + for pending in pending_edges { + let Some(child_thread_id) = pending.unresolved_spawn_thread_id else { + continue; + }; + if pending.kind != InteractionEdgeKind::SpawnAgent { + bail!( + "non-spawn interaction edge {} carried a spawn fallback target", + pending.edge_id + ); + } + if !self.rollout.threads.contains_key(&child_thread_id) { + continue; + } + + // Spawn normally resolves to the child task message because that is + // where the delegated work first becomes model-visible. A child can + // fail before that transcript item exists, but the spawned thread is + // still real and the spawning tool still created it. Preserve that + // relationship with the thread fallback instead of dropping the edge. + self.upsert_interaction_edge(InteractionEdge { + edge_id: pending.edge_id, + kind: pending.kind, + source: pending.source, + target: TraceAnchor::Thread { + thread_id: child_thread_id, + }, + started_at_unix_ms: pending.started_at_unix_ms, + ended_at_unix_ms: pending.ended_at_unix_ms, + carried_item_ids: Vec::new(), + carried_raw_payload_ids: pending.carried_raw_payload_ids, + })?; + } + Ok(()) + } + fn upsert_agent_interaction_edge_for_item( &mut self, pending: PendingAgentInteractionEdge, 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 8d4af37fae..deaf06dac7 100644 --- a/codex-rs/rollout-trace/src/reducer/tool/agents_tests.rs +++ b/codex-rs/rollout-trace/src/reducer/tool/agents_tests.rs @@ -10,6 +10,7 @@ use crate::model::ToolCallKind; use crate::model::ToolCallSummary; use crate::model::TraceAnchor; use crate::payload::RawPayloadKind; +use crate::payload::RawPayloadRef; use crate::raw_event::RawToolCallRequester; use crate::raw_event::RawTraceEventPayload; use crate::reducer::test_support::append_completed_inference; @@ -87,79 +88,12 @@ fn spawn_runtime_payload_targets_delivered_child_message() -> anyhow::Result<()> let temp = TempDir::new()?; let writer = create_started_agent_writer(&temp)?; start_agent_turn(&writer, "turn-1")?; - let invocation_payload = writer.write_json_payload( - RawPayloadKind::ToolInvocation, - &json!({ - "tool_name": "spawn_agent", - "payload": { - "type": "function", - "arguments": "{\"task_name\":\"repo_file_counter\",\"message\":\"count\"}" - } - }), - )?; - writer.append_with_context( - trace_context_for_agent("turn-1"), - RawTraceEventPayload::ToolCallStarted { - tool_call_id: "call-spawn".to_string(), - model_visible_call_id: Some("call-spawn".to_string()), - code_mode_runtime_tool_id: None, - requester: RawToolCallRequester::Model, - kind: ToolCallKind::SpawnAgent, - summary: ToolCallSummary::Generic { - label: "spawn_agent".to_string(), - input_preview: None, - output_preview: None, - }, - invocation_payload: Some(invocation_payload.clone()), - }, - )?; - let begin_payload = writer.write_json_payload( - RawPayloadKind::ToolRuntimeEvent, - &json!({ - "call_id": "call-spawn", - "sender_thread_id": "019d0000-0000-7000-8000-000000000001", - "prompt": "count" - }), - )?; - writer.append_with_context( - trace_context_for_agent("turn-1"), - RawTraceEventPayload::ToolCallRuntimeStarted { - tool_call_id: "call-spawn".to_string(), - runtime_payload: begin_payload.clone(), - }, - )?; - let end_payload = writer.write_json_payload( - RawPayloadKind::ToolRuntimeEvent, - &json!({ - "call_id": "call-spawn", - "sender_thread_id": "019d0000-0000-7000-8000-000000000001", - "new_thread_id": "019d0000-0000-7000-8000-000000000002", - "prompt": "count", - "model": "gpt-test", - "reasoning_effort": "medium", - "status": "running" - }), - )?; - writer.append_with_context( - trace_context_for_agent("turn-1"), - RawTraceEventPayload::ToolCallRuntimeEnded { - tool_call_id: "call-spawn".to_string(), - status: ExecutionStatus::Completed, - runtime_payload: end_payload.clone(), - }, - )?; - let result_payload = writer.write_json_payload( - RawPayloadKind::ToolResult, - &json!({"task_name": "/root/repo_file_counter"}), - )?; - writer.append_with_context( - trace_context_for_agent("turn-1"), - RawTraceEventPayload::ToolCallEnded { - tool_call_id: "call-spawn".to_string(), - status: ExecutionStatus::Completed, - result_payload: Some(result_payload.clone()), - }, - )?; + + let spawn_payloads = append_spawn_agent_tool_lifecycle(&writer, "turn-1")?; + + // Then record the child-side model-visible task message. This is the + // preferred target because it pinpoints where the delegated work entered + // the child timeline. start_thread( &writer, "019d0000-0000-7000-8000-000000000002", @@ -202,10 +136,58 @@ fn spawn_runtime_payload_targets_delivered_child_message() -> anyhow::Result<()> assert_eq!( edge.carried_raw_payload_ids, vec![ - invocation_payload.raw_payload_id, - begin_payload.raw_payload_id, - end_payload.raw_payload_id, - result_payload.raw_payload_id, + spawn_payloads.invocation.raw_payload_id, + spawn_payloads.begin.raw_payload_id, + spawn_payloads.end.raw_payload_id, + spawn_payloads.result.raw_payload_id, + ] + ); + + Ok(()) +} + +#[test] +fn spawn_runtime_payload_falls_back_to_child_thread_without_delivery_item() -> anyhow::Result<()> { + let temp = TempDir::new()?; + let writer = create_started_agent_writer(&temp)?; + start_agent_turn(&writer, "turn-1")?; + let spawn_payloads = append_spawn_agent_tool_lifecycle(&writer, "turn-1")?; + + // Deliberately start the child thread without appending an inference + // request containing the inter-agent task message. This reproduces the + // failure path where the child aborts before the reducer can target the + // precise child-side ConversationItem. + start_thread( + &writer, + "019d0000-0000-7000-8000-000000000002", + "/root/repo_file_counter", + )?; + + let replayed = replay_bundle(temp.path())?; + let edge = &replayed.interaction_edges["edge:spawn:019d0000-0000-7000-8000-000000000001:019d0000-0000-7000-8000-000000000002"]; + assert_eq!(edge.kind, InteractionEdgeKind::SpawnAgent); + assert_eq!( + edge.source, + TraceAnchor::ToolCall { + tool_call_id: "call-spawn".to_string() + } + ); + assert_eq!( + edge.target, + TraceAnchor::Thread { + thread_id: "019d0000-0000-7000-8000-000000000002".to_string() + } + ); + // No transcript item carried the task, so the fallback edge should not + // claim one. The raw payloads still preserve the tool evidence. + assert!(edge.carried_item_ids.is_empty()); + assert_eq!( + edge.carried_raw_payload_ids, + vec![ + spawn_payloads.invocation.raw_payload_id, + spawn_payloads.begin.raw_payload_id, + spawn_payloads.end.raw_payload_id, + spawn_payloads.result.raw_payload_id, ] ); @@ -528,6 +510,187 @@ fn agent_result_edge_links_child_result_to_parent_notification() -> anyhow::Resu Ok(()) } +#[test] +fn agent_result_edge_falls_back_to_child_thread_without_result_message() -> anyhow::Result<()> { + let temp = TempDir::new()?; + let writer = create_started_agent_writer(&temp)?; + + // The child thread and turn exist, but there is intentionally no completed + // assistant message for this turn. Failed child tasks can still notify the + // parent through AgentStatus, so the result edge must not require a final + // transcript item from the child. + start_thread( + &writer, + "019d0000-0000-7000-8000-000000000002", + "/root/child", + )?; + start_turn_for_thread( + &writer, + "019d0000-0000-7000-8000-000000000002", + "turn-child-1", + )?; + + let notification = r#"{"agent_path":"/root/child","status":{"failed":"boom"}}"#; + let carried_payload = writer.write_json_payload( + RawPayloadKind::AgentResult, + &json!({ + "child_agent_path": "/root/child", + "message": notification, + "status": {"failed": "boom"} + }), + )?; + writer.append_with_context( + trace_context_for_thread("019d0000-0000-7000-8000-000000000002", "turn-child-1"), + RawTraceEventPayload::AgentResultObserved { + edge_id: "edge:agent_result:thread-child:turn-child-1:thread-root".to_string(), + child_thread_id: "019d0000-0000-7000-8000-000000000002".to_string(), + child_codex_turn_id: "turn-child-1".to_string(), + parent_thread_id: "019d0000-0000-7000-8000-000000000001".to_string(), + message: notification.to_string(), + carried_payload: Some(carried_payload.clone()), + }, + )?; + + // The parent does receive the failure notification as a model-visible + // mailbox item. The target should remain that precise parent-side + // ConversationItem even though the source falls back to the child thread. + start_agent_turn(&writer, "turn-root-1")?; + let delivered = inter_agent_message( + "/root/child", + "/root", + notification, + /*trigger_turn*/ false, + ); + append_inference_request( + &writer, + "019d0000-0000-7000-8000-000000000001", + "turn-root-1", + "inference-root-1", + vec![message("assistant", &delivered)], + )?; + + let replayed = replay_bundle(temp.path())?; + let edge = + &replayed.interaction_edges["edge:agent_result:thread-child:turn-child-1:thread-root"]; + assert_eq!(edge.kind, InteractionEdgeKind::AgentResult); + assert_eq!( + edge.source, + TraceAnchor::Thread { + thread_id: "019d0000-0000-7000-8000-000000000002".to_string(), + } + ); + let target_item_id = target_conversation_item_id(&edge.target); + assert_eq!( + replayed.conversation_items[target_item_id].thread_id, + "019d0000-0000-7000-8000-000000000001" + ); + assert_eq!(edge.carried_item_ids, vec![target_item_id.clone()]); + assert_eq!( + edge.carried_raw_payload_ids, + vec![carried_payload.raw_payload_id] + ); + + Ok(()) +} + +struct SpawnAgentToolPayloads { + invocation: RawPayloadRef, + begin: RawPayloadRef, + end: RawPayloadRef, + result: RawPayloadRef, +} + +fn append_spawn_agent_tool_lifecycle( + writer: &TraceWriter, + turn_id: &str, +) -> anyhow::Result { + // Keep the parent-side tool lifecycle in one place so the spawn tests can + // focus on the child-side event that decides the edge target. + let invocation = writer.write_json_payload( + RawPayloadKind::ToolInvocation, + &json!({ + "tool_name": "spawn_agent", + "payload": { + "type": "function", + "arguments": r#"{"task_name":"repo_file_counter","message":"count"}"# + } + }), + )?; + writer.append_with_context( + trace_context_for_agent(turn_id), + RawTraceEventPayload::ToolCallStarted { + tool_call_id: "call-spawn".to_string(), + model_visible_call_id: Some("call-spawn".to_string()), + code_mode_runtime_tool_id: None, + requester: RawToolCallRequester::Model, + kind: ToolCallKind::SpawnAgent, + summary: ToolCallSummary::Generic { + label: "spawn_agent".to_string(), + input_preview: None, + output_preview: None, + }, + invocation_payload: Some(invocation.clone()), + }, + )?; + + let begin = writer.write_json_payload( + RawPayloadKind::ToolRuntimeEvent, + &json!({ + "call_id": "call-spawn", + "sender_thread_id": "019d0000-0000-7000-8000-000000000001", + "prompt": "count" + }), + )?; + writer.append_with_context( + trace_context_for_agent(turn_id), + RawTraceEventPayload::ToolCallRuntimeStarted { + tool_call_id: "call-spawn".to_string(), + runtime_payload: begin.clone(), + }, + )?; + + let end = writer.write_json_payload( + RawPayloadKind::ToolRuntimeEvent, + &json!({ + "call_id": "call-spawn", + "sender_thread_id": "019d0000-0000-7000-8000-000000000001", + "new_thread_id": "019d0000-0000-7000-8000-000000000002", + "prompt": "count", + "model": "gpt-test", + "reasoning_effort": "medium", + "status": "running" + }), + )?; + writer.append_with_context( + trace_context_for_agent(turn_id), + RawTraceEventPayload::ToolCallRuntimeEnded { + tool_call_id: "call-spawn".to_string(), + status: ExecutionStatus::Completed, + runtime_payload: end.clone(), + }, + )?; + + let result = writer.write_json_payload( + RawPayloadKind::ToolResult, + &json!({"task_name": "/root/repo_file_counter"}), + )?; + writer.append_with_context( + trace_context_for_agent(turn_id), + RawTraceEventPayload::ToolCallEnded { + tool_call_id: "call-spawn".to_string(), + status: ExecutionStatus::Completed, + result_payload: Some(result.clone()), + }, + )?; + + Ok(SpawnAgentToolPayloads { + invocation, + begin, + end, + result, + }) +} + fn inter_agent_message(author: &str, recipient: &str, content: &str, trigger_turn: bool) -> String { json!({ "author": author,