Relax MCP tool output validation

This commit is contained in:
Mohamed Boudra
2026-05-26 23:22:09 +07:00
parent 09bf981f13
commit a5b82d2a3b
4 changed files with 117 additions and 8 deletions
@@ -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<Record<string, unknown>> {
@@ -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) =>
+33 -8
View File
@@ -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<TConfig extends { outputSchema?: unknown }>(
config: TConfig,
): TConfig {
if (config.outputSchema === undefined) {
return config;
}
return {
...config,
outputSchema: relaxMcpOutputSchema(config.outputSchema),
} as TConfig;
}
type McpToolContext = RequestHandlerExtra<ServerRequest, ServerNotification>;
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(),
},
@@ -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"],
@@ -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(),