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 authored and copyberry committed 2026-08-28 21:40:23 +00:00
1 parent 03147407e3
commit f742dabc6f
38 files changed
+487 -224

No files matched your search

+14 -2
View File
@@ -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>,
}
+15
View File
@@ -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(
+1 -1
View File
@@ -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>,
}