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, }