diff --git a/services/agent-harness/src/tools/mcp.test.ts b/services/agent-harness/src/tools/mcp.test.ts new file mode 100644 index 0000000000..5c41eaf0df --- /dev/null +++ b/services/agent-harness/src/tools/mcp.test.ts @@ -0,0 +1,198 @@ +import { expect, it } from 'vitest'; +import type { Tool } from '@modelcontextprotocol/sdk/types.js'; +import { RunLimitsSchema } from '../commands'; +import { executeMcp } from './mcp'; +import { mcpTestFixture } from './mcp-test-fixture'; +import { toolDefinitions } from '@kilocode/agent-harness/tools'; +import { ToolCallSchema } from '@kilocode/agent-harness/contracts'; +import { evaluateDispatch } from '@kilocode/agent-harness/policy'; + +it('requires Ask approval despite a discovered read-only hint', async () => { + const request = await fixture().call(); + const definition = toolDefinitions.find(tool => tool.name === request.name); + const call = ToolCallSchema.parse({ + id: crypto.randomUUID(), + runId: crypto.randomUUID(), + name: request.name, + definitionVersion: definition?.version, + arguments: request.arguments, + context: { type: 'personal' }, + effect: definition?.effect, + executionTarget: { kind: 'backend' }, + approval: null, + state: 'pending', + result: null, + }); + expect( + evaluateDispatch(call, call, { + permissionMode: 'ask', + permissionRevision: 0, + expectedPermissionRevision: 0, + authorized: true, + available: true, + clientReady: true, + questionAnswered: true, + trustedRead: true, + }) + ).toBe('approval'); +}); + +function fixture() { + const { key, state, connection, fetchImpl } = mcpTestFixture(); + const run = (request: unknown = { name: 'mcp.discover', arguments: {} }, limits = {}) => + executeMcp( + request, + [connection], + { deadline: Date.now() + 1000, limits: RunLimitsSchema.parse(limits) }, + fetchImpl + ); + const call = async () => { + const result = await run(); + expect(result.status).toBe('succeeded'); + expect(JSON.stringify(result)).not.toContain('derived-secret'); + const [found] = (result as { output: any[] }).output; + return { + name: 'mcp.call', + arguments: { + serverId: found.serverId, + configurationVersion: found.configurationVersion, + definitionVersion: found.definitionVersion, + name: found.name, + arguments: { [key]: 2 }, + }, + }; + }; + return { key, state, connection, run, call }; +} + +it('discovers runtime-only schemas and returns validated business data without authorization', async () => { + const f = fixture(); + const call = await f.call(); + const result = await f.run(call); + expect(result).toEqual({ status: 'succeeded', output: f.state.result }); + expect(f.state.effects).toEqual([call.arguments.arguments]); + expect(JSON.stringify(result)).not.toContain('derived-secret'); +}); +it.each(['missing', 'wrong_type', 'minimum', 'extra'])( + 'rejects malformed arguments before an effect: %s', + async mode => { + const f = fixture(), + call = await f.call(); + const key = Object.keys(call.arguments.arguments)[0]; + const args = + mode === 'missing' + ? {} + : mode === 'extra' + ? { ...call.arguments.arguments, extra: true } + : { [key]: mode === 'minimum' ? 0 : 'two' }; + expect( + await f.run({ ...call, arguments: { ...call.arguments, arguments: args } }) + ).toMatchObject({ status: 'failed', error: { message: 'invalid_input' } }); + expect(f.state.effects).toEqual([]); + } +); +it.each(['inputSchema', 'outputSchema'] as const)( + 'rejects unsupported references and malformed %s before discovery or effects', + async field => { + for (const extra of [ + { $ref: 'https://evil.example/schema' }, + { $ref: '#/properties/value' }, + { required: 'wrong' }, + { required: ['value', 'value'] }, + { unevaluatedProperties: false }, + { properties: { value: { type: 'unsupported' } } }, + ]) { + const f = fixture(), + call = await f.call(); + f.state.tool[field] = { type: 'object', ...extra } as Tool['inputSchema']; + for (const request of [undefined, call]) + expect(await f.run(request), JSON.stringify(extra)).toMatchObject({ + status: 'failed', + error: { message: 'invalid_schema', retryable: false }, + }); + expect(f.state.effects).toEqual([]); + } + } +); +it.each([ + [401, 'reauthorization_required', true], + [403, 'reauthorization_required', true], + [503, 'unavailable_server', true], + [302, 'unsafe_destination', false], + [400, 'unsafe_destination', false], +] as const)('keeps HTTP failure %i distinct and sanitized', async (status, reason, retryable) => { + const f = fixture(); + f.state.status = status; + expect(await f.run()).toMatchObject({ status: 'failed', error: { message: reason, retryable } }); + expect(JSON.stringify(await f.run())).not.toContain('provider-secret'); +}); +it.each(['description', 'inputSchema', 'outputSchema'] as const)( + 'requires a new call after %s or configuration changes', + async field => { + const f = fixture(), + call = await f.call(); + expect((await f.call()).arguments).toEqual(call.arguments); + f.state.tool = { + ...f.state.tool, + [field]: + field === 'description' + ? 'Changed operation' + : { ...f.state.tool[field], additionalProperties: true }, + }; + expect(await f.run(call)).toMatchObject({ + status: 'failed', + error: { message: 'definition_changed' }, + }); + f.connection.configurationVersion = '2'; + expect(await f.run(call)).toMatchObject({ + status: 'failed', + error: { message: 'definition_changed' }, + }); + expect(f.state.effects).toEqual([]); + expect(await f.run(await f.call())).toMatchObject({ status: 'succeeded' }); + } +); +it.each([ + ['invalid', 'invalid_output'], + ['wrong_type', 'invalid_output'], + ['malformed', 'invalid_output'], + ['missing', 'invalid_output'], + ['error', 'invalid_output'], + ['lose', 'unavailable_server'], + ['overflow', 'limit_exceeded'], + ['stall', 'unavailable_server'], + ['businessLimit', 'limit_exceeded'], +])('retains unknown outcomes without mutation replay: %s', async (mode, reason) => { + const f = fixture(), + call = await f.call(); + if (mode === 'invalid') f.state.result = { content: [], structuredContent: { wrong: true } }; + if (mode === 'wrong_type') + f.state.result = { content: [], structuredContent: { [f.key]: 'two' } }; + if (mode === 'malformed') f.state.result = { content: 'provider-secret' }; + if (mode === 'missing') f.state.result = { content: [] }; + if (mode === 'error') f.state.result = { content: [], isError: true }; + if (mode === 'businessLimit') + f.state.result = { + ...(f.state.result as object), + content: [{ type: 'text', text: 'é'.repeat(1600) }], + }; + f.state.lose = mode === 'lose'; + f.state.overflow = mode === 'overflow'; + f.state.stall = mode === 'stall'; + const result = await f.run(call, { + httpResponseBytes: mode === 'businessLimit' ? 8192 : 2048, + toolOutputBytes: 2048, + toolAttemptMs: 100, + }); + expect(result).toEqual({ status: 'outcome_unknown', reason }); + expect(f.state.effects).toEqual([call.arguments.arguments]); + expect(JSON.stringify(result)).not.toContain('provider-secret'); +}); +it('returns no invented definitions without configured connections', async () => { + expect( + await executeMcp({ name: 'mcp.discover', arguments: {} }, [], { + deadline: Date.now() + 1000, + limits: RunLimitsSchema.parse({}), + }) + ).toEqual({ status: 'succeeded', output: [] }); +}); diff --git a/services/agent-harness/src/tools/mcp.ts b/services/agent-harness/src/tools/mcp.ts new file mode 100644 index 0000000000..aef5025367 --- /dev/null +++ b/services/agent-harness/src/tools/mcp.ts @@ -0,0 +1,199 @@ +import { z } from 'zod'; +import { ErrorCode, McpError } from '@modelcontextprotocol/sdk/types.js'; +import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import { CfWorkerJsonSchemaValidator } from '@modelcontextprotocol/sdk/validation/cfworker-provider.js'; +import type { JsonSchemaType } from '@modelcontextprotocol/sdk/validation/types.js'; +import { canonicalizeValidatedInput } from '@kilocode/agent-harness/commands'; +import type { ToolOutcome } from '@kilocode/agent-harness/contracts'; +import { toolDefinitions, ToolRequestSchema } from '@kilocode/agent-harness/tools'; +import { bytes, type RunLimits } from '../limits'; +import { createMcpTransportFactory, type McpConnection } from './mcp-transport'; + +export type { McpConnection } from './mcp-transport'; + +// Only this checked vocabulary is supported. References and unknown keywords fail closed. +const unique = (values: readonly unknown[]) => + new Set(values.map(canonicalizeValidatedInput)).size === values.length; +const kind = z.enum(['object', 'array', 'string', 'number', 'integer', 'boolean', 'null']); +const schema: z.ZodType = z.lazy(() => + z.strictObject({ + $schema: z.literal('https://json-schema.org/draft/2020-12/schema').optional(), + type: z.union([kind, z.array(kind).min(1).refine(unique)]).optional(), + properties: z.record(z.string(), schema).optional(), + required: z.array(z.string()).refine(unique).optional(), + additionalProperties: z.union([z.boolean(), schema]).optional(), + items: schema.optional(), + enum: z.array(z.json()).min(1).refine(unique).optional(), + const: z.json().optional(), + minimum: z.number().optional(), + maximum: z.number().optional(), + minLength: z.int().nonnegative().optional(), + maxLength: z.int().nonnegative().optional(), + minItems: z.int().nonnegative().optional(), + maxItems: z.int().nonnegative().optional(), + title: z.string().optional(), + description: z.string().optional(), + default: z.json().optional(), + }) +); +const codes = { + unavailable_server: 'unavailable_tool', + reauthorization_required: 'reauthorization_required', + invalid_schema: 'invalid_output', + unsafe_destination: 'invalid_input', + definition_changed: 'invalid_input', + invalid_input: 'invalid_input', + invalid_output: 'invalid_output', + limit_exceeded: 'limit_exceeded', +} as const; + +/** Connections contain ephemeral gateway authorization, never provider credentials or stored state. */ +export async function executeMcp( + input: unknown, + connections: readonly McpConnection[], + budget: { deadline: number; limits: RunLimits }, + fetchImpl: typeof fetch = fetch +): Promise { + let reason: keyof typeof codes | undefined; + let dispatched = false; + const abort = new AbortController(); + const stop = (value: keyof typeof codes): never => { + // Closing the transport on abort must not replace the original failure. + reason ??= value; + const error = new Error(reason); + abort.abort(error); + throw error; + }; + const timeout = Math.min(budget.limits.toolAttemptMs, budget.deadline - Date.now()); + const timer = setTimeout(() => abort.abort(), Math.max(0, timeout)); + const options = { + signal: abort.signal, + timeout: Math.max(1, timeout), + resetTimeoutOnProgress: false, + }; + const transportFor = createMcpTransportFactory( + abort.signal, + budget.limits.httpResponseBytes, + stop, + fetchImpl + ); + const validator = new CfWorkerJsonSchemaValidator(); + const checked = { + getValidator(value: JsonSchemaType) { + try { + if (!schema.safeParse(value).success) stop('invalid_schema'); + return validator.getValidator(value); + } catch { + return stop('invalid_schema'); + } + }, + }; + try { + if (timeout <= 0) stop('limit_exceeded'); + const parsed = ToolRequestSchema.safeParse(input); + if (!parsed.success) stop('invalid_input'); + const request = parsed.data; + if (request.name !== 'mcp.discover' && request.name !== 'mcp.call') stop('invalid_input'); + const contract = toolDefinitions + .filter(tool => tool.group === 'mcp') + .find(tool => tool.name === request.name); + if (!contract) stop('invalid_input'); + if (bytes(request.arguments) > budget.limits.toolInputBytes) stop('limit_exceeded'); + const definitions = []; + for (const connection of connections) { + if (request.name === 'mcp.call' && request.arguments.serverId !== connection.serverId) + continue; + if ( + request.name === 'mcp.call' && + request.arguments.configurationVersion !== connection.configurationVersion + ) + stop('definition_changed'); + const transport = transportFor(connection); + const client = new Client( + { name: 'kilo-agent-harness', version: '1' }, + { jsonSchemaValidator: checked } + ); + try { + await client.connect(transport, options); + let cursor: string | undefined; + const seen = new Set(); + do { + const page = await client.listTools(cursor ? { cursor } : {}, options); + for (const tool of page.tools) { + if (!tool.outputSchema || tool.execution?.taskSupport === 'required') + stop('invalid_schema'); + checked.getValidator(tool.inputSchema); + checked.getValidator(tool.outputSchema); + const digest = await crypto.subtle.digest( + 'SHA-256', + new TextEncoder().encode(canonicalizeValidatedInput(tool)) + ); + const definition = { + serverId: connection.serverId, + configurationVersion: connection.configurationVersion, + name: tool.name, + definitionVersion: Array.from(new Uint8Array(digest), byte => + byte.toString(16).padStart(2, '0') + ).join(''), + inputSchema: tool.inputSchema, + outputSchema: tool.outputSchema, + }; + definitions.push(definition); + if (bytes(definitions) > budget.limits.toolOutputBytes) stop('limit_exceeded'); + if (request.name !== 'mcp.call' || request.arguments.name !== tool.name) continue; + if (request.arguments.definitionVersion !== definition.definitionVersion) + stop('definition_changed'); + if (!checked.getValidator(tool.inputSchema)(request.arguments.arguments).valid) + stop('invalid_input'); + dispatched = true; + const result = await client.callTool( + { name: tool.name, arguments: request.arguments.arguments }, + undefined, + options + ); + if ( + result.isError || + !checked.getValidator(tool.outputSchema)(result.structuredContent).valid + ) + stop('invalid_output'); + const output = { content: result.content, structuredContent: result.structuredContent }; + if (bytes(output) > budget.limits.toolOutputBytes) stop('limit_exceeded'); + return { status: 'succeeded', output: contract.outputSchema.parse(output) }; + } + cursor = page.nextCursor; + if (cursor && seen.has(cursor)) stop('invalid_schema'); + if (cursor) seen.add(cursor); + } while (cursor); + } finally { + await client.close(); + } + } + if (request.name === 'mcp.call') stop('unavailable_server'); + return { status: 'succeeded', output: contract.outputSchema.parse(definitions) }; + } catch (error) { + if ( + (error instanceof McpError && + [ErrorCode.InvalidParams, ErrorCode.InvalidRequest, ErrorCode.ParseError].includes( + error.code + )) || + // SDK response parsing uses Zod Mini; the core error also covers Classic. + error instanceof z.core.$ZodError || + error instanceof SyntaxError + ) + reason ??= dispatched ? 'invalid_output' : 'invalid_schema'; + reason ??= 'unavailable_server'; + // No provider reconciliation contract exists here. A dispatched unknown effect is never replayed. + return dispatched + ? { status: 'outcome_unknown', reason } + : { + status: 'failed', + error: { + code: codes[reason], + message: reason, + retryable: reason === 'unavailable_server' || reason === 'reauthorization_required', + }, + }; + } finally { + clearTimeout(timer); + } +}