Skip to content

Commit 6bb9b15

Browse files
TheGreatAxiosTheGreatAxios
committed
fix(mcp): bound and scrub structured result evidence
1 parent 987966f commit 6bb9b15

6 files changed

Lines changed: 300 additions & 36 deletions

File tree

‎src/mcp/client.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -34,9 +34,9 @@ export interface MCPContentBlock {
3434
* protocol layer. The envelope carries all three so the plugin can surface
3535
* failures as errors and structured-only payloads as readable text.
3636
* `structuredContent` reaches the model JSON-serialized into the content
37-
* string under MCP_STRUCTURED_CONTENT_MARKER (see plugin.ts) with the raw
38-
* (policy-scrubbed) record preserved under ToolResult `detail` and in the
39-
* evidence archive — never raw.
37+
* string under MCP_STRUCTURED_CONTENT_MARKER (see plugin.ts). Small
38+
* policy-scrubbed records are preserved under ToolResult `detail`; the full
39+
* scrubbed record is retained in the evidence archive — never raw.
4040
*/
4141
export interface MCPToolResultEnvelope {
4242
blocks: MCPContentBlock[];

‎src/mcp/plugin.test.ts‎

Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
import { defined } from "../../tests/helpers/defined.js";
22
import { describe, test, expect } from "bun:test";
3+
import { mkdtempSync } from "node:fs";
4+
import { tmpdir } from "node:os";
5+
import { join } from "node:path";
6+
import type { ToolResult } from "@intx/types/runtime";
37
import { mcpClientToAgentTools } from "./plugin.js";
48
import { createPermissionGate } from "../permission/gate.js";
59
import {
@@ -8,6 +12,7 @@ import {
812
} from "../plugins/result-truncation-plugin.js";
913
import { toolOutputAbsolutePath } from "../plugins/tool-result-materialize.js";
1014
import { CREDENTIAL_REDACTION } from "../plugins/tool-result-secret-scrub.js";
15+
import { createCompactionArchive } from "../session/compaction-archive.js";
1116
import type { MCPClient, MCPContentBlock } from "./client.js";
1217

1318
interface ScriptedMcpEnvelope {
@@ -78,6 +83,39 @@ function fakeBlobStore() {
7883
};
7984
}
8085

86+
function memoryEvidenceArchive() {
87+
const blobs = new Map<string, Uint8Array>();
88+
const archive = createCompactionArchive({
89+
sessionId: "mcp-plugin-test",
90+
contextDir: mkdtempSync(join(tmpdir(), "mcp-plugin-archive-")),
91+
writeBlob: async (key, bytes) => {
92+
blobs.set(key, bytes);
93+
},
94+
readBlob: async (key) => {
95+
const bytes = blobs.get(key);
96+
if (bytes === undefined) throw new Error(`missing archive blob ${key}`);
97+
return bytes;
98+
},
99+
});
100+
return { archive, blobs };
101+
}
102+
103+
function serializePersistedToolResultTurn(result: ToolResult): string {
104+
return JSON.stringify({
105+
role: "user",
106+
content: [
107+
{
108+
type: "tool_result",
109+
callId: result.callId,
110+
content: [{ type: "text", text: String(result.content) }],
111+
...(result.detail !== undefined ? { detail: result.detail } : {}),
112+
...(result.isError !== undefined ? { isError: result.isError } : {}),
113+
},
114+
],
115+
timestamp: 0,
116+
});
117+
}
118+
81119
function skipGate() {
82120
return createPermissionGate({
83121
approvals: [],
@@ -292,6 +330,151 @@ describe("mcpClientToAgentTools", () => {
292330
expect(detailJson).not.toContain("sk-live-");
293331
});
294332

333+
test("scrubs structured keys from detail, archive bytes, and model content", async () => {
334+
const topLevelKey = ["sk-", "live-", "a".repeat(24)].join("");
335+
const nestedKey = ["sk-", "live-", "b".repeat(24)].join("");
336+
const { archive, blobs } = memoryEvidenceArchive();
337+
const result = await runEnvelopeTool(
338+
{
339+
blocks: [{ type: "resource", [topLevelKey]: "block-value" }],
340+
isError: false,
341+
structuredContent: {
342+
[topLevelKey]: "top-level",
343+
nested: { [nestedKey]: "nested" },
344+
},
345+
},
346+
"c-mcp-structured-key-secret",
347+
{ getEvidenceArchive: () => archive },
348+
);
349+
350+
const detail = JSON.stringify(result.detail);
351+
const modelTurn = serializePersistedToolResultTurn(result);
352+
const archiveBytes = [...blobs.values()].map((bytes) =>
353+
new TextDecoder().decode(bytes),
354+
);
355+
const surfaces = [
356+
detail,
357+
String(result.content),
358+
modelTurn,
359+
...archiveBytes,
360+
];
361+
for (const surface of surfaces) {
362+
expect(surface).not.toContain(topLevelKey);
363+
expect(surface).not.toContain(nestedKey);
364+
}
365+
expect(detail).toContain(CREDENTIAL_REDACTION);
366+
expect(String(result.content)).toContain(CREDENTIAL_REDACTION);
367+
expect(modelTurn).toContain(CREDENTIAL_REDACTION);
368+
expect(archiveBytes.join("\n")).toContain(CREDENTIAL_REDACTION);
369+
});
370+
371+
test("keeps oversized structured content full only in the evidence archive", async () => {
372+
const { archive } = memoryEvidenceArchive();
373+
const hugeValue = "x".repeat(MAX_RESULT_CHARS * 4);
374+
const result = await runEnvelopeTool(
375+
{
376+
blocks: [],
377+
isError: false,
378+
structuredContent: { hugeValue },
379+
},
380+
"c-mcp-oversized-detail",
381+
{ getEvidenceArchive: () => archive },
382+
);
383+
384+
expect(result.detail).toBeUndefined();
385+
expect(JSON.stringify(result).length).toBeLessThanOrEqual(
386+
MAX_RESULT_CHARS + 256,
387+
);
388+
const serializedTurn = serializePersistedToolResultTurn(result);
389+
expect(serializedTurn.length).toBeLessThanOrEqual(MAX_RESULT_CHARS + 512);
390+
expect(serializedTurn).not.toContain(hugeValue);
391+
392+
const [occurrence] = await archive.listOccurrences();
393+
if (occurrence === undefined) throw new Error("missing archive occurrence");
394+
const archived = JSON.parse(
395+
await archive.readAuthorizedPayload(occurrence.occurrenceId),
396+
) as { structuredContent: { hugeValue: string } };
397+
expect(archived.structuredContent.hugeValue).toBe(hugeValue);
398+
});
399+
400+
test("omits unserializable structured detail without failing text content", async () => {
401+
const result = await runEnvelopeTool(
402+
{
403+
blocks: [{ type: "text", text: "usable text" }],
404+
isError: false,
405+
structuredContent: { unsupported: 1n },
406+
},
407+
"c-mcp-unserializable-detail",
408+
);
409+
410+
expect(result.isError).toBeUndefined();
411+
expect(result.content).toBe("usable text");
412+
expect(result.detail).toBeUndefined();
413+
expect(JSON.stringify(result).length).toBeLessThan(MAX_RESULT_CHARS);
414+
});
415+
416+
test("archives identical success and failure payloads with distinct isError", async () => {
417+
const { archive } = memoryEvidenceArchive();
418+
const envelope = {
419+
blocks: [{ type: "text", text: "same payload" }],
420+
structuredContent: { answer: 42 },
421+
};
422+
423+
await runEnvelopeTool(
424+
{ ...envelope, isError: false },
425+
"c-mcp-archive-success",
426+
{ getEvidenceArchive: () => archive },
427+
);
428+
await runEnvelopeTool(
429+
{ ...envelope, isError: true },
430+
"c-mcp-archive-failure",
431+
{ getEvidenceArchive: () => archive },
432+
);
433+
434+
const occurrences = await archive.listOccurrences();
435+
expect(occurrences).toHaveLength(2);
436+
const payloads = await Promise.all(
437+
occurrences.map(async (occurrence) =>
438+
JSON.parse(
439+
await archive.readAuthorizedPayload(occurrence.occurrenceId),
440+
),
441+
),
442+
);
443+
expect(payloads.map((payload) => payload.isError)).toEqual([false, true]);
444+
});
445+
446+
test("falls back to legacy call when block and envelope methods are absent", async () => {
447+
let calls = 0;
448+
const client: MCPClient = {
449+
serverName: "legacy",
450+
tools: [
451+
{
452+
name: "echo",
453+
description: "returns legacy text",
454+
inputSchema: { type: "object", properties: {} },
455+
},
456+
],
457+
call: async () => {
458+
calls++;
459+
return "legacy response";
460+
},
461+
close: async () => undefined,
462+
};
463+
const [tool] = mcpClientToAgentTools(client, skipGate());
464+
if (tool?.kind !== "full") throw new Error("expected full tool");
465+
466+
const result = await tool.handler(
467+
{ id: "c-mcp-legacy", name: "mcp__legacy__echo", arguments: {} },
468+
new AbortController().signal,
469+
);
470+
471+
expect(calls).toBe(1);
472+
expect(result).toEqual({
473+
callId: "c-mcp-legacy",
474+
content: "legacy response",
475+
});
476+
});
477+
295478
test("thrown transport failures still surface as error results", async () => {
296479
const gate = skipGate();
297480
const client: MCPClient = {

‎src/mcp/plugin.ts‎

Lines changed: 39 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import {
77
scrubSecretShapedValue,
88
} from "../plugins/tool-result-secret-scrub.js";
99
import {
10+
MAX_RESULT_CHARS,
1011
truncateToolResultContent,
1112
type SpillBlobWriter,
1213
} from "../plugins/result-truncation-plugin.js";
@@ -43,21 +44,9 @@ export interface McpSpillOptions {
4344
}
4445

4546
function applyPolicyToBlocks(blocks: MCPContentBlock[]): MCPContentBlock[] {
46-
return blocks.map((block) => {
47-
const next = { ...block };
48-
if (typeof next.text === "string") {
49-
next.text = scrubSecretShapedContent(next.text);
50-
}
51-
for (const [key, value] of Object.entries(next)) {
52-
if (key === "type" || key === "text") continue;
53-
if (typeof value === "string") {
54-
next[key] = scrubSecretShapedContent(value);
55-
} else if (value !== null && typeof value === "object") {
56-
next[key] = scrubSecretShapedValue(value);
57-
}
58-
}
59-
return next;
60-
});
47+
return blocks.map(
48+
(block) => scrubSecretShapedValue(block) as MCPContentBlock,
49+
);
6150
}
6251

6352
// MCP results never reach the posix runner, so the secret-scrub and truncation
@@ -75,6 +64,20 @@ function sanitizeMcpResultContent(
7564
);
7665
}
7766

67+
function serializeStructuredContent(
68+
value: Record<string, unknown>,
69+
): { serialized: string; detail?: Record<string, unknown> } | undefined {
70+
try {
71+
const serialized = JSON.stringify(value);
72+
return {
73+
serialized,
74+
...(serialized.length <= MAX_RESULT_CHARS ? { detail: value } : {}),
75+
};
76+
} catch {
77+
return undefined;
78+
}
79+
}
80+
7881
export function mcpClientTools(
7982
client: MCPClient,
8083
spillOptions: McpSpillOptions = {},
@@ -134,18 +137,23 @@ export function mcpClientTools(
134137
string,
135138
unknown
136139
>);
140+
const isError = envelope.isError === true;
141+
const serializedStructured =
142+
scrubbedStructured === undefined
143+
? undefined
144+
: serializeStructuredContent(scrubbedStructured);
137145
const archive = getEvidenceArchive?.();
138146
if (archive !== undefined) {
139147
try {
140148
await archive.recordAuthorizedPayload({
141149
kind: "tool_result",
142-
payload:
143-
scrubbedStructured === undefined
144-
? { blocks: authorizedBlocks }
145-
: {
146-
blocks: authorizedBlocks,
147-
structuredContent: scrubbedStructured,
148-
},
150+
payload: {
151+
blocks: authorizedBlocks,
152+
isError,
153+
...(scrubbedStructured !== undefined
154+
? { structuredContent: scrubbedStructured }
155+
: {}),
156+
},
149157
callId: call.id,
150158
provenance: "mcp:post-policy-pre-flatten",
151159
});
@@ -154,15 +162,16 @@ export function mcpClientTools(
154162
}
155163
}
156164
const flattened = unwrapToolContent(authorizedBlocks);
157-
const isError = envelope.isError === true;
158165
const baseContent =
159166
flattened !== ""
160167
? flattened
161-
: scrubbedStructured !== undefined
162-
? `${MCP_STRUCTURED_CONTENT_MARKER}\n${JSON.stringify(scrubbedStructured)}`
163-
: isError
164-
? `MCP tool ${client.serverName}/${tool.name} reported an error with empty content.`
165-
: flattened;
168+
: serializedStructured !== undefined
169+
? `${MCP_STRUCTURED_CONTENT_MARKER}\n${serializedStructured.serialized}`
170+
: scrubbedStructured !== undefined
171+
? `${MCP_STRUCTURED_CONTENT_MARKER}\n[structured content unavailable]`
172+
: isError
173+
? `MCP tool ${client.serverName}/${tool.name} reported an error with empty content.`
174+
: flattened;
166175
const writeBlob = getBlobWriter?.();
167176
const contextDir = getContextDir?.();
168177
const spill =
@@ -178,8 +187,8 @@ export function mcpClientTools(
178187
callId: call.id,
179188
content,
180189
...(isError ? { isError: true as const } : {}),
181-
...(scrubbedStructured !== undefined
182-
? { detail: scrubbedStructured }
190+
...(serializedStructured?.detail !== undefined
191+
? { detail: serializedStructured.detail }
183192
: {}),
184193
};
185194
} catch (err) {

‎src/plugins/tool-result-secret-scrub.test.ts‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,46 @@ describe("scrubSecretShapedValue", () => {
6262
});
6363
});
6464

65+
describe("scrubSecretShapedValue key handling", () => {
66+
test("scrubs credential-shaped keys recursively", () => {
67+
const topLevelKey = ["sk-", "live-", "a".repeat(24)].join("");
68+
const nestedKey = `prefix-${["sk-", "live-", "b".repeat(24)].join("")}`;
69+
70+
const out = scrubSecretShapedValue({
71+
[topLevelKey]: "top-level",
72+
nested: { [nestedKey]: "nested" },
73+
});
74+
const serialized = JSON.stringify(out);
75+
76+
expect(out).toEqual({
77+
[CREDENTIAL_REDACTION]: "top-level",
78+
nested: { [`prefix-${CREDENTIAL_REDACTION}`]: "nested" },
79+
});
80+
expect(serialized).not.toContain(topLevelKey);
81+
expect(serialized).not.toContain(nestedKey);
82+
});
83+
84+
test("resolves scrubbed key collisions deterministically without raw keys", () => {
85+
const firstKey = ["sk-", "live-", "a".repeat(24)].join("");
86+
const secondKey = ["sk-", "live-", "b".repeat(24)].join("");
87+
88+
const out = scrubSecretShapedValue({
89+
[firstKey]: "first",
90+
[secondKey]: "second",
91+
[CREDENTIAL_REDACTION]: "already-redacted",
92+
});
93+
const serialized = JSON.stringify(out);
94+
95+
expect(out).toEqual({
96+
[CREDENTIAL_REDACTION]: "first",
97+
[`${CREDENTIAL_REDACTION} [2]`]: "second",
98+
[`${CREDENTIAL_REDACTION} [3]`]: "already-redacted",
99+
});
100+
expect(serialized).not.toContain(firstKey);
101+
expect(serialized).not.toContain(secondKey);
102+
});
103+
});
104+
65105
describe("toolResultSecretScrubPlugin", () => {
66106
const next =
67107
(content: ToolResult["content"], isError = false) =>

0 commit comments

Comments
 (0)