diff --git a/packages/mcp-utils/src/index.ts b/packages/mcp-utils/src/index.ts index a9fc30b1..fb5b00d2 100644 --- a/packages/mcp-utils/src/index.ts +++ b/packages/mcp-utils/src/index.ts @@ -1,3 +1,4 @@ export * from './server.js'; export * from './stream-transport.js'; +export * from './tool-policy.js'; export * from './types.js'; diff --git a/packages/mcp-utils/src/server.test.ts b/packages/mcp-utils/src/server.test.ts index 6a89b16d..1db3e3c2 100644 --- a/packages/mcp-utils/src/server.test.ts +++ b/packages/mcp-utils/src/server.test.ts @@ -1,16 +1,23 @@ -import { Client } from '@modelcontextprotocol/client'; +import { + Client, + StreamableHTTPClientTransport, +} from '@modelcontextprotocol/client'; import type { CallToolRequestParams } from '@modelcontextprotocol/client'; -import type { Server } from '@modelcontextprotocol/server'; -import { describe, expect, test, vi } from 'vitest'; +import { createMcpHandler, inputRequired } from '@modelcontextprotocol/server'; +import type { Server, ServerOptions } from '@modelcontextprotocol/server'; +import { afterEach, describe, expect, test, vi } from 'vitest'; import { z } from 'zod/v4'; import { createMcpServer, + type McpServerOptions, resource, resources, resourceTemplate, tool, -} from './server.js'; + type Tool, + type ToolPolicy, +} from './index.js'; import { StreamTransport } from './stream-transport.js'; export const MCP_CLIENT_NAME = 'test-client'; @@ -78,6 +85,71 @@ async function setup(options: SetupOptions) { return { client, clientTransport, callTool, server, serverTransport }; } +// https://blog.modelcontextprotocol.io/posts/2026-07-28-release-candidate/ +const MODERN_PROTOCOL_VERSION = '2026-07-28'; +const MCP_ENDPOINT = new URL('https://mcp.test'); +const telemetry = { + policyId: 'test-policy', + policyVersion: 1, + outcome: 'allowed', +}; +const cleanups: Array<() => Promise> = []; + +/** + * A policy that always proceeds. Attaching any policy is what opts a tool + * into structured results, so this is the minimal opted-in configuration. + */ +const passThroughPolicy: ToolPolicy<{ value: string }, undefined> = { + resolve: async () => ({ + type: 'execute', + resolution: undefined, + telemetry, + }), +}; + +afterEach(async () => { + for (const cleanup of cleanups.splice(0).reverse()) { + await cleanup(); + } +}); + +/** + * Connects a client on the modern revision, the serving path that carries the + * per-request `_meta` envelope and the multi-round-trip `requestState` + * vocabulary. + */ +async function setupModernClient( + options: Omit +) { + const handler = createMcpHandler( + () => + createMcpServer({ + name: 'test-server', + version: '0.0.0', + ...options, + }), + { legacy: 'reject' } + ); + const transport = new StreamableHTTPClientTransport(MCP_ENDPOINT, { + fetch: (url, init) => handler.fetch(new Request(url, init)), + }); + const client = new Client( + { name: MCP_CLIENT_NAME, version: MCP_CLIENT_VERSION }, + { + capabilities: {}, + versionNegotiation: { mode: { pin: MODERN_PROTOCOL_VERSION } }, + } + ); + + await client.connect(transport); + cleanups.push( + () => client.close(), + () => handler.close() + ); + + return client; +} + describe('tools', () => { test('parameter set to default value when omitted by caller', async () => { const server = createMcpServer({ @@ -298,6 +370,401 @@ describe('tools', () => { }); }); +describe('structured tool results', () => { + // Expressible exactly as it was before this package grew structured + // results: no policy, no formatResult. + const plainTools = () => ({ + fixture: tool({ + description: 'Fixture', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + execute: async ({ value }: { value: string }) => ({ value }), + }), + }); + + const policyTools = ( + formatResult?: (result: { value: string }) => string + ) => ({ + fixture: tool({ + description: 'Fixture', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: passThroughPolicy, + execute: async ({ value }: { value: string }) => ({ value }), + ...(formatResult ? { formatResult } : {}), + }), + }); + + // Both strings below were measured by running this exact fixture against + // base fc54ea2's `server.ts`, not written by hand. + const BASE_LEGACY_DISCOVERY = + '{"tools":[{"name":"fixture","description":"Fixture","inputSchema":' + + '{"type":"object","properties":{"value":{"type":"string"}},' + + '"required":["value"],"$schema":"http://json-schema.org/draft-07/schema#",' + + '"additionalProperties":false}}]}'; + const BASE_LEGACY_CALL = + '{"content":[{"type":"text","text":"{\\"value\\":\\"hi\\"}"}]}'; + const BASE_MODERN_DISCOVERY = + '{"_meta":{"io.modelcontextprotocol/serverInfo":{"name":"test-server",' + + '"version":"0.0.0"}},"ttlMs":0,"cacheScope":"private","tools":[' + + '{"name":"fixture","description":"Fixture","inputSchema":' + + '{"$schema":"http://json-schema.org/draft-07/schema#","type":"object",' + + '"properties":{"value":{"type":"string"}},"required":["value"],' + + '"additionalProperties":false}}]}'; + const BASE_MODERN_CALL = + '{"_meta":{"io.modelcontextprotocol/serverInfo":{"name":"test-server",' + + '"version":"0.0.0"}},"content":[{"type":"text",' + + '"text":"{\\"value\\":\\"hi\\"}"}]}'; + + /** + * A policy attached on every context that normalizes only the modern era, + * suppressing structured results elsewhere. This is the shape PR C needs: + * one policy, routed inside the policy. + */ + const suppressingTools = () => ({ + fixture: tool({ + description: 'Fixture', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + outputSchema: (schema, ctx) => + ctx.era === 'modern' ? schema : undefined, + resolve: passThroughPolicy.resolve, + } satisfies ToolPolicy<{ value: string }, undefined>, + execute: async ({ value }: { value: string }) => ({ value }), + formatResult: ({ value }: { value: string }) => `PREFIX:${value}`, + }), + }); + + test('a policy-free tool keeps base bytes on the legacy path', async () => { + const server = createMcpServer({ + name: 'test-server', + version: '0.0.0', + tools: plainTools(), + }); + + const { client } = await setup({ server }); + + expect(JSON.stringify(await client.listTools())).toBe( + BASE_LEGACY_DISCOVERY + ); + expect( + JSON.stringify( + await client.callTool({ name: 'fixture', arguments: { value: 'hi' } }) + ) + ).toBe(BASE_LEGACY_CALL); + }); + + test('a policy-free tool keeps base bytes on the modern path', async () => { + const client = await setupModernClient({ tools: plainTools() }); + + expect(JSON.stringify(await client.listTools())).toBe( + BASE_MODERN_DISCOVERY + ); + expect( + JSON.stringify( + await client.callTool({ name: 'fixture', arguments: { value: 'hi' } }) + ) + ).toBe(BASE_MODERN_CALL); + }); + + test('a policy advertises outputSchema and returns structuredContent with single-encoded text', async () => { + const server = createMcpServer({ + name: 'test-server', + version: '0.0.0', + tools: policyTools(), + }); + + const { client } = await setup({ server }); + + const { tools } = await client.listTools(); + expect(tools[0]?.outputSchema).toMatchObject({ + type: 'object', + properties: { value: { type: 'string' } }, + }); + + const result = await client.callTool({ + name: 'fixture', + arguments: { value: 'hi' }, + }); + + expect(result.structuredContent).toEqual({ value: 'hi' }); + expect(result.content).toEqual([ + { type: 'text', text: JSON.stringify({ value: 'hi' }) }, + ]); + }); + + test('a policy without an output-schema hook normalizes on both paths', async () => { + const legacyServer = createMcpServer({ + name: 'test-server', + version: '0.0.0', + tools: policyTools(), + }); + + const legacy = await setup({ server: legacyServer }); + const modern = await setupModernClient({ tools: policyTools() }); + + for (const client of [legacy.client, modern]) { + const { tools } = await client.listTools(); + expect(tools[0]?.outputSchema).toMatchObject({ type: 'object' }); + + const result = await client.callTool({ + name: 'fixture', + arguments: { value: 'hi' }, + }); + + expect(result.structuredContent).toEqual({ value: 'hi' }); + } + }); + + test('an output-schema hook returning undefined restores the whole base result', async () => { + const legacyServer = createMcpServer({ + name: 'test-server', + version: '0.0.0', + tools: suppressingTools(), + }); + + const legacy = await setup({ server: legacyServer }); + const modern = await setupModernClient({ tools: suppressingTools() }); + + // Suppressed lane: byte-identical to the policy-free fixture's measured + // base bytes, including `formatResult` being skipped, even though this + // tool both attaches a policy and declares `formatResult`. + expect(JSON.stringify(await legacy.client.listTools())).toBe( + BASE_LEGACY_DISCOVERY + ); + expect( + JSON.stringify( + await legacy.client.callTool({ + name: 'fixture', + arguments: { value: 'hi' }, + }) + ) + ).toBe(BASE_LEGACY_CALL); + + // Normalized lane: same policy, same tool, structured results and the + // tool's own `formatResult`. + const modernDiscovery = await modern.listTools(); + expect(modernDiscovery.tools[0]?.outputSchema).toMatchObject({ + type: 'object', + properties: { value: { type: 'string' } }, + }); + + const modernResult = await modern.callTool({ + name: 'fixture', + arguments: { value: 'hi' }, + }); + + expect(modernResult.structuredContent).toEqual({ value: 'hi' }); + expect(modernResult.content).toEqual([{ type: 'text', text: 'PREFIX:hi' }]); + }); + + test('formatResult changes text only, never discovery or structuredContent', async () => { + const plainServer = createMcpServer({ + name: 'test-server', + version: '0.0.0', + tools: policyTools(), + }); + const formattedServer = createMcpServer({ + name: 'test-server', + version: '0.0.0', + tools: policyTools(({ value }) => `PREFIX:${value}`), + }); + + const plain = await setup({ server: plainServer }); + const formatted = await setup({ server: formattedServer }); + + // Discovery is byte-identical with and without `formatResult`. + expect(JSON.stringify(await formatted.client.listTools())).toBe( + JSON.stringify(await plain.client.listTools()) + ); + + const result = await formatted.client.callTool({ + name: 'fixture', + arguments: { value: 'hi' }, + }); + + expect(result.structuredContent).toEqual({ value: 'hi' }); + expect(result.content).toEqual([{ type: 'text', text: 'PREFIX:hi' }]); + }); + + test('formatResult on a policy-free tool still emits no structuredContent', async () => { + const server = createMcpServer({ + name: 'test-server', + version: '0.0.0', + tools: { + fixture: tool({ + description: 'Fixture', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + execute: async ({ value }) => ({ value }), + formatResult: ({ value }) => `PREFIX:${value}`, + }), + }, + }); + + const { client } = await setup({ server }); + + // Discovery still carries no `outputSchema` key. + expect(JSON.stringify(await client.listTools())).toBe( + BASE_LEGACY_DISCOVERY + ); + + const result = await client.callTool({ + name: 'fixture', + arguments: { value: 'hi' }, + }); + + expect(result.structuredContent).toBeUndefined(); + expect(result.content).toEqual([{ type: 'text', text: 'PREFIX:hi' }]); + }); + + test('a null tool result produces no content', async () => { + const server = createMcpServer({ + name: 'test-server', + version: '0.0.0', + tools: { + empty: tool({ + description: 'Empty', + parameters: z.object({}), + outputSchema: z.object({}), + execute: async () => null as unknown as Record, + }), + }, + }); + + const { client } = await setup({ server }); + + const result = await client.callTool({ name: 'empty', arguments: {} }); + + expect(result.content).toEqual([]); + expect(result.structuredContent).toBeUndefined(); + }); +}); + +describe('SDK request state pass-through', () => { + test('the verifier runs before dispatch and its value reaches the handler', async () => { + const order: string[] = []; + let round = 0; + // Typed as the whole SDK option object, so narrowing it to `verify` or + // republishing its fields would fail to compile. + const requestState: ServerOptions['requestState'] = { + verify: async (state) => { + order.push(`verify:${state}`); + return { approved: true }; + }, + }; + const client = await setupModernClient({ + requestState, + tools: { + guarded: tool({ + description: 'Guarded', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + resolve: async (_params, ctx) => { + round += 1; + order.push( + `resolve:${round}:${JSON.stringify(ctx.server.mcpReq.requestState())}` + ); + + if (round === 1) { + return { + type: 'result', + result: inputRequired({ requestState: 'minted-state' }), + telemetry, + }; + } + + return { type: 'execute', resolution: undefined, telemetry }; + }, + }, + execute: async ({ value }) => ({ value }), + }), + }, + }); + + const result = await client.callTool({ + name: 'guarded', + arguments: { value: 'hi' }, + }); + + expect(order).toEqual([ + 'resolve:1:undefined', + 'verify:minted-state', + 'resolve:2:{"approved":true}', + ]); + expect(result.structuredContent).toEqual({ value: 'hi' }); + }); + + test('verifier rejection propagates the SDK-owned -32602 and skips the handler', async () => { + const resolve = vi.fn(async () => ({ + type: 'result' as const, + result: inputRequired({ requestState: 'minted-state' }), + telemetry, + })); + const client = await setupModernClient({ + requestState: { + verify: async () => { + throw new Error('bad signature'); + }, + }, + tools: { + guarded: tool({ + description: 'Guarded', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { resolve }, + execute: async ({ value }) => ({ value }), + }), + }, + }); + + const call = client.callTool({ + name: 'guarded', + arguments: { value: 'hi' }, + }); + + // The SDK owns the code, the frozen message, and the reason. This package + // neither wraps nor translates them. + await expect(call).rejects.toMatchObject({ + code: -32602, + message: 'Invalid or expired requestState', + data: { reason: 'invalid_request_state' }, + }); + // Only the first round reached the handler. + expect(resolve.mock.calls).toHaveLength(1); + }); +}); + +describe('public package surface', () => { + test('a tool with no formatResult and no policy stays assignable', async () => { + // Mirrors how consumers wrap `Tool` today, without going through the + // `tool()` helper: no `formatResult`, no policy, one-argument `execute`. + const consumerTool: Tool< + z.ZodObject<{ value: z.ZodString }>, + z.ZodObject<{ value: z.ZodString }> + > = { + description: 'Consumer', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + execute: async ({ value }) => ({ value }), + }; + + const server = createMcpServer({ + name: 'test-server', + version: '0.0.0', + tools: { consumer: consumerTool }, + }); + + const { callTool } = await setup({ server }); + + await expect( + callTool({ name: 'consumer', arguments: { value: 'hi' } }) + ).resolves.toEqual({ value: 'hi' }); + }); +}); + describe('resources helper', () => { test('should add scheme to resource URIs', () => { const output = resources('my-scheme', [ diff --git a/packages/mcp-utils/src/server.ts b/packages/mcp-utils/src/server.ts index 01328e50..924caeb7 100644 --- a/packages/mcp-utils/src/server.ts +++ b/packages/mcp-utils/src/server.ts @@ -1,4 +1,9 @@ -import { Server } from '@modelcontextprotocol/server'; +import { + CLIENT_CAPABILITIES_META_KEY, + CLIENT_INFO_META_KEY, + PROTOCOL_VERSION_META_KEY, + Server, +} from '@modelcontextprotocol/server'; import type { ClientCapabilities, Implementation, @@ -8,10 +13,17 @@ import type { Tool as McpTool, ReadResourceResult, ServerCapabilities, + ServerContext, + ServerOptions, } from '@modelcontextprotocol/server'; import { z } from 'zod/v4'; import type { ExtractParams } from './types.js'; +import type { + ToolPolicy, + ToolPolicyTelemetry, + ToolRequestContext, +} from './tool-policy.js'; import { assertValidUri, compareUris, matchUriTemplate } from './util.js'; export type Scheme = string; @@ -45,14 +57,45 @@ export type Tool< // MCP spec restricts outputSchema to type "object" at the root level: // https://modelcontextprotocol.io/specification/2025-11-25/schema#tool-outputschema OutputSchema extends z.ZodObject = z.ZodObject, + Resolution = never, > = { description: Prop; annotations?: Annotations; parameters: Params; outputSchema: OutputSchema; - /** If true, excludes the tool from `tools/list` while keeping it callable via `tools/call`. */ - hidden?: boolean; - execute(params: z.infer): Promise>; + /** + * If true, excludes the tool from `tools/list` while keeping it callable via + * `tools/call`. A function decides the same thing per request, from the + * normalized request context. + * + * The function runs inside `tools/list`, once per tool, and must not throw: + * a throw fails the whole discovery response, not just this entry. + */ + hidden?: boolean | ((ctx: ToolRequestContext) => boolean); + /** + * Pre-execution policy consulted for discovery schemas, argument + * normalization, and the decision to execute or answer directly. + */ + policy?: ToolPolicy, Resolution>; + execute( + params: z.infer, + ...resolution: [Resolution] extends [never] ? [] : [Resolution] + ): Promise>; + /** + * Renders the tool result as MCP text content. + * + * Defaults to `JSON.stringify`, which keeps the text content a + * single-encoded rendering of the business result. Setting it never + * changes discovery bytes and never decides whether `structuredContent` + * is emitted. + * + * It is skipped on a request whose policy suppressed structured results + * via `policy.outputSchema` returning `undefined`, because that request + * reproduces the whole pre-normalization result. On a policy-free tool it + * always applies: setting it there is an explicit authoring choice, + * unrelated to suppression. + */ + formatResult?: (result: z.infer) => string; }; /** @@ -162,7 +205,8 @@ export function jsonResourceResponse( export function tool< Params extends z.ZodObject, OutputSchema extends z.ZodObject, ->(tool: Tool) { + Resolution = never, +>(tool: Tool) { return tool; } @@ -189,8 +233,38 @@ type ToolCallErrorDetails = ToolCallBaseDetails & { export type ToolCallDetails = ToolCallSuccessDetails | ToolCallErrorDetails; +/** Safe, product-neutral report of one pre-execution policy decision. */ +type ToolPolicyCallDetails = { + /** Name of the tool whose policy produced the decision. */ + name: string; + /** + * Whether the decision short-circuited business execution. + * + * `'rejected'` is this package's own outcome for a decision it could not + * recognize, never the policy's raw string. On `'rejected'` the + * policy-reported `telemetry.outcome` is unreliable by construction, so + * this label is the authoritative one. + */ + decision: 'execute' | 'result' | 'rejected'; + /** Wall-clock duration of policy resolution, in milliseconds. */ + durationMs: number; + /** Client identity, when the request carried it. */ + clientInfo?: Implementation; + /** + * Allowlisted telemetry reported by the policy. + * + * Partial because types erase at runtime: a policy written in plain + * JavaScript can omit the required identity fields, and the record then + * carries only what survived the allowlist. + */ + telemetry: Partial; +}; + export type InitCallback = (initData: InitData) => void | Promise; export type ToolCallCallback = (details: ToolCallDetails) => void; +type ToolPolicyCallCallback = ( + details: ToolPolicyCallDetails +) => void | Promise; export type PropCallback = () => T | Promise; export type Prop = T | PropCallback; @@ -234,6 +308,27 @@ export type McpServerOptions = { */ onToolCall?: ToolCallCallback; + /** + * Callback for each pre-execution policy decision. Receives allowlisted + * telemetry only; a failure here never changes a tool result. + * + * The callback is not awaited. Its completion is not ordered before the + * response, and on serving paths that stop work once the response is sent + * it is not guaranteed to run to completion at all. A failing or slow sink + * never changes or delays the tool result. + */ + onToolPolicyCall?: ToolPolicyCallCallback; + + /** + * Multi-round-trip request-state integrity hook, forwarded unchanged to the + * SDK server. + * + * The SDK owns the option shape, when `verify` runs, what a resolved value + * means, and the JSON-RPC error a rejection produces. This package adds one + * typed pass-through and no behavior of its own. + */ + requestState?: ServerOptions['requestState']; + /** * Resources to be served by the server. These can be defined as a static * object or as a function that dynamically returns the object synchronously @@ -256,9 +351,233 @@ export type McpServerOptions = { * asks for the list of tools or invokes a tool. This allows for dynamic tools * that can change after the server has started. */ - tools?: Prop>; + tools?: Prop>>; }; +/** + * Whether a value is a plain object, which is what MCP permits at the root of + * structured output. + * + * Prototype-exact on purpose. A `Date` or anything else carrying `toJSON` + * serializes to a non-object root, and a `Map` or `Set` serializes to `{}`, + * silently dropping its contents. Both would defeat a root check that only + * asked `typeof`. The cost is that a class instance is rejected even though + * it serializes to its own fields: this seam advertises a data shape, so a + * tool that wants to emit one returns the data. + */ +function isPlainObject(value: unknown): value is Record { + if (typeof value !== 'object' || value === null) { + return false; + } + + const prototype = Object.getPrototypeOf(value); + + return prototype === Object.prototype || prototype === null; +} + +/** Names a non-plain-object value in an error a tool author can act on. */ +function describeNonPlainObject(value: unknown): string { + if (value === null) { + return 'null'; + } + + if (Array.isArray(value)) { + return 'an array'; + } + + if (typeof value !== 'object') { + return typeof value; + } + + return `an instance of ${value.constructor?.name ?? 'an anonymous class'}`; +} + +/** + * JSON-value equality, not a general-purpose deep equal. + * + * Object keys compare order-insensitively, array elements order-sensitively, + * and everything else by identity. It knows nothing about `Date`, `Map`, + * cycles or `NaN`, and it does not need to: the sole caller compares a value + * that survived a JSON round trip against the raw result, so anything JSON + * cannot represent faithfully has already become something else by the time + * it arrives here and simply compares unequal. Divergence is what rejects a + * non-JSON value, not a precondition. Only the root is checked for + * plain-objectness before the walk, so arbitrary values do reach it. + */ +function jsonValueEquals(a: unknown, b: unknown): boolean { + if (a === b) { + return true; + } + + if (Array.isArray(a) || Array.isArray(b)) { + return ( + Array.isArray(a) && + Array.isArray(b) && + a.length === b.length && + a.every((element, index) => jsonValueEquals(element, b[index])) + ); + } + + if (!isPlainObject(a) || !isPlainObject(b)) { + return false; + } + + const keys = Object.keys(a); + + return ( + keys.length === Object.keys(b).length && + keys.every( + (key) => Object.hasOwn(b, key) && jsonValueEquals(a[key], b[key]) + ) + ); +} + +/** + * The result shape one request gets. Discovery and the call path resolve this + * once per request from the same function, so they always agree. + */ +type RequestResultShape = + /** Policy-free: pre-normalization bytes, with `formatResult` still applied. */ + | { kind: 'plain' } + /** Policy suppressed this request: the base result, `formatResult` skipped. */ + | { kind: 'suppressed' } + /** Structured results advertised and emitted against this schema. */ + | { kind: 'normalized'; outputSchema: z.ZodType }; + +/** + * Decides whether a request carries structured results. + * + * Structured results follow the policy seam: `policy.outputSchema` is the + * only hook that can contextualize an advertised output schema, so a + * policy-free tool never advertises. A policy that defines the hook decides + * per request and can suppress; a policy without one always normalizes. + */ +function resolveResultShape( + tool: Tool, + ctx: ToolRequestContext +): RequestResultShape { + if (!tool.policy) { + return { kind: 'plain' }; + } + + if (!tool.policy.outputSchema) { + return { kind: 'normalized', outputSchema: tool.outputSchema }; + } + + const outputSchema = tool.policy.outputSchema(tool.outputSchema, ctx); + + return outputSchema === undefined + ? { kind: 'suppressed' } + : { kind: 'normalized', outputSchema }; +} + +/** + * Resolves the input schema one request advertises and enforces. + * + * `tools/list` and `tools/call` both resolve it here, from the same hook + * call and in the same strict form, so the advertised schema and the schema + * strict parsing enforces cannot disagree. A hook returning a loose or + * catchall object still advertises and rejects unknown keys: strictness at + * this seam belongs to the server, not to the policy. + */ +function resolveParameters( + tool: Tool, + ctx: ToolRequestContext +): z.ZodObject { + return ( + tool.policy?.inputSchema?.(tool.parameters, ctx) ?? tool.parameters + ).strict(); +} + +/** + * Renders one `tools/list` entry for one request. + */ +async function describeTool( + name: string, + tool: Tool, + ctx: ToolRequestContext +): Promise { + const inputSchema = z.toJSONSchema(resolveParameters(tool, ctx), { + target: 'draft-7', + }); + const entry = { + name, + description: + typeof tool.description === 'function' + ? await tool.description() + : tool.description, + annotations: tool.annotations, + // Casting the same as the SDK does: + // https://github.com/modelcontextprotocol/typescript-sdk/blob/fb07af810b51003c338dc4885a9e42f54519f9af/src/server/mcp.ts#L154 + inputSchema: inputSchema as McpTool['inputSchema'], + }; + + // A request advertises only when the shared decision normalizes it: a + // policy-free tool never does, and a policy can suppress per request. + // Suppressed and policy-free entries keep the discovery bytes they had + // before structured results existed. + const shape = resolveResultShape(tool, ctx); + + if (shape.kind !== 'normalized') { + return entry; + } + + const outputSchema = z.toJSONSchema(shape.outputSchema, { + target: 'draft-7', + }); + + // MCP restricts a tool's structured output to an object root, so a + // resolved schema that declares anything else can never be advertised. + // Fail the whole discovery response rather than dropping the key, which + // would advertise the suppressed shape while the call path still emits + // `structuredContent`, or dropping the entry, which would silently hide a + // tool on a policy bug. The error is deterministic: every `tools/list` + // fails identically, so an authoring error surfaces in development. + if (outputSchema.type !== 'object') { + throw new Error( + `Tool "${name}" resolved an output schema that is not object-rooted: ` + + `root type ${JSON.stringify(outputSchema.type ?? null)}. ` + + 'MCP restricts structured output to an object root, so model ' + + 'variants beneath an object root instead of at the root.' + ); + } + + return { + ...entry, + outputSchema: outputSchema as McpTool['outputSchema'], + }; +} + +/** + * Derives the product-neutral request context from SDK-owned metadata. + * Called once per `tools/list` and `tools/call`, so every policy hook and + * visibility filter in one request sees the same facts. + */ +function normalizeToolRequestContext( + requestContext: ServerContext, + server: Server +): ToolRequestContext { + const envelope = requestContext.mcpReq.envelope as + | Record + | undefined; + + return { + server: requestContext, + // The per-request `_meta` envelope exists only on the modern revision. + era: + envelope?.[PROTOCOL_VERSION_META_KEY] === undefined ? 'legacy' : 'modern', + // The modern per-request envelope is authoritative. The legacy path + // carries no envelope, so fall back to what initialization captured. + clientInfo: + (envelope?.[CLIENT_INFO_META_KEY] as Implementation | undefined) ?? + server.getClientVersion(), + clientCapabilities: + (envelope?.[CLIENT_CAPABILITIES_META_KEY] as + | ClientCapabilities + | undefined) ?? server.getClientCapabilities(), + }; +} + /** * Creates an MCP server with the given options. * @@ -285,6 +604,7 @@ export function createMcpServer(options: McpServerOptions) { { capabilities, instructions: options.instructions, + requestState: options.requestState, } ); @@ -440,36 +760,29 @@ export function createMcpServer(options: McpServerOptions) { if (options.tools) { server.setRequestHandler( 'tools/list', - async (): Promise => { + async (_request, serverContext): Promise => { const tools = await getTools(); + const context = normalizeToolRequestContext(serverContext, server); + const visibleTools = Object.entries(tools).filter( + ([, tool]) => + !(typeof tool.hidden === 'function' + ? tool.hidden(context) + : tool.hidden) + ); return { tools: await Promise.all( - Object.entries(tools) - .filter(([, tool]) => !tool.hidden) - .map(async ([name, { description, annotations, parameters }]) => { - const inputSchema = z.toJSONSchema(parameters, { - target: 'draft-7', - }); - - return { - name, - description: - typeof description === 'function' - ? await description() - : description, - annotations, - // Casting the same as the SDK does: - // https://github.com/modelcontextprotocol/typescript-sdk/blob/fb07af810b51003c338dc4885a9e42f54519f9af/src/server/mcp.ts#L154 - inputSchema: inputSchema as McpTool['inputSchema'], - }; - }) + visibleTools.map(([name, tool]) => + describeTool(name, tool, context) + ) ), } satisfies ListToolsResult; } ); - server.setRequestHandler('tools/call', async (request) => { + server.setRequestHandler('tools/call', async (request, serverContext) => { + const context = normalizeToolRequestContext(serverContext, server); + try { const tools = await getTools(); const toolName = request.params.name; @@ -483,14 +796,97 @@ export function createMcpServer(options: McpServerOptions) { if (!tool) { throw new Error('tool not found'); } - const args = tool.parameters - .strict() - .parse(request.params.arguments ?? {}); - const executeWithCallback = async (tool: Tool) => { + const rawArguments = request.params.arguments ?? {}; + // Check for the hook, not for a nullish result: a policy that + // normalizes arguments away must not silently fall back to the raw + // arguments, which strict parsing would then reject. + const normalizedArguments = tool.policy?.normalizeArguments + ? tool.policy.normalizeArguments(rawArguments, context) + : rawArguments; + // Already strict, and the exact schema discovery advertised. + const args = resolveParameters(tool, context).parse( + normalizedArguments + ) as Record; + + // Resolved once per request, before the policy runs, so the hook + // cannot observe anything `resolve` changed and discovery and the + // call path can never disagree. + const shape = resolveResultShape(tool, context); + + let resolution: unknown; + if (tool.policy) { + const policyStartedAt = performance.now(); + const decision = await tool.policy.resolve(args, context); + const durationMs = performance.now() - policyStartedAt; + + // Guard first, then narrow, so the audit record is built from a + // recognized decision instead of whatever the policy returned. + // Types erase at runtime: a plain JavaScript policy or a cast can + // return a nullish value, a non-object, or an unknown `type`, and + // all three are unrecognizable the same way. + const rawDecision = decision as + | { type?: unknown; telemetry?: ToolPolicyTelemetry } + | null + | undefined; + const recognized = + typeof decision === 'object' && + decision !== null && + (decision.type === 'execute' || decision.type === 'result') + ? decision + : undefined; + + // Fire-and-forget, exactly once per request and after narrowing, so + // every outcome reaches the sink under a label this package owns. + // The callback is an audit sink that the JSDoc promises cannot + // change the result, so a slow or never-settling one must not stall + // the request. `Promise.resolve().then` also captures a synchronous + // throw from a plain JavaScript callback. + void Promise.resolve() + .then(() => + options.onToolPolicyCall?.({ + name: toolName, + decision: recognized?.type ?? 'rejected', + clientInfo: context.clientInfo, + durationMs, + telemetry: safeToolPolicyTelemetry(rawDecision?.telemetry), + }) + ) + .catch((error) => { + // Don't fail the tool call if the callback fails + console.error('Failed to run tool policy callback', error); + }); + + if (!recognized) { + // Fail closed instead of falling through and executing the + // guarded tool with no resolution. The throw lands in this + // handler's catch, which is how a policy that throws already + // behaves. + throw new Error( + `Unrecognized tool policy decision type: ${String(rawDecision?.type)}` + ); + } + + switch (recognized.type) { + case 'execute': + resolution = recognized.resolution; + break; + case 'result': + return recognized.result; + } + } + + const executeWithCallback = async () => { + // Policy-free tools keep the existing one-argument execute call. + const executeResult = tool.policy + ? tool.execute(args, resolution) + : ( + tool.execute as ( + args: Record + ) => Promise + )(args); // Wrap success or error in a result value - const res = await tool - .execute(args) + const res = await executeResult .then((data: unknown) => ({ success: true as const, data })) .catch((error) => ({ success: false as const, error })); @@ -513,15 +909,89 @@ export function createMcpServer(options: McpServerOptions) { return res.data; }; - const result = await executeWithCallback(tool); + const result = await executeWithCallback(); + + if (shape.kind !== 'normalized') { + // These lanes advertise nothing, so they keep the answers they gave + // before structured results existed, a nullish result included. + if (result == null) { + return { content: [] }; + } + + const base = result as Record; + // A suppressed request reproduces the whole pre-normalization + // result, which means the default single-encoded text: + // `formatResult` is skipped. A policy-free request applies it. + const text = + shape.kind === 'suppressed' || !tool.formatResult + ? JSON.stringify(base) + : tool.formatResult(base); + + return { content: [{ type: 'text', text }] }; + } + + // This request advertised a schema, so it owes the caller output that + // conforms to it. Two independent checks, because neither implies the + // other: `outputSchema` hooks may resolve a permissive schema, under + // which `null`, a scalar or a `Date` parses successfully and is still + // output MCP does not permit at the root. + if (!isPlainObject(result)) { + throw new Error( + `Tool "${toolName}" advertised an output schema but produced ` + + `${describeNonPlainObject(result)}. ` + + 'Structured output must be a plain object.' + ); + } + + // Compare and validate what the wire will actually carry. A + // permissive leaf (`z.any()`, `z.unknown()`) hands its input straight + // back, so parsing the raw result would let a `Map`, `Date` or + // `Infinity` one level down pass by identity and then be mangled by + // serialization. The round trip makes any such value diverge from the + // raw result below, or fail this parse outright. + const roundTripped = JSON.parse(JSON.stringify(result)) as unknown; + const parsed = shape.outputSchema.safeParse(roundTripped); + + if (!parsed.success) { + throw new Error( + `Tool "${toolName}" produced output that does not conform to its ` + + 'advertised output schema: ' + + parsed.error.issues + .map( + (issue) => + `${issue.path.join('.') || '(root)'}: ${issue.message}` + ) + .join('; ') + ); + } - const content = - result != null - ? [{ type: 'text' as const, text: JSON.stringify(result) }] - : []; + // A successful parse is not yet agreement. A zod object strips + // undeclared keys and coercions rewrite values, both while reporting + // success, and the advertised JSON for those same schemas says + // `additionalProperties: false`. So require the schema to have + // accepted the result as-is, comparing against the raw result: + // anything the schema or serialization had to change is mismatched + // output, answered rather than silently normalized away. + if (!jsonValueEquals(parsed.data, result)) { + throw new Error( + `Tool "${toolName}" produced output its advertised output schema ` + + 'did not accept as-is: the schema stripped or coerced fields, ' + + 'or the result carries values JSON cannot represent. ' + + 'Return exactly the advertised shape using JSON values, and do ' + + 'not advertise a transforming or coercing schema.' + ); + } + + // Equal to `parsed.data`, so emitting the business result keeps A1's + // "structuredContent equals the business result" literally true. + const structuredContent = result; + const text = tool.formatResult + ? tool.formatResult(structuredContent) + : JSON.stringify(structuredContent); return { - content, + structuredContent, + content: [{ type: 'text', text }], }; } catch (error) { return { @@ -540,6 +1010,55 @@ export function createMcpServer(options: McpServerOptions) { return server; } +/** + * Copies only the allowlisted scalar telemetry fields. + * + * `ToolPolicyTelemetry` is a compile-time contract and types erase at + * runtime, so a policy that spreads a wider object, or any policy written in + * plain JavaScript, could otherwise push raw arguments or request state into + * a telemetry sink. Widening the allowlist means changing both the type and + * this function. + * + * A decision missing the object entirely (again, a plain JavaScript policy or + * a cast) yields empty telemetry rather than a `TypeError`, which would + * otherwise drop the audit record and log it as a callback failure. + */ +function safeToolPolicyTelemetry( + telemetry: ToolPolicyTelemetry | undefined +): Partial { + if (!telemetry || typeof telemetry !== 'object') { + return {}; + } + + const safe: Partial = {}; + + if (typeof telemetry.interactionId === 'string') { + safe.interactionId = telemetry.interactionId; + } + + if (typeof telemetry.authorityPath === 'string') { + safe.authorityPath = telemetry.authorityPath; + } + + if (typeof telemetry.outcome === 'string') { + safe.outcome = telemetry.outcome; + } + + if (typeof telemetry.reason === 'string') { + safe.reason = telemetry.reason; + } + + if (typeof telemetry.policyId === 'string') { + safe.policyId = telemetry.policyId; + } + + if (typeof telemetry.policyVersion === 'number') { + safe.policyVersion = telemetry.policyVersion; + } + + return safe; +} + function enumerateError(error: unknown) { if (!error) { return error; diff --git a/packages/mcp-utils/src/tool-policy.test.ts b/packages/mcp-utils/src/tool-policy.test.ts new file mode 100644 index 00000000..a8b27cfc --- /dev/null +++ b/packages/mcp-utils/src/tool-policy.test.ts @@ -0,0 +1,1106 @@ +import { + Client, + StreamableHTTPClientTransport, +} from '@modelcontextprotocol/client'; +import { + createMcpHandler, + type ClientCapabilities, +} from '@modelcontextprotocol/server'; +import { afterEach, describe, expect, test, vi } from 'vitest'; +import { z } from 'zod/v4'; + +import { createMcpServer, type McpServerOptions, tool } from './server.js'; +import { StreamTransport } from './stream-transport.js'; +import type { + ToolPolicy, + ToolPolicyDecision, + ToolPolicyTelemetry, + ToolRequestContext, +} from './tool-policy.js'; + +// https://blog.modelcontextprotocol.io/posts/2026-07-28-release-candidate/ +const MODERN_PROTOCOL_VERSION = '2026-07-28'; +const MCP_ENDPOINT = new URL('https://mcp.test'); + +const telemetry: ToolPolicyTelemetry = { + policyId: 'test-policy', + policyVersion: 1, + outcome: 'allowed', +}; + +const cleanups: Array<() => Promise> = []; + +afterEach(async () => { + for (const cleanup of cleanups.splice(0).reverse()) { + await cleanup(); + } +}); + +/** + * Connects a client that negotiates the modern protocol revision, so the + * server sees the per-request `_meta` envelope. + */ +async function setupModernClient( + options: Omit, + capabilities: ClientCapabilities = {} +) { + const handler = createMcpHandler( + () => + createMcpServer({ + name: 'policy-test-server', + version: '0.0.0', + ...options, + }), + { legacy: 'reject' } + ); + const transport = new StreamableHTTPClientTransport(MCP_ENDPOINT, { + fetch: (url, init) => handler.fetch(new Request(url, init)), + }); + const client = new Client( + { name: 'policy-test-client', version: '1.2.3' }, + { + capabilities, + versionNegotiation: { mode: { pin: MODERN_PROTOCOL_VERSION } }, + } + ); + + await client.connect(transport); + cleanups.push( + () => client.close(), + () => handler.close() + ); + + return client; +} + +/** + * Connects a client over the legacy in-process transport, which carries no + * per-request `_meta` envelope. + */ +async function setupLegacyClient( + options: Omit +) { + const server = createMcpServer({ + name: 'policy-test-server', + version: '0.0.0', + ...options, + }); + const clientTransport = new StreamTransport(); + const serverTransport = new StreamTransport(); + + clientTransport.readable.pipeTo(serverTransport.writable); + serverTransport.readable.pipeTo(clientTransport.writable); + + const client = new Client( + { name: 'policy-test-client', version: '1.2.3' }, + { capabilities: {} } + ); + + await server.connect(serverTransport); + await client.connect(clientTransport); + + return client; +} + +describe('tool request context', () => { + test('a modern request carries the SDK-owned era, client info and capabilities', async () => { + const contexts: ToolRequestContext[] = []; + const client = await setupModernClient( + { + tools: { + capture: tool({ + description: 'Capture', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + resolve: async (_params, ctx) => { + contexts.push(ctx); + return { type: 'execute', resolution: undefined, telemetry }; + }, + }, + execute: async ({ value }) => ({ value }), + }), + }, + }, + { elicitation: {} } + ); + + await client.callTool({ name: 'capture', arguments: { value: 'ok' } }); + + expect(contexts).toHaveLength(1); + expect(contexts[0]?.era).toBe('modern'); + expect(contexts[0]?.clientInfo).toEqual({ + name: 'policy-test-client', + version: '1.2.3', + }); + expect(contexts[0]?.clientCapabilities).toMatchObject({ elicitation: {} }); + expect(contexts[0]?.server.mcpReq).toBeDefined(); + }); + + test('a legacy request falls back to the metadata captured at initialization', async () => { + const contexts: ToolRequestContext[] = []; + const client = await setupLegacyClient({ + tools: { + capture: tool({ + description: 'Capture', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + resolve: async (_params, ctx) => { + contexts.push(ctx); + return { type: 'execute', resolution: undefined, telemetry }; + }, + }, + execute: async ({ value }) => ({ value }), + }), + }, + }); + + await client.callTool({ name: 'capture', arguments: { value: 'ok' } }); + + expect(contexts).toHaveLength(1); + expect(contexts[0]?.era).toBe('legacy'); + expect(contexts[0]?.clientInfo).toEqual({ + name: 'policy-test-client', + version: '1.2.3', + }); + expect(contexts[0]?.clientCapabilities).toEqual({}); + }); +}); + +describe('pre-execution tool policy', () => { + test('discovery applies contextual visibility and policy schemas', async () => { + const policy: ToolPolicy< + { value: string }, + { era: ToolRequestContext['era'] } + > = { + inputSchema: (schema, ctx) => + ctx.era === 'modern' + ? schema.extend({ confirmation: z.string() }) + : schema, + outputSchema: (schema, ctx) => + ctx.era === 'modern' + ? schema.extend({ confirmed: z.boolean() }) + : schema, + // The resolution carries the era, so `execute` can satisfy the schema + // this request advertised without re-deriving the context. + resolve: async (_args, ctx) => ({ + type: 'execute', + resolution: { era: ctx.era }, + telemetry, + }), + }; + const tools = { + contextual: tool({ + description: 'Contextual', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + hidden: (ctx) => ctx.era !== 'modern', + policy, + execute: async ({ value }, { era }) => + era === 'modern' ? { value, confirmed: true } : { value }, + }), + }; + + const modernClient = await setupModernClient({ tools }); + const legacyClient = await setupLegacyClient({ tools }); + + const modernDiscovery = await modernClient.listTools(); + const legacyDiscovery = await legacyClient.listTools(); + + expect(modernDiscovery.tools[0]?.inputSchema).toHaveProperty( + 'properties.confirmation' + ); + expect(modernDiscovery.tools[0]?.outputSchema).toHaveProperty( + 'properties.confirmed' + ); + expect(legacyDiscovery.tools).toEqual([]); + + // The call path must resolve its input schema from the same hook, not + // from `tool.parameters`: the hook-supplied `confirmation` key would + // otherwise be rejected by strict parsing. + const modernResult = await modernClient.callTool({ + name: 'contextual', + arguments: { value: 'ok', confirmation: 'yes' }, + }); + + expect(modernResult.structuredContent).toEqual({ + value: 'ok', + confirmed: true, + }); + + // Hidden from `tools/list`, still alive via `tools/call`. + const legacyResult = await legacyClient.callTool({ + name: 'contextual', + arguments: { value: 'ok' }, + }); + + expect(legacyResult.isError).not.toBe(true); + expect(legacyResult.structuredContent).toEqual({ value: 'ok' }); + }); + + test('normalizeArguments runs before strict parsing', async () => { + const execute = vi.fn(async ({ value }: { value: string }) => ({ value })); + const client = await setupModernClient({ + tools: { + normalized: tool({ + description: 'Normalized', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + normalizeArguments: (raw) => { + const { legacy: _legacy, ...rest } = raw as Record< + string, + unknown + >; + return rest; + }, + resolve: async () => ({ + type: 'execute', + resolution: undefined, + telemetry, + }), + }, + execute, + }), + }, + }); + + const accepted = await client.callTool({ + name: 'normalized', + arguments: { value: 'ok', legacy: true }, + }); + const rejected = await client.callTool({ + name: 'normalized', + arguments: { value: 'no', unknown: true }, + }); + + expect(accepted.isError).not.toBe(true); + expect(rejected.isError).toBe(true); + // Only the normalized call reaches execution; the unknown field is fatal. + // A tool with a policy always receives the resolution as its second + // argument, `undefined` here. + expect(execute.mock.calls).toEqual([[{ value: 'ok' }, undefined]]); + }); + + test('a result decision short-circuits business execution', async () => { + const execute = vi.fn(); + const client = await setupModernClient({ + tools: { + guarded: tool({ + description: 'Guarded', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + // The request is normalized, so it advertises an output schema. + // A terminal result on a normalized request must be one the + // protocol exempts from `structuredContent`: an error here. + // This package never synthesizes structured content for it. + resolve: async () => ({ + type: 'result', + result: { + isError: true, + content: [ + { + type: 'text' as const, + text: 'Not permitted: request approval before retrying.', + }, + ], + }, + telemetry, + }), + }, + execute, + }), + }, + }); + + const result = await client.callTool({ + name: 'guarded', + arguments: { value: 'ignored' }, + }); + + expect(execute).not.toHaveBeenCalled(); + expect(result.isError).toBe(true); + expect(result.content).toEqual([ + { + type: 'text', + text: 'Not permitted: request approval before retrying.', + }, + ]); + expect(result.structuredContent).toBeUndefined(); + }); + + test('a result decision on a suppressed request stays content-only', async () => { + const execute = vi.fn(); + const client = await setupModernClient({ + tools: { + guarded: tool({ + description: 'Guarded', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + // Suppressed: the request advertises no output schema, so a + // content-only terminal result is the shape its discovery entry + // promised. + outputSchema: () => undefined, + resolve: async () => ({ + type: 'result', + result: { + content: [{ type: 'text' as const, text: 'intercepted' }], + }, + telemetry, + }), + }, + execute, + }), + }, + }); + + const discovery = await client.listTools(); + + expect(discovery.tools[0]?.outputSchema).toBeUndefined(); + + const result = await client.callTool({ + name: 'guarded', + arguments: { value: 'ignored' }, + }); + + expect(execute).not.toHaveBeenCalled(); + expect(result.content).toEqual([{ type: 'text', text: 'intercepted' }]); + expect(result.structuredContent).toBeUndefined(); + }); + + test('an unrecognized decision type fails closed', async () => { + const execute = vi.fn(); + const onToolPolicyCall = vi.fn(); + const client = await setupModernClient({ + onToolPolicyCall, + tools: { + guarded: tool({ + description: 'Guarded', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + // A policy written in plain JavaScript, or one casting its + // decision, can return a type this package does not recognize. + resolve: async () => + ({ + type: 'deny', + telemetry, + }) as unknown as ToolPolicyDecision, + }, + execute, + }), + }, + }); + + const result = await client.callTool({ + name: 'guarded', + arguments: { value: 'ignored' }, + }); + + // Fail closed: the guarded tool never runs and the caller gets an error. + expect(result.isError).toBe(true); + expect(execute).not.toHaveBeenCalled(); + // The sink sees this package's own label, never the policy's raw string. + expect(onToolPolicyCall).toHaveBeenCalledWith( + expect.objectContaining({ name: 'guarded', decision: 'rejected' }) + ); + }); + + test('a nullish decision fails closed and still records an audit call', async () => { + const execute = vi.fn(); + const onToolPolicyCall = vi.fn(); + const consoleError = vi + .spyOn(console, 'error') + .mockImplementation(() => {}); + const client = await setupModernClient({ + onToolPolicyCall, + tools: { + guarded: tool({ + description: 'Guarded', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + // A plain JavaScript policy can resolve to nothing at all. + resolve: async () => + undefined as unknown as ToolPolicyDecision, + }, + execute, + }), + }, + }); + + const result = await client.callTool({ + name: 'guarded', + arguments: { value: 'ignored' }, + }); + + expect(result.isError).toBe(true); + expect(execute).not.toHaveBeenCalled(); + expect(onToolPolicyCall).toHaveBeenCalledWith( + expect.objectContaining({ name: 'guarded', decision: 'rejected' }) + ); + // The record is built before the throw, so nothing is misattributed to + // the callback. + expect(consoleError).not.toHaveBeenCalledWith( + 'Failed to run tool policy callback', + expect.any(Error) + ); + consoleError.mockRestore(); + }); + + test('an execute decision hands the parsed arguments and the resolution to execute', async () => { + const execute = vi.fn(async ({ value }: { value: string }) => ({ value })); + const resolve = vi.fn(async () => ({ + type: 'execute' as const, + resolution: { grant: 'approved' }, + telemetry, + })); + const client = await setupModernClient({ + tools: { + guarded: tool({ + description: 'Guarded', + parameters: z.object({ + value: z.string(), + flag: z.boolean().default(false), + }), + outputSchema: z.object({ value: z.string() }), + policy: { resolve }, + execute, + }), + }, + }); + + await client.callTool({ name: 'guarded', arguments: { value: 'input' } }); + + // Defaults are applied by strict parsing before the policy sees them. + expect(resolve).toHaveBeenCalledWith( + { value: 'input', flag: false }, + expect.anything() + ); + expect(execute).toHaveBeenCalledWith( + { value: 'input', flag: false }, + { grant: 'approved' } + ); + }); + + test('a policy-free tool still executes with exactly one argument', async () => { + const execute = vi.fn(async ({ value }: { value: string }) => ({ value })); + const client = await setupModernClient({ + tools: { + plain: tool({ + description: 'Plain', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + execute, + }), + }, + }); + + await client.callTool({ name: 'plain', arguments: { value: 'ok' } }); + + expect(execute).toHaveBeenCalledWith({ value: 'ok' }); + }); + + test('policy resolution runs after parsing and before business execution', async () => { + const order: string[] = []; + const client = await setupModernClient({ + tools: { + ordered: tool({ + description: 'Ordered', + parameters: z.object({ value: z.string() }).refine((parsed) => { + order.push('parse'); + return parsed.value.length > 0; + }) as unknown as z.ZodObject<{ value: z.ZodString }>, + outputSchema: z.object({ value: z.string() }), + policy: { + resolve: async () => { + order.push('resolve'); + return { type: 'execute', resolution: undefined, telemetry }; + }, + }, + execute: async ({ value }) => { + order.push('execute'); + return { value }; + }, + }), + }, + }); + + await client.callTool({ name: 'ordered', arguments: { value: 'ok' } }); + + expect(order).toEqual(['parse', 'resolve', 'execute']); + }); + + test('a field outside the allowlist never reaches the policy callback', async () => { + const onToolPolicyCall = vi.fn(); + const client = await setupModernClient({ + onToolPolicyCall, + tools: { + leaky: tool({ + description: 'Leaky', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + resolve: async (params) => ({ + type: 'execute', + resolution: undefined, + // A policy compiled from plain JavaScript, or one spreading a + // wider object, can carry fields the type does not allow. + telemetry: { + ...telemetry, + rawArguments: params, + token: 'super-secret', + } as ToolPolicyTelemetry, + }), + }, + execute: async ({ value }) => ({ value }), + }), + }, + }); + + await client.callTool({ name: 'leaky', arguments: { value: 'ok' } }); + + expect(onToolPolicyCall).toHaveBeenCalledWith( + expect.objectContaining({ telemetry }) + ); + }); + + test('the policy callback receives only the allowlisted telemetry fields', async () => { + const onToolPolicyCall = vi.fn(); + const client = await setupModernClient({ + onToolPolicyCall, + tools: { + observed: tool({ + description: 'Observed', + parameters: z.object({ secret: z.string() }), + outputSchema: z.object({ secret: z.string() }), + policy: { + resolve: async () => ({ + type: 'execute', + resolution: undefined, + telemetry: { + interactionId: 'interaction-1', + authorityPath: 'test-authority', + outcome: 'allowed', + reason: 'internal-reason', + policyId: 'test-policy', + policyVersion: 2, + }, + }), + }, + execute: async ({ secret }) => ({ secret }), + }), + }, + }); + + await client.callTool({ + name: 'observed', + arguments: { secret: 'do-not-log-me' }, + }); + + // An exact call tuple: raw arguments never reach the telemetry callback. + expect(onToolPolicyCall.mock.calls).toEqual([ + [ + { + name: 'observed', + decision: 'execute', + clientInfo: { name: 'policy-test-client', version: '1.2.3' }, + durationMs: expect.any(Number), + telemetry: { + interactionId: 'interaction-1', + authorityPath: 'test-authority', + outcome: 'allowed', + reason: 'internal-reason', + policyId: 'test-policy', + policyVersion: 2, + }, + }, + ], + ]); + }); + + test('a failing policy callback cannot change the tool result', async () => { + const onToolPolicyCall = vi.fn(async () => { + throw new Error('callback failed'); + }); + const consoleError = vi + .spyOn(console, 'error') + .mockImplementation(() => {}); + const client = await setupModernClient({ + onToolPolicyCall, + tools: { + observed: tool({ + description: 'Observed', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + resolve: async () => ({ + type: 'execute', + resolution: undefined, + telemetry, + }), + }, + execute: async ({ value }) => ({ value }), + }), + }, + }); + + const result = await client.callTool({ + name: 'observed', + arguments: { value: 'ok' }, + }); + + expect(result.isError).not.toBe(true); + expect(result.structuredContent).toEqual({ value: 'ok' }); + expect(consoleError).toHaveBeenCalledWith( + 'Failed to run tool policy callback', + expect.any(Error) + ); + consoleError.mockRestore(); + }); + + test('a decision without telemetry still records an audit call', async () => { + const onToolPolicyCall = vi.fn(); + const client = await setupModernClient({ + onToolPolicyCall, + tools: { + observed: tool({ + description: 'Observed', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + // A plain JavaScript policy can omit the telemetry object. + resolve: async () => + ({ + type: 'execute', + resolution: undefined, + }) as unknown as ToolPolicyDecision, + }, + execute: async ({ value }) => ({ value }), + }), + }, + }); + + const result = await client.callTool({ + name: 'observed', + arguments: { value: 'ok' }, + }); + + expect(result.structuredContent).toEqual({ value: 'ok' }); + // The record survives the missing telemetry object: the sanitizer + // returns an empty allowlist instead of throwing. + expect(onToolPolicyCall).toHaveBeenCalledWith( + expect.objectContaining({ decision: 'execute', telemetry: {} }) + ); + }); + + test('a never-settling policy callback cannot delay the tool result', async () => { + const client = await setupModernClient({ + // The callback is not awaited, so a sink that never completes must not + // hold the response. + onToolPolicyCall: () => new Promise(() => {}), + tools: { + observed: tool({ + description: 'Observed', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + resolve: async () => ({ + type: 'execute', + resolution: undefined, + telemetry, + }), + }, + execute: async ({ value }) => ({ value }), + }), + }, + }); + + const result = await client.callTool({ + name: 'observed', + arguments: { value: 'ok' }, + }); + + expect(result.structuredContent).toEqual({ value: 'ok' }); + }); +}); + +describe('policy schema contracts', () => { + test('a loose input schema hook advertises the strict schema parsing enforces', async () => { + const client = await setupModernClient({ + tools: { + open: tool({ + description: 'Open', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + // A hook is free to hand back a loose or catchall object. The + // call path parses strictly either way, so discovery must not + // advertise unknown keys as acceptable. + inputSchema: (schema) => schema.loose(), + resolve: async () => ({ + type: 'execute', + resolution: undefined, + telemetry, + }), + }, + execute: async ({ value }) => ({ value }), + }), + }, + }); + + const discovery = await client.listTools(); + + expect(discovery.tools[0]?.inputSchema).toMatchObject({ + additionalProperties: false, + }); + + const result = await client.callTool({ + name: 'open', + arguments: { value: 'ok', extra: 'rejected' }, + }); + + expect(result.isError).toBe(true); + }); + + test('output that contradicts the advertised schema fails closed', async () => { + const client = await setupModernClient({ + tools: { + drifting: tool({ + description: 'Drifting', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + outputSchema: (schema) => schema.extend({ confirmed: z.boolean() }), + resolve: async () => ({ + type: 'execute', + resolution: undefined, + telemetry, + }), + }, + // Drifts from what this request advertised: `confirmed` is missing. + execute: async ({ value }) => + ({ value }) as unknown as { value: string; confirmed: boolean }, + }), + }, + }); + + const result = await client.callTool({ + name: 'drifting', + arguments: { value: 'ok' }, + }); + + expect(result.isError).toBe(true); + expect(result.structuredContent).toBeUndefined(); + expect(result.content).toEqual([ + { type: 'text', text: expect.stringContaining('confirmed') }, + ]); + }); + + test('a null result on a normalized request fails closed', async () => { + const client = await setupModernClient({ + tools: { + empty: tool({ + description: 'Empty', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + resolve: async () => ({ + type: 'execute', + resolution: undefined, + telemetry, + }), + }, + // Reachable from plain JavaScript: the cast is only needed to get + // past `execute`'s declared return type. + execute: async () => null as unknown as { value: string }, + }), + }, + }); + + const result = await client.callTool({ + name: 'empty', + arguments: { value: 'ok' }, + }); + + expect(result.isError).toBe(true); + expect(result.content).not.toEqual([]); + }); + + test('a scalar result on a normalized request fails closed', async () => { + const client = await setupModernClient({ + tools: { + scalar: tool({ + description: 'Scalar', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + resolve: async () => ({ + type: 'execute', + resolution: undefined, + telemetry, + }), + }, + execute: async () => 7 as unknown as { value: string }, + }), + }, + }); + + const result = await client.callTool({ + name: 'scalar', + arguments: { value: 'ok' }, + }); + + expect(result.isError).toBe(true); + expect(result.structuredContent).toBeUndefined(); + }); + + test('a non-object-rooted output schema fails the whole discovery response', async () => { + const client = await setupModernClient({ + tools: { + healthy: tool({ + description: 'Healthy', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + execute: async ({ value }) => ({ value }), + }), + broken: tool({ + description: 'Broken', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + // MCP restricts structured output to an object root, so this + // resolved schema can never be advertised. + outputSchema: () => z.string(), + resolve: async () => ({ + type: 'execute', + resolution: undefined, + telemetry, + }), + }, + execute: async ({ value }) => ({ value }), + }), + }, + }); + + // The whole list fails, deliberately: an authoring error must not hide + // behind a healthy neighbour entry. + await expect(client.listTools()).rejects.toThrow(/broken/); + }); + + test('output the advertised schema only accepts after stripping fails closed', async () => { + const client = await setupModernClient({ + tools: { + leaky: tool({ + description: 'Leaky', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + resolve: async () => ({ + type: 'execute', + resolution: undefined, + telemetry, + }), + }, + // A zod object strips undeclared keys and still reports success, + // while the advertised JSON for that same schema says + // `additionalProperties: false`. The extra key is mismatched + // output, so it must be answered, not quietly removed. + execute: async ({ value }) => + ({ value, extra: 'unadvertised' }) as unknown as { + value: string; + }, + }), + }, + }); + + const result = await client.callTool({ + name: 'leaky', + arguments: { value: 'ok' }, + }); + + expect(result.isError).toBe(true); + expect(result.structuredContent).toBeUndefined(); + expect(result.content).toEqual([ + { type: 'text', text: expect.stringContaining('did not accept') }, + ]); + }); + + test('a conforming result is emitted exactly as execute returned it', async () => { + const businessResult = { value: 'ok', count: 2 }; + const client = await setupModernClient({ + tools: { + conforming: tool({ + description: 'Conforming', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ count: z.number(), value: z.string() }), + policy: { + resolve: async () => ({ + type: 'execute', + resolution: undefined, + telemetry, + }), + }, + execute: async () => businessResult, + }), + }, + }); + + const result = await client.callTool({ + name: 'conforming', + arguments: { value: 'ok' }, + }); + + // The validation lane must be byte-transparent for a conforming result: + // same keys, same order, same single-encoded text. + expect(result.isError).not.toBe(true); + expect(JSON.stringify(result.structuredContent)).toBe( + JSON.stringify(businessResult) + ); + expect(result.content).toEqual([ + { type: 'text', text: JSON.stringify(businessResult) }, + ]); + }); + + test('a non-plain object on a normalized request fails closed', async () => { + const client = await setupModernClient({ + tools: { + dated: tool({ + description: 'Dated', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + // A permissive resolved schema is what makes this reachable: a + // `Date` parses here and then serializes to a JSON string, so + // the advertised object root would be contradicted on the wire. + outputSchema: () => z.unknown() as unknown as z.ZodType, + resolve: async () => ({ + type: 'execute', + resolution: undefined, + telemetry, + }), + }, + execute: async () => new Date(0) as unknown as { value: string }, + }), + }, + }); + + const result = await client.callTool({ + name: 'dated', + arguments: { value: 'ok' }, + }); + + expect(result.isError).toBe(true); + expect(result.structuredContent).toBeUndefined(); + }); + + test('a Map under a permissive output leaf fails closed', async () => { + const client = await setupModernClient({ + tools: { + mapping: tool({ + description: 'Mapping', + parameters: z.object({ value: z.string() }), + // A permissive leaf advertises `{}` and parses any value by + // identity, so the root plain-object check never sees the `Map`. + // JSON serializes it to `{}`, silently dropping its contents. + outputSchema: z.object({ value: z.unknown() }), + policy: { + resolve: async () => ({ + type: 'execute', + resolution: undefined, + telemetry, + }), + }, + execute: async () => + ({ value: new Map([['a', 1]]) }) as unknown as { value: unknown }, + }), + }, + }); + + const result = await client.callTool({ + name: 'mapping', + arguments: { value: 'ok' }, + }); + + expect(result.isError).toBe(true); + expect(result.structuredContent).toBeUndefined(); + }); + + test('an own `toJSON` property that replaces the root fails closed', async () => { + const client = await setupModernClient({ + tools: { + hijacking: tool({ + description: 'Hijacking', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + // Loose keeps the extra key from being stripped, which is what + // isolates the hazard: `toJSON` is an own data property, so the + // prototype-exact root check passes it, and serialization then + // replaces the whole object with a string. + outputSchema: (schema) => schema.loose(), + resolve: async () => ({ + type: 'execute', + resolution: undefined, + telemetry, + }), + }, + execute: async ({ value }) => + ({ value, toJSON: () => 'hijacked' }) as unknown as { + value: string; + }, + }), + }, + }); + + const result = await client.callTool({ + name: 'hijacking', + arguments: { value: 'ok' }, + }); + + expect(result.isError).toBe(true); + expect(result.structuredContent).toBeUndefined(); + }); + + test('a loose output schema passes conforming extra keys through', async () => { + const businessResult = { value: 'ok', extra: 1 }; + const client = await setupModernClient({ + tools: { + spacious: tool({ + description: 'Spacious', + parameters: z.object({ value: z.string() }), + outputSchema: z.object({ value: z.string() }), + policy: { + // The documented escape hatch: a tool that owes callers more + // than it can declare advertises `additionalProperties: {}` and + // keeps emitting the undeclared keys. + outputSchema: (schema) => schema.loose(), + resolve: async () => ({ + type: 'execute', + resolution: undefined, + telemetry, + }), + }, + execute: async () => businessResult as unknown as { value: string }, + }), + }, + }); + + const discovery = await client.listTools(); + + expect(discovery.tools[0]?.outputSchema).toMatchObject({ + additionalProperties: {}, + }); + + const result = await client.callTool({ + name: 'spacious', + arguments: { value: 'ok' }, + }); + + expect(result.isError).not.toBe(true); + expect(result.structuredContent).toEqual(businessResult); + }); +}); diff --git a/packages/mcp-utils/src/tool-policy.ts b/packages/mcp-utils/src/tool-policy.ts new file mode 100644 index 00000000..eefbed7c --- /dev/null +++ b/packages/mcp-utils/src/tool-policy.ts @@ -0,0 +1,156 @@ +import type { + CallToolResult, + ClientCapabilities, + Implementation, + InputRequiredResult, + ServerContext, +} from '@modelcontextprotocol/server'; +import type { z } from 'zod/v4'; + +/** + * Product-neutral facts about one tool request, normalized once per + * `tools/list` and `tools/call` from SDK-owned request metadata. + * + * This context deliberately carries no product concern: it does not decide + * elicitation support, parse serving-path URLs, or model per-connection + * opt-out. A consumer that needs those builds them on top. + */ +export type ToolRequestContext = { + /** The SDK request context, including the multi-round-trip accessors. */ + server: ServerContext; + /** + * `modern` when the request carried the per-request `_meta` envelope + * introduced with the multi-round-trip protocol revision, `legacy` + * otherwise. + */ + era: 'legacy' | 'modern'; + /** Client identity, when the request envelope carried it. */ + clientInfo?: Implementation; + /** Client capabilities, when the request envelope carried them. */ + clientCapabilities?: ClientCapabilities; +}; + +/** + * Closed allowlist of fields a policy may report about its own decision. + * + * Deliberately has no index signature: every field is a named scalar, so raw + * arguments, request state, and response content cannot reach a telemetry sink + * through this type. Widening the allowlist is a deliberate, reviewed change + * rather than a smuggled open record. + * + * Policy identity and the outcome are required: an audit record that cannot + * name the policy that produced it, its contract version, and how the request + * ended is not auditable. + * + * The values are opaque to this package. A downstream policy chooses them. + */ +export type ToolPolicyTelemetry = { + /** Stable identifier of the policy that produced the decision. */ + policyId: string; + /** Version of the policy contract that produced the decision. */ + policyVersion: number; + /** Terminal classification of the decision. */ + outcome: string; + /** Correlates every round and repeated attempt of one logical interaction. */ + interactionId?: string; + /** Which authority path granted the action. */ + authorityPath?: string; + /** Internal, non-user-facing explanation of the outcome. */ + reason?: string; +}; + +/** + * The outcome of consulting a policy before business execution: either + * proceed with a resolution, or answer the caller directly. + */ +export type ToolPolicyDecision = + | { + type: 'execute'; + resolution: Resolution; + telemetry: ToolPolicyTelemetry; + } + | { + type: 'result'; + /** + * The terminal result the caller receives instead of business output. + * + * On a normalized request (one whose discovery entry advertises an + * output schema) this must be an `InputRequiredResult`, set `isError`, + * or carry `structuredContent` conforming to the advertised schema. + * mcp-utils passes the result through verbatim and never synthesizes + * `structuredContent`, so a content-only success on a normalized + * request contradicts what that request advertised. + */ + result: CallToolResult | InputRequiredResult; + telemetry: ToolPolicyTelemetry; + }; + +/** + * A pre-execution guard around one tool. + * + * The server applies the hooks in a fixed order: contextual discovery and + * schema selection, argument normalization, strict parsing, then `resolve`. + * Business execution runs only after `resolve` returns an `execute` decision. + */ +export type ToolPolicy = { + /** + * Replaces the advertised and enforced input schema for this request. + * + * It runs inside `tools/list`, once per tool, and must not throw: a throw + * fails the whole discovery response, not just this tool's entry. + */ + inputSchema?( + schema: z.ZodObject, + ctx: ToolRequestContext + ): z.ZodObject; + + /** + * Chooses the output schema this request advertises, and by doing so + * decides whether the request carries structured results at all. + * + * Returning a schema normalizes the request: its `tools/list` entry + * advertises that schema, its call result carries `structuredContent`, and + * the tool's `formatResult` renders the text. + * + * Returning `undefined` suppresses structured results for that request, + * restoring the whole pre-normalization result: no `outputSchema` key in + * its `tools/list` entry, no `structuredContent`, and single-encoded + * `JSON.stringify` text with `formatResult` skipped. Use it to hold one + * serving path or client generation on byte-exact legacy output while + * another is normalized. + * + * Defining no hook at all normalizes every request against the tool's own + * `outputSchema`. Discovery and the call path evaluate this once per + * request and always agree. + * + * Like `inputSchema`, this runs inside `tools/list`, once per tool, and must + * not throw: a throw fails the whole discovery response, not just this + * tool's entry. + * + * The author owes two things on a normalized request: the resolved schema + * must be object-rooted, because MCP restricts structured output to an + * object root, and `execute`'s output must conform to it as-is. mcp-utils + * checks both and answers `isError` rather than emitting output that + * contradicts what the request advertised. Conformance is identity, so a + * transforming, coercing or stripping schema is unsupported by + * construction: it would diverge from the business result on every call + * and therefore always error. + * + * A resolved schema that is not object-rooted fails the entire discovery + * response, by design: it is an authoring error, not a runtime condition, + * so every `tools/list` fails identically and it cannot reach production. + */ + outputSchema?( + schema: z.ZodObject, + ctx: ToolRequestContext + ): z.ZodType | undefined; + + /** Adjusts raw arguments before strict parsing rejects unknown fields. */ + normalizeArguments?(raw: unknown, ctx: ToolRequestContext): unknown; + + /** Decides whether the tool executes, and with what resolution. */ + resolve( + params: Params, + ctx: ToolRequestContext + ): Promise>; +};