Relax MCP tool output validation
This commit is contained in:
@@ -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) =>
|
||||
|
||||
@@ -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(),
|
||||
|
||||
Reference in New Issue
Block a user