mirror of
https://github.com/openai/codex.git
synced 2026-10-28 18:58:46 +00:00
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
This commit is contained in:
committed by
copyberry
parent
6d8f311b00
commit
f2df3e3b56
@@ -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()
|
||||
|
||||
@@ -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::<std::io::Result<Vec<_>>>()?;
|
||||
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::<serde_json::Value>)
|
||||
.collect::<serde_json::Result<Vec<_>>>()?;
|
||||
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());
|
||||
}
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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<ToolRuntimeTr
|
||||
},
|
||||
payload: ToolRuntimePayload::McpToolCallBegin(event),
|
||||
}),
|
||||
EventMsg::McpToolCallEnd(event) => 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,
|
||||
|
||||
@@ -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)?;
|
||||
|
||||
@@ -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<CallToolResult, String>,
|
||||
) {
|
||||
let ThreadTraceContextState::Enabled(context) = &self.state else {
|
||||
return;
|
||||
};
|
||||
#[derive(Serialize)]
|
||||
struct Payload<'a> {
|
||||
result: &'a Result<CallToolResult, String>,
|
||||
}
|
||||
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,
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user