mirror of
https://github.com/openai/codex.git
synced 2026-11-11 04:42:10 +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:
1 parent
03147407e3
commit
f742dabc6f
38 files changed
+487
-224
No files matched your search
@@ -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>,
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user