From aa1e88fbf3bd11eb1d7797aeee0fb6a5b82c8d7d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fortissolu=C3=A7=C3=B5es?= Date: Fri, 17 Jul 2026 16:34:51 -0300 Subject: [PATCH] fix(agent-runtime): keep MCP input schemas JSON-serializable Store MCP tool input schemas as plain JSON Schema objects instead of converting them to Zod schemas before storage. Zod schemas contain circular internal references that cause JSON.stringify to fail with 'cannot serialize cyclic structures' when the agent state is serialized. The runtime now preserves the original JSON Schema and converts through ensureZodSchema only at the consumption boundary (getToolSet). Tests: 16 new tests covering schema storage, serialization roundtrip, multi-server scenarios, circular rejection defense, and full pipeline validation. --- .../__tests__/mcp-schema-integration.test.ts | 227 ++++++++++++ .../mcp-schema-serialization.test.ts | 329 ++++++++++++++++++ packages/agent-runtime/src/mcp.ts | 3 +- 3 files changed, 557 insertions(+), 2 deletions(-) create mode 100644 packages/agent-runtime/src/__tests__/mcp-schema-integration.test.ts create mode 100644 packages/agent-runtime/src/__tests__/mcp-schema-serialization.test.ts diff --git a/packages/agent-runtime/src/__tests__/mcp-schema-integration.test.ts b/packages/agent-runtime/src/__tests__/mcp-schema-integration.test.ts new file mode 100644 index 0000000000..b1bbfdf325 --- /dev/null +++ b/packages/agent-runtime/src/__tests__/mcp-schema-integration.test.ts @@ -0,0 +1,227 @@ +/** + * End-to-end integration test for the MCP schema serialization fix. + * + * Validates the full pipeline: + * MCP tools/list → getMCPToolData → getToolSet → JSON.stringify(agentState) + * + * Run with: bun test packages/agent-runtime/src/__tests__/mcp-schema-integration.test.ts + */ +import { describe, test, expect } from 'bun:test' + +import { getMCPToolData } from '../mcp' +import { ensureZodSchema, getToolSet } from '../tools/prompts' +// @ts-expect-error - tree-sitter-wasm will be resolved at runtime in the test env +import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime' + +import type { AgentTemplate } from '../templates/types' +import type { MCPConfig } from '@codebuff/common/types/mcp' +import type { CustomToolDefinitions } from '@codebuff/common/util/file' + +// --------------------------------------------------------------------------- +// Realistic MCP tool schemas (based on actual Engram tool definitions) +// --------------------------------------------------------------------------- + +/** A realistic Engram mem_save tool schema (complex nested JSON Schema) */ +const engramMemSaveSchema = { + type: 'object', + properties: { + title: { type: 'string', description: 'Short, searchable title' }, + type: { + type: 'string', + enum: ['bugfix', 'decision', 'architecture', 'discovery', 'pattern', 'config', 'preference'], + description: 'Category', + }, + scope: { type: 'string', enum: ['project', 'personal'] }, + topic_key: { type: 'string', description: 'Optional topic identifier for upserts' }, + capture_prompt: { type: 'boolean', description: 'Capture the current user prompt' }, + content: { type: 'string', description: 'Structured content with What/Why/Where/Learned' }, + project: { type: 'string', description: 'Optional explicit project' }, + session_id: { type: 'string' }, + }, + required: ['title'], + additionalProperties: false, +} as Record + +/** A realistic Engram mem_search tool schema */ +const engramMemSearchSchema = { + type: 'object', + properties: { + query: { type: 'string', description: 'Search query — natural language or keywords' }, + limit: { type: 'number', minimum: 1, maximum: 20 }, + type: { type: 'string' }, + scope: { type: 'string' }, + project: { type: 'string' }, + all_projects: { type: 'boolean' }, + match_mode: { type: 'string', enum: ['all', 'any'] }, + }, + required: ['query'], + additionalProperties: false, +} as Record + +/** A realistic Engram mem_context tool schema */ +const engramMemContextSchema = { + type: 'object', + properties: { + project: { type: 'string' }, + scope: { type: 'string' }, + }, +} as Record + +/** Codegraph explore schema */ +const codegraphExploreSchema = { + type: 'object', + properties: { + query: { type: 'string' }, + maxFiles: { type: 'number' }, + projectPath: { type: 'string' }, + }, + required: ['query'], +} as Record + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +describe('MCP schema serialization — integration tests', () => { + describe('getMCPToolData produces serializable CustomToolDefinitions', () => { + test('multiple MCP servers with complex real-world schemas', async () => { + const engramTools = [ + { name: 'mem_save', description: 'Save memory', inputSchema: engramMemSaveSchema }, + { name: 'mem_search', description: 'Search memory', inputSchema: engramMemSearchSchema }, + { name: 'mem_context', description: 'Get memory context', inputSchema: engramMemContextSchema }, + ] + + const codegraphTools = [ + { name: 'codegraph_explore', description: 'Explore code', inputSchema: codegraphExploreSchema }, + ] + + const result = await getMCPToolData({ + toolNames: [ + 'engram/mem_save', + 'engram/mem_search', + 'engram/mem_context', + 'codegraph/codegraph_explore', + ], + mcpServers: { + engram: { command: 'engram', args: ['stdio'] } as any, + codegraph: { command: 'codegraph', args: ['stdio'] } as any, + }, + requestMcpToolData: async ({ mcpConfig }) => { + const name = (mcpConfig as any).command as string + return name === 'engram' ? engramTools : codegraphTools + }, + }) + + // All tools should be present + expect(Object.keys(result).length).toBe(4) + + // EVERY inputSchema must be a plain object, NOT a Zod schema + for (const [name, def] of Object.entries(result)) { + expect(typeof (def.inputSchema as any)?.safeParse).not.toBe('function') + } + + // JSON.stringify must work + expect(() => JSON.stringify(result)).not.toThrow() + + // Roundtrip preserves data + const json = JSON.stringify(result) + const parsed = JSON.parse(json) as CustomToolDefinitions + expect(Object.keys(parsed).length).toBe(4) + expect(parsed['engram__mem_save'].endsAgentStep).toBe(true) + expect(parsed['engram__mem_search'].description).toBe('Search memory') + }) + }) + + describe('full pipeline: getMCPToolData → getToolSet → JSON.stringify', () => { + test('tool set built from MCP definitions is JSON-serializable', async () => { + // Step 1: Get MCP tool definitions (simulating real Engram tools) + const mcpDefs = await getMCPToolData({ + toolNames: ['engram/mem_search', 'engram/mem_context'], + mcpServers: { + engram: { command: 'engram', args: ['stdio'] } as any, + }, + requestMcpToolData: async () => [ + { name: 'mem_search', description: 'Search memory', inputSchema: engramMemSearchSchema }, + { name: 'mem_context', description: 'Get memory context', inputSchema: engramMemContextSchema }, + ], + }) + + // Step 2: Build the full tool set (mimics what run-agent-step.ts does) + const toolSet = await getToolSet({ + toolNames: ['read_files'], + additionalToolDefinitions: async () => mcpDefs, + agentTools: {}, + skills: {}, + }) + + // The tool set should contain built-in tools + MCP tools + expect(toolSet['read_files']).toBeDefined() + expect(toolSet['engram__mem_search']).toBeDefined() + expect(toolSet['engram__mem_context']).toBeDefined() + + // Step 3: Convert tool definitions to agent state format (mimics run-agent-step.ts:910-913) + const toolDefinitions = Object.fromEntries( + Object.entries(toolSet).map(([name, tool]) => [ + name, + { description: tool.description, inputSchema: tool.inputSchema }, + ]), + ) + + // Step 4: THE DECISIVE TEST — serialization of the agent state payload + expect(() => JSON.stringify(toolDefinitions)).not.toThrow() + + const json = JSON.stringify(toolDefinitions) + const parsed = JSON.parse(json) + + // MCP tools survived roundtrip in the serialized payload + expect(parsed['engram__mem_search']).toBeDefined() + expect(parsed['engram__mem_context']).toBeDefined() + // Built-in tools are Zod schemas — they may or may not survive serialization + // (that's a different concern; our fix is about MCP schemas) + + // Step 5: ensureZodSchema works on the stored MCP schemas + for (const [name, def] of Object.entries(mcpDefs)) { + const zod = ensureZodSchema(def.inputSchema) + expect(typeof zod.safeParse).toBe('function') + + // Verify schema can parse valid input + if (name === 'engram__mem_search') { + const result = zod.safeParse({ query: 'test query' }) + expect(result.success).toBe(true) + } + } + }) + }) + + describe('checks: no Zod, no closures, no sockets in definitions', () => { + test('stored definitions are pure data — no functions or instances', async () => { + const result = await getMCPToolData({ + toolNames: ['test/safe'], + mcpServers: { test: { command: 'test', args: [] } as any }, + requestMcpToolData: async () => [ + { + name: 'safe', + description: 'Safe tool', + inputSchema: { type: 'object', properties: { x: { type: 'number' } } }, + }, + ], + }) + + const def = result['test__safe']! + + // No functions + for (const [key, value] of Object.entries(def)) { + if (key === 'inputSchema') continue // checked below + expect(typeof (value as any)).not.toBe('function') + } + + // No Zod + expect(typeof (def.inputSchema as any)?.safeParse).not.toBe('function') + + // No internal state references + expect(def.inputSchema).not.toHaveProperty('_def') + expect(def.inputSchema).not.toHaveProperty('_type') + expect(def.inputSchema).not.toHaveProperty('_cached') + }) + }) +}) diff --git a/packages/agent-runtime/src/__tests__/mcp-schema-serialization.test.ts b/packages/agent-runtime/src/__tests__/mcp-schema-serialization.test.ts new file mode 100644 index 0000000000..efb6739129 --- /dev/null +++ b/packages/agent-runtime/src/__tests__/mcp-schema-serialization.test.ts @@ -0,0 +1,329 @@ +/** + * Tests for MCP tool schema serialization fix. + * + * The core invariant: MCP tool inputSchemas stored in customToolDefinitions + * MUST be plain JSON Schema objects, not Zod schemas, so that + * JSON.stringify() never fails with "cannot serialize cyclic structures". + * + * Zod conversion happens at the consumption boundary via ensureZodSchema(). + */ +import { describe, test, expect } from 'bun:test' + +import { getMCPToolData } from '../mcp' +import { ensureZodSchema } from '../tools/prompts' + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +/** Create a logger that tracks warnings. */ +function makeLogger(): { + debug: () => void + info: () => void + warn: (ctx: unknown, msg: string) => void + error: () => void + _warnings: string[] +} { + const warnings: string[] = [] + return { + debug: () => {}, + info: () => {}, + warn: (_ctx: unknown, msg: string) => { + warnings.push(msg) + }, + error: () => {}, + _warnings: warnings, + } +} + +/** Minimal valid JSON Schema objects covering typical MCP tool shapes. */ +const jsonSchemas = { + empty: {} as Record, + booleanTrue: true as unknown as Record, + booleanFalse: false as unknown as Record, + simple: { + type: 'object', + properties: { name: { type: 'string' } }, + } as Record, + nested: { + type: 'object', + properties: { + items: { + type: 'array', + items: { type: 'object', properties: { id: { type: 'number' } } }, + }, + }, + } as Record, + withEnum: { + type: 'object', + properties: { color: { type: 'string', enum: ['red', 'green', 'blue'] } }, + } as Record, + withRequired: { + type: 'object', + properties: { name: { type: 'string' }, age: { type: 'number' } }, + required: ['name'], + } as Record, + withAdditional: { + type: 'object', + properties: { key: { type: 'string' } }, + additionalProperties: false, + } as Record, +} + +// A truly circular object (simulates what someone might accidentally inject) +function makeCircular(): Record { + const obj: Record = { name: 'circular' } + obj.self = obj + return obj +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +describe('MCP schema serialization', () => { + // --- Storage invariants (the actual fix) --- + + describe('stored schemas are JSON-serializable', () => { + const toolKinds = Object.entries(jsonSchemas) + + for (const [kind, inputSchema] of toolKinds) { + test(`stores "${kind}" schema as plain object, not Zod`, async () => { + const logger = makeLogger() + const result = await getMCPToolData({ + toolNames: ['test/' + kind], + mcpServers: { test: { command: 'fake', args: [] } as any }, + requestMcpToolData: async () => [ + { name: kind, description: `Schema kind: ${kind}`, inputSchema }, + ], + logger, + }) + + const key = 'test__' + kind + expect(result[key]).toBeDefined() + const def = result[key]! + + // Invariant: inputSchema must NOT be a Zod schema + expect(typeof (def.inputSchema as any)?.safeParse).not.toBe('function') + + // Invariant: JSON.stringify + JSON.parse roundtrip must succeed + expect(() => JSON.stringify(result)).not.toThrow() + const roundtripped = JSON.parse(JSON.stringify(result)) + expect(roundtripped[key]).toBeDefined() + }) + } + }) + + // --- Roundtrip for multiple tools / multiple servers --- + + test('multiple tools from multiple MCP servers serialize correctly', async () => { + const logger = makeLogger() + const result = await getMCPToolData({ + toolNames: ['srv1/alpha', 'srv1/beta', 'srv2/gamma'], + mcpServers: { + srv1: { command: 'srv1', args: [] } as any, + srv2: { command: 'srv2', args: [] } as any, + }, + requestMcpToolData: async ({ mcpConfig }) => { + const prefix = (mcpConfig as any).command as string + return [ + { + name: prefix === 'srv1' ? 'alpha' : 'gamma', + description: `Tool from ${prefix}`, + inputSchema: { type: 'object', properties: { x: { type: 'number' } } }, + }, + ...(prefix === 'srv1' + ? [ + { + name: 'beta', + description: 'Another tool', + inputSchema: { type: 'object', properties: {} }, + }, + ] + : []), + ] + }, + logger, + }) + + expect(Object.keys(result).length).toBe(3) + expect(() => JSON.stringify(result)).not.toThrow() + + // ensureZodSchema should convert every stored schema back to Zod + for (const [name, def] of Object.entries(result)) { + const zod = ensureZodSchema(def.inputSchema) + expect(typeof zod.safeParse).toBe('function') + const parsed = zod.safeParse({}) + expect(parsed.success).toBe(true) + } + }) + + // --- Circular rejection (defense-in-depth) --- + + test('truly circular object still fails JSON.stringify (defense-in-depth)', async () => { + // If someone bypasses getMCPToolData and injects a circular object + // directly into customToolDefinitions, JSON.stringify should still + // fail — we must NOT silently swallow circular references. + const circular = makeCircular() + expect(() => JSON.stringify(circular)).toThrow() + + // But our path should never produce circular objects + const logger = makeLogger() + const result = await getMCPToolData({ + toolNames: ['test/safe'], + mcpServers: { test: { command: 'safe', args: [] } as any }, + requestMcpToolData: async () => [ + { + name: 'safe', + description: 'Safe tool', + inputSchema: { type: 'object', properties: {} }, + }, + ], + logger, + }) + + expect(() => JSON.stringify(result)).not.toThrow() + }) + + // --- Error with cause --- + + test('MCP server failure is caught and logged, not thrown', async () => { + const logger = makeLogger() + const cause = new Error('connection refused') + + const result = await getMCPToolData({ + toolNames: ['bad/missing'], + mcpServers: { bad: { command: 'bad', args: [] } as any }, + requestMcpToolData: async () => { + const err = new Error('MCP server failed') + err.cause = cause + throw err + }, + logger, + }) + + // Should return empty (no tools from failed server) + expect(Object.keys(result).length).toBe(0) + // Should have logged a warning + expect(logger._warnings.length).toBeGreaterThan(0) + expect(logger._warnings[0]).toContain('bad') + }) + + // --- JSON.stringify of the final agent payload --- + + test('final agent payload (customToolDefinitions) survives JSON.stringify', () => { + // Simulate what the agent runtime stores in fileContext.customToolDefinitions + // after getMCPToolData returns. This is the payload that eventually gets + // serialized via saveChatState → JSON.stringify(runState). + + const customToolDefinitions: Record; description?: string; endsAgentStep?: boolean }> = { + 'engram__mem_save': { + inputSchema: { + type: 'object', + properties: { + title: { type: 'string' }, + content: { type: 'string' }, + type: { type: 'string' }, + scope: { type: 'string' }, + }, + required: ['title'], + }, + description: 'Save a memory', + endsAgentStep: true, + }, + 'engram__mem_search': { + inputSchema: { + type: 'object', + properties: { + query: { type: 'string' }, + limit: { type: 'number' }, + }, + }, + description: 'Search memory', + endsAgentStep: true, + }, + 'engram__mem_context': { + inputSchema: { + type: 'object', + properties: { + project: { type: 'string' }, + scope: { type: 'string' }, + }, + }, + description: 'Get memory context', + endsAgentStep: true, + }, + 'codegraph__codegraph_explore': { + inputSchema: { + type: 'object', + properties: { + query: { type: 'string' }, + maxFiles: { type: 'number' }, + projectPath: { type: 'string' }, + }, + required: ['query'], + }, + description: 'Explore code', + endsAgentStep: true, + }, + 'builtin__read': { + inputSchema: { + type: 'object', + properties: { + filePath: { type: 'string' }, + offset: { type: 'number' }, + limit: { type: 'number' }, + }, + required: ['filePath'], + }, + description: 'Read a file', + endsAgentStep: false, + }, + } + + // THE DECISIVE TEST: JSON.stringify must not throw + expect(() => JSON.stringify(customToolDefinitions)).not.toThrow() + + const json = JSON.stringify(customToolDefinitions) + const parsed = JSON.parse(json) + + // Verify all tools survived the roundtrip + expect(Object.keys(parsed).length).toBe(5) + expect(parsed['engram__mem_save'].inputSchema.required).toEqual(['title']) + expect(parsed['engram__mem_search'].inputSchema.properties.query.type).toBe('string') + expect(parsed['codegraph__codegraph_explore'].inputSchema.required).toEqual(['query']) + expect(parsed['builtin__read'].inputSchema.properties.filePath.type).toBe('string') + + // ensureZodSchema should work on all of them + for (const [name, def] of Object.entries(customToolDefinitions)) { + const zod = ensureZodSchema(def.inputSchema) + expect(typeof zod.safeParse).toBe('function') + } + }) + + // --- Tool call result doesn't leak internals --- + + test('tool call result does not contain Zod or circular refs', () => { + // After a tool executes, the result is stored in the message history. + // This must also be JSON-serializable. + + const toolCall = { + toolCallId: 'call-1', + toolName: 'engram__mem_search', + input: { query: 'PR #876', limit: 5 }, + } + + const toolResult = { + role: 'tool' as const, + toolCallId: 'call-1', + content: JSON.stringify([ + { id: 1, title: 'Memory about PR #876' }, + { id: 2, title: 'Another memory' }, + ]), + } + + // Both must be serializable + expect(() => JSON.stringify(toolCall)).not.toThrow() + expect(() => JSON.stringify(toolResult)).not.toThrow() + }) +}) diff --git a/packages/agent-runtime/src/mcp.ts b/packages/agent-runtime/src/mcp.ts index a7390f219c..a2845e11ea 100644 --- a/packages/agent-runtime/src/mcp.ts +++ b/packages/agent-runtime/src/mcp.ts @@ -1,5 +1,4 @@ import { getErrorObject } from '@codebuff/common/util/error' -import { convertJsonSchemaToZod } from 'zod-from-json-schema' import { MCP_TOOL_SEPARATOR } from './mcp-constants' @@ -56,7 +55,7 @@ export async function getMCPToolData( for (const { name, description, inputSchema } of mcpData) { writeTo[mcpName + MCP_TOOL_SEPARATOR + name] = { - inputSchema: convertJsonSchemaToZod(inputSchema as any) as any, + inputSchema: inputSchema as Record, endsAgentStep: true, description, }