Skip to content

feat(mcp): bound successful tool results #4916

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Astro-Han merged 1 commit into from
Sep 7, 2026
Merged

feat(mcp): bound successful tool results #4916

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
167 changes: 167 additions & 0 deletions packages/mcp/src/__tests__/manager.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -971,6 +971,133 @@ describe('McpClientManager E2E', { concurrency: false }, () => {
);
});

test('accepts the content block limit and rejects the next block', async () => {
const fixture = await createRemoteFixture('streamable-http', {
advertiseTools: false,
legacyToolCallResult: {
content: Array.from({ length: 256 }, () => ({ type: 'text', text: 'x' })),
},
});
const manager = createManager();
await manager.sync(remoteConfig(fixture.url));

const accepted = await manager.callTool(bindingFor(manager, 'remote', 'echo'), {});
assert.equal(accepted.content.length, 256);

await assert.rejects(
callToolResultFixture({
content: Array.from({ length: 257 }, () => ({ type: 'text', text: 'x' })),
}),
/exceeds the content block limit/u,
);
});

test('accepts JSON depth 32 and rejects depth 33', async () => {
const fixture = await createRemoteFixture('streamable-http', {
advertiseTools: false,
legacyToolCallResult: {
content: [{ type: 'text', text: 'deep' }],
structuredContent: nestedStructuredContent(30),
},
});
const manager = createManager();
await manager.sync(remoteConfig(fixture.url));

const accepted = await manager.callTool(bindingFor(manager, 'remote', 'echo'), {});
assert.deepEqual(accepted.structuredContent, nestedStructuredContent(30));

await assert.rejects(
callToolResultFixture({
content: [{ type: 'text', text: 'deep' }],
structuredContent: nestedStructuredContent(31),
}),
/exceeds the JSON depth limit/u,
);
});

test('accepts 8192 JSON nodes and rejects node 8193', async () => {
const fixture = await createRemoteFixture('streamable-http', {
advertiseTools: false,
legacyToolCallResult: nodeBoundaryToolResult(8_192),
});
const manager = createManager();
await manager.sync(remoteConfig(fixture.url));

const accepted = await manager.callTool(bindingFor(manager, 'remote', 'echo'), {});
assert.equal((accepted.structuredContent as { items: unknown[] }).items.length, 8_185);

await assert.rejects(
callToolResultFixture(nodeBoundaryToolResult(8_193)),
/exceeds the JSON node limit/u,
);
});

test('counts exact serialized UTF-8 bytes and rejects the next byte', async () => {
const fixture = await createRemoteFixture('streamable-http', {
advertiseTools: false,
legacyToolCallResult: exactByteToolResult(24 * 1024 * 1024, '😀'),
});
const manager = createManager();
await manager.sync(remoteConfig(fixture.url));

const accepted = await manager.callTool(bindingFor(manager, 'remote', 'echo'), {});
assert.equal(accepted.content[0]?.type, 'text');

const oversized = exactByteToolResult(24 * 1024 * 1024 + 1, '\0');
assert.equal(Buffer.byteLength(JSON.stringify(oversized), 'utf8'), 24 * 1024 * 1024 + 1);
await assert.rejects(callToolResultFixture(oversized), /exceeds the byte limit/u);
});

test('applies the JSON budget to metadata on known content blocks', async () => {
const result = {
content: [
{
type: 'text',
text: 'safe',
_meta: nestedStructuredContent(31),
},
],
};

await assert.rejects(callToolResultFixture(result), /exceeds the JSON depth limit/u);
});

test('rechecks the budget after credential scrubbing expands a result', async () => {
const fixture = await createRemoteFixture('streamable-http', {
advertiseTools: false,
legacyToolCallResult: exactByteToolResult(24 * 1024 * 1024, 'abcd'),
});
const manager = createManager();
await manager.sync({
version: MCP_CONFIG_VERSION,
mcpServers: {
remote: {
url: fixture.url,
transport: 'streamable-http',
headers: { Authorization: 'abcd' },
},
},
});

await assert.rejects(
manager.callTool(bindingFor(manager, 'remote', 'echo'), {}),
/exceeds the byte limit/u,
);
});

test('checks the success budget before output-schema validation', async () => {
const fixture = await createRemoteFixture('streamable-http');
const manager = createManager();
await manager.sync(remoteConfig(fixture.url));

await assert.rejects(
manager.callTool(bindingFor(manager, 'remote', 'invalid-output'), {
mode: 'oversized-success',
}),
/exceeds the byte limit/u,
);
});

