mirror of
https://github.com/openai/codex.git
synced 2026-10-26 18:38:47 +00:00
Support per-tool MCP output limits (#41421)
## What changed - Add a positive `output_token_limit` setting to each entry under an MCP server's `tools` configuration. - Apply the most restrictive limit when plugin and user policies overlap, while keeping approval policy independent. - Carry the effective MCP output budget in conversation history so tool output, post-tool hook responses, and resumed sessions use the same truncation limit. ## Testing - Cover configuration parsing, serialization, schema validation, and plugin policy merging. - Cover MCP output below and above the configured limit, post-tool hook responses, and session resume. GitOrigin-RevId: d0beb4fca9ba6055d9e1d31c137373b465d50d61
This commit is contained in:
pakrym-oai
BIN
Binary file not shown.
@@ -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<usize> {
|
||||
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()
|
||||
}
|
||||
|
||||
@@ -51,6 +51,7 @@ fn server(url: &str) -> McpServerConfig {
|
||||
"read".to_string(),
|
||||
McpServerToolConfig {
|
||||
approval_mode: Some(AppToolApproval::Approve),
|
||||
..Default::default()
|
||||
},
|
||||
)]),
|
||||
}
|
||||
|
||||
@@ -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));
|
||||
|
||||
@@ -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<AppToolApproval>,
|
||||
|
||||
/// 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<NonZeroUsize>,
|
||||
}
|
||||
|
||||
impl McpServerToolConfig {
|
||||
/// Applies the stricter explicit output budget without changing approval policy.
|
||||
pub fn restrict_output_token_limit(&mut self, limit: Option<NonZeroUsize>) {
|
||||
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<String>,
|
||||
|
||||
/// 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<String, McpServerToolConfig>,
|
||||
}
|
||||
|
||||
@@ -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::<McpServerConfig>(&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(
|
||||
|
||||
@@ -879,7 +879,7 @@ pub struct PluginMcpServerConfig {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub disabled_tools: Option<Vec<String>>,
|
||||
|
||||
/// 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<String, McpServerToolConfig>,
|
||||
}
|
||||
|
||||
@@ -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),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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),
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
},
|
||||
|
||||
@@ -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()
|
||||
}),
|
||||
},
|
||||
);
|
||||
|
||||
@@ -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] {
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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<String, McpServerConfig>) -> bool {
|
||||
fn replace_mcp_servers(
|
||||
&mut self,
|
||||
servers: &BTreeMap<String, McpServerConfig>,
|
||||
) -> anyhow::Result<bool> {
|
||||
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<String> = 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 {
|
||||
|
||||
@@ -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<TomlTable> {
|
||||
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<TomlItem> {
|
||||
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<TomlItem> {
|
||||
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<InlineTable> {
|
||||
Ok(serialize_mcp_server_table(config)?.into_inline_table())
|
||||
}
|
||||
|
||||
pub(super) fn merge_inline_table(existing: &mut InlineTable, replacement: InlineTable) {
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
|
||||
@@ -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).
|
||||
|
||||
@@ -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<usize>,
|
||||
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]
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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]"));
|
||||
|
||||
@@ -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<ResponseItemEnvelope>,
|
||||
) {
|
||||
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)
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -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<BoxFuture<'static, CodexResult<ResponseInputItem>>>,
|
||||
in_flight: &mut FuturesOrdered<InFlightFuture<'static>>,
|
||||
sess: Arc<Session>,
|
||||
turn_context: Arc<TurnContext>,
|
||||
) -> 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<BoxFuture<'static, CodexResult<ResponseInputItem>>> =
|
||||
FuturesOrdered::new();
|
||||
let mut in_flight: FuturesOrdered<InFlightFuture<'static>> = FuturesOrdered::new();
|
||||
let mut needs_follow_up = false;
|
||||
let mut last_agent_message: Option<String> = None;
|
||||
let mut active_item: Option<TurnItem> = None;
|
||||
|
||||
@@ -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<Box<dyn Future<Output = Result<ResponseInputItem>> + Send + 'f>>;
|
||||
Pin<Box<dyn Future<Output = Result<ResponseItemEnvelope>> + Send + 'f>>;
|
||||
|
||||
#[derive(Default)]
|
||||
pub(crate) struct OutputItemResult {
|
||||
|
||||
@@ -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<usize> {
|
||||
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
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<Output = Result<ResponseInputItem, CodexErr>> {
|
||||
) -> impl std::future::Future<Output = Result<ResponseItemEnvelope, CodexErr>> {
|
||||
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()
|
||||
|
||||
@@ -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<usize> {
|
||||
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)
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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")
|
||||
};
|
||||
|
||||
@@ -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>,
|
||||
|
||||
@@ -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<usize>,
|
||||
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<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_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(())
|
||||
|
||||
@@ -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()
|
||||
},
|
||||
),
|
||||
]),
|
||||
|
||||
@@ -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<usize>,
|
||||
}
|
||||
|
||||
impl ResponseItemEnvelope {
|
||||
|
||||
@@ -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()))?;
|
||||
|
||||
@@ -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<usize> {
|
||||
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<usize> {
|
||||
(**self).fallback_token_limit_override()
|
||||
}
|
||||
|
||||
fn to_response_item(&self, call_id: &str, payload: &ToolPayload) -> ResponseInputItem {
|
||||
(**self).to_response_item(call_id, payload)
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user