diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 24e4f109df..e7c700939b 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -4471,6 +4471,7 @@ dependencies = [ "codex-protocol", "codex-state", "codex-utils-absolute-path", + "codex-utils-output-truncation", "codex-utils-path", "codex-utils-process", "codex-utils-string", @@ -5090,6 +5091,7 @@ dependencies = [ "codex-protocol", "codex-utils-string", "pretty_assertions", + "serde_json", ] [[package]] diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 26d31b12ce..1fe0016e6e 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -87,6 +87,7 @@ use codex_rollout::state_db; use codex_tools::ToolName; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::truncate_mcp_tool_result; use codex_utils_output_truncation::truncate_text; use codex_utils_path_uri::PathUri; use codex_utils_pty::DEFAULT_OUTPUT_BYTES_CAP; @@ -982,39 +983,11 @@ fn truncate_mcp_tool_result_for_event( result: &Result, ) -> Result { match result { - Ok(call_tool_result) => { - // The app-server rebuilds `ThreadItem::McpToolCall` from this item, - // so avoid persisting multi-megabyte results in rollout storage. - let Ok(serialized) = serde_json::to_string(call_tool_result) else { - return Ok(call_tool_result.clone()); - }; - if serialized.len() <= MCP_TOOL_CALL_EVENT_RESULT_MAX_BYTES { - return Ok(call_tool_result.clone()); - } - - // A huge MCP result can put bytes in `content`, `structuredContent`, - // or `_meta`. Collapse the event copy to a text preview of the whole - // serialized result so the UI still has useful context without - // preserving a multi-megabyte structured payload. - // - // This budget applies to the preview text, not the final event JSON. - // The preview is itself serialized into a JSON string, so quotes and - // backslashes can be escaped again and the stored event may end up - // somewhat larger than this byte budget. - let truncated = truncate_text( - &serialized, - TruncationPolicy::Bytes(MCP_TOOL_CALL_EVENT_RESULT_MAX_BYTES), - ); - Ok(CallToolResult { - content: vec![serde_json::json!({ - "type": "text", - "text": truncated, - })], - structured_content: None, - is_error: call_tool_result.is_error, - meta: None, - }) - } + Ok(call_tool_result) => Ok(truncate_mcp_tool_result( + call_tool_result, + MCP_TOOL_CALL_EVENT_RESULT_MAX_BYTES, + ) + .into_owned()), Err(message) => Err(truncate_text( message, TruncationPolicy::Bytes(MCP_TOOL_CALL_EVENT_RESULT_MAX_BYTES), diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index a1c0549575..8cace5296b 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -1071,61 +1071,6 @@ fn sanitize_mcp_tool_result_for_model_preserves_supported_media() { assert_eq!(got, original); } -#[test] -fn truncate_mcp_tool_result_for_event_preserves_small_result() { - let original = CallToolResult { - content: vec![serde_json::json!({ - "type": "text", - "text": "hello", - })], - structured_content: Some(serde_json::json!({"x": 1})), - is_error: Some(false), - meta: Some(serde_json::json!({"k": "v"})), - }; - - let got = truncate_mcp_tool_result_for_event(&Ok(original.clone())) - .expect("small result should remain successful"); - - assert_eq!(got, original); -} - -#[test] -fn truncate_mcp_tool_result_for_event_bounds_large_result() { - let original = CallToolResult { - content: vec![serde_json::json!({ - "type": "text", - "text": "long-message-with-newlines-\n".repeat(200_000), - })], - structured_content: Some(serde_json::json!({ - "structured": "structured-value-".repeat(200_000), - })), - is_error: Some(false), - meta: Some(serde_json::json!({ - "meta": "meta-value-".repeat(200_000), - })), - }; - - let got = truncate_mcp_tool_result_for_event(&Ok(original)) - .expect("large result should remain successful"); - let serialized = serde_json::to_string(&got).expect("truncated result should serialize"); - - // The truncated preview is embedded as a JSON string, so quotes and - // backslashes can be escaped again. That can roughly double the preview - // bytes in the worst case. The extra buffer covers the small result wrapper - // and marker. - assert!(serialized.len() < MCP_TOOL_CALL_EVENT_RESULT_MAX_BYTES * 2 + 1024); - assert_eq!(got.structured_content, None); - assert_eq!(got.meta, None); - assert_eq!(got.is_error, Some(false)); - assert!( - got.content[0] - .get("text") - .and_then(serde_json::Value::as_str) - .is_some_and(|text| text.contains("truncated")), - "large event result should contain a truncation marker: {got:?}" - ); -} - #[test] fn truncate_mcp_tool_result_for_event_bounds_large_error() { let got = truncate_mcp_tool_result_for_event(&Err("error-message-".repeat(200_000))) diff --git a/codex-rs/rollout/Cargo.toml b/codex-rs/rollout/Cargo.toml index 6ce9637f65..6bb676c9b6 100644 --- a/codex-rs/rollout/Cargo.toml +++ b/codex-rs/rollout/Cargo.toml @@ -24,6 +24,7 @@ codex-otel = { workspace = true } codex-protocol = { workspace = true } codex-state = { workspace = true } codex-utils-absolute-path = { workspace = true } +codex-utils-output-truncation = { workspace = true } codex-utils-path = { workspace = true } codex-utils-string = { workspace = true } memchr = { workspace = true } diff --git a/codex-rs/rollout/src/policy.rs b/codex-rs/rollout/src/policy.rs index 968bfbe73c..a650768231 100644 --- a/codex-rs/rollout/src/policy.rs +++ b/codex-rs/rollout/src/policy.rs @@ -4,13 +4,16 @@ use crate::RolloutItem; use crate::protocol::EventMsg; use codex_extension_items::ExtensionItem; use codex_protocol::items::CommandExecutionItem; +use codex_protocol::items::McpToolCallItem; use codex_protocol::items::TurnItem; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::ItemCompletedEvent; use codex_protocol::protocol::SubAgentActivityKind; use codex_protocol::protocol::ThreadHistoryMode; +use codex_utils_output_truncation::truncate_mcp_tool_result; use codex_utils_string::truncate_middle_with_marker; +const PERSISTED_MCP_RESULT_MAX_BYTES: usize = 64 * 1024; pub(super) const PERSISTED_COMMAND_OUTPUT_MAX_BYTES: usize = 64 * 1024; const PERSISTED_COMMAND_OUTPUT_TRUNCATION_MARKER: &str = "\n... command output truncated for persistence ...\n"; @@ -259,6 +262,21 @@ fn persisted_item_completed_event<'a>( } Some(Cow::Owned(EventMsg::ItemCompleted(event))) } + TurnItem::McpToolCall(McpToolCallItem { + result: Some(result), + .. + }) => { + let Cow::Owned(result) = + truncate_mcp_tool_result(result, PERSISTED_MCP_RESULT_MAX_BYTES) + else { + return Some(Cow::Borrowed(ev)); + }; + let mut event = event.clone(); + if let TurnItem::McpToolCall(tool_call) = &mut event.item { + tool_call.result = Some(result); + } + Some(Cow::Owned(EventMsg::ItemCompleted(event))) + } _ => Some(Cow::Borrowed(ev)), }, } diff --git a/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs b/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs index e7a663ffaf..058cc29a12 100644 --- a/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs +++ b/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs @@ -4,14 +4,18 @@ use std::path::Path; use std::time::Duration; use chrono::Utc; +use codex_app_server_protocol::McpToolCallResult; use codex_app_server_protocol::ThreadItem; use codex_protocol::ThreadId; use codex_protocol::items::AgentMessageContent; use codex_protocol::items::AgentMessageItem; use codex_protocol::items::CommandExecutionItem; use codex_protocol::items::CommandExecutionStatus; +use codex_protocol::items::McpToolCallItem; +use codex_protocol::items::McpToolCallStatus; use codex_protocol::items::TurnItem; use codex_protocol::items::UserMessageItem; +use codex_protocol::mcp::CallToolResult; use codex_protocol::models::BaseInstructions; use codex_protocol::models::ContentItem; use codex_protocol::models::MessagePhase; @@ -505,6 +509,96 @@ async fn paginated_command_history_caps_aggregated_output() { assert!(output.contains("command output truncated for persistence")); } +#[tokio::test] +async fn paginated_mcp_history_caps_large_results() { + let home = TempDir::new().expect("temp dir"); + let store = projection_store(home.path()).await; + let thread_id = ThreadId::default(); + create_paginated_thread(&store, thread_id).await; + let arguments = serde_json::json!({"query": "large result"}); + let tool_call = McpToolCallItem { + id: "mcp-1".to_string(), + server: "test".to_string(), + tool: "large_result".to_string(), + arguments: arguments.clone(), + connector_id: None, + mcp_app_resource_uri: None, + mcp_app_ui: None, + link_id: None, + app_name: None, + action_name: None, + plugin_id: None, + read_only_hint: None, + status: McpToolCallStatus::Failed, + result: Some(CallToolResult { + content: vec![serde_json::json!({ + "type": "text", + "text": format!("head\n{}\ntail", "x".repeat(200_000)), + })], + structured_content: None, + is_error: Some(true), + meta: None, + }), + error: None, + duration: Some(Duration::from_millis(25)), + }; + + store + .append_items(AppendThreadItemsParams { + thread_id, + items: vec![ + turn_started("turn-1"), + completed_item(thread_id, "turn-1", TurnItem::McpToolCall(tool_call)), + turn_completed("turn-1"), + ], + }) + .await + .expect("append MCP history"); + + let pool = codex_state::open_thread_history_db(&codex_state::SqliteConfig::new_for_testing( + home.path().abs(), + )) + .await + .expect("open thread history db"); + let item_json = sqlx::query_scalar::<_, String>( + "SELECT item_json FROM thread_items WHERE thread_id = ? AND item_id = ?", + ) + .bind(thread_id.to_string()) + .bind("mcp-1") + .fetch_one(&pool) + .await + .expect("read projected MCP tool call"); + let projected: ThreadItem = + serde_json::from_str(&item_json).expect("deserialize projected MCP tool call"); + let ThreadItem::McpToolCall { + arguments: projected_arguments, + result: Some(result), + .. + } = projected + else { + panic!("expected projected MCP tool call"); + }; + assert_eq!(projected_arguments, arguments); + let preview = result.content[0]["text"] + .as_str() + .expect("persisted MCP result preview"); + assert_eq!( + *result, + McpToolCallResult { + content: vec![serde_json::json!({ + "type": "text", + "text": preview, + })], + structured_content: None, + meta: None, + } + ); + assert!(preview.len() < 65 * 1024); + assert!(preview.contains("head")); + assert!(preview.contains("chars truncated")); + assert!(preview.contains("tail")); +} + #[tokio::test] async fn paginated_realtime_items_materialize_separately_in_rollout_order() { let home = TempDir::new().expect("temp dir"); diff --git a/codex-rs/utils/output-truncation/Cargo.toml b/codex-rs/utils/output-truncation/Cargo.toml index 17176c9f56..3f83bd5fb1 100644 --- a/codex-rs/utils/output-truncation/Cargo.toml +++ b/codex-rs/utils/output-truncation/Cargo.toml @@ -10,6 +10,7 @@ workspace = true [dependencies] codex-protocol = { workspace = true } codex-utils-string = { workspace = true } +serde_json = { workspace = true } [dev-dependencies] pretty_assertions = { workspace = true } diff --git a/codex-rs/utils/output-truncation/src/lib.rs b/codex-rs/utils/output-truncation/src/lib.rs index 1cae07741c..cee4a744ae 100644 --- a/codex-rs/utils/output-truncation/src/lib.rs +++ b/codex-rs/utils/output-truncation/src/lib.rs @@ -1,5 +1,8 @@ //! Shared byte/token truncation for tool and exec output. +use std::borrow::Cow; + +use codex_protocol::mcp::CallToolResult; use codex_protocol::models::FunctionCallOutputBody; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputPayload; @@ -37,6 +40,30 @@ pub fn truncate_text(content: &str, policy: TruncationPolicy) -> String { } } +/// Replaces an oversized serialized MCP result with a text preview while preserving its error +/// status and supported bounded metadata. +pub fn truncate_mcp_tool_result( + result: &CallToolResult, + max_bytes: usize, +) -> Cow<'_, CallToolResult> { + let Ok(serialized) = serde_json::to_string(result) else { + return Cow::Borrowed(result); + }; + if serialized.len() <= max_bytes { + return Cow::Borrowed(result); + } + + Cow::Owned(CallToolResult { + content: vec![serde_json::json!({ + "type": "text", + "text": truncate_text(&serialized, TruncationPolicy::Bytes(max_bytes)), + })], + structured_content: None, + is_error: result.is_error, + meta: None, + }) +} + /// Applies the existing byte/token budget without changing success metadata or media ordering. pub fn truncate_function_output_payload( output: &mut FunctionCallOutputPayload, diff --git a/codex-rs/utils/output-truncation/src/truncate_tests.rs b/codex-rs/utils/output-truncation/src/truncate_tests.rs index 736a283563..df34be4f02 100644 --- a/codex-rs/utils/output-truncation/src/truncate_tests.rs +++ b/codex-rs/utils/output-truncation/src/truncate_tests.rs @@ -4,12 +4,69 @@ use crate::approx_tokens_from_byte_count_i64; use crate::formatted_truncate_text; use crate::formatted_truncate_text_content_items_with_policy; use crate::truncate_function_output_items_with_policy; +use crate::truncate_mcp_tool_result; use crate::truncate_text; +use codex_protocol::mcp::CallToolResult; use codex_protocol::models::DEFAULT_IMAGE_DETAIL; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::ImageReference; use pretty_assertions::assert_eq; +#[test] +fn truncate_mcp_tool_result_borrows_small_result() { + let result = CallToolResult { + content: vec![serde_json::json!({ + "type": "text", + "text": "hello", + })], + structured_content: Some(serde_json::json!({"x": 1})), + is_error: Some(false), + meta: Some(serde_json::json!({"k": "v"})), + }; + + let truncated = truncate_mcp_tool_result(&result, /*max_bytes*/ 1024); + + assert!(matches!(truncated, std::borrow::Cow::Borrowed(_))); + assert_eq!(truncated.as_ref(), &result); +} + +#[test] +fn truncate_mcp_tool_result_compacts_large_result() { + let result = CallToolResult { + content: vec![serde_json::json!({ + "type": "text", + "text": format!("head\n{}\ntail", "x".repeat(5_000)), + })], + structured_content: Some(serde_json::json!({"x": "y".repeat(5_000)})), + is_error: Some(true), + meta: Some(serde_json::json!({ + "private": format!("meta-head{}meta-tail", "z".repeat(5_000)), + })), + }; + + let truncated = truncate_mcp_tool_result(&result, /*max_bytes*/ 1024).into_owned(); + + let preview = truncated.content[0]["text"] + .as_str() + .expect("truncated result should contain a text preview"); + assert_eq!( + truncated, + CallToolResult { + content: vec![serde_json::json!({ + "type": "text", + "text": preview, + })], + structured_content: None, + is_error: Some(true), + meta: None, + } + ); + assert!(preview.len() < 2 * 1024); + assert!(preview.contains("head")); + assert!(preview.contains("chars truncated")); + assert!(preview.contains("meta-tail")); +} + #[test] fn truncate_bytes_less_than_placeholder_returns_placeholder() { let content = "example output";