From a5b82d2a3b0eda47207f542901c8b2b396e67dbc Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Tue, 26 May 2026 23:22:09 +0700 Subject: [PATCH] Relax MCP tool output validation --- .../src/server/agent/mcp-server.test.ts | 82 +++++++++++++++++++ .../server/src/server/agent/mcp-server.ts | 41 ++++++++-- .../server/src/server/schedule/service.ts | 1 + packages/server/src/server/schedule/types.ts | 1 + 4 files changed, 117 insertions(+), 8 deletions(-) diff --git a/packages/server/src/server/agent/mcp-server.test.ts b/packages/server/src/server/agent/mcp-server.test.ts index 6a7beb68b..20342f297 100644 --- a/packages/server/src/server/agent/mcp-server.test.ts +++ b/packages/server/src/server/agent/mcp-server.test.ts @@ -4,7 +4,9 @@ import { realpathSync } from "node:fs"; import { mkdtemp, mkdir, rm, writeFile } from "node:fs/promises"; import { join, resolve as resolvePath } from "node:path"; import { tmpdir } from "node:os"; +import Ajv from "ajv"; import { z } from "zod"; +import { zodToJsonSchema } from "zod-to-json-schema"; import { createTestLogger } from "../../test-utils/test-logger.js"; import { createAgentMcpServer } from "./mcp-server.js"; @@ -51,6 +53,7 @@ interface LooseStructuredContent { interface RegisteredMcpTool { inputSchema: LooseInputSchema; + outputSchema?: unknown; callback?: ( input: unknown, extra?: unknown, @@ -103,6 +106,14 @@ async function invokeToolWithParsedInput( return tool.handler(parsed.data); } +function expectOutputSchemaAccepts(tool: RegisteredMcpTool, data: unknown): void { + expect(tool.outputSchema).toBeDefined(); + const jsonSchema = zodToJsonSchema(tool.outputSchema as z.ZodTypeAny); + const ajv = new Ajv({ allErrors: true, strict: false }); + const validate = ajv.compile(jsonSchema); + expect(validate(data), JSON.stringify(validate.errors, null, 2)).toBe(true); +} + function agentsOf(response: { structuredContent: LooseStructuredContent; }): Array> { @@ -545,6 +556,38 @@ describe("create_agent MCP tool", () => { ); }); + it("advertises create_agent output schema that accepts full provider modes", async () => { + const { agentManager, agentStorage, spies } = createTestDeps(); + spies.agentManager.createAgent.mockResolvedValue({ + id: "mode-agent", + cwd: REPO_CWD, + lifecycle: "idle", + currentModeId: "build", + availableModes: [ + { + id: "build", + label: "Build", + description: null, + icon: "hammer", + colorTier: "dangerous", + }, + ], + config: { title: "Mode test" }, + } as ManagedAgent); + + const server = await createAgentMcpServer({ agentManager, agentStorage, logger }); + const tool = registeredTool(server, "create_agent"); + const response = await tool.handler({ + cwd: existingCwd, + title: "Mode test", + provider: "codex/gpt-5.4", + initialPrompt: "Do work", + background: true, + }); + + expectOutputSchemaAccepts(tool, response.structuredContent); + }); + it("requires provider as provider/model and rejects the old model field", async () => { const { agentManager, agentStorage } = createTestDeps(); const server = await createAgentMcpServer({ agentManager, agentStorage, logger }); @@ -1762,6 +1805,45 @@ describe("create_schedule MCP tool", () => { ); }); + it("advertises create_schedule output schema that accepts inherited feature values", async () => { + const { agentManager, agentStorage, spies } = createTestDeps(); + spies.agentManager.getAgent.mockReturnValue({ + id: "parent-agent", + provider: "opencode", + cwd: REPO_CWD, + lifecycle: "idle", + currentModeId: "build", + availableModes: [], + config: { + title: "Parent agent", + model: "openai/gpt-5.5", + featureValues: { auto_accept: true }, + }, + } as ManagedAgent); + const create = vi.fn(async (input: CreateScheduleInput) => createStoredSchedule(input)); + const server = await createAgentMcpServer({ + agentManager, + agentStorage, + scheduleService: { create } as unknown as ScheduleService, + callerAgentId: "parent-agent", + logger, + }); + const tool = registeredTool(server, "create_schedule"); + + const response = await tool.handler({ + prompt: "say hello", + every: "5m", + target: "new-agent", + provider: "opencode/openai/gpt-5.5", + }); + + expect(response.structuredContent.target).toMatchObject({ + type: "new-agent", + config: { featureValues: { auto_accept: true } }, + }); + expectOutputSchemaAccepts(tool, response.structuredContent); + }); + it("accepts a blank cron field when every is provided", async () => { const { agentManager, agentStorage } = createTestDeps(); const create = vi.fn(async (scheduleInput: CreateScheduleInput) => diff --git a/packages/server/src/server/agent/mcp-server.ts b/packages/server/src/server/agent/mcp-server.ts index 59c037b0c..a8f334c19 100644 --- a/packages/server/src/server/agent/mcp-server.ts +++ b/packages/server/src/server/agent/mcp-server.ts @@ -246,6 +246,37 @@ function formatStructuredContentForModel(structuredContent: unknown): string { return summary.length > 0 ? `${summary.join("\n")}\n\n${json}` : json; } +function isZodSchema(value: unknown): value is z.ZodTypeAny { + return ( + typeof value === "object" && value !== null && "_def" in value && "safeParseAsync" in value + ); +} + +function relaxMcpOutputSchema(outputSchema: unknown): unknown { + if (!outputSchema) { + return outputSchema; + } + + if (isZodSchema(outputSchema)) { + return outputSchema instanceof z.ZodObject ? outputSchema.passthrough() : outputSchema; + } + + return z.object(outputSchema as z.ZodRawShape).passthrough(); +} + +function relaxMcpToolOutputSchema( + config: TConfig, +): TConfig { + if (config.outputSchema === undefined) { + return config; + } + + return { + ...config, + outputSchema: relaxMcpOutputSchema(config.outputSchema), + } as TConfig; +} + type McpToolContext = RequestHandlerExtra; function parseTimestamp(value: string | null | undefined): number { @@ -564,7 +595,7 @@ export async function createAgentMcpServer(options: AgentMcpServerOptions): Prom }); const registerRawTool = server.registerTool.bind(server); const registerTool: McpServer["registerTool"] = (name, config, handler) => - registerRawTool(name, config, (async (args: never, extra: never) => + registerRawTool(name, relaxMcpToolOutputSchema(config), (async (args: never, extra: never) => addModelVisibleStructuredContent(await handler(args, extra))) as typeof handler); const resolveCallerAgent = () => { @@ -1122,13 +1153,7 @@ export async function createAgentMcpServer(options: AgentMcpServerOptions): Prom status: AgentStatusEnum, cwd: z.string(), currentModeId: z.string().nullable(), - availableModes: z.array( - z.object({ - id: z.string(), - label: z.string(), - description: z.string().nullable().optional(), - }), - ), + availableModes: z.array(ProviderModeSchema), lastMessage: z.string().nullable().optional(), permission: AgentPermissionRequestPayloadSchema.nullable().optional(), }, diff --git a/packages/server/src/server/schedule/service.ts b/packages/server/src/server/schedule/service.ts index b7449ebe2..39b647bdc 100644 --- a/packages/server/src/server/schedule/service.ts +++ b/packages/server/src/server/schedule/service.ts @@ -536,6 +536,7 @@ export class ScheduleService { sandboxMode: schedule.target.config.sandboxMode, networkAccess: schedule.target.config.networkAccess, webSearch: schedule.target.config.webSearch, + featureValues: schedule.target.config.featureValues, extra: schedule.target.config.extra, systemPrompt: schedule.target.config.systemPrompt, mcpServers: schedule.target.config.mcpServers as AgentSessionConfig["mcpServers"], diff --git a/packages/server/src/server/schedule/types.ts b/packages/server/src/server/schedule/types.ts index dd24e175b..1178f31f1 100644 --- a/packages/server/src/server/schedule/types.ts +++ b/packages/server/src/server/schedule/types.ts @@ -34,6 +34,7 @@ export const ScheduleTargetSchema = z.discriminatedUnion("type", [ sandboxMode: z.string().trim().min(1).optional(), networkAccess: z.boolean().optional(), webSearch: z.boolean().optional(), + featureValues: z.record(z.unknown()).optional(), extra: z .object({ codex: z.record(z.unknown()).optional(),