diff --git a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst index 6c39f60fd6..bec977fae9 100644 Binary files a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst and b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst differ diff --git a/codex-rs/codex-mcp/src/binding.rs b/codex-rs/codex-mcp/src/binding.rs index c6d79fea0d..4fbe33a433 100644 --- a/codex-rs/codex-mcp/src/binding.rs +++ b/codex-rs/codex-mcp/src/binding.rs @@ -268,6 +268,18 @@ impl PreparedMcpCall { .tool_approval_mode(&self.tool_info.tool.name) } + /// Returns the explicit output budget captured with this call's effective server config. + pub fn output_token_limit(&self) -> Option { + self.config + .mcp_server_catalog + .server(&self.server_name)? + .config() + .tools + .get(self.tool_info.tool.name.as_ref())? + .output_token_limit + .map(std::num::NonZeroUsize::get) + } + pub fn plugin_id(&self) -> Option<&str> { self.plugin_id.as_deref() } diff --git a/codex-rs/codex-mcp/src/catalog_tests.rs b/codex-rs/codex-mcp/src/catalog_tests.rs index ef41e110a8..896b2d4522 100644 --- a/codex-rs/codex-mcp/src/catalog_tests.rs +++ b/codex-rs/codex-mcp/src/catalog_tests.rs @@ -51,6 +51,7 @@ fn server(url: &str) -> McpServerConfig { "read".to_string(), McpServerToolConfig { approval_mode: Some(AppToolApproval::Approve), + ..Default::default() }, )]), } diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index ee20c57598..8388d49c73 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -3966,6 +3966,7 @@ fn server_metadata_preserves_tool_approval_policy() { "search".to_string(), McpServerToolConfig { approval_mode: Some(AppToolApproval::Approve), + ..Default::default() }, ); let metadata = McpServerMetadata::from(&EffectiveMcpServer::configured(config)); diff --git a/codex-rs/config/src/mcp_types.rs b/codex-rs/config/src/mcp_types.rs index 4716f03494..5d18433be1 100644 --- a/codex-rs/config/src/mcp_types.rs +++ b/codex-rs/config/src/mcp_types.rs @@ -3,6 +3,7 @@ use std::borrow::Cow; use std::collections::HashMap; use std::fmt; +use std::num::NonZeroUsize; use std::time::Duration; use base64::Engine; @@ -71,13 +72,24 @@ impl fmt::Display for McpServerDisabledReason { } } -/// Per-tool approval settings for a single MCP server tool. +/// Per-tool settings for a single MCP server tool. #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Default, JsonSchema)] #[schemars(deny_unknown_fields)] pub struct McpServerToolConfig { /// Approval mode for this tool. #[serde(default, skip_serializing_if = "Option::is_none")] pub approval_mode: Option, + + /// Token budget for this tool's output, before the standard 20% serialization allowance. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub output_token_limit: Option, +} + +impl McpServerToolConfig { + /// Applies the stricter explicit output budget without changing approval policy. + pub fn restrict_output_token_limit(&mut self, limit: Option) { + self.output_token_limit = self.output_token_limit.into_iter().chain(limit).min(); + } } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema)] @@ -250,7 +262,7 @@ pub struct McpServerConfig { #[serde(default, skip_serializing_if = "Option::is_none")] pub oauth_resource: Option, - /// Per-tool approval settings keyed by tool name. + /// Per-tool settings keyed by tool name. #[serde(default, skip_serializing_if = "HashMap::is_empty")] pub tools: HashMap, } diff --git a/codex-rs/config/src/mcp_types_tests.rs b/codex-rs/config/src/mcp_types_tests.rs index d0aafaf54d..f7a70787c5 100644 --- a/codex-rs/config/src/mcp_types_tests.rs +++ b/codex-rs/config/src/mcp_types_tests.rs @@ -480,6 +480,7 @@ fn deserialize_server_config_with_default_tool_approval_mode() { [tools.search] approval_mode = "prompt" + output_token_limit = 30000 "#, ) .expect("should deserialize default tool approval mode"); @@ -492,17 +493,31 @@ fn deserialize_server_config_with_default_tool_approval_mode() { cfg.tools.get("search"), Some(&McpServerToolConfig { approval_mode: Some(AppToolApproval::Prompt), + output_token_limit: std::num::NonZeroUsize::new(30_000), }) ); let serialized = toml::to_string(&cfg).expect("should serialize MCP config"); assert!(serialized.contains("default_tools_approval_mode = \"approve\"")); + assert!(serialized.contains("output_token_limit = 30000")); let round_tripped: McpServerConfig = toml::from_str(&serialized).expect("should deserialize serialized MCP config"); assert_eq!(round_tripped, cfg); } +#[test] +fn deserialize_rejects_nonpositive_mcp_tool_output_limits() { + for output_token_limit in [0, -1] { + let config = format!( + "command = \"echo\"\n[tools.search]\noutput_token_limit = {output_token_limit}\n" + ); + let error = toml::from_str::(&config) + .expect_err("MCP tool output limit must be positive"); + assert!(error.to_string().contains("output_token_limit")); + } +} + #[test] fn serialize_round_trips_server_config_with_parallel_tool_calls() { let cfg: McpServerConfig = toml::from_str( diff --git a/codex-rs/config/src/types.rs b/codex-rs/config/src/types.rs index 4e01557be5..2eb73e577a 100644 --- a/codex-rs/config/src/types.rs +++ b/codex-rs/config/src/types.rs @@ -879,7 +879,7 @@ pub struct PluginMcpServerConfig { #[serde(default, skip_serializing_if = "Option::is_none")] pub disabled_tools: Option>, - /// Per-tool approval settings keyed by tool name. + /// Per-tool policy settings keyed by tool name. #[serde(default, skip_serializing_if = "HashMap::is_empty")] pub tools: HashMap, } diff --git a/codex-rs/core-plugins/src/loader.rs b/codex-rs/core-plugins/src/loader.rs index 0c6ce406b6..752795e668 100644 --- a/codex-rs/core-plugins/src/loader.rs +++ b/codex-rs/core-plugins/src/loader.rs @@ -974,6 +974,7 @@ fn apply_plugin_mcp_server_policy(config: &mut McpServerConfig, policy: &PluginM if let Some(approval_mode) = tool_policy.approval_mode { tool_config.approval_mode = Some(approval_mode); } + tool_config.restrict_output_token_limit(tool_policy.output_token_limit); } } @@ -1449,7 +1450,7 @@ pub fn apply_configured_plugin_mcp_server_policies( } } for (tool_name, tool_policy) in &policy.tools { - if tool_policy.approval_mode.is_some() { + if tool_policy.approval_mode.is_some() || tool_policy.output_token_limit.is_some() { server.tools.entry(tool_name.clone()).or_default(); } } @@ -1467,6 +1468,12 @@ pub fn apply_configured_plugin_mcp_server_policies( .restrict_to(approval_mode), ); } + tool_config.restrict_output_token_limit( + policy + .tools + .get(tool_name) + .and_then(|tool_policy| tool_policy.output_token_limit), + ); } } } diff --git a/codex-rs/core-plugins/src/manager_tests.rs b/codex-rs/core-plugins/src/manager_tests.rs index 008ada15f4..8e1bf449ef 100644 --- a/codex-rs/core-plugins/src/manager_tests.rs +++ b/codex-rs/core-plugins/src/manager_tests.rs @@ -1113,7 +1113,7 @@ async fn load_plugins_applies_plugin_mcp_server_policy() { "default_tools_approval_mode": "prompt", "enabled_tools": ["read", "search"], "tools": { - "search": { "approval_mode": "prompt" } + "search": { "approval_mode": "prompt", "output_token_limit": 8000 } } } } @@ -1134,6 +1134,7 @@ disabled_tools = ["delete"] [plugins."sample@test".mcp_servers.sample.tools.search] approval_mode = "approve" +output_token_limit = 12000 "#; let outcome = @@ -1154,6 +1155,7 @@ approval_mode = "approve" server.tools.get("search"), Some(&McpServerToolConfig { approval_mode: Some(AppToolApproval::Approve), + output_token_limit: std::num::NonZeroUsize::new(8_000), }) ); } diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 1d8a9f4e1f..9b453c0db9 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -1966,7 +1966,7 @@ }, "McpServerToolConfig": { "additionalProperties": false, - "description": "Per-tool approval settings for a single MCP server tool.", + "description": "Per-tool settings for a single MCP server tool.", "properties": { "approval_mode": { "allOf": [ @@ -1975,6 +1975,12 @@ } ], "description": "Approval mode for this tool." + }, + "output_token_limit": { + "description": "Token budget for this tool's output, before the standard 20% serialization allowance.", + "format": "uint", + "minimum": 1.0, + "type": "integer" } }, "type": "object" @@ -2990,7 +2996,7 @@ "additionalProperties": { "$ref": "#/definitions/McpServerToolConfig" }, - "description": "Per-tool approval settings keyed by tool name.", + "description": "Per-tool policy settings keyed by tool name.", "type": "object" } }, diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 331751c716..adc4b36c0d 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -882,6 +882,7 @@ mod tests { None, Some(CodexHarnessMetadata { client_authored: true, + ..Default::default() }), Some(CodexHarnessMetadata::default()), None, @@ -905,6 +906,7 @@ mod tests { item: client.clone(), metadata: Some(CodexHarnessMetadata { client_authored: true, + ..Default::default() }), }, ); diff --git a/codex-rs/core/src/compact_remote_v2_image_budget_tests.rs b/codex-rs/core/src/compact_remote_v2_image_budget_tests.rs index 30f44b952e..ffc3a6f85a 100644 --- a/codex-rs/core/src/compact_remote_v2_image_budget_tests.rs +++ b/codex-rs/core/src/compact_remote_v2_image_budget_tests.rs @@ -147,6 +147,7 @@ fn image_treatment_preserves_client_developer_boundary_behavior() { item, metadata: Some(CodexHarnessMetadata { client_authored: true, + ..Default::default() }), }; for max_tokens in [1871, 1879, 1880, 1950, 2200] { diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 3a4554456f..113686e856 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -6773,6 +6773,7 @@ approval_mode = "approve" server.tools.get("search"), Some(&McpServerToolConfig { approval_mode: Some(AppToolApproval::Approve), + ..Default::default() }) ); } @@ -6816,6 +6817,7 @@ approval_mode = "approve" tool, &McpServerToolConfig { approval_mode: Some(AppToolApproval::Approve), + ..Default::default() } ); } diff --git a/codex-rs/core/src/config/edit.rs b/codex-rs/core/src/config/edit.rs index 6d13d9a755..52447014fe 100644 --- a/codex-rs/core/src/config/edit.rs +++ b/codex-rs/core/src/config/edit.rs @@ -314,7 +314,7 @@ impl ConfigDocument { &[NOTICE_TABLE_KEY, "model_migrations", from.as_str()], value(to.clone()), )), - ConfigEdit::ReplaceMcpServers(servers) => Ok(self.replace_mcp_servers(servers)), + ConfigEdit::ReplaceMcpServers(servers) => self.replace_mcp_servers(servers), ConfigEdit::AddToolSuggestDisabledTool(disabled_tool) => { Ok(self.add_tool_suggest_disabled_tool(disabled_tool)) } @@ -432,9 +432,12 @@ impl ConfigDocument { self.remove(segments) } - fn replace_mcp_servers(&mut self, servers: &BTreeMap) -> bool { + fn replace_mcp_servers( + &mut self, + servers: &BTreeMap, + ) -> anyhow::Result { if servers.is_empty() { - return self.clear(&["mcp_servers"]); + return Ok(self.clear(&["mcp_servers"])); } let root = self.doc.as_table_mut(); @@ -446,7 +449,7 @@ impl ConfigDocument { } let Some(item) = root.get_mut("mcp_servers") else { - return false; + return Ok(false); }; if document_helpers::ensure_table_for_write(item).is_none() { @@ -454,7 +457,7 @@ impl ConfigDocument { } let Some(table) = item.as_table_mut() else { - return false; + return Ok(false); }; let keys_to_remove: Vec = table @@ -472,17 +475,17 @@ impl ConfigDocument { if let TomlItem::Value(value) = existing && let Some(inline) = value.as_inline_table_mut() { - let replacement = document_helpers::serialize_mcp_server_inline(config); + let replacement = document_helpers::serialize_mcp_server_inline(config)?; document_helpers::merge_inline_table(inline, replacement); } else { - *existing = document_helpers::serialize_mcp_server(config); + *existing = document_helpers::serialize_mcp_server(config)?; } } else { - table.insert(name, document_helpers::serialize_mcp_server(config)); + table.insert(name, document_helpers::serialize_mcp_server(config)?); } } - true + Ok(true) } fn set_skill_config(&mut self, selector: SkillConfigSelector, enabled: bool) -> bool { diff --git a/codex-rs/core/src/config/edit/document_helpers.rs b/codex-rs/core/src/config/edit/document_helpers.rs index 60883137d0..acf598bca1 100644 --- a/codex-rs/core/src/config/edit/document_helpers.rs +++ b/codex-rs/core/src/config/edit/document_helpers.rs @@ -1,3 +1,4 @@ +use anyhow::Context; use codex_config::types::AppToolApproval; use codex_config::types::McpServerAuth; use codex_config::types::McpServerConfig; @@ -45,7 +46,7 @@ pub(super) fn ensure_table_for_read(item: &mut TomlItem) -> Option<&mut TomlTabl } } -fn serialize_mcp_server_table(config: &McpServerConfig) -> TomlTable { +fn serialize_mcp_server_table(config: &McpServerConfig) -> anyhow::Result { let mut entry = TomlTable::new(); entry.set_implicit(false); @@ -175,15 +176,15 @@ fn serialize_mcp_server_table(config: &McpServerConfig) -> TomlTable { let mut tool_entries: Vec<_> = config.tools.iter().collect(); tool_entries.sort_by_key(|(name, _)| *name); for (name, tool_config) in tool_entries { - tools.insert(name, serialize_mcp_server_tool(tool_config)); + tools.insert(name, serialize_mcp_server_tool(tool_config)?); } entry.insert("tools", TomlItem::Table(tools)); } - entry + Ok(entry) } -fn serialize_mcp_server_tool(config: &McpServerToolConfig) -> TomlItem { +fn serialize_mcp_server_tool(config: &McpServerToolConfig) -> anyhow::Result { let mut entry = TomlTable::new(); entry.set_implicit(false); if let Some(approval_mode) = config.approval_mode { @@ -194,15 +195,21 @@ fn serialize_mcp_server_tool(config: &McpServerToolConfig) -> TomlItem { AppToolApproval::Approve => "approve", }); } - TomlItem::Table(entry) + if let Some(output_token_limit) = config.output_token_limit { + entry["output_token_limit"] = value( + i64::try_from(output_token_limit.get()) + .context("output_token_limit exceeds the TOML integer range")?, + ); + } + Ok(TomlItem::Table(entry)) } -pub(super) fn serialize_mcp_server(config: &McpServerConfig) -> TomlItem { - TomlItem::Table(serialize_mcp_server_table(config)) +pub(super) fn serialize_mcp_server(config: &McpServerConfig) -> anyhow::Result { + Ok(TomlItem::Table(serialize_mcp_server_table(config)?)) } -pub(super) fn serialize_mcp_server_inline(config: &McpServerConfig) -> InlineTable { - serialize_mcp_server_table(config).into_inline_table() +pub(super) fn serialize_mcp_server_inline(config: &McpServerConfig) -> anyhow::Result { + Ok(serialize_mcp_server_table(config)?.into_inline_table()) } pub(super) fn merge_inline_table(existing: &mut InlineTable, replacement: InlineTable) { diff --git a/codex-rs/core/src/config/edit_tests.rs b/codex-rs/core/src/config/edit_tests.rs index 7271ce78b5..19434c39f2 100644 --- a/codex-rs/core/src/config/edit_tests.rs +++ b/codex-rs/core/src/config/edit_tests.rs @@ -1161,8 +1161,9 @@ B = \"2\" assert_eq!(raw, expected); } -#[test] -fn blocking_replace_mcp_servers_serializes_tool_approval_overrides() { +#[test_case::test_case(30_000; "positive limit")] +#[test_case::test_case(usize::MAX; "largest native limit")] +fn blocking_replace_mcp_servers_serializes_tool_approval_overrides(output_token_limit: usize) { let tmp = tempdir().expect("tmpdir"); let codex_home = tmp.path(); @@ -1196,22 +1197,36 @@ fn blocking_replace_mcp_servers_serializes_tool_approval_overrides() { "search".to_string(), McpServerToolConfig { approval_mode: Some(AppToolApproval::Approve), + output_token_limit: std::num::NonZeroUsize::new(output_token_limit), }, )]), }, ); - apply_blocking(codex_home, &[ConfigEdit::ReplaceMcpServers(servers)]).expect("persist"); + let result = apply_blocking(codex_home, &[ConfigEdit::ReplaceMcpServers(servers)]); + if i64::try_from(output_token_limit).is_err() { + assert!( + result + .unwrap_err() + .to_string() + .contains("output_token_limit exceeds the TOML integer range") + ); + return; + } + result.expect("persist"); let raw = std::fs::read_to_string(codex_home.join(CONFIG_TOML_FILE)).expect("read config"); - let expected = "\ + let expected = format!( + "\ [mcp_servers.docs] command = \"docs-server\" default_tools_approval_mode = \"prompt\" [mcp_servers.docs.tools.search] approval_mode = \"approve\" -"; +output_token_limit = {output_token_limit} +" + ); assert_eq!(raw, expected); } diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 954e185283..1b62e756b8 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -21,7 +21,6 @@ use codex_protocol::models::BaseInstructions; use codex_protocol::models::ContentItem; use codex_protocol::models::FunctionCallOutputBody; use codex_protocol::models::FunctionCallOutputContentItem; -use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ImageDetail; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::InputModality; @@ -37,8 +36,7 @@ use codex_utils_output_truncation::TruncationPolicy; use codex_utils_output_truncation::approx_bytes_for_tokens; use codex_utils_output_truncation::approx_token_count; use codex_utils_output_truncation::approx_tokens_from_byte_count_i64; -use codex_utils_output_truncation::truncate_function_output_items_with_policy; -use codex_utils_output_truncation::truncate_text; +use codex_utils_output_truncation::truncate_function_output_payload; use std::num::NonZeroUsize; use std::ops::Deref; use std::sync::Arc; @@ -171,7 +169,7 @@ impl ContextManager { self.record_items_with_metadata(items.into_iter().map(|item| (item, None)), policy); } - /// Records history envelopes while preserving their history-only metadata. + /// Records output while preserving its history-only metadata. pub(crate) fn record_annotated_items( &mut self, items: &[ResponseItemEnvelope], @@ -196,10 +194,20 @@ impl ContextManager { continue; } - let processed = ResponseItemEnvelope { - item: Self::process_item(item, policy), + let mut processed = ResponseItemEnvelope { + item: item.clone(), metadata: metadata.cloned(), }; + if let ResponseItem::FunctionCallOutput { output, .. } + | ResponseItem::CustomToolCallOutput { output, .. } = &mut processed.item + { + // The override already includes the tool's serialization allowance. + let policy = metadata + .and_then(|metadata| metadata.fallback_token_limit_override) + .map(TruncationPolicy::Tokens) + .unwrap_or(policy * 1.2); + truncate_function_output_payload(output, policy, estimate_audio_token_count); + } Arc::make_mut(&mut self.items).push(processed); } } @@ -473,55 +481,6 @@ impl ContextManager { normalize::strip_audio_when_unsupported(input_modalities, items); } - fn process_item(item: &ResponseItem, policy: TruncationPolicy) -> ResponseItem { - let policy_with_serialization_budget = policy * 1.2; - match item { - ResponseItem::FunctionCallOutput { - id, - call_id, - name, - namespace, - output, - internal_chat_message_metadata_passthrough: metadata, - } => ResponseItem::FunctionCallOutput { - id: id.clone(), - call_id: call_id.clone(), - name: name.clone(), - namespace: namespace.clone(), - output: truncate_function_output_payload(output, policy_with_serialization_budget), - internal_chat_message_metadata_passthrough: metadata.clone(), - }, - ResponseItem::CustomToolCallOutput { - id, - call_id, - name, - output, - internal_chat_message_metadata_passthrough: metadata, - } => ResponseItem::CustomToolCallOutput { - id: id.clone(), - call_id: call_id.clone(), - name: name.clone(), - output: truncate_function_output_payload(output, policy_with_serialization_budget), - internal_chat_message_metadata_passthrough: metadata.clone(), - }, - ResponseItem::AdditionalTools { .. } - | ResponseItem::Message { .. } - | ResponseItem::AgentMessage { .. } - | ResponseItem::Reasoning { .. } - | ResponseItem::LocalShellCall { .. } - | ResponseItem::FunctionCall { .. } - | ResponseItem::ToolSearchCall { .. } - | ResponseItem::ToolSearchOutput { .. } - | ResponseItem::WebSearchCall { .. } - | ResponseItem::ImageGenerationCall { .. } - | ResponseItem::CustomToolCall { .. } - | ResponseItem::Compaction { .. } - | ResponseItem::CompactionTrigger { .. } - | ResponseItem::ContextCompaction { .. } - | ResponseItem::Other => item.clone(), - } - } - /// Walk backward from a rollback cut and trim contiguous pre-turn context-update items. /// /// Returns the adjusted cut index after removing contextual developer/user items immediately @@ -570,25 +529,6 @@ impl ContextManager { } } -pub(crate) fn truncate_function_output_payload( - output: &FunctionCallOutputPayload, - policy: TruncationPolicy, -) -> FunctionCallOutputPayload { - let body = match &output.body { - FunctionCallOutputBody::Text(content) => { - FunctionCallOutputBody::Text(truncate_text(content, policy)) - } - FunctionCallOutputBody::ContentItems(items) => FunctionCallOutputBody::ContentItems( - truncate_function_output_items_with_policy(items, policy, estimate_audio_token_count), - ), - }; - - FunctionCallOutputPayload { - body, - success: output.success, - } -} - /// API messages include every non-system item (user/assistant messages, reasoning, /// tool calls, tool outputs, shell calls, web-search calls, and image-generation /// calls). diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index cdf5b3da97..2c661588f6 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -41,6 +41,7 @@ use image::Luma; use image::Rgba; use pretty_assertions::assert_eq; use regex_lite::Regex; +use test_case::test_case; const EXEC_FORMAT_MAX_BYTES: usize = 10_000; const EXEC_FORMAT_MAX_TOKENS: usize = 2_500; @@ -526,8 +527,16 @@ fn annotated_history_apis_preserve_envelopes() { assert_eq!(history.into_raw_items(), vec![first_item]); } -#[test] -fn record_annotated_items_preserves_metadata_while_processing_item() { +#[test_case(None, 100, 5, true; "model policy")] +#[test_case(Some(200), 100, 200, false; "configured override")] +#[test_case(Some(100), 85, 100, true; "saved limit has no additional allowance")] +#[test_case(Some(30_000), 20_000, 30_000, false; "large explicit budget")] +fn record_annotated_items_preserves_metadata_while_processing_item( + fallback_token_limit_override: Option, + repeat_count: usize, + expected_token_limit: usize, + expected_truncation: bool, +) { let envelope = ResponseItemEnvelope { item: ResponseItem::FunctionCallOutput { id: None, @@ -535,23 +544,36 @@ fn record_annotated_items_preserves_metadata_while_processing_item() { name: None, namespace: None, output: FunctionCallOutputPayload { - body: FunctionCallOutputBody::Text("word ".repeat(100)), + body: FunctionCallOutputBody::Text("word ".repeat(repeat_count)), success: Some(true), }, internal_chat_message_metadata_passthrough: None, }, - metadata: Some(CodexHarnessMetadata::default()), + metadata: Some(CodexHarnessMetadata { + fallback_token_limit_override, + ..Default::default() + }), }; let mut history = ContextManager::new(); history.record_annotated_items(std::slice::from_ref(&envelope), TruncationPolicy::Tokens(4)); assert_eq!(history.annotated_items().len(), 1); + assert_eq!(history.annotated_items()[0].metadata, envelope.metadata); assert_eq!( - history.annotated_items()[0].metadata, - Some(CodexHarnessMetadata::default()) + history.annotated_items()[0].item != envelope.item, + expected_truncation + ); + let ResponseItem::FunctionCallOutput { output, .. } = &history.annotated_items()[0].item else { + panic!("expected function call output"); + }; + assert_eq!( + output.body, + FunctionCallOutputBody::Text(truncate_text( + &"word ".repeat(repeat_count), + TruncationPolicy::Tokens(expected_token_limit), + )) ); - assert_ne!(history.annotated_items()[0].item, envelope.item); } #[test] diff --git a/codex-rs/core/src/context_manager/mod.rs b/codex-rs/core/src/context_manager/mod.rs index 97cd7c2109..e67eb7113f 100644 --- a/codex-rs/core/src/context_manager/mod.rs +++ b/codex-rs/core/src/context_manager/mod.rs @@ -6,4 +6,3 @@ pub(crate) use history::ContextManager; pub(crate) use history::estimate_image_bytes; pub(crate) use history::estimate_item_token_count; pub(crate) use history::is_user_turn_boundary; -pub(crate) use history::truncate_function_output_payload; diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 4833da53f3..80b0794b7a 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -2102,6 +2102,7 @@ async fn persist_custom_mcp_tool_approval_writes_tool_override() { tool, &McpServerToolConfig { approval_mode: Some(AppToolApproval::Approve), + ..Default::default() } ); assert!(contents.contains("[mcp_servers.docs.tools.search]")); @@ -2319,6 +2320,7 @@ async fn maybe_persist_mcp_tool_approval_reloads_session_config_for_custom_serve tool, &McpServerToolConfig { approval_mode: Some(AppToolApproval::Approve), + ..Default::default() } ); assert_eq!(mcp_tool_approval_is_remembered(&session, &key).await, true); @@ -2368,6 +2370,7 @@ enabled = true tool, &McpServerToolConfig { approval_mode: Some(AppToolApproval::Approve), + ..Default::default() } ); assert!(contents.contains(r#"[plugins."sample@test".mcp_servers.sample.tools.search]"#)); @@ -2423,6 +2426,7 @@ async fn maybe_persist_mcp_tool_approval_writes_project_config_for_project_serve tool, &McpServerToolConfig { approval_mode: Some(AppToolApproval::Approve), + ..Default::default() } ); assert!(contents.contains("[mcp_servers.docs.tools.search]")); diff --git a/codex-rs/core/src/session/inject.rs b/codex-rs/core/src/session/inject.rs index 6e51aa40e2..8c17a8ae5e 100644 --- a/codex-rs/core/src/session/inject.rs +++ b/codex-rs/core/src/session/inject.rs @@ -72,6 +72,7 @@ impl Session { && matches!(&item, ResponseItem::Message { role, .. } if role == "developer")) .then_some(CodexHarnessMetadata { client_authored: true, + ..Default::default() }); ResponseItemEnvelope { item, metadata } @@ -82,9 +83,7 @@ impl Session { turn_context: &TurnContext, items: Vec, ) { - if !self.enabled(Feature::RetainClientDeveloperMessages) - || items.iter().all(|item| item.metadata.is_none()) - { + if items.iter().all(|item| item.metadata.is_none()) { let items = items .into_iter() .map(ResponseItemEnvelope::into_item) diff --git a/codex-rs/core/src/session/input_queue.rs b/codex-rs/core/src/session/input_queue.rs index e7bd6642bd..6e71857aa3 100644 --- a/codex-rs/core/src/session/input_queue.rs +++ b/codex-rs/core/src/session/input_queue.rs @@ -401,6 +401,7 @@ mod tests { item: ResponseItem::Other, metadata: Some(CodexHarnessMetadata { client_authored: true, + ..Default::default() }), }); assert!(serde_json::to_value(annotated).is_err()); diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 1583f2d9b0..8636a4bd8f 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -40,6 +40,7 @@ use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::skills::emit_explicit_skill_invocations; use crate::stream_events_utils::HandleOutputCtx; +use crate::stream_events_utils::InFlightFuture; use crate::stream_events_utils::TurnItemContributorPolicy; use crate::stream_events_utils::finalize_non_tool_response_item; use crate::stream_events_utils::handle_non_tool_response_item; @@ -90,7 +91,6 @@ use codex_protocol::items::build_hook_prompt_message; use codex_protocol::models::BaseInstructions; use codex_protocol::models::ContentItem; use codex_protocol::models::MessagePhase; -use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::AgentMessageContentDeltaEvent; use codex_protocol::protocol::AgentReasoningSectionBreakEvent; @@ -124,7 +124,6 @@ use codex_utils_stream_parser::AssistantTextStreamParser; use codex_utils_stream_parser::ProposedPlanSegment; use codex_utils_stream_parser::extract_proposed_plan_text; use codex_utils_stream_parser::strip_citations; -use futures::future::BoxFuture; use futures::prelude::*; use futures::stream::FuturesOrdered; use tokio_util::sync::CancellationToken; @@ -2154,22 +2153,21 @@ async fn handle_assistant_item_done_in_plan_mode( #[instrument(level = "trace", skip_all)] async fn drain_in_flight( - in_flight: &mut FuturesOrdered>>, + in_flight: &mut FuturesOrdered>, sess: Arc, turn_context: Arc, ) -> CodexResult<()> { while let Some(res) = in_flight.next().await { match res { - Ok(response_input) => { - let response_item = response_input.into(); - sess.record_conversation_items(&turn_context, std::slice::from_ref(&response_item)) - .await; + Ok(envelope) => { mark_thread_memory_mode_polluted_if_external_context( sess.as_ref(), turn_context.as_ref(), - &response_item, + &envelope.item, ) .await; + sess.record_annotated_conversation_items(&turn_context, vec![envelope]) + .await; } Err(err) => { error_or_panic(format!("in-flight tool future failed during drain: {err}")); @@ -2247,8 +2245,7 @@ async fn try_run_sampling_request( .instrument(trace_span!("stream_request")) .or_cancel(&cancellation_token) .await??; - let mut in_flight: FuturesOrdered>> = - FuturesOrdered::new(); + let mut in_flight: FuturesOrdered> = FuturesOrdered::new(); let mut needs_follow_up = false; let mut last_agent_message: Option = None; let mut active_item: Option = None; diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 5445313656..a75dbe2af8 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -2,6 +2,7 @@ use std::pin::Pin; use std::sync::Arc; use codex_extension_api::ExtensionData; +use codex_history::ResponseItemEnvelope; use codex_protocol::ResponseItemId; use codex_protocol::config_types::ModeKind; use codex_protocol::items::TurnItem; @@ -191,7 +192,7 @@ async fn record_stage1_output_usage_for_memory_citation( /// queuing any tool execution futures. This records items immediately so /// history and rollout stay in sync even if the turn is later cancelled. pub(crate) type InFlightFuture<'f> = - Pin> + Send + 'f>>; + Pin> + Send + 'f>>; #[derive(Default)] pub(crate) struct OutputItemResult { diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index 378f947058..72f9481ff3 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -1,4 +1,3 @@ -use crate::context_manager::truncate_function_output_payload; use crate::original_image_detail::sanitize_original_image_detail; use crate::session::session::Session; use crate::session::step_context::StepContext; @@ -16,9 +15,11 @@ use codex_protocol::models::ResponseItem; use codex_protocol::models::function_call_output_content_items_to_text; use codex_tools::LoadableToolSpec; use codex_tools::ToolName; +use codex_utils_audio::estimate_audio_token_count; use codex_utils_output_truncation::TruncationPolicy; use codex_utils_output_truncation::approx_token_count; use codex_utils_output_truncation::formatted_truncate_text; +use codex_utils_output_truncation::truncate_function_output_payload; use codex_utils_output_truncation::truncate_text; use serde::Serialize; use serde_json::Value as JsonValue; @@ -122,6 +123,10 @@ impl ToolOutput for McpToolOutput { self.result.success() } + fn fallback_token_limit_override(&self) -> Option { + Some((self.truncation_policy * 1.2).token_budget()) + } + fn to_response_item(&self, call_id: &str, _payload: &ToolPayload) -> ResponseInputItem { ResponseInputItem::FunctionCallOutput { call_id: call_id.to_string(), @@ -165,13 +170,13 @@ impl McpToolOutput { } } - // This is the context-injection form, so keep it aligned with the - // function-call output truncation that conversation history already - // applies. Code-mode consumers still get the raw `CallToolResult`. - // - // The text is serialized again inside the Responses payload, so allow - // a small buffer for JSON escaping and wrapper overhead. - truncate_function_output_payload(&payload, self.truncation_policy * 1.2) + // History receives this budget in tokens. Code Mode keeps the raw result. + truncate_function_output_payload( + &mut payload, + self.truncation_policy * 1.2, + estimate_audio_token_count, + ); + payload } } diff --git a/codex-rs/core/src/tools/context_tests.rs b/codex-rs/core/src/tools/context_tests.rs index 3011b2ef51..391b8b0c03 100644 --- a/codex-rs/core/src/tools/context_tests.rs +++ b/codex-rs/core/src/tools/context_tests.rs @@ -230,8 +230,11 @@ fn mcp_tool_output_response_item_preserves_content_items() { } } -#[test] -fn mcp_tool_output_code_mode_result_preserves_content_without_private_metadata() { +#[test_case::test_case(TruncationPolicy::Bytes(64); "byte budget")] +#[test_case::test_case(TruncationPolicy::Tokens(1); "token budget")] +fn mcp_tool_output_code_mode_result_preserves_content_without_private_metadata( + truncation_policy: TruncationPolicy, +) { let large_content = "large structured value ".repeat(1_000); let output = McpToolOutput { result: CallToolResult { @@ -250,12 +253,13 @@ fn mcp_tool_output_code_mode_result_preserves_content_without_private_metadata() tool_input: json!({}), wall_time: std::time::Duration::from_millis(1250), original_image_detail_supported: false, - truncation_policy: TruncationPolicy::Bytes(64), + truncation_policy, }; - let result = output.code_mode_result(&ToolPayload::Function { + let payload = ToolPayload::Function { arguments: "{}".to_string(), - }); + }; + let result = output.code_mode_result(&payload); assert_eq!( result, diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 9d84ddc36b..4d45c64349 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -37,6 +37,7 @@ use codex_tools::agent_plugin_mcp_tool_to_responses_api_tool; use codex_tools::mcp_tool_to_responses_api_tool; use codex_utils_image::PromptImageMode; use codex_utils_image::load_data_url_for_prompt_uncached; +use codex_utils_output_truncation::TruncationPolicy; use codex_utils_string::take_bytes_at_char_boundary; use futures::future::BoxFuture; use serde_json::Map; @@ -217,6 +218,12 @@ impl McpHandler { } }; + // Capture presentation policy from the same config snapshot used for execution. + let truncation_policy = prepared_mcp_call + .as_ref() + .and_then(codex_mcp::PreparedMcpCall::output_token_limit) + .map(TruncationPolicy::Tokens) + .unwrap_or(turn.model_info().truncation_policy.into()); let started = Instant::now(); let result = handle_mcp_tool_call( Arc::clone(&session), @@ -237,7 +244,7 @@ impl McpHandler { tool_input: result.tool_input, wall_time: started.elapsed(), original_image_detail_supported: can_request_original_image_detail(turn.model_info()), - truncation_policy: turn.model_info().truncation_policy.into(), + truncation_policy, })) } } diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index 6b92c89270..67c065faee 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -25,6 +25,7 @@ use crate::tools::registry::AnyToolResult; use crate::tools::registry::ToolArgumentDiffConsumer; use crate::tools::router::ToolCall; use crate::tools::router::ToolCallSource; +use codex_history::ResponseItemEnvelope; use codex_protocol::error::CodexErr; use codex_protocol::models::ResponseInputItem; @@ -74,7 +75,7 @@ impl ToolCallRuntime { self, call: ToolCall, cancellation_token: CancellationToken, - ) -> impl std::future::Future> { + ) -> impl std::future::Future> { let error_call = call.clone(); let source = call.direct_source(); let future = self.handle_tool_call_with_source(call, source, cancellation_token); @@ -82,7 +83,9 @@ impl ToolCallRuntime { match future.await { Ok(response) => Ok(response.into_response()), Err(FunctionCallError::Fatal(message)) => Err(CodexErr::Fatal(message)), - Err(other) => Ok(Self::failure_response(error_call, other)), + Err(other) => Ok(ResponseItemEnvelope::new( + Self::failure_response(error_call, other).into(), + )), } } .in_current_span() @@ -655,7 +658,10 @@ mod tests { success: Some(true), }, }; - assert_eq!(expected_response, response); + assert_eq!( + ResponseItemEnvelope::new(expected_response.into()), + response + ); let actual = records .lock() diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 15c2cf27e5..6564a22d1f 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -30,6 +30,8 @@ use crate::tools::tool_dispatch_trace::ToolDispatchTrace; use crate::util::error_or_panic; use codex_analytics::ControlToolCallStatus; use codex_extension_api::ToolCallOutcome; +use codex_history::CodexHarnessMetadata; +use codex_history::ResponseItemEnvelope; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseInputItem; use codex_protocol::parse_command::ParsedCommand; @@ -191,14 +193,22 @@ pub(crate) struct AnyToolResult { } impl AnyToolResult { - pub(crate) fn into_response(self) -> ResponseInputItem { + pub(crate) fn into_response(self) -> ResponseItemEnvelope { let Self { call_id, payload, result, .. } = self; - result.to_response_item(&call_id, &payload) + ResponseItemEnvelope { + item: result.to_response_item(&call_id, &payload).into(), + metadata: result + .fallback_token_limit_override() + .map(|limit| CodexHarnessMetadata { + fallback_token_limit_override: Some(limit), + ..Default::default() + }), + } } pub(crate) fn code_mode_result(self) -> serde_json::Value { @@ -223,6 +233,10 @@ impl ToolOutput for PostToolUseFeedbackOutput { self.original.success_for_logging() } + fn fallback_token_limit_override(&self) -> Option { + self.original.fallback_token_limit_override() + } + fn to_response_item(&self, call_id: &str, payload: &ToolPayload) -> ResponseInputItem { self.model_visible.to_response_item(call_id, payload) } diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index b862988bec..65c19f6ae3 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/codex-rs/core/src/tools/registry_tests.rs @@ -2,6 +2,8 @@ use super::*; use crate::session::session::Session; use crate::session::step_context::StepContext; use codex_protocol::DEFAULT_FUNCTION_NAMESPACE; +use codex_protocol::models::ResponseItem; +use codex_utils_output_truncation::TruncationPolicy; use futures::future::BoxFuture; use pretty_assertions::assert_eq; use std::sync::atomic::AtomicUsize; @@ -587,6 +589,53 @@ async fn write_stdin_does_not_expose_default_pre_tool_use_payload() { assert_eq!(write_stdin.pre_tool_use_payload(&invocation), None); } +#[test_case::test_case(TruncationPolicy::Tokens(1), 2; "token budget")] +#[test_case::test_case(TruncationPolicy::Bytes(401), 121; "scale bytes before converting to tokens")] +fn post_tool_use_feedback_output_preserves_fallback_token_limit_override( + truncation_policy: TruncationPolicy, + expected_token_limit: usize, +) { + let result = AnyToolResult { + call_id: "call-1".to_string(), + payload: ToolPayload::Function { + arguments: "{}".to_string(), + }, + result: Box::new(PostToolUseFeedbackOutput { + original: Box::new(crate::tools::context::McpToolOutput { + result: codex_protocol::mcp::CallToolResult { + content: Vec::new(), + structured_content: None, + is_error: None, + meta: None, + }, + tool_input: serde_json::json!({}), + wall_time: Duration::ZERO, + original_image_detail_supported: false, + truncation_policy, + }), + model_visible: crate::tools::context::FunctionToolOutput::from_text( + "hook feedback".to_string(), + /*success*/ None, + ), + }), + post_tool_use_payload: None, + }; + + assert_eq!( + result.into_response(), + ResponseItemEnvelope { + item: ResponseItem::from(ResponseInputItem::FunctionCallOutput { + call_id: "call-1".to_string(), + output: FunctionCallOutputPayload::from_text("hook feedback".to_string()), + }), + metadata: Some(CodexHarnessMetadata { + fallback_token_limit_override: Some(expected_token_limit), + ..Default::default() + }), + } + ); +} + #[test] fn post_tool_use_feedback_output_keeps_code_mode_result_typed() { let result = AnyToolResult { @@ -608,12 +657,15 @@ fn post_tool_use_feedback_output_keeps_code_mode_result_typed() { assert_eq!( result.into_response(), - ResponseInputItem::FunctionCallOutput { - call_id: "call-1".to_string(), - output: codex_protocol::models::FunctionCallOutputPayload::from_text( - "hook feedback".to_string() - ), - } + ResponseItemEnvelope::new( + ResponseInputItem::FunctionCallOutput { + call_id: "call-1".to_string(), + output: codex_protocol::models::FunctionCallOutputPayload::from_text( + "hook feedback".to_string() + ), + } + .into() + ) ); let result = AnyToolResult { diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index d01de04d99..01935232b4 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -26,7 +26,6 @@ use codex_protocol::dynamic_tools::DynamicToolNamespaceTool; use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::models::ContentItem; use codex_protocol::models::FunctionCallOutputBody; -use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; use codex_tools::ResponsesApiNamespace; use codex_tools::ResponsesApiNamespaceTool; @@ -604,10 +603,12 @@ async fn extension_tool_executors_are_model_visible_and_dispatchable() -> anyhow ) .await?; - let response = result.into_response(); + let response = result.into_response().item; match response { - ResponseInputItem::FunctionCallOutput { call_id, output } => { - assert_eq!(call_id, "call-extension"); + ResponseItem::FunctionCallOutput { + call_id, output, .. + } => { + assert_eq!(call_id.as_deref(), Some("call-extension")); let FunctionCallOutputBody::Text(text) = output.body else { panic!("expected text function call output") }; diff --git a/codex-rs/core/tests/suite/hooks_mcp.rs b/codex-rs/core/tests/suite/hooks_mcp.rs index 5c81b00af1..2243b15778 100644 --- a/codex-rs/core/tests/suite/hooks_mcp.rs +++ b/codex-rs/core/tests/suite/hooks_mcp.rs @@ -233,7 +233,7 @@ print({python_output_literal}) Ok(()) } -fn write_mcp_tool_hook( +pub(super) fn write_mcp_tool_hook( home: &Path, event_name: &str, matcher: Option<&str>, diff --git a/codex-rs/core/tests/suite/truncation.rs b/codex-rs/core/tests/suite/truncation.rs index 974f59b9ff..ed1c77f0f7 100644 --- a/codex-rs/core/tests/suite/truncation.rs +++ b/codex-rs/core/tests/suite/truncation.rs @@ -26,7 +26,10 @@ use core_test_support::responses::mount_sse_once; use core_test_support::responses::sse; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; +use core_test_support::skip_if_wine_exec; use core_test_support::stdio_server_bin; +use core_test_support::test_codex::TestCodex; +use core_test_support::test_codex::TestCodexBuilder; use core_test_support::test_codex::local_selections; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; @@ -35,6 +38,11 @@ use serde_json::Value; use serde_json::json; use std::collections::HashMap; use std::time::Duration; +use test_case::test_case; +use wiremock::MockServer; + +use super::rmcp_client::remote_aware_environment_id; +use super::rmcp_client::remote_aware_stdio_server_bin; fn assert_wall_time_header(output: &str) { let (wall_time, marker) = output @@ -745,21 +753,19 @@ async fn exec_command_output_not_truncated_with_custom_limit() -> Result<()> { Ok(()) } -// MCP server output should also remain intact when the config increases the token limit. -#[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn mcp_tool_call_output_not_truncated_with_custom_limit() -> Result<()> { - skip_if_no_network!(Ok(())); - - let server = start_mock_server().await; - - let call_id = "rmcp-untruncated"; +async fn call_mcp_echo( + server: &MockServer, + builder: TestCodexBuilder, + output_token_limit: Option, + message_bytes: usize, +) -> Result<(TestCodex, String)> { + let call_id = "rmcp-output"; let server_name = "rmcp"; let namespace = format!("mcp__{server_name}"); - let large_msg = "a".repeat(80_000); - let args_json = serde_json::json!({ "message": large_msg }); + let args_json = json!({ "message": "a".repeat(message_bytes) }); mount_sse_once( - &server, + server, sse(vec![ responses::ev_response_created("resp-1"), responses::ev_function_call_with_namespace( @@ -772,8 +778,8 @@ async fn mcp_tool_call_output_not_truncated_with_custom_limit() -> Result<()> { ]), ) .await; - let mock2 = mount_sse_once( - &server, + let response = mount_sse_once( + server, sse(vec![ responses::ev_assistant_message("msg-1", "rmcp echo tool completed."), responses::ev_completed("resp-2"), @@ -781,67 +787,139 @@ async fn mcp_tool_call_output_not_truncated_with_custom_limit() -> Result<()> { ) .await; - let rmcp_test_server_bin = stdio_server_bin()?; - - let mut builder = test_codex().with_config(move |config| { - config.tool_output_token_limit = Some(50_000); - let mut servers = config.mcp_servers.get().clone(); - servers.insert( - server_name.to_string(), - codex_config::types::McpServerConfig { - auth: Default::default(), - transport: codex_config::types::McpServerTransportConfig::Stdio { - command: rmcp_test_server_bin, - args: Vec::new(), - env: None, - env_vars: Vec::new(), - cwd: None, - }, - environment_id: "local".to_string(), - enabled: true, - required: false, - supports_parallel_tool_calls: false, - omit_tools_from: None, - disabled_reason: None, - startup_timeout_sec: Some(std::time::Duration::from_secs(10)), - tool_timeout_sec: None, - default_tools_approval_mode: None, - enabled_tools: None, - disabled_tools: None, - scopes: None, - oauth: None, - oauth_resource: None, - tools: HashMap::new(), - }, - ); + let mcp_server = serde_json::from_value(json!({ + "command": remote_aware_stdio_server_bin()?, + "environment_id": remote_aware_environment_id(), + "startup_timeout_sec": 10, + "tools": { "echo": { "output_token_limit": output_token_limit } }, + }))?; + let mut builder = builder.with_config(move |config| { config .mcp_servers - .set(servers) + .set(HashMap::from([(server_name.to_string(), mcp_server)])) .expect("test mcp servers should accept any configuration"); }); - let fixture = builder.build(&server).await?; + let fixture = builder.build_with_auto_env(server).await?; wait_for_mcp_server(&fixture.codex, server_name).await?; + fixture.submit_text_turn("call the MCP echo tool").await?; - fixture - .submit_turn_with_permission_profile( - "call the rmcp echo tool with a very large message", - PermissionProfile::read_only(), - ) - .await?; - - let output = mock2 + let output = response .single_request() .function_call_output_text(call_id) - .context("function_call_output present for rmcp call")?; + .context("model-facing MCP output text")?; + Ok((fixture, output)) +} +#[test_case(3_000, 13_000; "serialization allowance")] +#[test_case(30_000, 116_000; "large override")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn mcp_tool_output_limit_preserves_output_that_fits( + output_token_limit: usize, + message_bytes: usize, +) -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_wine_exec!(Ok(()), "requires a Windows test_stdio_server binary"); + + let server = start_mock_server().await; + let builder = test_codex().with_config(|config| config.tool_output_token_limit = Some(50)); + let (_fixture, output) = + call_mcp_echo(&server, builder, Some(output_token_limit), message_bytes).await?; + + assert!(output.contains(&"a".repeat(message_bytes))); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn mcp_tool_output_limit_truncates_oversized_output() -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_wine_exec!(Ok(()), "requires a Windows test_stdio_server binary"); + + let server = start_mock_server().await; + let builder = test_codex().with_config(|config| config.tool_output_token_limit = Some(50)); + let (_fixture, output) = call_mcp_echo( + &server, + builder, + Some(30_000), + /*message_bytes*/ 150_000, + ) + .await?; + + assert!(output.contains("truncated")); + // 30k tokens plus the serialization allowance leaves about 144k bytes. + assert!((140_000..145_000).contains(&output.len())); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn mcp_tool_output_limit_applies_to_hook_feedback() -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_wine_exec!(Ok(()), "requires a Windows test_stdio_server binary"); + + let server = start_mock_server().await; + let builder = test_codex() + .with_pre_build_hook(|home| { + super::hooks_mcp::write_mcp_tool_hook( + home, + "PostToolUse", + Some("^mcp__rmcp__echo$"), + "rmcp", + &json!({ "continue": false, "stopReason": "hook feedback ".repeat(100) }) + .to_string(), + ) + .expect("write MCP post-tool hook"); + }) + .with_config(|config| { + core_test_support::hooks::trust_discovered_hooks(config); + config.tool_output_token_limit = Some(50); + }); + let (_fixture, output) = + call_mcp_echo(&server, builder, Some(100), /*message_bytes*/ 0).await?; + + assert!(output.starts_with("hook feedback ")); + assert!(output.contains("truncated")); + // The tool's 120-token budget applies, not the 60-token global budget. + assert!((400..600).contains(&output.len())); + Ok(()) +} + +#[test_case(None; "model default")] +#[test_case(Some(30_000); "tool override")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn mcp_tool_output_limit_survives_resume(output_token_limit: Option) -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_wine_exec!(Ok(()), "requires a Windows test_stdio_server binary"); + + let server = start_mock_server().await; + let builder = test_codex().with_config(|config| config.tool_output_token_limit = Some(50_000)); + let (fixture, output) = call_mcp_echo( + &server, + builder, + output_token_limit, + /*message_bytes*/ 150_000, + ) + .await?; + + fixture.codex.ensure_rollout_materialized().await; + fixture.codex.flush_rollout().await?; + let resumed_response = mount_sse_once( + &server, + sse(vec![ + responses::ev_assistant_message("msg-2", "resumed"), + responses::ev_completed("resp-3"), + ]), + ) + .await; + let mut resume_builder = test_codex().with_config(|config| { + config.tool_output_token_limit = Some(50); + }); + let resumed = resume_builder.restart(&server, &fixture).await?; + resumed.submit_turn("continue").await?; assert_eq!( - output.len(), - 80065, - "MCP output should retain its serialized length plus wall-time header" - ); - assert!( - !output.contains("truncated"), - "output should not include truncation markers when limit is raised: {output}" + resumed_response + .single_request() + .function_call_output_text("rmcp-output") + .context("resumed MCP output")?, + output ); Ok(()) diff --git a/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs b/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs index f37f3411ec..3ab8077d70 100644 --- a/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs +++ b/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs @@ -231,8 +231,8 @@ async fn high_level_discovery_matches_the_existing_plugin_provider() -> TestResu "enabled_tools": ["read", "deploy", "trusted", "package-only"], "disabled_tools": ["package-denied"], "tools": { - "read": {"approval_mode": "prompt"}, - "deploy": {"approval_mode": "approve"}, + "read": {"approval_mode": "prompt", "output_token_limit": 12000}, + "deploy": {"approval_mode": "approve", "output_token_limit": 4000}, "trusted": {"approval_mode": "approve"} } }, @@ -255,6 +255,10 @@ disabled_tools = ["write"] [plugins."selected-root".mcp_servers.first.tools.read] approval_mode = "approve" +output_token_limit = 8000 + +[plugins."selected-root".mcp_servers.first.tools.deploy] +output_token_limit = 9000 [plugins."selected-root".mcp_servers.first.tools.trusted] approval_mode = "approve" @@ -308,18 +312,21 @@ default_tools_approval_mode = "auto" "read".to_string(), McpServerToolConfig { approval_mode: Some(AppToolApproval::Prompt), + output_token_limit: std::num::NonZeroUsize::new(8_000), }, ), ( "deploy".to_string(), McpServerToolConfig { approval_mode: Some(AppToolApproval::Prompt), + output_token_limit: std::num::NonZeroUsize::new(4_000), }, ), ( "trusted".to_string(), McpServerToolConfig { approval_mode: Some(AppToolApproval::Approve), + ..Default::default() }, ), ]), diff --git a/codex-rs/history/src/lib.rs b/codex-rs/history/src/lib.rs index 6e15c01bc8..0241394fcb 100644 --- a/codex-rs/history/src/lib.rs +++ b/codex-rs/history/src/lib.rs @@ -48,6 +48,11 @@ pub struct CodexHarnessMetadata { /// Whether a developer message was supplied by an app-server client. #[serde(default)] pub client_authored: bool, + + /// Overrides history's fallback truncation budget, including on resume. + /// Measured in tokens, with any tool-specific allowance already included. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub fallback_token_limit_override: Option, } impl ResponseItemEnvelope { diff --git a/codex-rs/history/src/tests.rs b/codex-rs/history/src/tests.rs index 97b61cf511..cca1050cd7 100644 --- a/codex-rs/history/src/tests.rs +++ b/codex-rs/history/src/tests.rs @@ -72,6 +72,7 @@ fn response_item_envelope_stores_metadata_beside_rollout_payload() -> Result<()> item: response_item.clone(), metadata: Some(CodexHarnessMetadata { client_authored: true, + fallback_token_limit_override: Some(20_000), }), }), }; @@ -84,7 +85,7 @@ fn response_item_envelope_stores_metadata_beside_rollout_payload() -> Result<()> "ordinal": 7, "type": "response_item", "payload": response_item, - "metadata": { "client_authored": true }, + "metadata": { "client_authored": true, "fallback_token_limit_override": 20_000 }, }) ); assert_eq!(serialized["payload"].get("metadata"), None); @@ -97,6 +98,7 @@ fn response_item_envelope_stores_metadata_beside_rollout_payload() -> Result<()> envelope.metadata, Some(CodexHarnessMetadata { client_authored: true, + fallback_token_limit_override: Some(20_000), }) ); Ok(()) @@ -184,6 +186,7 @@ fn compacted_replacement_history_stores_metadata_in_an_aligned_sidecar() -> Resu item: developer_message.clone(), metadata: Some(CodexHarnessMetadata { client_authored: true, + ..Default::default() }), }, ResponseItemEnvelope::new(compaction_item.clone()), @@ -216,6 +219,7 @@ fn compacted_replacement_history_stores_metadata_in_an_aligned_sidecar() -> Resu item: developer_message, metadata: Some(CodexHarnessMetadata { client_authored: true, + ..Default::default() }), }, ResponseItemEnvelope { @@ -277,6 +281,7 @@ fn compacted_metadata_remains_compatible_with_legacy_response_item_readers() -> item: response_item.clone(), metadata: Some(CodexHarnessMetadata { client_authored: true, + ..Default::default() }), }; let response_line = serde_json::to_value(RolloutItem::ResponseItem(envelope.clone()))?; diff --git a/codex-rs/tools/src/tool_output.rs b/codex-rs/tools/src/tool_output.rs index a8cfa569d3..0aa22b626d 100644 --- a/codex-rs/tools/src/tool_output.rs +++ b/codex-rs/tools/src/tool_output.rs @@ -23,6 +23,12 @@ pub trait ToolOutput: Send { false } + /// Overrides history's fallback token limit after tool-specific truncation. + /// Include any serialization allowance; history uses this limit unchanged. + fn fallback_token_limit_override(&self) -> Option { + None + } + fn to_response_item(&self, call_id: &str, payload: &ToolPayload) -> ResponseInputItem; /// Returns the tool call id exposed to `PostToolUse` hooks for this output. @@ -67,6 +73,10 @@ where (**self).contains_external_context() } + fn fallback_token_limit_override(&self) -> Option { + (**self).fallback_token_limit_override() + } + fn to_response_item(&self, call_id: &str, payload: &ToolPayload) -> ResponseInputItem { (**self).to_response_item(call_id, payload) } diff --git a/codex-rs/utils/output-truncation/src/lib.rs b/codex-rs/utils/output-truncation/src/lib.rs index fc7d43bfc2..e612457396 100644 --- a/codex-rs/utils/output-truncation/src/lib.rs +++ b/codex-rs/utils/output-truncation/src/lib.rs @@ -1,6 +1,8 @@ -//! Helpers for truncating tool and exec output using [`TruncationPolicy`](codex_protocol::protocol::TruncationPolicy). +//! Shared byte/token truncation for tool and exec output. +use codex_protocol::models::FunctionCallOutputBody; use codex_protocol::models::FunctionCallOutputContentItem; +use codex_protocol::models::FunctionCallOutputPayload; pub use codex_utils_string::approx_bytes_for_tokens; pub use codex_utils_string::approx_token_count; pub use codex_utils_string::approx_tokens_from_byte_count; @@ -29,6 +31,24 @@ pub fn truncate_text(content: &str, policy: TruncationPolicy) -> String { } } +/// Applies the existing byte/token budget without changing success metadata or media ordering. +pub fn truncate_function_output_payload( + output: &mut FunctionCallOutputPayload, + policy: TruncationPolicy, + estimate_audio_token_count: impl Fn(&str) -> usize, +) { + match &mut output.body { + FunctionCallOutputBody::Text(text) => *text = truncate_text(text, policy), + FunctionCallOutputBody::ContentItems(items) => { + *items = truncate_function_output_items_with_policy( + items, + policy, + estimate_audio_token_count, + ); + } + } +} + pub fn formatted_truncate_text_content_items_with_policy( items: &[FunctionCallOutputContentItem], policy: TruncationPolicy,