Make MCP and Code Mode input schema budgets configurable (#47936)

## Why

Fixed schema budgets can strip parameter descriptions from large MCP tools or render their Code Mode input types as `unknown`. Allow larger budgets to preserve that guidance.

## What changed

- Add `mcp_servers.<name>.tool_input_schema_max_bytes` to configure the UTF-8 byte threshold for best-effort input schema compaction, retaining the 5,000-byte default.
- Add `features.code_mode.tool_input_schema_max_bytes` for rendered input types. Use the largest of the 16,000-byte default, the Code Mode setting, and the ordinary MCP server's explicit setting.
- Apply rendering budgets to prompt declarations and the runtime tool catalog, and avoid reusing cached declarations after the Code Mode budget changes.

## Testing

Add coverage for UTF-8 compaction thresholds, budget precedence, configuration round trips, and preservation of parameter guidance in both the Code Mode prompt and `ALL_TOOLS`.

GitOrigin-RevId: 0ef2823a614d4794e636f518b09c0b03a5239b45
This commit is contained in:
Vivian Fang
2026-09-24 20:39:37 +00:00
committed by copyberry
parent b62fcd3c4b
commit 339e981ba7
63 changed files with 670 additions and 86 deletions
+21 -1
View File
@@ -10,6 +10,7 @@ use std::collections::BTreeMap;
use crate::PUBLIC_TOOL_NAME;
use crate::json_schema_types::render_json_schema_to_typescript;
use crate::json_schema_types::render_json_schema_to_typescript_with_budget;
const MAX_JS_SAFE_INTEGER: u64 = (1_u64 << 53) - 1;
const DEFERRED_NESTED_TOOLS_GUIDANCE: &str = r#"Some deferred nested tools may be omitted from this description. They are still available on the global `tools` object and listed in `ALL_TOOLS`.
@@ -143,6 +144,9 @@ pub struct ToolDefinition {
pub description: String,
pub kind: CodeModeToolKind,
pub input_schema: Option<JsonValue>,
/// Internal budget used while rendering the input declaration, before sending it to the host.
#[serde(skip)]
pub input_schema_max_bytes: Option<usize>,
pub output_schema: Option<JsonValue>,
}
@@ -433,7 +437,10 @@ fn render_code_mode_sample_for_definition(definition: &ToolDefinition) -> String
CodeModeToolKind::Function => definition
.input_schema
.as_ref()
.map(render_json_schema_to_typescript)
.map(|schema| match definition.input_schema_max_bytes {
Some(max_bytes) => render_json_schema_to_typescript_with_budget(schema, max_bytes),
None => render_json_schema_to_typescript(schema),
})
.unwrap_or_else(|| "unknown".to_string()),
CodeModeToolKind::Freeform => "string".to_string(),
};
@@ -610,6 +617,7 @@ mod tests {
"required": ["city"],
"additionalProperties": false
})),
input_schema_max_bytes: None,
output_schema: Some(json!({
"type": "object",
"properties": { "ok": { "type": "boolean" } },
@@ -650,6 +658,7 @@ mod tests {
},
"required": ["weather"]
})),
input_schema_max_bytes: None,
output_schema: Some(json!({
"type": "object",
"properties": {
@@ -686,6 +695,7 @@ mod tests {
"properties": {},
"additionalProperties": false
})),
input_schema_max_bytes: None,
output_schema: Some(mcp_call_tool_result_schema(json!({
"type": "object",
"properties": {
@@ -725,6 +735,7 @@ mod tests {
description: "bar".to_string(),
kind: CodeModeToolKind::Function,
input_schema: None,
input_schema_max_bytes: None,
output_schema: None,
}],
&[],
@@ -778,6 +789,7 @@ bar"
"properties": {},
"additionalProperties": false
})),
input_schema_max_bytes: None,
output_schema: Some(mcp_call_tool_result_schema(json!({
"type": "object",
"properties": {},
@@ -794,6 +806,7 @@ bar"
"properties": {},
"additionalProperties": false
})),
input_schema_max_bytes: None,
output_schema: Some(mcp_call_tool_result_schema(json!({
"type": "object",
"properties": {},
@@ -838,6 +851,7 @@ bar"
"properties": {},
"additionalProperties": false
})),
input_schema_max_bytes: None,
output_schema: Some(mcp_call_tool_result_schema(json!({
"type": "object",
"properties": {},
@@ -868,6 +882,7 @@ bar"
"properties": {},
"additionalProperties": false
})),
input_schema_max_bytes: None,
output_schema: Some(json!({
"type": "object",
"properties": {
@@ -902,6 +917,7 @@ bar"
"properties": {},
"additionalProperties": false
})),
input_schema_max_bytes: None,
output_schema: Some(json!({
"type": "object",
"properties": {
@@ -935,6 +951,7 @@ bar"
description: "First tool".to_string(),
kind: first_tool.kind,
input_schema: first_tool.input_schema,
input_schema_max_bytes: None,
output_schema: first_tool.output_schema,
},
ToolDefinition {
@@ -943,6 +960,7 @@ bar"
description: "Second tool".to_string(),
kind: second_tool.kind,
input_schema: second_tool.input_schema,
input_schema_max_bytes: None,
output_schema: second_tool.output_schema,
},
],
@@ -975,6 +993,7 @@ bar"
"properties": {},
"additionalProperties": false
})),
input_schema_max_bytes: None,
output_schema: Some(mcp_call_tool_result_schema(json!({
"type": "object",
"properties": {},
@@ -1007,6 +1026,7 @@ bar"
description: "Deferred tool".to_string(),
kind: CodeModeToolKind::Function,
input_schema: None,
input_schema_max_bytes: None,
output_schema: None,
}],
&BTreeMap::new(),
@@ -88,6 +88,7 @@ fn exec_override_preserves_runtime_sections() {
description: "First tool".to_string(),
kind: CodeModeToolKind::Function,
input_schema: None,
input_schema_max_bytes: None,
output_schema: None,
}];
let deferred_tools = [ToolDefinition {
@@ -96,6 +97,7 @@ fn exec_override_preserves_runtime_sections() {
description: "Deferred tool".to_string(),
kind: CodeModeToolKind::Function,
input_schema: None,
input_schema_max_bytes: None,
output_schema: Some(json!({
"type": "object",
"properties": {
@@ -166,6 +166,7 @@ impl From<WireToolDefinition> for ToolDefinition {
description: value.description,
kind: value.kind.into(),
input_schema: value.input_schema,
input_schema_max_bytes: None,
output_schema: value.output_schema,
}
}
@@ -10,14 +10,21 @@ const MAX_LOCAL_REF_EXPANSIONS_PER_PATH: usize = 2;
// compact schema cannot expand into an arbitrarily large model-visible item.
const MAX_TOTAL_LOCAL_REF_EXPANSIONS: usize = 32;
// Bound individual schema rendering before assembling the final declaration.
const MAX_RENDERED_SCHEMA_BYTES: usize = 16_000;
pub const DEFAULT_INPUT_SCHEMA_MAX_BYTES: usize = 16_000;
// Charge intermediate render strings as they are built so repeated local refs
// cannot allocate unbounded expanded copies before the final schema cap runs.
const MAX_RENDER_WORK_BYTES: usize = MAX_RENDERED_SCHEMA_BYTES * 4;
const RENDER_WORK_MULTIPLIER: usize = 4;
pub fn render_json_schema_to_typescript(schema: &JsonValue) -> String {
let rendered = JsonSchemaTypeRenderer::new(schema).render(schema);
if rendered.len() > MAX_RENDERED_SCHEMA_BYTES {
render_json_schema_to_typescript_with_budget(schema, DEFAULT_INPUT_SCHEMA_MAX_BYTES)
}
pub(crate) fn render_json_schema_to_typescript_with_budget(
schema: &JsonValue,
max_bytes: usize,
) -> String {
let rendered = JsonSchemaTypeRenderer::new(schema, max_bytes).render(schema);
if rendered.len() > max_bytes {
"unknown".to_string()
} else {
rendered
@@ -34,13 +41,13 @@ struct JsonSchemaTypeRenderer<'a> {
}
impl<'a> JsonSchemaTypeRenderer<'a> {
fn new(root: &'a JsonValue) -> Self {
fn new(root: &'a JsonValue, max_bytes: usize) -> Self {
Self {
root,
nested_schema_resource_depth: 0,
active_local_ref_expansions: BTreeMap::new(),
remaining_local_ref_expansions: MAX_TOTAL_LOCAL_REF_EXPANSIONS,
remaining_render_work_bytes: MAX_RENDER_WORK_BYTES,
remaining_render_work_bytes: max_bytes.saturating_mul(RENDER_WORK_MULTIPLIER),
render_work_budget_exhausted: false,
}
}
@@ -158,14 +158,14 @@ fn repeated_large_ref_expansions_exhaust_render_work_budget() {
"properties": {
"value": {
"type": "string",
"description": "x".repeat(MAX_RENDERED_SCHEMA_BYTES / 2)
"description": "x".repeat(DEFAULT_INPUT_SCHEMA_MAX_BYTES / 2)
}
}
}
}
});
let mut renderer = JsonSchemaTypeRenderer::new(&schema);
let mut renderer = JsonSchemaTypeRenderer::new(&schema, DEFAULT_INPUT_SCHEMA_MAX_BYTES);
assert_eq!(renderer.render(&schema), "unknown");
assert!(renderer.render_work_budget_exhausted);
}
@@ -176,19 +176,19 @@ fn oversized_ref_literal_exhausts_render_work_budget() {
"$ref": "#/$defs/Value",
"$defs": {
"Value": {
"const": "x".repeat(MAX_RENDER_WORK_BYTES)
"const": "x".repeat(DEFAULT_INPUT_SCHEMA_MAX_BYTES * RENDER_WORK_MULTIPLIER)
}
}
});
let mut renderer = JsonSchemaTypeRenderer::new(&schema);
let mut renderer = JsonSchemaTypeRenderer::new(&schema, DEFAULT_INPUT_SCHEMA_MAX_BYTES);
assert_eq!(renderer.render(&schema), "unknown");
assert!(renderer.render_work_budget_exhausted);
}
#[test]
fn rendered_schema_has_a_hard_size_cap() {
let description = "x".repeat(MAX_RENDERED_SCHEMA_BYTES);
let description = "x".repeat(DEFAULT_INPUT_SCHEMA_MAX_BYTES);
let schema = json!({
"type": "object",
"properties": {
+1
View File
@@ -20,6 +20,7 @@ pub use description::is_code_mode_nested_tool;
pub use description::normalize_code_mode_identifier;
pub use description::parse_exec_source;
pub use description::render_code_mode_sample;
pub use json_schema_types::DEFAULT_INPUT_SCHEMA_MAX_BYTES;
pub use json_schema_types::render_json_schema_to_typescript;
pub use response::DEFAULT_IMAGE_DETAIL;
pub use response::FunctionCallOutputContentItem;