test('rejects an invalid binding before the wire', async () => {
const fixture = await createRemoteFixture('streamable-http');
const manager = createManager();
Expand Down Expand Up @@ -1722,6 +1849,12 @@ function createProtocolServer(options: {
content: Array.from({ length: 101 }, () => ({ type: 'text' as const, text: 'x' })),
};
}
if (args.mode === 'oversized-success') {
return {
content: [{ type: 'text', text: 'x'.repeat(24 * 1024 * 1024) }],
structuredContent: { wrong: true },
};
}
return {
content: [{ type: 'text', text: 'invalid' }],
structuredContent: { wrong: true },
Expand Down Expand Up @@ -1774,6 +1907,40 @@ function remoteToolDefinition(
};
}

function nestedStructuredContent(depth: number): Record<string, unknown> {
let value: Record<string, unknown> = { value: 'done' };
for (let index = 0; index < depth; index += 1) value = { child: value };
return value;
}

function nodeBoundaryToolResult(nodes: number): Record<string, unknown> {
// result + content + block + block.type + block.text + structuredContent + items = 7 nodes.
return {
content: [{ type: 'text', text: 'many' }],
structuredContent: { items: Array.from({ length: nodes - 7 }, (_, index) => index) },
};
}

function exactByteToolResult(bytes: number, suffix: string): Record<string, unknown> {
const result = { content: [{ type: 'text', text: suffix }] };
const overhead = Buffer.byteLength(JSON.stringify(result), 'utf8') - Buffer.byteLength(suffix);
const fillBytes = bytes - overhead - Buffer.byteLength(suffix);
assert.ok(fillBytes >= 0);
result.content[0]!.text = `${'x'.repeat(fillBytes)}${suffix}`;
assert.equal(Buffer.byteLength(JSON.stringify(result), 'utf8'), bytes);
return result;
}

async function callToolResultFixture(result: unknown) {
const fixture = await createRemoteFixture('streamable-http', {
advertiseTools: false,
legacyToolCallResult: result,
});
const manager = createManager();
await manager.sync(remoteConfig(fixture.url));
return manager.callTool(bindingFor(manager, 'remote', 'echo'), {});
}

function reorderedRemoteEchoDefinition() {
return {
inputSchema: {
Expand Down
162 changes: 161 additions & 1 deletion packages/mcp/src/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ import {
} from '@modelcontextprotocol/client';
import { StdioClientTransport } from '@modelcontextprotocol/client/stdio';
import { redactSecrets } from '@maka/core/redaction';
import { serializedByteLength } from '@maka/core/serialized-byte-length';
import {
deepScrubMcpSecrets,
EMPTY_MCP_SECRET_INVENTORY,
Expand Down Expand Up @@ -133,6 +134,10 @@ const STDERR_OVERSIZED_LINE = '[stderr line omitted: exceeds diagnostic limit]';
const STDERR_CONTINUATION = '[stderr continuation omitted]';
const MAX_SUMMARIZED_ERROR_BLOCKS = 100;
const OVERSIZED_TOOL_ERROR_CONTENT = 'server returned oversized error content';
const MAX_SUCCESS_TOOL_RESULT_CONTENT_BLOCKS = 256;
const MAX_SUCCESS_TOOL_RESULT_BYTES = 24 * 1024 * 1024;
const MAX_SUCCESS_TOOL_RESULT_JSON_DEPTH = 32;
const MAX_SUCCESS_TOOL_RESULT_JSON_NODES = 8_192;
const MAX_TOOL_REFRESH_PASSES = 3;
const TOOL_REFRESH_BURST_IDLE_MS = 1_000;
// Recency-bounded per-server scrub material: enough to cover every value a
Expand Down Expand Up @@ -855,6 +860,7 @@ export class McpClientManager {
scrubKnownSecrets(redactSecrets(summarizeErrorContent(result.content)), inventory),
);
}
assertSuccessfulToolResultBudget(serverId, toolName, result, { raw: true });
const validateOutput = preparation.value.validateOutput;
if (validateOutput) {
if (result.structuredContent === undefined) {
Expand All @@ -876,10 +882,12 @@ export class McpClientManager {
}
// Success payloads cross toward the renderer and the transcript too; a
// server can embed the credential it was just sent into a result.
return {
const published = {
content: deepScrub(result.content.map(normalizeContent), inventory),
structuredContent: deepScrub(result.structuredContent, inventory),
};
assertSuccessfulToolResultBudget(serverId, toolName, published);
return published;
}

async test(serverId: string): Promise<McpTestResult> {
Expand Down Expand Up @@ -2390,6 +2398,158 @@ function normalizeContent(value: unknown): McpContentBlock {
return { type: 'unknown', value };
}

function assertSuccessfulToolResultBudget(
serverId: string,
toolName: string,
result: { content: unknown[]; structuredContent?: unknown },
options: { raw?: boolean } = {},
): void {
if (result.content.length > MAX_SUCCESS_TOOL_RESULT_CONTENT_BLOCKS) {
throw successfulToolResultBudgetError(serverId, toolName, 'content block');
}
const budget: SuccessfulToolResultBudget = {
bytes: 0,
nodes: 0,
active: new WeakSet<object>(),
};
try {
// Bound the untrusted content before normalization and then the exact
// result that crosses the public boundary. Known-block metadata may be
// discarded later, but it must not provide a path around the first pass.
countSuccessfulToolResultValue(
result.structuredContent === undefined
? { content: result.content }
: { content: result.content, structuredContent: result.structuredContent },
budget,
0,
options.raw === true,
);
} catch (error) {
if (error instanceof SuccessfulToolResultBudgetViolation) {
throw successfulToolResultBudgetError(serverId, toolName, error.limit);
}
throw new McpToolCallError(serverId, toolName, 'server returned an invalid tool result');
}
}

type SuccessfulToolResultBudgetLimit = 'byte' | 'content block' | 'JSON depth' | 'JSON node';

class SuccessfulToolResultBudgetViolation extends Error {
constructor(readonly limit: Exclude<SuccessfulToolResultBudgetLimit, 'content block'>) {
super(limit);
}
}

interface SuccessfulToolResultBudget {
bytes: number;
nodes: number;
active: WeakSet<object>;
}

function countSuccessfulToolResultValue(
value: unknown,
budget: SuccessfulToolResultBudget,
depth: number,
rejectUndefined = false,
): void {
budget.nodes += 1;
if (budget.nodes > MAX_SUCCESS_TOOL_RESULT_JSON_NODES) {
throw new SuccessfulToolResultBudgetViolation('JSON node');
}
if (depth > MAX_SUCCESS_TOOL_RESULT_JSON_DEPTH) {
throw new SuccessfulToolResultBudgetViolation('JSON depth');
}
if (value === null) {
countSuccessfulToolResultBytes(budget, 4);
return;
}
if (typeof value === 'string') {
countSuccessfulToolResultJsonStringBytes(budget, value);
return;
}
if (typeof value === 'boolean') {
countSuccessfulToolResultBytes(budget, value ? 4 : 5);
return;
}
if (typeof value === 'number') {
if (!Number.isFinite(value)) throw new Error('non-finite number');
countSuccessfulToolResultBytes(budget, JSON.stringify(value).length);
return;
}
if (typeof value !== 'object') throw new Error('unsupported JSON value');
if (budget.active.has(value)) throw new Error('cyclic JSON value');
if (typeof (value as { toJSON?: unknown }).toJSON === 'function') {
throw new Error('custom JSON serializer');
}
if (!Array.isArray(value)) {
const prototype = Object.getPrototypeOf(value);
if (prototype !== Object.prototype && prototype !== null) {
throw new Error('non-plain JSON object');
}
}

budget.active.add(value);
try {
if (Array.isArray(value)) {
countSuccessfulToolResultBytes(budget, 1);
for (let index = 0; index < value.length; index += 1) {
if (index > 0) countSuccessfulToolResultBytes(budget, 1);
countSuccessfulToolResultValue(value[index], budget, depth + 1, rejectUndefined);
}
countSuccessfulToolResultBytes(budget, 1);
return;
}

countSuccessfulToolResultBytes(budget, 1);
let emitted = 0;
for (const key in value) {
if (!Object.hasOwn(value, key)) continue;
const item = (value as Record<string, unknown>)[key];
if (item === undefined) {
if (rejectUndefined) throw new Error('unsupported JSON value');
continue;
}
if (emitted > 0) countSuccessfulToolResultBytes(budget, 1);
countSuccessfulToolResultJsonStringBytes(budget, key);
countSuccessfulToolResultBytes(budget, 1);
countSuccessfulToolResultValue(item, budget, depth + 1, rejectUndefined);
emitted += 1;
}
countSuccessfulToolResultBytes(budget, 1);
} finally {
budget.active.delete(value);
}
}

function countSuccessfulToolResultJsonStringBytes(
budget: SuccessfulToolResultBudget,
value: string,
): void {
const remainingBytes = MAX_SUCCESS_TOOL_RESULT_BYTES - budget.bytes;
const bytes = serializedByteLength(value, remainingBytes);
if (bytes > remainingBytes) throw new SuccessfulToolResultBudgetViolation('byte');
budget.bytes += bytes;
}

function countSuccessfulToolResultBytes(budget: SuccessfulToolResultBudget, bytes: number): void {
if (budget.bytes > MAX_SUCCESS_TOOL_RESULT_BYTES - bytes) {
throw new SuccessfulToolResultBudgetViolation('byte');
}
budget.bytes += bytes;
}

function successfulToolResultBudgetError(
serverId: string,
toolName: string,
limit: SuccessfulToolResultBudgetLimit,
): McpToolCallError {
return new McpToolCallError(
serverId,
toolName,
`server returned a tool result that exceeds the ${limit} limit`,
);
}

function summarizeErrorContent(content: unknown[]): string {
if (content.length > MAX_SUMMARIZED_ERROR_BLOCKS) return OVERSIZED_TOOL_ERROR_CONTENT;
const fragments: string[] = [];
Expand Down