From f2df3e3b56351c87e893e223f3a408b0ed5e6ee8 Mon Sep 17 00:00:00 2001 From: cassirer-openai Date: Mon, 26 Oct 2026 20:12:27 +0000 Subject: [PATCH] Preserve full MCP results in rollout traces (#57280) ## Why MCP completion events can truncate large results and omit content removed for model consumption. Rollout traces need the original response, including content and metadata, as well as results for calls denied before reaching the backend. ## What changed - Record MCP runtime results before callbacks, content filtering, and UI truncation, and record skipped calls as failures. - Stop deriving MCP runtime completion records from protocol events to avoid duplicate records. - Treat MCP runtime payloads as opaque results when reducing agent interactions. - Buffer JSON trace writes and explicitly flush them. ## Testing Extend app-server tests to verify that traces retain oversized MCP text, audio, structured content, and metadata while UI results remain truncated, and that denied calls produce exactly one failed runtime result with the denial reason. GitOrigin-RevId: 57880770545eb5869420d7dc6fa4a15e481cae9d --- .../app-server/tests/suite/v2/guardian_v2.rs | 39 ++++++++++++++- .../app-server/tests/suite/v2/mcp_tool.rs | 48 ++++++++++++++++++- codex-rs/core/src/mcp_tool_call.rs | 22 +++++++-- codex-rs/rollout-trace/src/protocol_event.rs | 20 ++------ .../rollout-trace/src/reducer/tool/agents.rs | 4 ++ codex-rs/rollout-trace/src/thread.rs | 36 ++++++++++++++ codex-rs/rollout-trace/src/writer.rs | 8 +++- 7 files changed, 154 insertions(+), 23 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/guardian_v2.rs b/codex-rs/app-server/tests/suite/v2/guardian_v2.rs index 9e1bd0a63c..c15253a118 100644 --- a/codex-rs/app-server/tests/suite/v2/guardian_v2.rs +++ b/codex-rs/app-server/tests/suite/v2/guardian_v2.rs @@ -1003,9 +1003,15 @@ async fn guardian_v2_routes_scoped_tool_approvals( Some(thread_id) } }; + let traces = TempDir::new()?; + let trace_root = matches!(review_outcome, ReviewOutcome::Deny) + .then(|| traces.path().to_str().expect("trace root path")); let mut app_server = TestAppServer::builder() .with_codex_home(codex_home.path()) - .with_env_overrides(&[("OTEL_METRIC_EXPORT_INTERVAL", Some("25"))]) + .with_env_overrides(&[ + ("OTEL_METRIC_EXPORT_INTERVAL", Some("25")), + ("CODEX_ROLLOUT_TRACE_ROOT", trace_root), + ]) .build_initialized_with_timeout(TIMEOUT) .await?; let thread = match lifecycle { @@ -2066,6 +2072,37 @@ async fn guardian_v2_routes_scoped_tool_approvals( }) .await?; timeout(TIMEOUT, app_server.shutdown_gracefully()).await??; + if matches!(review_outcome, ReviewOutcome::Deny) { + // Denied calls have no backend response, but their runtime result must + // still be available to trace consumers alongside executed calls. + let mut denied_results = Vec::new(); + for entry in std::fs::read_dir(traces.path())? { + let bundle = entry?.path(); + for line in std::fs::read_to_string(bundle.join("trace.jsonl"))?.lines() { + let event: Value = serde_json::from_str(line)?; + let payload = &event["payload"]; + if event["thread_id"] == reviewed_thread_id + && payload["type"] == "tool_call_runtime_ended" + && payload["tool_call_id"] == "guardian-action-0" + { + assert_eq!(payload["status"], "failed"); + let path = payload["runtime_payload"]["path"] + .as_str() + .expect("MCP runtime payload path"); + let result: Value = + serde_json::from_reader(std::fs::File::open(bundle.join(path))?)?; + denied_results.push(result); + } + } + } + assert_eq!(denied_results.len(), 1); + assert!( + denied_results[0]["result"]["Err"] + .as_str() + .expect("denied MCP result") + .contains("The destination is not authorized.") + ); + } let reviewer_id = responses_state .guardian_requests .lock() diff --git a/codex-rs/app-server/tests/suite/v2/mcp_tool.rs b/codex-rs/app-server/tests/suite/v2/mcp_tool.rs index 847bfbe8d3..5820ab23c8 100644 --- a/codex-rs/app-server/tests/suite/v2/mcp_tool.rs +++ b/codex-rs/app-server/tests/suite/v2/mcp_tool.rs @@ -999,6 +999,7 @@ async fn mcp_server_tool_call_forwards_url_elicitation() -> Result<()> { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn mcp_tool_call_completion_notification_contains_truncated_large_result() -> Result<()> { + let traces = TempDir::new()?; let call_id = "call-large-mcp"; let namespace = format!("mcp__{TEST_SERVER_NAME}"); let responses = vec![ @@ -1028,6 +1029,7 @@ async fn mcp_tool_call_completion_notification_contains_truncated_large_result() let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) + .with_env_overrides(&[("CODEX_ROLLOUT_TRACE_ROOT", traces.path().to_str())]) .build_initialized() .await?; let ThreadStartResponse { thread, .. } = mcp @@ -1105,6 +1107,47 @@ async fn mcp_tool_call_completion_notification_contains_truncated_large_result() ) .await??; + assert!(mcp.shutdown_gracefully().await?.success()); + + // The UI receives a truncated result, but the trace retains the original + // content and metadata, including audio removed during model sanitization. + let bundles = std::fs::read_dir(traces.path())?.collect::>>()?; + assert_eq!(bundles.len(), 1); + let bundle = bundles[0].path(); + let events = std::fs::read_to_string(bundle.join("trace.jsonl"))? + .lines() + .map(serde_json::from_str::) + .collect::>>()?; + let completions: Vec<_> = events + .iter() + .map(|event| &event["payload"]) + .filter(|event| { + event["type"] == "tool_call_runtime_ended" && event["tool_call_id"] == call_id + }) + .collect(); + assert_eq!(completions.len(), 1); + let payload_path = completions[0]["runtime_payload"]["path"] + .as_str() + .expect("MCP runtime payload path"); + let payload: serde_json::Value = + serde_json::from_reader(std::fs::File::open(bundle.join(payload_path))?)?; + let expected = json!({ + "content": [ + {"type": "text", "text": "large-mcp-content-".repeat(DEFAULT_OUTPUT_BYTES_CAP / 8)}, + {"type": "audio", "data": "UklGRg==", "mimeType": "audio/wav"} + ], + "structuredContent": {"large": "structured-value-".repeat(DEFAULT_OUTPUT_BYTES_CAP / 8)}, + "isError": false, + "_meta": { + "calledBy": "mcp-app", + } + }); + // Avoid constructing a text diff of the intentionally oversized response. + assert!( + payload["result"]["Ok"] == expected, + "trace did not retain the full MCP response" + ); + mcp_server_handle.abort(); let _ = mcp_server_handle.await; @@ -1443,7 +1486,10 @@ impl ServerHandler for ToolAppsMcpServer { let mut result = CallToolResult::structured(json!({ "large": "structured-value-".repeat(DEFAULT_OUTPUT_BYTES_CAP / 8), })); - result.content = vec![ContentBlock::text(large_text)]; + result.content = vec![ + ContentBlock::text(large_text), + ContentBlock::audio("UklGRg==", "audio/wav"), + ]; result.meta = Some(meta); return Ok(result.into()); } diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 6bf51b9ab1..8c4ce1f5d0 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -489,7 +489,7 @@ async fn handle_approved_mcp_tool_call( let mut request_error_code = None; let result = async { let result = async { - let mut result = prepared_call + let result = prepared_call .call_with_preparation(/*requested_timeout*/ None, || async { if let McpToolApprovalApplication::Apply { decision, policy } = &approval_application @@ -572,7 +572,14 @@ async fn handle_approved_mcp_tool_call( .map_err(|error| { request_error_code = codex_rmcp_client::request_error_code(&error); format!("tool call error: {error:?}") - })?; + }); + // Keep the backend response intact, before callbacks or model/UI limits can change it. + sess.services.rollout_thread_trace.record_mcp_tool_result( + &turn_context.sub_id, + call_id, + &result, + ); + let mut result = result?; crate::tools::record_confirmed_code_mode_send( sess, &turn_context.sub_id, @@ -2588,6 +2595,13 @@ async fn notify_mcp_tool_call_skip( .await; } + let result = Err(message); + // Skipped calls never reach the backend-result recorder. + sess.services.rollout_thread_trace.record_mcp_tool_result( + &turn_context.sub_id, + call_id, + &result, + ); notify_mcp_tool_call_completed( sess, turn_context, @@ -2595,10 +2609,10 @@ async fn notify_mcp_tool_call_skip( invocation, item_metadata, Duration::ZERO, - truncate_mcp_tool_result_for_event(&Err(message.clone())), + truncate_mcp_tool_result_for_event(&result), ) .await; - Err(message) + result } #[cfg(test)] diff --git a/codex-rs/rollout-trace/src/protocol_event.rs b/codex-rs/rollout-trace/src/protocol_event.rs index 5253a07cb4..ed49f8a943 100644 --- a/codex-rs/rollout-trace/src/protocol_event.rs +++ b/codex-rs/rollout-trace/src/protocol_event.rs @@ -2,9 +2,9 @@ //! //! The session layer already emits protocol events for turn lifecycle, terminal //! sessions, patch application, MCP calls, and collaboration tools. Rollout -//! tracing reuses those observations instead of adding another set of hooks in -//! `codex-core`: this module translates the protocol surface into the smaller -//! trace vocabulary and keeps the mapping isolated inside `codex-rollout-trace`. +//! tracing reuses those observations where they preserve the original payload. +//! MCP results are recorded before protocol-event truncation instead. This module +//! translates protocol events into the smaller trace vocabulary. //! //! The long explicit `EventMsg` matches are intentional. Most protocol events //! are not trace runtime boundaries, but spelling them out makes new protocol @@ -17,7 +17,6 @@ use codex_protocol::protocol::ExecCommandEndEvent; use codex_protocol::protocol::ExecCommandSource; use codex_protocol::protocol::ExecCommandStatus; use codex_protocol::protocol::McpToolCallBeginEvent; -use codex_protocol::protocol::McpToolCallEndEvent; use codex_protocol::protocol::PatchApplyBeginEvent; use codex_protocol::protocol::PatchApplyEndEvent; use codex_protocol::protocol::PatchApplyStatus; @@ -116,7 +115,6 @@ pub(crate) enum ToolRuntimePayload<'a> { PatchApplyBegin(&'a PatchApplyBeginEvent), PatchApplyEnd(&'a PatchApplyEndEvent), McpToolCallBegin(&'a McpToolCallBeginEvent), - McpToolCallEnd(&'a McpToolCallEndEvent), CollabAgentSpawnBegin(&'a codex_protocol::protocol::CollabAgentSpawnBeginEvent), CollabAgentSpawnEnd(&'a codex_protocol::protocol::CollabAgentSpawnEndEvent), CollabAgentInteractionBegin(&'a codex_protocol::protocol::CollabAgentInteractionBeginEvent), @@ -143,7 +141,6 @@ impl Serialize for ToolRuntimePayload<'_> { ToolRuntimePayload::PatchApplyBegin(event) => event.serialize(serializer), ToolRuntimePayload::PatchApplyEnd(event) => event.serialize(serializer), ToolRuntimePayload::McpToolCallBegin(event) => event.serialize(serializer), - ToolRuntimePayload::McpToolCallEnd(event) => event.serialize(serializer), ToolRuntimePayload::CollabAgentSpawnBegin(event) => event.serialize(serializer), ToolRuntimePayload::CollabAgentSpawnEnd(event) => event.serialize(serializer), ToolRuntimePayload::CollabAgentInteractionBegin(event) => event.serialize(serializer), @@ -316,15 +313,8 @@ pub(crate) fn tool_runtime_trace_event(event: &EventMsg) -> Option Some(ToolRuntimeTraceEvent::Ended { - tool_call_id: &event.call_id, - status: if event.result.is_ok() { - ExecutionStatus::Completed - } else { - ExecutionStatus::Failed - }, - payload: ToolRuntimePayload::McpToolCallEnd(event), - }), + // The full backend result is recorded before this UI event can truncate it. + EventMsg::McpToolCallEnd(_) => None, EventMsg::CollabAgentSpawnBegin(event) => Some(ToolRuntimeTraceEvent::Started { tool_call_id: &event.call_id, kind: ToolCallKind::SpawnAgent, diff --git a/codex-rs/rollout-trace/src/reducer/tool/agents.rs b/codex-rs/rollout-trace/src/reducer/tool/agents.rs index 7c71ddef7d..c7274825aa 100644 --- a/codex-rs/rollout-trace/src/reducer/tool/agents.rs +++ b/codex-rs/rollout-trace/src/reducer/tool/agents.rs @@ -133,6 +133,10 @@ impl TraceReducer { runtime_payload: &RawPayloadRef, ) -> Result<()> { let kind = self.rollout.tool_calls[tool_call_id].kind.clone(); + // MCP results are opaque evidence; they cannot describe agent interactions. + if matches!(kind, ToolCallKind::Mcp { .. }) { + return Ok(()); + } let runtime_payload_json = self.read_payload_json(runtime_payload)?; if runtime_payload_json.get("agent_thread_id").is_some() { let payload: SubAgentActivityEvent = serde_json::from_value(runtime_payload_json)?; diff --git a/codex-rs/rollout-trace/src/thread.rs b/codex-rs/rollout-trace/src/thread.rs index 969bdc776c..b8e752ec07 100644 --- a/codex-rs/rollout-trace/src/thread.rs +++ b/codex-rs/rollout-trace/src/thread.rs @@ -5,6 +5,7 @@ //! thread-local event methods here avoids repeatedly plumbing `thread_id` //! through session code. +use codex_protocol::mcp::CallToolResult; use codex_protocol::protocol::AgentStatus; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::SessionSource; @@ -21,6 +22,7 @@ use crate::CodeCellTraceContext; use crate::CodexTurnId; use crate::CompactionId; use crate::CompactionTraceContext; +use crate::ExecutionStatus; use crate::InferenceTraceContext; use crate::McpCallTraceContext; use crate::RawPayloadKind; @@ -247,6 +249,40 @@ impl ThreadTraceContext { context.append_with_context_best_effort(codex_turn_id.into(), payload); } + /// Records the full MCP backend result before callbacks, sanitization, or UI truncation. + /// The dispatch trace separately records the result returned to the caller. + pub fn record_mcp_tool_result( + &self, + codex_turn_id: &str, + tool_call_id: &str, + result: &Result, + ) { + let ThreadTraceContextState::Enabled(context) = &self.state else { + return; + }; + #[derive(Serialize)] + struct Payload<'a> { + result: &'a Result, + } + let Some(runtime_payload) = context + .write_json_payload_best_effort(RawPayloadKind::ToolRuntimeEvent, &Payload { result }) + else { + return; + }; + context.append_with_context_best_effort( + codex_turn_id.to_owned(), + RawTraceEventPayload::ToolCallRuntimeEnded { + tool_call_id: tool_call_id.to_owned(), + status: if result.is_ok() { + ExecutionStatus::Completed + } else { + ExecutionStatus::Failed + }, + runtime_payload, + }, + ); + } + /// Records the exec actually selected by stdin, since its process handle can be reused. pub fn record_stdin_target( &self, diff --git a/codex-rs/rollout-trace/src/writer.rs b/codex-rs/rollout-trace/src/writer.rs index 925f7bab7a..fe09cf5b03 100644 --- a/codex-rs/rollout-trace/src/writer.rs +++ b/codex-rs/rollout-trace/src/writer.rs @@ -143,8 +143,12 @@ impl TraceWriter { fn write_json_file(path: &Path, value: &impl Serialize) -> Result<()> { let file = File::create(path).with_context(|| format!("create {}", path.display()))?; - serde_json::to_writer_pretty(file, value) - .with_context(|| format!("write JSON {}", path.display())) + let mut writer = BufWriter::new(file); + serde_json::to_writer_pretty(&mut writer, value) + .with_context(|| format!("write JSON {}", path.display()))?; + writer + .flush() + .with_context(|| format!("flush JSON {}", path.display())) } pub(crate) fn unix_time_ms() -> i64 {