diff --git a/src/server/index.test.ts b/src/server/index.test.ts index 2957b872..24543393 100644 --- a/src/server/index.test.ts +++ b/src/server/index.test.ts @@ -15,21 +15,10 @@ import { ListToolsRequestSchema, SetLevelRequestSchema, ErrorCode, - ListToolsResultSchema, - CallToolResultSchema, - ListResourcesResultSchema, - ListResourceTemplatesResultSchema, - ReadResourceResultSchema, - ListPromptsResultSchema, - GetPromptResultSchema, - CompleteResultSchema, } from "../types.js"; import { Transport } from "../shared/transport.js"; import { InMemoryTransport } from "../inMemory.js"; import { Client } from "../client/index.js"; -import { UriTemplate } from "../shared/uriTemplate.js"; -import { ResourceTemplate } from "./index.js"; -import { completable } from "./completable.js"; test("should accept latest protocol version", async () => { let sendPromiseResolve: (value: unknown) => void; @@ -557,1338 +546,3 @@ test("should handle request timeout", async () => { code: ErrorCode.RequestTimeout, }); }); - -describe("ResourceTemplate", () => { - test("should create ResourceTemplate with string pattern", () => { - const template = new ResourceTemplate("test://{category}/{id}", { - list: undefined, - }); - expect(template.uriTemplate.toString()).toBe("test://{category}/{id}"); - expect(template.listCallback).toBeUndefined(); - }); - - test("should create ResourceTemplate with UriTemplate", () => { - const uriTemplate = new UriTemplate("test://{category}/{id}"); - const template = new ResourceTemplate(uriTemplate, { list: undefined }); - expect(template.uriTemplate).toBe(uriTemplate); - expect(template.listCallback).toBeUndefined(); - }); - - test("should create ResourceTemplate with list callback", async () => { - const list = jest.fn().mockResolvedValue({ - resources: [{ name: "Test", uri: "test://example" }], - }); - - const template = new ResourceTemplate("test://{id}", { list }); - expect(template.listCallback).toBe(list); - - const abortController = new AbortController(); - const result = await template.listCallback?.({ - signal: abortController.signal, - }); - expect(result?.resources).toHaveLength(1); - expect(list).toHaveBeenCalled(); - }); -}); - -describe("Server.tool", () => { - test("should register zero-argument tool", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - const client = new Client({ - name: "test client", - version: "1.0", - }); - - server.tool("test", async () => ({ - content: [ - { - type: "text", - text: "Test response", - }, - ], - })); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - const result = await client.request( - { - method: "tools/list", - }, - ListToolsResultSchema, - ); - - expect(result.tools).toHaveLength(1); - expect(result.tools[0].name).toBe("test"); - expect(result.tools[0].inputSchema).toEqual({ - type: "object", - }); - }); - - test("should register tool with args schema", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - const client = new Client({ - name: "test client", - version: "1.0", - }); - - server.tool( - "test", - { - name: z.string(), - value: z.number(), - }, - async ({ name, value }) => ({ - content: [ - { - type: "text", - text: `${name}: ${value}`, - }, - ], - }), - ); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - const result = await client.request( - { - method: "tools/list", - }, - ListToolsResultSchema, - ); - - expect(result.tools).toHaveLength(1); - expect(result.tools[0].name).toBe("test"); - expect(result.tools[0].inputSchema).toMatchObject({ - type: "object", - properties: { - name: { type: "string" }, - value: { type: "number" }, - }, - }); - }); - - test("should register tool with description", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - const client = new Client({ - name: "test client", - version: "1.0", - }); - - server.tool("test", "Test description", async () => ({ - content: [ - { - type: "text", - text: "Test response", - }, - ], - })); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - const result = await client.request( - { - method: "tools/list", - }, - ListToolsResultSchema, - ); - - expect(result.tools).toHaveLength(1); - expect(result.tools[0].name).toBe("test"); - expect(result.tools[0].description).toBe("Test description"); - }); - - test("should validate tool args", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - - const client = new Client( - { - name: "test client", - version: "1.0", - }, - { - capabilities: { - tools: {}, - }, - }, - ); - - server.tool( - "test", - { - name: z.string(), - value: z.number(), - }, - async ({ name, value }) => ({ - content: [ - { - type: "text", - text: `${name}: ${value}`, - }, - ], - }), - ); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - await expect( - client.request( - { - method: "tools/call", - params: { - name: "test", - arguments: { - name: "test", - value: "not a number", - }, - }, - }, - CallToolResultSchema, - ), - ).rejects.toThrow(/Invalid arguments/); - }); - - test("should prevent duplicate tool registration", () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - - server.tool("test", async () => ({ - content: [ - { - type: "text", - text: "Test response", - }, - ], - })); - - expect(() => { - server.tool("test", async () => ({ - content: [ - { - type: "text", - text: "Test response 2", - }, - ], - })); - }).toThrow(/already registered/); - }); - - test("should allow client to call server tools", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - - const client = new Client( - { - name: "test client", - version: "1.0", - }, - { - capabilities: { - tools: {}, - }, - }, - ); - - server.tool( - "test", - "Test tool", - { - input: z.string(), - }, - async ({ input }) => ({ - content: [ - { - type: "text", - text: `Processed: ${input}`, - }, - ], - }), - ); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - const result = await client.request( - { - method: "tools/call", - params: { - name: "test", - arguments: { - input: "hello", - }, - }, - }, - CallToolResultSchema, - ); - - expect(result.content).toEqual([ - { - type: "text", - text: "Processed: hello", - }, - ]); - }); - - test("should handle server tool errors gracefully", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - - const client = new Client( - { - name: "test client", - version: "1.0", - }, - { - capabilities: { - tools: {}, - }, - }, - ); - - server.tool("error-test", async () => { - throw new Error("Tool execution failed"); - }); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - const result = await client.request( - { - method: "tools/call", - params: { - name: "error-test", - }, - }, - CallToolResultSchema, - ); - - expect(result.isError).toBe(true); - expect(result.content).toEqual([ - { - type: "text", - text: "Tool execution failed", - }, - ]); - }); - - test("should throw McpError for invalid tool name", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - - const client = new Client( - { - name: "test client", - version: "1.0", - }, - { - capabilities: { - tools: {}, - }, - }, - ); - - server.tool("test-tool", async () => ({ - content: [ - { - type: "text", - text: "Test response", - }, - ], - })); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - await expect( - client.request( - { - method: "tools/call", - params: { - name: "nonexistent-tool", - }, - }, - CallToolResultSchema, - ), - ).rejects.toThrow(/Tool nonexistent-tool not found/); - }); -}); - -describe("Server.resource", () => { - test("should register resource with uri and readCallback", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - const client = new Client({ - name: "test client", - version: "1.0", - }); - - server.resource("test", "test://resource", async () => ({ - contents: [ - { - uri: "test://resource", - text: "Test content", - }, - ], - })); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - const result = await client.request( - { - method: "resources/list", - }, - ListResourcesResultSchema, - ); - - expect(result.resources).toHaveLength(1); - expect(result.resources[0].name).toBe("test"); - expect(result.resources[0].uri).toBe("test://resource"); - }); - - test("should register resource with metadata", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - const client = new Client({ - name: "test client", - version: "1.0", - }); - - server.resource( - "test", - "test://resource", - { - description: "Test resource", - mimeType: "text/plain", - }, - async () => ({ - contents: [ - { - uri: "test://resource", - text: "Test content", - }, - ], - }), - ); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - const result = await client.request( - { - method: "resources/list", - }, - ListResourcesResultSchema, - ); - - expect(result.resources).toHaveLength(1); - expect(result.resources[0].description).toBe("Test resource"); - expect(result.resources[0].mimeType).toBe("text/plain"); - }); - - test("should register resource template", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - const client = new Client({ - name: "test client", - version: "1.0", - }); - - server.resource( - "test", - new ResourceTemplate("test://resource/{id}", { list: undefined }), - async () => ({ - contents: [ - { - uri: "test://resource/123", - text: "Test content", - }, - ], - }), - ); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - const result = await client.request( - { - method: "resources/templates/list", - }, - ListResourceTemplatesResultSchema, - ); - - expect(result.resourceTemplates).toHaveLength(1); - expect(result.resourceTemplates[0].name).toBe("test"); - expect(result.resourceTemplates[0].uriTemplate).toBe( - "test://resource/{id}", - ); - }); - - test("should register resource template with listCallback", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - const client = new Client({ - name: "test client", - version: "1.0", - }); - - server.resource( - "test", - new ResourceTemplate("test://resource/{id}", { - list: async () => ({ - resources: [ - { - name: "Resource 1", - uri: "test://resource/1", - }, - { - name: "Resource 2", - uri: "test://resource/2", - }, - ], - }), - }), - async (uri) => ({ - contents: [ - { - uri: uri.href, - text: "Test content", - }, - ], - }), - ); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - const result = await client.request( - { - method: "resources/list", - }, - ListResourcesResultSchema, - ); - - expect(result.resources).toHaveLength(2); - expect(result.resources[0].name).toBe("Resource 1"); - expect(result.resources[0].uri).toBe("test://resource/1"); - expect(result.resources[1].name).toBe("Resource 2"); - expect(result.resources[1].uri).toBe("test://resource/2"); - }); - - test("should pass template variables to readCallback", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - const client = new Client({ - name: "test client", - version: "1.0", - }); - - server.resource( - "test", - new ResourceTemplate("test://resource/{category}/{id}", { - list: undefined, - }), - async (uri, { category, id }) => ({ - contents: [ - { - uri: uri.href, - text: `Category: ${category}, ID: ${id}`, - }, - ], - }), - ); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - const result = await client.request( - { - method: "resources/read", - params: { - uri: "test://resource/books/123", - }, - }, - ReadResourceResultSchema, - ); - - expect(result.contents[0].text).toBe("Category: books, ID: 123"); - }); - - test("should prevent duplicate resource registration", () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - - server.resource("test", "test://resource", async () => ({ - contents: [ - { - uri: "test://resource", - text: "Test content", - }, - ], - })); - - expect(() => { - server.resource("test2", "test://resource", async () => ({ - contents: [ - { - uri: "test://resource", - text: "Test content 2", - }, - ], - })); - }).toThrow(/already registered/); - }); - - test("should prevent duplicate resource template registration", () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - - server.resource( - "test", - new ResourceTemplate("test://resource/{id}", { list: undefined }), - async () => ({ - contents: [ - { - uri: "test://resource/123", - text: "Test content", - }, - ], - }), - ); - - expect(() => { - server.resource( - "test", - new ResourceTemplate("test://resource/{id}", { list: undefined }), - async () => ({ - contents: [ - { - uri: "test://resource/123", - text: "Test content 2", - }, - ], - }), - ); - }).toThrow(/already registered/); - }); - - test("should handle resource read errors gracefully", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - const client = new Client({ - name: "test client", - version: "1.0", - }); - - server.resource("error-test", "test://error", async () => { - throw new Error("Resource read failed"); - }); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - await expect( - client.request( - { - method: "resources/read", - params: { - uri: "test://error", - }, - }, - ReadResourceResultSchema, - ), - ).rejects.toThrow(/Resource read failed/); - }); - - test("should throw McpError for invalid resource URI", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - const client = new Client({ - name: "test client", - version: "1.0", - }); - - server.resource("test", "test://resource", async () => ({ - contents: [ - { - uri: "test://resource", - text: "Test content", - }, - ], - })); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - await expect( - client.request( - { - method: "resources/read", - params: { - uri: "test://nonexistent", - }, - }, - ReadResourceResultSchema, - ), - ).rejects.toThrow(/Resource test:\/\/nonexistent not found/); - }); - - test("should support completion of resource template parameters", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - - const client = new Client( - { - name: "test client", - version: "1.0", - }, - { - capabilities: { - resources: {}, - }, - }, - ); - - server.resource( - "test", - new ResourceTemplate("test://resource/{category}", { - list: undefined, - complete: { - category: () => ["books", "movies", "music"], - }, - }), - async () => ({ - contents: [ - { - uri: "test://resource/test", - text: "Test content", - }, - ], - }), - ); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - const result = await client.request( - { - method: "completion/complete", - params: { - ref: { - type: "ref/resource", - uri: "test://resource/{category}", - }, - argument: { - name: "category", - value: "", - }, - }, - }, - CompleteResultSchema, - ); - - expect(result.completion.values).toEqual(["books", "movies", "music"]); - expect(result.completion.total).toBe(3); - }); - - test("should support filtered completion of resource template parameters", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - - const client = new Client( - { - name: "test client", - version: "1.0", - }, - { - capabilities: { - resources: {}, - }, - }, - ); - - server.resource( - "test", - new ResourceTemplate("test://resource/{category}", { - list: undefined, - complete: { - category: (test) => - ["books", "movies", "music"].filter((value) => - value.startsWith(test), - ), - }, - }), - async () => ({ - contents: [ - { - uri: "test://resource/test", - text: "Test content", - }, - ], - }), - ); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - const result = await client.request( - { - method: "completion/complete", - params: { - ref: { - type: "ref/resource", - uri: "test://resource/{category}", - }, - argument: { - name: "category", - value: "m", - }, - }, - }, - CompleteResultSchema, - ); - - expect(result.completion.values).toEqual(["movies", "music"]); - expect(result.completion.total).toBe(2); - }); -}); - -describe("Server.prompt", () => { - test("should register zero-argument prompt", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - const client = new Client({ - name: "test client", - version: "1.0", - }); - - server.prompt("test", async () => ({ - messages: [ - { - role: "assistant", - content: { - type: "text", - text: "Test response", - }, - }, - ], - })); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - const result = await client.request( - { - method: "prompts/list", - }, - ListPromptsResultSchema, - ); - - expect(result.prompts).toHaveLength(1); - expect(result.prompts[0].name).toBe("test"); - expect(result.prompts[0].arguments).toBeUndefined(); - }); - - test("should register prompt with args schema", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - const client = new Client({ - name: "test client", - version: "1.0", - }); - - server.prompt( - "test", - { - name: z.string(), - value: z.string(), - }, - async ({ name, value }) => ({ - messages: [ - { - role: "assistant", - content: { - type: "text", - text: `${name}: ${value}`, - }, - }, - ], - }), - ); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - const result = await client.request( - { - method: "prompts/list", - }, - ListPromptsResultSchema, - ); - - expect(result.prompts).toHaveLength(1); - expect(result.prompts[0].name).toBe("test"); - expect(result.prompts[0].arguments).toEqual([ - { name: "name", required: true }, - { name: "value", required: true }, - ]); - }); - - test("should register prompt with description", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - const client = new Client({ - name: "test client", - version: "1.0", - }); - - server.prompt("test", "Test description", async () => ({ - messages: [ - { - role: "assistant", - content: { - type: "text", - text: "Test response", - }, - }, - ], - })); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - const result = await client.request( - { - method: "prompts/list", - }, - ListPromptsResultSchema, - ); - - expect(result.prompts).toHaveLength(1); - expect(result.prompts[0].name).toBe("test"); - expect(result.prompts[0].description).toBe("Test description"); - }); - - test("should validate prompt args", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - - const client = new Client( - { - name: "test client", - version: "1.0", - }, - { - capabilities: { - prompts: {}, - }, - }, - ); - - server.prompt( - "test", - { - name: z.string(), - value: z.string().min(3), - }, - async ({ name, value }) => ({ - messages: [ - { - role: "assistant", - content: { - type: "text", - text: `${name}: ${value}`, - }, - }, - ], - }), - ); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - await expect( - client.request( - { - method: "prompts/get", - params: { - name: "test", - arguments: { - name: "test", - value: "ab", // Too short - }, - }, - }, - GetPromptResultSchema, - ), - ).rejects.toThrow(/Invalid arguments/); - }); - - test("should prevent duplicate prompt registration", () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - - server.prompt("test", async () => ({ - messages: [ - { - role: "assistant", - content: { - type: "text", - text: "Test response", - }, - }, - ], - })); - - expect(() => { - server.prompt("test", async () => ({ - messages: [ - { - role: "assistant", - content: { - type: "text", - text: "Test response 2", - }, - }, - ], - })); - }).toThrow(/already registered/); - }); - - test("should throw McpError for invalid prompt name", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - - const client = new Client( - { - name: "test client", - version: "1.0", - }, - { - capabilities: { - prompts: {}, - }, - }, - ); - - server.prompt("test-prompt", async () => ({ - messages: [ - { - role: "assistant", - content: { - type: "text", - text: "Test response", - }, - }, - ], - })); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - await expect( - client.request( - { - method: "prompts/get", - params: { - name: "nonexistent-prompt", - }, - }, - GetPromptResultSchema, - ), - ).rejects.toThrow(/Prompt nonexistent-prompt not found/); - }); - test("should support completion of prompt arguments", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - - const client = new Client( - { - name: "test client", - version: "1.0", - }, - { - capabilities: { - prompts: {}, - }, - }, - ); - - server.prompt( - "test-prompt", - { - name: completable(z.string(), () => ["Alice", "Bob", "Charlie"]), - }, - async ({ name }) => ({ - messages: [ - { - role: "assistant", - content: { - type: "text", - text: `Hello ${name}`, - }, - }, - ], - }), - ); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - const result = await client.request( - { - method: "completion/complete", - params: { - ref: { - type: "ref/prompt", - name: "test-prompt", - }, - argument: { - name: "name", - value: "", - }, - }, - }, - CompleteResultSchema, - ); - - expect(result.completion.values).toEqual(["Alice", "Bob", "Charlie"]); - expect(result.completion.total).toBe(3); - }); - - test("should support filtered completion of prompt arguments", async () => { - const server = new Server({ - name: "test server", - version: "1.0", - }); - - const client = new Client( - { - name: "test client", - version: "1.0", - }, - { - capabilities: { - prompts: {}, - }, - }, - ); - - server.prompt( - "test-prompt", - { - name: completable(z.string(), (test) => - ["Alice", "Bob", "Charlie"].filter((value) => value.startsWith(test)), - ), - }, - async ({ name }) => ({ - messages: [ - { - role: "assistant", - content: { - type: "text", - text: `Hello ${name}`, - }, - }, - ], - }), - ); - - const [clientTransport, serverTransport] = - InMemoryTransport.createLinkedPair(); - - await Promise.all([ - client.connect(clientTransport), - server.connect(serverTransport), - ]); - - const result = await client.request( - { - method: "completion/complete", - params: { - ref: { - type: "ref/prompt", - name: "test-prompt", - }, - argument: { - name: "name", - value: "A", - }, - }, - }, - CompleteResultSchema, - ); - - expect(result.completion.values).toEqual(["Alice"]); - expect(result.completion.total).toBe(1); - }); -}); diff --git a/src/server/index.ts b/src/server/index.ts index 1810883a..bf8cd6b2 100644 --- a/src/server/index.ts +++ b/src/server/index.ts @@ -1,60 +1,25 @@ -import z, { - AnyZodObject, - ZodObject, - ZodOptional, - ZodRawShape, - ZodString, - ZodType, - ZodTypeAny, - ZodTypeDef, -} from "zod"; -import { zodToJsonSchema } from "zod-to-json-schema"; import { mergeCapabilities, Protocol, ProtocolOptions, - RequestHandlerExtra, RequestOptions, } from "../shared/protocol.js"; import { - CallToolRequestSchema, - CallToolResult, ClientCapabilities, - CompleteRequest, - CompleteRequestSchema, - CompleteResult, CreateMessageRequest, CreateMessageResultSchema, EmptyResultSchema, - ErrorCode, - GetPromptRequestSchema, - GetPromptResult, Implementation, InitializedNotificationSchema, InitializeRequest, InitializeRequestSchema, InitializeResult, LATEST_PROTOCOL_VERSION, - ListPromptsRequestSchema, - ListPromptsResult, - ListResourcesRequestSchema, - ListResourcesResult, - ListResourceTemplatesRequestSchema, ListRootsRequest, ListRootsResultSchema, - ListToolsRequestSchema, - ListToolsResult, LoggingMessageNotification, - McpError, Notification, - Prompt, - PromptArgument, - PromptReference, - ReadResourceRequestSchema, - ReadResourceResult, Request, - Resource, - ResourceReference, ResourceUpdatedNotification, Result, ServerCapabilities, @@ -62,10 +27,7 @@ import { ServerRequest, ServerResult, SUPPORTED_PROTOCOL_VERSIONS, - Tool, } from "../types.js"; -import { UriTemplate, Variables } from "../shared/uriTemplate.js"; -import { Completable, CompletableDef } from "./completable.js"; export type ServerOptions = ProtocolOptions & { /** @@ -111,12 +73,6 @@ export class Server< private _clientCapabilities?: ClientCapabilities; private _clientVersion?: Implementation; private _capabilities: ServerCapabilities; - private _registeredResources: { [uri: string]: RegisteredResource } = {}; - private _registeredResourceTemplates: { - [name: string]: RegisteredResourceTemplate; - } = {}; - private _registeredTools: { [name: string]: RegisteredTool } = {}; - private _registeredPrompts: { [name: string]: RegisteredPrompt } = {}; /** * Callback for when initialization has fully completed (i.e., the client has sent an `initialized` notification). @@ -365,701 +321,4 @@ export class Server< async sendPromptListChanged() { return this.notification({ method: "notifications/prompts/list_changed" }); } - - private setToolRequestHandlers() { - this.assertCanSetRequestHandler(ListToolsRequestSchema.shape.method.value); - this.assertCanSetRequestHandler(CallToolRequestSchema.shape.method.value); - - this.registerCapabilities({ - tools: {}, - }); - - this.setRequestHandler( - ListToolsRequestSchema, - (): ListToolsResult => ({ - tools: Object.entries(this._registeredTools).map( - ([name, tool]): Tool => { - return { - name, - description: tool.description, - inputSchema: tool.inputSchema - ? (zodToJsonSchema(tool.inputSchema) as Tool["inputSchema"]) - : EMPTY_OBJECT_JSON_SCHEMA, - }; - }, - ), - }), - ); - - this.setRequestHandler( - CallToolRequestSchema, - async (request, extra): Promise => { - const tool = this._registeredTools[request.params.name]; - if (!tool) { - throw new McpError( - ErrorCode.InvalidParams, - `Tool ${request.params.name} not found`, - ); - } - - if (tool.inputSchema) { - const parseResult = await tool.inputSchema.safeParseAsync( - request.params.arguments, - ); - if (!parseResult.success) { - throw new McpError( - ErrorCode.InvalidParams, - `Invalid arguments for tool ${request.params.name}: ${parseResult.error.message}`, - ); - } - - const args = parseResult.data; - const cb = tool.callback as ToolCallback; - try { - return await Promise.resolve(cb(args, extra)); - } catch (error) { - return { - content: [ - { - type: "text", - text: error instanceof Error ? error.message : String(error), - }, - ], - isError: true, - }; - } - } else { - const cb = tool.callback as ToolCallback; - try { - return await Promise.resolve(cb(extra)); - } catch (error) { - return { - content: [ - { - type: "text", - text: error instanceof Error ? error.message : String(error), - }, - ], - isError: true, - }; - } - } - }, - ); - } - - /** - * Registers a zero-argument tool `name`, which will run the given function when the client calls it. - */ - tool(name: string, cb: ToolCallback): void; - - /** - * Registers a zero-argument tool `name` (with a description) which will run the given function when the client calls it. - */ - tool(name: string, description: string, cb: ToolCallback): void; - - /** - * Registers a tool `name` accepting the given arguments, which must be an object containing named properties associated with Zod schemas. When the client calls it, the function will be run with the parsed and validated arguments. - */ - tool( - name: string, - paramsSchema: Args, - cb: ToolCallback, - ): void; - - /** - * Registers a tool `name` (with a description) accepting the given arguments, which must be an object containing named properties associated with Zod schemas. When the client calls it, the function will be run with the parsed and validated arguments. - */ - tool( - name: string, - description: string, - paramsSchema: Args, - cb: ToolCallback, - ): void; - - tool(name: string, ...rest: unknown[]): void { - if (this._registeredTools[name]) { - throw new Error(`Tool ${name} is already registered`); - } - - let description: string | undefined; - if (typeof rest[0] === "string") { - description = rest.shift() as string; - } - - let paramsSchema: ZodRawShape | undefined; - if (rest.length > 1) { - paramsSchema = rest.shift() as ZodRawShape; - } - - const cb = rest[0] as ToolCallback; - this._registeredTools[name] = { - description, - inputSchema: - paramsSchema === undefined ? undefined : z.object(paramsSchema), - callback: cb, - }; - - this.setToolRequestHandlers(); - } - - private setCompletionRequestHandler() { - this.assertCanSetRequestHandler(CompleteRequestSchema.shape.method.value); - - this.setRequestHandler( - CompleteRequestSchema, - async (request): Promise => { - switch (request.params.ref.type) { - case "ref/prompt": - return this.handlePromptCompletion(request, request.params.ref); - - case "ref/resource": - return this.handleResourceCompletion(request, request.params.ref); - - default: - throw new McpError( - ErrorCode.InvalidParams, - `Invalid completion reference: ${request.params.ref}`, - ); - } - }, - ); - } - - private async handlePromptCompletion( - request: CompleteRequest, - ref: PromptReference, - ): Promise { - const prompt = this._registeredPrompts[ref.name]; - if (!prompt) { - throw new McpError( - ErrorCode.InvalidParams, - `Prompt ${request.params.ref.name} not found`, - ); - } - - if (!prompt.argsSchema) { - return EMPTY_COMPLETION_RESULT; - } - - const field = prompt.argsSchema.shape[request.params.argument.name]; - if (!(field instanceof Completable)) { - return EMPTY_COMPLETION_RESULT; - } - - const def: CompletableDef = field._def; - const suggestions = await def.complete(request.params.argument.value); - return createCompletionResult(suggestions); - } - - private async handleResourceCompletion( - request: CompleteRequest, - ref: ResourceReference, - ): Promise { - const template = Object.values(this._registeredResourceTemplates).find( - (t) => t.resourceTemplate.uriTemplate.toString() === ref.uri, - ); - - if (!template) { - if (this._registeredResources[ref.uri]) { - // Attempting to autocomplete a fixed resource URI is not an error in the spec (but probably should be). - return EMPTY_COMPLETION_RESULT; - } - - throw new McpError( - ErrorCode.InvalidParams, - `Resource template ${request.params.ref.uri} not found`, - ); - } - - const completer = template.resourceTemplate.completeCallback( - request.params.argument.name, - ); - if (!completer) { - return EMPTY_COMPLETION_RESULT; - } - - const suggestions = await completer(request.params.argument.value); - return createCompletionResult(suggestions); - } - - private setResourceRequestHandlers() { - this.assertCanSetRequestHandler( - ListResourcesRequestSchema.shape.method.value, - ); - this.assertCanSetRequestHandler( - ListResourceTemplatesRequestSchema.shape.method.value, - ); - this.assertCanSetRequestHandler( - ReadResourceRequestSchema.shape.method.value, - ); - - this.registerCapabilities({ - resources: {}, - }); - - this.setRequestHandler( - ListResourcesRequestSchema, - async (request, extra) => { - const resources = Object.entries(this._registeredResources).map( - ([uri, resource]) => ({ - uri, - name: resource.name, - ...resource.metadata, - }), - ); - - const templateResources: Resource[] = []; - for (const template of Object.values( - this._registeredResourceTemplates, - )) { - if (!template.resourceTemplate.listCallback) { - continue; - } - - const result = await template.resourceTemplate.listCallback(extra); - for (const resource of result.resources) { - templateResources.push({ - ...resource, - ...template.metadata, - }); - } - } - - return { resources: [...resources, ...templateResources] }; - }, - ); - - this.setRequestHandler(ListResourceTemplatesRequestSchema, async () => { - const resourceTemplates = Object.entries( - this._registeredResourceTemplates, - ).map(([name, template]) => ({ - name, - uriTemplate: template.resourceTemplate.uriTemplate.toString(), - ...template.metadata, - })); - - return { resourceTemplates }; - }); - - this.setRequestHandler( - ReadResourceRequestSchema, - async (request, extra) => { - const uri = new URL(request.params.uri); - - // First check for exact resource match - const resource = this._registeredResources[uri.toString()]; - if (resource) { - return resource.readCallback(uri, extra); - } - - // Then check templates - for (const template of Object.values( - this._registeredResourceTemplates, - )) { - const variables = template.resourceTemplate.uriTemplate.match( - uri.toString(), - ); - if (variables) { - return template.readCallback(uri, variables, extra); - } - } - - throw new McpError( - ErrorCode.InvalidParams, - `Resource ${uri} not found`, - ); - }, - ); - - this.setCompletionRequestHandler(); - } - - /** - * Registers a resource `name` at a fixed URI, which will use the given callback to respond to read requests. - */ - resource(name: string, uri: string, readCallback: ReadResourceCallback): void; - - /** - * Registers a resource `name` at a fixed URI with metadata, which will use the given callback to respond to read requests. - */ - resource( - name: string, - uri: string, - metadata: ResourceMetadata, - readCallback: ReadResourceCallback, - ): void; - - /** - * Registers a resource `name` with a template pattern, which will use the given callback to respond to read requests. - */ - resource( - name: string, - template: ResourceTemplate, - readCallback: ReadResourceTemplateCallback, - ): void; - - /** - * Registers a resource `name` with a template pattern and metadata, which will use the given callback to respond to read requests. - */ - resource( - name: string, - template: ResourceTemplate, - metadata: ResourceMetadata, - readCallback: ReadResourceTemplateCallback, - ): void; - - resource( - name: string, - uriOrTemplate: string | ResourceTemplate, - ...rest: unknown[] - ): void { - let metadata: ResourceMetadata | undefined; - if (typeof rest[0] === "object") { - metadata = rest.shift() as ResourceMetadata; - } - - const readCallback = rest[0] as - | ReadResourceCallback - | ReadResourceTemplateCallback; - - if (typeof uriOrTemplate === "string") { - this.registerResource({ - name, - uri: uriOrTemplate, - metadata, - readCallback: readCallback as ReadResourceCallback, - }); - } else { - this.registerResourceTemplate({ - name, - resourceTemplate: uriOrTemplate, - metadata, - readCallback: readCallback as ReadResourceTemplateCallback, - }); - } - } - - private registerResource({ - name, - uri, - metadata, - readCallback, - }: { - name: string; - uri: string; - metadata?: ResourceMetadata; - readCallback: ReadResourceCallback; - }): void { - if (this._registeredResources[uri]) { - throw new Error(`Resource ${uri} is already registered`); - } - - this._registeredResources[uri] = { - name, - metadata, - readCallback, - }; - - this.setResourceRequestHandlers(); - } - - private registerResourceTemplate({ - name, - resourceTemplate, - metadata, - readCallback, - }: { - name: string; - resourceTemplate: ResourceTemplate; - metadata?: ResourceMetadata; - readCallback: ReadResourceTemplateCallback; - }): void { - if (this._registeredResourceTemplates[name]) { - throw new Error(`Resource template ${name} is already registered`); - } - - this._registeredResourceTemplates[name] = { - resourceTemplate, - metadata, - readCallback, - }; - - this.setResourceRequestHandlers(); - } - - private setPromptRequestHandlers() { - this.assertCanSetRequestHandler( - ListPromptsRequestSchema.shape.method.value, - ); - this.assertCanSetRequestHandler(GetPromptRequestSchema.shape.method.value); - - this.registerCapabilities({ - prompts: {}, - }); - - this.setRequestHandler( - ListPromptsRequestSchema, - (): ListPromptsResult => ({ - prompts: Object.entries(this._registeredPrompts).map( - ([name, prompt]): Prompt => { - return { - name, - description: prompt.description, - arguments: prompt.argsSchema - ? promptArgumentsFromSchema(prompt.argsSchema) - : undefined, - }; - }, - ), - }), - ); - - this.setRequestHandler( - GetPromptRequestSchema, - async (request, extra): Promise => { - const prompt = this._registeredPrompts[request.params.name]; - if (!prompt) { - throw new McpError( - ErrorCode.InvalidParams, - `Prompt ${request.params.name} not found`, - ); - } - - if (prompt.argsSchema) { - const parseResult = await prompt.argsSchema.safeParseAsync( - request.params.arguments, - ); - if (!parseResult.success) { - throw new McpError( - ErrorCode.InvalidParams, - `Invalid arguments for prompt ${request.params.name}: ${parseResult.error.message}`, - ); - } - - const args = parseResult.data; - const cb = prompt.callback as PromptCallback; - return await Promise.resolve(cb(args, extra)); - } else { - const cb = prompt.callback as PromptCallback; - return await Promise.resolve(cb(extra)); - } - }, - ); - - this.setCompletionRequestHandler(); - } - - prompt(name: string, cb: PromptCallback): void; - prompt(name: string, description: string, cb: PromptCallback): void; - prompt( - name: string, - argsSchema: Args, - cb: PromptCallback, - ): void; - prompt( - name: string, - description: string, - argsSchema: Args, - cb: PromptCallback, - ): void; - - prompt(name: string, ...rest: unknown[]): void { - if (this._registeredPrompts[name]) { - throw new Error(`Prompt ${name} is already registered`); - } - - let description: string | undefined; - if (typeof rest[0] === "string") { - description = rest.shift() as string; - } - - let argsSchema: PromptArgsRawShape | undefined; - if (rest.length > 1) { - argsSchema = rest.shift() as PromptArgsRawShape; - } - - const cb = rest[0] as PromptCallback; - this._registeredPrompts[name] = { - description, - argsSchema: argsSchema === undefined ? undefined : z.object(argsSchema), - callback: cb, - }; - - this.setPromptRequestHandlers(); - } } - -/** - * A callback to complete one variable within a resource template's URI template. - */ -export type CompleteResourceTemplateCallback = ( - value: string, -) => string[] | Promise; - -/** - * A resource template combines a URI pattern with optional functionality to enumerate - * all resources matching that pattern. - */ -export class ResourceTemplate { - private _uriTemplate: UriTemplate; - - constructor( - uriTemplate: string | UriTemplate, - private _callbacks: { - /** - * A callback to list all resources matching this template. This is required to specified, even if `undefined`, to avoid accidentally forgetting resource listing. - */ - list: ListResourcesCallback | undefined; - - /** - * An optional callback to autocomplete variables within the URI template. Useful for clients and users to discover possible values. - */ - complete?: { - [variable: string]: CompleteResourceTemplateCallback; - }; - }, - ) { - this._uriTemplate = - typeof uriTemplate === "string" - ? new UriTemplate(uriTemplate) - : uriTemplate; - } - - /** - * Gets the URI template pattern. - */ - get uriTemplate(): UriTemplate { - return this._uriTemplate; - } - - /** - * Gets the list callback, if one was provided. - */ - get listCallback(): ListResourcesCallback | undefined { - return this._callbacks.list; - } - - /** - * Gets the callback for completing a specific URI template variable, if one was provided. - */ - completeCallback( - variable: string, - ): CompleteResourceTemplateCallback | undefined { - return this._callbacks.complete?.[variable]; - } -} - -/** - * Callback for a tool handler registered with Server.tool(). - * - * Parameters will include tool arguments, if applicable, as well as other request handler context. - */ -export type ToolCallback = - Args extends ZodRawShape - ? ( - args: z.objectOutputType, - extra: RequestHandlerExtra, - ) => CallToolResult | Promise - : (extra: RequestHandlerExtra) => CallToolResult | Promise; - -type RegisteredTool = { - description?: string; - inputSchema?: AnyZodObject; - callback: ToolCallback; -}; - -const EMPTY_OBJECT_JSON_SCHEMA = { - type: "object" as const, -}; - -/** - * Additional, optional information for annotating a resource. - */ -export type ResourceMetadata = Omit; - -/** - * Callback to list all resources matching a given template. - */ -export type ListResourcesCallback = ( - extra: RequestHandlerExtra, -) => ListResourcesResult | Promise; - -/** - * Callback to read a resource at a given URI. - */ -export type ReadResourceCallback = ( - uri: URL, - extra: RequestHandlerExtra, -) => ReadResourceResult | Promise; - -type RegisteredResource = { - name: string; - metadata?: ResourceMetadata; - readCallback: ReadResourceCallback; -}; - -/** - * Callback to read a resource at a given URI, following a filled-in URI template. - */ -export type ReadResourceTemplateCallback = ( - uri: URL, - variables: Variables, - extra: RequestHandlerExtra, -) => ReadResourceResult | Promise; - -type RegisteredResourceTemplate = { - resourceTemplate: ResourceTemplate; - metadata?: ResourceMetadata; - readCallback: ReadResourceTemplateCallback; -}; - -type PromptArgsRawShape = { - [k: string]: - | ZodType - | ZodOptional>; -}; - -export type PromptCallback< - Args extends undefined | PromptArgsRawShape = undefined, -> = Args extends PromptArgsRawShape - ? ( - args: z.objectOutputType, - extra: RequestHandlerExtra, - ) => GetPromptResult | Promise - : (extra: RequestHandlerExtra) => GetPromptResult | Promise; - -type RegisteredPrompt = { - description?: string; - argsSchema?: ZodObject; - callback: PromptCallback; -}; - -function promptArgumentsFromSchema( - schema: ZodObject, -): PromptArgument[] { - return Object.entries(schema.shape).map( - ([name, field]): PromptArgument => ({ - name, - description: field.description, - required: !field.isOptional(), - }), - ); -} - -function createCompletionResult(suggestions: string[]): CompleteResult { - return { - completion: { - values: suggestions.slice(0, 100), - total: suggestions.length, - hasMore: suggestions.length > 100, - }, - }; -} - -const EMPTY_COMPLETION_RESULT: CompleteResult = { - completion: { - values: [], - hasMore: false, - }, -}; diff --git a/src/server/mcp.test.ts b/src/server/mcp.test.ts new file mode 100644 index 00000000..33fdeb91 --- /dev/null +++ b/src/server/mcp.test.ts @@ -0,0 +1,1395 @@ +import { McpServer } from "./mcp.js"; +import { Client } from "../client/index.js"; +import { InMemoryTransport } from "../inMemory.js"; +import { z } from "zod"; +import { + ListToolsResultSchema, + CallToolResultSchema, + ListResourcesResultSchema, + ListResourceTemplatesResultSchema, + ReadResourceResultSchema, + ListPromptsResultSchema, + GetPromptResultSchema, + CompleteResultSchema, +} from "../types.js"; +import { ResourceTemplate } from "./mcp.js"; +import { completable } from "./completable.js"; +import { UriTemplate } from "../shared/uriTemplate.js"; + +describe("McpServer", () => { + test("should expose underlying Server instance", () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + + expect(mcpServer.server).toBeDefined(); + }); + + test("should allow sending notifications via Server", async () => { + const mcpServer = new McpServer( + { + name: "test server", + version: "1.0", + }, + { capabilities: { logging: {} } }, + ); + + const client = new Client({ + name: "test client", + version: "1.0", + }); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + // This should work because we're using the underlying server + await expect( + mcpServer.server.sendLoggingMessage({ + level: "info", + data: "Test log message", + }), + ).resolves.not.toThrow(); + }); +}); + +describe("ResourceTemplate", () => { + test("should create ResourceTemplate with string pattern", () => { + const template = new ResourceTemplate("test://{category}/{id}", { + list: undefined, + }); + expect(template.uriTemplate.toString()).toBe("test://{category}/{id}"); + expect(template.listCallback).toBeUndefined(); + }); + + test("should create ResourceTemplate with UriTemplate", () => { + const uriTemplate = new UriTemplate("test://{category}/{id}"); + const template = new ResourceTemplate(uriTemplate, { list: undefined }); + expect(template.uriTemplate).toBe(uriTemplate); + expect(template.listCallback).toBeUndefined(); + }); + + test("should create ResourceTemplate with list callback", async () => { + const list = jest.fn().mockResolvedValue({ + resources: [{ name: "Test", uri: "test://example" }], + }); + + const template = new ResourceTemplate("test://{id}", { list }); + expect(template.listCallback).toBe(list); + + const abortController = new AbortController(); + const result = await template.listCallback?.({ + signal: abortController.signal, + }); + expect(result?.resources).toHaveLength(1); + expect(list).toHaveBeenCalled(); + }); +}); + +describe("tool()", () => { + test("should register zero-argument tool", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const client = new Client({ + name: "test client", + version: "1.0", + }); + + mcpServer.tool("test", async () => ({ + content: [ + { + type: "text", + text: "Test response", + }, + ], + })); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "tools/list", + }, + ListToolsResultSchema, + ); + + expect(result.tools).toHaveLength(1); + expect(result.tools[0].name).toBe("test"); + expect(result.tools[0].inputSchema).toEqual({ + type: "object", + }); + }); + + test("should register tool with args schema", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const client = new Client({ + name: "test client", + version: "1.0", + }); + + mcpServer.tool( + "test", + { + name: z.string(), + value: z.number(), + }, + async ({ name, value }) => ({ + content: [ + { + type: "text", + text: `${name}: ${value}`, + }, + ], + }), + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "tools/list", + }, + ListToolsResultSchema, + ); + + expect(result.tools).toHaveLength(1); + expect(result.tools[0].name).toBe("test"); + expect(result.tools[0].inputSchema).toMatchObject({ + type: "object", + properties: { + name: { type: "string" }, + value: { type: "number" }, + }, + }); + }); + + test("should register tool with description", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const client = new Client({ + name: "test client", + version: "1.0", + }); + + mcpServer.tool("test", "Test description", async () => ({ + content: [ + { + type: "text", + text: "Test response", + }, + ], + })); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "tools/list", + }, + ListToolsResultSchema, + ); + + expect(result.tools).toHaveLength(1); + expect(result.tools[0].name).toBe("test"); + expect(result.tools[0].description).toBe("Test description"); + }); + + test("should validate tool args", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + + const client = new Client( + { + name: "test client", + version: "1.0", + }, + { + capabilities: { + tools: {}, + }, + }, + ); + + mcpServer.tool( + "test", + { + name: z.string(), + value: z.number(), + }, + async ({ name, value }) => ({ + content: [ + { + type: "text", + text: `${name}: ${value}`, + }, + ], + }), + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + await expect( + client.request( + { + method: "tools/call", + params: { + name: "test", + arguments: { + name: "test", + value: "not a number", + }, + }, + }, + CallToolResultSchema, + ), + ).rejects.toThrow(/Invalid arguments/); + }); + + test("should prevent duplicate tool registration", () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + + mcpServer.tool("test", async () => ({ + content: [ + { + type: "text", + text: "Test response", + }, + ], + })); + + expect(() => { + mcpServer.tool("test", async () => ({ + content: [ + { + type: "text", + text: "Test response 2", + }, + ], + })); + }).toThrow(/already registered/); + }); + + test("should allow client to call server tools", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + + const client = new Client( + { + name: "test client", + version: "1.0", + }, + { + capabilities: { + tools: {}, + }, + }, + ); + + mcpServer.tool( + "test", + "Test tool", + { + input: z.string(), + }, + async ({ input }) => ({ + content: [ + { + type: "text", + text: `Processed: ${input}`, + }, + ], + }), + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "tools/call", + params: { + name: "test", + arguments: { + input: "hello", + }, + }, + }, + CallToolResultSchema, + ); + + expect(result.content).toEqual([ + { + type: "text", + text: "Processed: hello", + }, + ]); + }); + + test("should handle server tool errors gracefully", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + + const client = new Client( + { + name: "test client", + version: "1.0", + }, + { + capabilities: { + tools: {}, + }, + }, + ); + + mcpServer.tool("error-test", async () => { + throw new Error("Tool execution failed"); + }); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "tools/call", + params: { + name: "error-test", + }, + }, + CallToolResultSchema, + ); + + expect(result.isError).toBe(true); + expect(result.content).toEqual([ + { + type: "text", + text: "Tool execution failed", + }, + ]); + }); + + test("should throw McpError for invalid tool name", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + + const client = new Client( + { + name: "test client", + version: "1.0", + }, + { + capabilities: { + tools: {}, + }, + }, + ); + + mcpServer.tool("test-tool", async () => ({ + content: [ + { + type: "text", + text: "Test response", + }, + ], + })); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + await expect( + client.request( + { + method: "tools/call", + params: { + name: "nonexistent-tool", + }, + }, + CallToolResultSchema, + ), + ).rejects.toThrow(/Tool nonexistent-tool not found/); + }); +}); + +describe("resource()", () => { + test("should register resource with uri and readCallback", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const client = new Client({ + name: "test client", + version: "1.0", + }); + + mcpServer.resource("test", "test://resource", async () => ({ + contents: [ + { + uri: "test://resource", + text: "Test content", + }, + ], + })); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "resources/list", + }, + ListResourcesResultSchema, + ); + + expect(result.resources).toHaveLength(1); + expect(result.resources[0].name).toBe("test"); + expect(result.resources[0].uri).toBe("test://resource"); + }); + + test("should register resource with metadata", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const client = new Client({ + name: "test client", + version: "1.0", + }); + + mcpServer.resource( + "test", + "test://resource", + { + description: "Test resource", + mimeType: "text/plain", + }, + async () => ({ + contents: [ + { + uri: "test://resource", + text: "Test content", + }, + ], + }), + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "resources/list", + }, + ListResourcesResultSchema, + ); + + expect(result.resources).toHaveLength(1); + expect(result.resources[0].description).toBe("Test resource"); + expect(result.resources[0].mimeType).toBe("text/plain"); + }); + + test("should register resource template", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const client = new Client({ + name: "test client", + version: "1.0", + }); + + mcpServer.resource( + "test", + new ResourceTemplate("test://resource/{id}", { list: undefined }), + async () => ({ + contents: [ + { + uri: "test://resource/123", + text: "Test content", + }, + ], + }), + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "resources/templates/list", + }, + ListResourceTemplatesResultSchema, + ); + + expect(result.resourceTemplates).toHaveLength(1); + expect(result.resourceTemplates[0].name).toBe("test"); + expect(result.resourceTemplates[0].uriTemplate).toBe( + "test://resource/{id}", + ); + }); + + test("should register resource template with listCallback", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const client = new Client({ + name: "test client", + version: "1.0", + }); + + mcpServer.resource( + "test", + new ResourceTemplate("test://resource/{id}", { + list: async () => ({ + resources: [ + { + name: "Resource 1", + uri: "test://resource/1", + }, + { + name: "Resource 2", + uri: "test://resource/2", + }, + ], + }), + }), + async (uri) => ({ + contents: [ + { + uri: uri.href, + text: "Test content", + }, + ], + }), + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "resources/list", + }, + ListResourcesResultSchema, + ); + + expect(result.resources).toHaveLength(2); + expect(result.resources[0].name).toBe("Resource 1"); + expect(result.resources[0].uri).toBe("test://resource/1"); + expect(result.resources[1].name).toBe("Resource 2"); + expect(result.resources[1].uri).toBe("test://resource/2"); + }); + + test("should pass template variables to readCallback", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const client = new Client({ + name: "test client", + version: "1.0", + }); + + mcpServer.resource( + "test", + new ResourceTemplate("test://resource/{category}/{id}", { + list: undefined, + }), + async (uri, { category, id }) => ({ + contents: [ + { + uri: uri.href, + text: `Category: ${category}, ID: ${id}`, + }, + ], + }), + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "resources/read", + params: { + uri: "test://resource/books/123", + }, + }, + ReadResourceResultSchema, + ); + + expect(result.contents[0].text).toBe("Category: books, ID: 123"); + }); + + test("should prevent duplicate resource registration", () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + + mcpServer.resource("test", "test://resource", async () => ({ + contents: [ + { + uri: "test://resource", + text: "Test content", + }, + ], + })); + + expect(() => { + mcpServer.resource("test2", "test://resource", async () => ({ + contents: [ + { + uri: "test://resource", + text: "Test content 2", + }, + ], + })); + }).toThrow(/already registered/); + }); + + test("should prevent duplicate resource template registration", () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + + mcpServer.resource( + "test", + new ResourceTemplate("test://resource/{id}", { list: undefined }), + async () => ({ + contents: [ + { + uri: "test://resource/123", + text: "Test content", + }, + ], + }), + ); + + expect(() => { + mcpServer.resource( + "test", + new ResourceTemplate("test://resource/{id}", { list: undefined }), + async () => ({ + contents: [ + { + uri: "test://resource/123", + text: "Test content 2", + }, + ], + }), + ); + }).toThrow(/already registered/); + }); + + test("should handle resource read errors gracefully", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const client = new Client({ + name: "test client", + version: "1.0", + }); + + mcpServer.resource("error-test", "test://error", async () => { + throw new Error("Resource read failed"); + }); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + await expect( + client.request( + { + method: "resources/read", + params: { + uri: "test://error", + }, + }, + ReadResourceResultSchema, + ), + ).rejects.toThrow(/Resource read failed/); + }); + + test("should throw McpError for invalid resource URI", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const client = new Client({ + name: "test client", + version: "1.0", + }); + + mcpServer.resource("test", "test://resource", async () => ({ + contents: [ + { + uri: "test://resource", + text: "Test content", + }, + ], + })); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + await expect( + client.request( + { + method: "resources/read", + params: { + uri: "test://nonexistent", + }, + }, + ReadResourceResultSchema, + ), + ).rejects.toThrow(/Resource test:\/\/nonexistent not found/); + }); + + test("should support completion of resource template parameters", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + + const client = new Client( + { + name: "test client", + version: "1.0", + }, + { + capabilities: { + resources: {}, + }, + }, + ); + + mcpServer.resource( + "test", + new ResourceTemplate("test://resource/{category}", { + list: undefined, + complete: { + category: () => ["books", "movies", "music"], + }, + }), + async () => ({ + contents: [ + { + uri: "test://resource/test", + text: "Test content", + }, + ], + }), + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "completion/complete", + params: { + ref: { + type: "ref/resource", + uri: "test://resource/{category}", + }, + argument: { + name: "category", + value: "", + }, + }, + }, + CompleteResultSchema, + ); + + expect(result.completion.values).toEqual(["books", "movies", "music"]); + expect(result.completion.total).toBe(3); + }); + + test("should support filtered completion of resource template parameters", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + + const client = new Client( + { + name: "test client", + version: "1.0", + }, + { + capabilities: { + resources: {}, + }, + }, + ); + + mcpServer.resource( + "test", + new ResourceTemplate("test://resource/{category}", { + list: undefined, + complete: { + category: (test: string) => + ["books", "movies", "music"].filter((value) => + value.startsWith(test), + ), + }, + }), + async () => ({ + contents: [ + { + uri: "test://resource/test", + text: "Test content", + }, + ], + }), + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "completion/complete", + params: { + ref: { + type: "ref/resource", + uri: "test://resource/{category}", + }, + argument: { + name: "category", + value: "m", + }, + }, + }, + CompleteResultSchema, + ); + + expect(result.completion.values).toEqual(["movies", "music"]); + expect(result.completion.total).toBe(2); + }); +}); + +describe("prompt()", () => { + test("should register zero-argument prompt", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const client = new Client({ + name: "test client", + version: "1.0", + }); + + mcpServer.prompt("test", async () => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: "Test response", + }, + }, + ], + })); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "prompts/list", + }, + ListPromptsResultSchema, + ); + + expect(result.prompts).toHaveLength(1); + expect(result.prompts[0].name).toBe("test"); + expect(result.prompts[0].arguments).toBeUndefined(); + }); + + test("should register prompt with args schema", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const client = new Client({ + name: "test client", + version: "1.0", + }); + + mcpServer.prompt( + "test", + { + name: z.string(), + value: z.string(), + }, + async ({ name, value }) => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: `${name}: ${value}`, + }, + }, + ], + }), + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "prompts/list", + }, + ListPromptsResultSchema, + ); + + expect(result.prompts).toHaveLength(1); + expect(result.prompts[0].name).toBe("test"); + expect(result.prompts[0].arguments).toEqual([ + { name: "name", required: true }, + { name: "value", required: true }, + ]); + }); + + test("should register prompt with description", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const client = new Client({ + name: "test client", + version: "1.0", + }); + + mcpServer.prompt("test", "Test description", async () => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: "Test response", + }, + }, + ], + })); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "prompts/list", + }, + ListPromptsResultSchema, + ); + + expect(result.prompts).toHaveLength(1); + expect(result.prompts[0].name).toBe("test"); + expect(result.prompts[0].description).toBe("Test description"); + }); + + test("should validate prompt args", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + + const client = new Client( + { + name: "test client", + version: "1.0", + }, + { + capabilities: { + prompts: {}, + }, + }, + ); + + mcpServer.prompt( + "test", + { + name: z.string(), + value: z.string().min(3), + }, + async ({ name, value }) => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: `${name}: ${value}`, + }, + }, + ], + }), + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + await expect( + client.request( + { + method: "prompts/get", + params: { + name: "test", + arguments: { + name: "test", + value: "ab", // Too short + }, + }, + }, + GetPromptResultSchema, + ), + ).rejects.toThrow(/Invalid arguments/); + }); + + test("should prevent duplicate prompt registration", () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + + mcpServer.prompt("test", async () => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: "Test response", + }, + }, + ], + })); + + expect(() => { + mcpServer.prompt("test", async () => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: "Test response 2", + }, + }, + ], + })); + }).toThrow(/already registered/); + }); + + test("should throw McpError for invalid prompt name", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + + const client = new Client( + { + name: "test client", + version: "1.0", + }, + { + capabilities: { + prompts: {}, + }, + }, + ); + + mcpServer.prompt("test-prompt", async () => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: "Test response", + }, + }, + ], + })); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + await expect( + client.request( + { + method: "prompts/get", + params: { + name: "nonexistent-prompt", + }, + }, + GetPromptResultSchema, + ), + ).rejects.toThrow(/Prompt nonexistent-prompt not found/); + }); + + test("should support completion of prompt arguments", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + + const client = new Client( + { + name: "test client", + version: "1.0", + }, + { + capabilities: { + prompts: {}, + }, + }, + ); + + mcpServer.prompt( + "test-prompt", + { + name: completable(z.string(), () => ["Alice", "Bob", "Charlie"]), + }, + async ({ name }) => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: `Hello ${name}`, + }, + }, + ], + }), + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "completion/complete", + params: { + ref: { + type: "ref/prompt", + name: "test-prompt", + }, + argument: { + name: "name", + value: "", + }, + }, + }, + CompleteResultSchema, + ); + + expect(result.completion.values).toEqual(["Alice", "Bob", "Charlie"]); + expect(result.completion.total).toBe(3); + }); + + test("should support filtered completion of prompt arguments", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + + const client = new Client( + { + name: "test client", + version: "1.0", + }, + { + capabilities: { + prompts: {}, + }, + }, + ); + + mcpServer.prompt( + "test-prompt", + { + name: completable(z.string(), (test) => + ["Alice", "Bob", "Charlie"].filter((value) => value.startsWith(test)), + ), + }, + async ({ name }) => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: `Hello ${name}`, + }, + }, + ], + }), + ); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "completion/complete", + params: { + ref: { + type: "ref/prompt", + name: "test-prompt", + }, + argument: { + name: "name", + value: "A", + }, + }, + }, + CompleteResultSchema, + ); + + expect(result.completion.values).toEqual(["Alice"]); + expect(result.completion.total).toBe(1); + }); +}); diff --git a/src/server/mcp.ts b/src/server/mcp.ts new file mode 100644 index 00000000..db270994 --- /dev/null +++ b/src/server/mcp.ts @@ -0,0 +1,767 @@ +import { Server, ServerOptions } from "./index.js"; +import { zodToJsonSchema } from "zod-to-json-schema"; +import { + z, + ZodRawShape, + ZodObject, + ZodString, + AnyZodObject, + ZodTypeAny, + ZodType, + ZodTypeDef, + ZodOptional, +} from "zod"; +import { + Implementation, + Tool, + ListToolsResult, + CallToolResult, + McpError, + ErrorCode, + CompleteRequest, + CompleteResult, + PromptReference, + ResourceReference, + Resource, + ListResourcesResult, + ListResourceTemplatesRequestSchema, + ReadResourceRequestSchema, + ListToolsRequestSchema, + CallToolRequestSchema, + ListResourcesRequestSchema, + ListPromptsRequestSchema, + GetPromptRequestSchema, + CompleteRequestSchema, + ListPromptsResult, + Prompt, + PromptArgument, + GetPromptResult, + ReadResourceResult, + CallToolResultSchema, +} from "../types.js"; +import { Completable, CompletableDef } from "./completable.js"; +import { UriTemplate, Variables } from "../shared/uriTemplate.js"; +import { RequestHandlerExtra } from "../shared/protocol.js"; +import { Transport } from "../shared/transport.js"; + +/** + * High-level MCP server that provides a simpler API for working with resources, tools, and prompts. + * For advanced usage (like sending notifications or setting custom request handlers), use the underlying + * Server instance available via the `server` property. + */ +export class McpServer { + /** + * The underlying Server instance, useful for advanced operations like sending notifications. + */ + public readonly server: Server; + + private _registeredResources: { [uri: string]: RegisteredResource } = {}; + private _registeredResourceTemplates: { + [name: string]: RegisteredResourceTemplate; + } = {}; + private _registeredTools: { [name: string]: RegisteredTool } = {}; + private _registeredPrompts: { [name: string]: RegisteredPrompt } = {}; + + constructor(serverInfo: Implementation, options?: ServerOptions) { + this.server = new Server(serverInfo, options); + } + + /** + * Attaches to the given transport, starts it, and starts listening for messages. + * + * The `server` object assumes ownership of the Transport, replacing any callbacks that have already been set, and expects that it is the only user of the Transport instance going forward. + */ + async connect(transport: Transport): Promise { + return await this.server.connect(transport); + } + + /** + * Closes the connection. + */ + async close(): Promise { + await this.server.close(); + } + + private setToolRequestHandlers() { + this.server.assertCanSetRequestHandler( + ListToolsRequestSchema.shape.method.value, + ); + this.server.assertCanSetRequestHandler( + CallToolRequestSchema.shape.method.value, + ); + + this.server.registerCapabilities({ + tools: {}, + }); + + this.server.setRequestHandler( + ListToolsRequestSchema, + (): ListToolsResult => ({ + tools: Object.entries(this._registeredTools).map( + ([name, tool]): Tool => { + return { + name, + description: tool.description, + inputSchema: tool.inputSchema + ? (zodToJsonSchema(tool.inputSchema) as Tool["inputSchema"]) + : EMPTY_OBJECT_JSON_SCHEMA, + }; + }, + ), + }), + ); + + this.server.setRequestHandler( + CallToolRequestSchema, + async (request, extra): Promise => { + const tool = this._registeredTools[request.params.name]; + if (!tool) { + throw new McpError( + ErrorCode.InvalidParams, + `Tool ${request.params.name} not found`, + ); + } + + if (tool.inputSchema) { + const parseResult = await tool.inputSchema.safeParseAsync( + request.params.arguments, + ); + if (!parseResult.success) { + throw new McpError( + ErrorCode.InvalidParams, + `Invalid arguments for tool ${request.params.name}: ${parseResult.error.message}`, + ); + } + + const args = parseResult.data; + const cb = tool.callback as ToolCallback; + try { + return await Promise.resolve(cb(args, extra)); + } catch (error) { + return { + content: [ + { + type: "text", + text: error instanceof Error ? error.message : String(error), + }, + ], + isError: true, + }; + } + } else { + const cb = tool.callback as ToolCallback; + try { + return await Promise.resolve(cb(extra)); + } catch (error) { + return { + content: [ + { + type: "text", + text: error instanceof Error ? error.message : String(error), + }, + ], + isError: true, + }; + } + } + }, + ); + } + + private setCompletionRequestHandler() { + this.server.assertCanSetRequestHandler( + CompleteRequestSchema.shape.method.value, + ); + + this.server.setRequestHandler( + CompleteRequestSchema, + async (request): Promise => { + switch (request.params.ref.type) { + case "ref/prompt": + return this.handlePromptCompletion(request, request.params.ref); + + case "ref/resource": + return this.handleResourceCompletion(request, request.params.ref); + + default: + throw new McpError( + ErrorCode.InvalidParams, + `Invalid completion reference: ${request.params.ref}`, + ); + } + }, + ); + } + + private async handlePromptCompletion( + request: CompleteRequest, + ref: PromptReference, + ): Promise { + const prompt = this._registeredPrompts[ref.name]; + if (!prompt) { + throw new McpError( + ErrorCode.InvalidParams, + `Prompt ${request.params.ref.name} not found`, + ); + } + + if (!prompt.argsSchema) { + return EMPTY_COMPLETION_RESULT; + } + + const field = prompt.argsSchema.shape[request.params.argument.name]; + if (!(field instanceof Completable)) { + return EMPTY_COMPLETION_RESULT; + } + + const def: CompletableDef = field._def; + const suggestions = await def.complete(request.params.argument.value); + return createCompletionResult(suggestions); + } + + private async handleResourceCompletion( + request: CompleteRequest, + ref: ResourceReference, + ): Promise { + const template = Object.values(this._registeredResourceTemplates).find( + (t) => t.resourceTemplate.uriTemplate.toString() === ref.uri, + ); + + if (!template) { + if (this._registeredResources[ref.uri]) { + // Attempting to autocomplete a fixed resource URI is not an error in the spec (but probably should be). + return EMPTY_COMPLETION_RESULT; + } + + throw new McpError( + ErrorCode.InvalidParams, + `Resource template ${request.params.ref.uri} not found`, + ); + } + + const completer = template.resourceTemplate.completeCallback( + request.params.argument.name, + ); + if (!completer) { + return EMPTY_COMPLETION_RESULT; + } + + const suggestions = await completer(request.params.argument.value); + return createCompletionResult(suggestions); + } + + private setResourceRequestHandlers() { + this.server.assertCanSetRequestHandler( + ListResourcesRequestSchema.shape.method.value, + ); + this.server.assertCanSetRequestHandler( + ListResourceTemplatesRequestSchema.shape.method.value, + ); + this.server.assertCanSetRequestHandler( + ReadResourceRequestSchema.shape.method.value, + ); + + this.server.registerCapabilities({ + resources: {}, + }); + + this.server.setRequestHandler( + ListResourcesRequestSchema, + async (request, extra) => { + const resources = Object.entries(this._registeredResources).map( + ([uri, resource]) => ({ + uri, + name: resource.name, + ...resource.metadata, + }), + ); + + const templateResources: Resource[] = []; + for (const template of Object.values( + this._registeredResourceTemplates, + )) { + if (!template.resourceTemplate.listCallback) { + continue; + } + + const result = await template.resourceTemplate.listCallback(extra); + for (const resource of result.resources) { + templateResources.push({ + ...resource, + ...template.metadata, + }); + } + } + + return { resources: [...resources, ...templateResources] }; + }, + ); + + this.server.setRequestHandler( + ListResourceTemplatesRequestSchema, + async () => { + const resourceTemplates = Object.entries( + this._registeredResourceTemplates, + ).map(([name, template]) => ({ + name, + uriTemplate: template.resourceTemplate.uriTemplate.toString(), + ...template.metadata, + })); + + return { resourceTemplates }; + }, + ); + + this.server.setRequestHandler( + ReadResourceRequestSchema, + async (request, extra) => { + const uri = new URL(request.params.uri); + + // First check for exact resource match + const resource = this._registeredResources[uri.toString()]; + if (resource) { + return resource.readCallback(uri, extra); + } + + // Then check templates + for (const template of Object.values( + this._registeredResourceTemplates, + )) { + const variables = template.resourceTemplate.uriTemplate.match( + uri.toString(), + ); + if (variables) { + return template.readCallback(uri, variables, extra); + } + } + + throw new McpError( + ErrorCode.InvalidParams, + `Resource ${uri} not found`, + ); + }, + ); + + this.setCompletionRequestHandler(); + } + + private setPromptRequestHandlers() { + this.server.assertCanSetRequestHandler( + ListPromptsRequestSchema.shape.method.value, + ); + this.server.assertCanSetRequestHandler( + GetPromptRequestSchema.shape.method.value, + ); + + this.server.registerCapabilities({ + prompts: {}, + }); + + this.server.setRequestHandler( + ListPromptsRequestSchema, + (): ListPromptsResult => ({ + prompts: Object.entries(this._registeredPrompts).map( + ([name, prompt]): Prompt => { + return { + name, + description: prompt.description, + arguments: prompt.argsSchema + ? promptArgumentsFromSchema(prompt.argsSchema) + : undefined, + }; + }, + ), + }), + ); + + this.server.setRequestHandler( + GetPromptRequestSchema, + async (request, extra): Promise => { + const prompt = this._registeredPrompts[request.params.name]; + if (!prompt) { + throw new McpError( + ErrorCode.InvalidParams, + `Prompt ${request.params.name} not found`, + ); + } + + if (prompt.argsSchema) { + const parseResult = await prompt.argsSchema.safeParseAsync( + request.params.arguments, + ); + if (!parseResult.success) { + throw new McpError( + ErrorCode.InvalidParams, + `Invalid arguments for prompt ${request.params.name}: ${parseResult.error.message}`, + ); + } + + const args = parseResult.data; + const cb = prompt.callback as PromptCallback; + return await Promise.resolve(cb(args, extra)); + } else { + const cb = prompt.callback as PromptCallback; + return await Promise.resolve(cb(extra)); + } + }, + ); + + this.setCompletionRequestHandler(); + } + + /** + * Registers a resource `name` at a fixed URI, which will use the given callback to respond to read requests. + */ + resource(name: string, uri: string, readCallback: ReadResourceCallback): void; + + /** + * Registers a resource `name` at a fixed URI with metadata, which will use the given callback to respond to read requests. + */ + resource( + name: string, + uri: string, + metadata: ResourceMetadata, + readCallback: ReadResourceCallback, + ): void; + + /** + * Registers a resource `name` with a template pattern, which will use the given callback to respond to read requests. + */ + resource( + name: string, + template: ResourceTemplate, + readCallback: ReadResourceTemplateCallback, + ): void; + + /** + * Registers a resource `name` with a template pattern and metadata, which will use the given callback to respond to read requests. + */ + resource( + name: string, + template: ResourceTemplate, + metadata: ResourceMetadata, + readCallback: ReadResourceTemplateCallback, + ): void; + + resource( + name: string, + uriOrTemplate: string | ResourceTemplate, + ...rest: unknown[] + ): void { + let metadata: ResourceMetadata | undefined; + if (typeof rest[0] === "object") { + metadata = rest.shift() as ResourceMetadata; + } + + const readCallback = rest[0] as + | ReadResourceCallback + | ReadResourceTemplateCallback; + + if (typeof uriOrTemplate === "string") { + if (this._registeredResources[uriOrTemplate]) { + throw new Error(`Resource ${uriOrTemplate} is already registered`); + } + + this._registeredResources[uriOrTemplate] = { + name, + metadata, + readCallback: readCallback as ReadResourceCallback, + }; + } else { + if (this._registeredResourceTemplates[name]) { + throw new Error(`Resource template ${name} is already registered`); + } + + this._registeredResourceTemplates[name] = { + resourceTemplate: uriOrTemplate, + metadata, + readCallback: readCallback as ReadResourceTemplateCallback, + }; + } + + this.setResourceRequestHandlers(); + } + + /** + * Registers a zero-argument tool `name`, which will run the given function when the client calls it. + */ + tool(name: string, cb: ToolCallback): void; + + /** + * Registers a zero-argument tool `name` (with a description) which will run the given function when the client calls it. + */ + tool(name: string, description: string, cb: ToolCallback): void; + + /** + * Registers a tool `name` accepting the given arguments, which must be an object containing named properties associated with Zod schemas. When the client calls it, the function will be run with the parsed and validated arguments. + */ + tool( + name: string, + paramsSchema: Args, + cb: ToolCallback, + ): void; + + /** + * Registers a tool `name` (with a description) accepting the given arguments, which must be an object containing named properties associated with Zod schemas. When the client calls it, the function will be run with the parsed and validated arguments. + */ + tool( + name: string, + description: string, + paramsSchema: Args, + cb: ToolCallback, + ): void; + + tool(name: string, ...rest: unknown[]): void { + if (this._registeredTools[name]) { + throw new Error(`Tool ${name} is already registered`); + } + + let description: string | undefined; + if (typeof rest[0] === "string") { + description = rest.shift() as string; + } + + let paramsSchema: ZodRawShape | undefined; + if (rest.length > 1) { + paramsSchema = rest.shift() as ZodRawShape; + } + + const cb = rest[0] as ToolCallback; + this._registeredTools[name] = { + description, + inputSchema: + paramsSchema === undefined ? undefined : z.object(paramsSchema), + callback: cb, + }; + + this.setToolRequestHandlers(); + } + + /** + * Registers a zero-argument prompt `name`, which will run the given function when the client calls it. + */ + prompt(name: string, cb: PromptCallback): void; + + /** + * Registers a zero-argument prompt `name` (with a description) which will run the given function when the client calls it. + */ + prompt(name: string, description: string, cb: PromptCallback): void; + + /** + * Registers a prompt `name` accepting the given arguments, which must be an object containing named properties associated with Zod schemas. When the client calls it, the function will be run with the parsed and validated arguments. + */ + prompt( + name: string, + argsSchema: Args, + cb: PromptCallback, + ): void; + + /** + * Registers a prompt `name` (with a description) accepting the given arguments, which must be an object containing named properties associated with Zod schemas. When the client calls it, the function will be run with the parsed and validated arguments. + */ + prompt( + name: string, + description: string, + argsSchema: Args, + cb: PromptCallback, + ): void; + + prompt(name: string, ...rest: unknown[]): void { + if (this._registeredPrompts[name]) { + throw new Error(`Prompt ${name} is already registered`); + } + + let description: string | undefined; + if (typeof rest[0] === "string") { + description = rest.shift() as string; + } + + let argsSchema: PromptArgsRawShape | undefined; + if (rest.length > 1) { + argsSchema = rest.shift() as PromptArgsRawShape; + } + + const cb = rest[0] as PromptCallback; + this._registeredPrompts[name] = { + description, + argsSchema: argsSchema === undefined ? undefined : z.object(argsSchema), + callback: cb, + }; + + this.setPromptRequestHandlers(); + } +} + +/** + * A callback to complete one variable within a resource template's URI template. + */ +export type CompleteResourceTemplateCallback = ( + value: string, +) => string[] | Promise; + +/** + * A resource template combines a URI pattern with optional functionality to enumerate + * all resources matching that pattern. + */ +export class ResourceTemplate { + private _uriTemplate: UriTemplate; + + constructor( + uriTemplate: string | UriTemplate, + private _callbacks: { + /** + * A callback to list all resources matching this template. This is required to specified, even if `undefined`, to avoid accidentally forgetting resource listing. + */ + list: ListResourcesCallback | undefined; + + /** + * An optional callback to autocomplete variables within the URI template. Useful for clients and users to discover possible values. + */ + complete?: { + [variable: string]: CompleteResourceTemplateCallback; + }; + }, + ) { + this._uriTemplate = + typeof uriTemplate === "string" + ? new UriTemplate(uriTemplate) + : uriTemplate; + } + + /** + * Gets the URI template pattern. + */ + get uriTemplate(): UriTemplate { + return this._uriTemplate; + } + + /** + * Gets the list callback, if one was provided. + */ + get listCallback(): ListResourcesCallback | undefined { + return this._callbacks.list; + } + + /** + * Gets the callback for completing a specific URI template variable, if one was provided. + */ + completeCallback( + variable: string, + ): CompleteResourceTemplateCallback | undefined { + return this._callbacks.complete?.[variable]; + } +} + +/** + * Callback for a tool handler registered with Server.tool(). + * + * Parameters will include tool arguments, if applicable, as well as other request handler context. + */ +export type ToolCallback = + Args extends ZodRawShape + ? ( + args: z.objectOutputType, + extra: RequestHandlerExtra, + ) => CallToolResult | Promise + : (extra: RequestHandlerExtra) => CallToolResult | Promise; + +type RegisteredTool = { + description?: string; + inputSchema?: AnyZodObject; + callback: ToolCallback; +}; + +const EMPTY_OBJECT_JSON_SCHEMA = { + type: "object" as const, +}; + +/** + * Additional, optional information for annotating a resource. + */ +export type ResourceMetadata = Omit; + +/** + * Callback to list all resources matching a given template. + */ +export type ListResourcesCallback = ( + extra: RequestHandlerExtra, +) => ListResourcesResult | Promise; + +/** + * Callback to read a resource at a given URI. + */ +export type ReadResourceCallback = ( + uri: URL, + extra: RequestHandlerExtra, +) => ReadResourceResult | Promise; + +type RegisteredResource = { + name: string; + metadata?: ResourceMetadata; + readCallback: ReadResourceCallback; +}; + +/** + * Callback to read a resource at a given URI, following a filled-in URI template. + */ +export type ReadResourceTemplateCallback = ( + uri: URL, + variables: Variables, + extra: RequestHandlerExtra, +) => ReadResourceResult | Promise; + +type RegisteredResourceTemplate = { + resourceTemplate: ResourceTemplate; + metadata?: ResourceMetadata; + readCallback: ReadResourceTemplateCallback; +}; + +type PromptArgsRawShape = { + [k: string]: + | ZodType + | ZodOptional>; +}; + +export type PromptCallback< + Args extends undefined | PromptArgsRawShape = undefined, +> = Args extends PromptArgsRawShape + ? ( + args: z.objectOutputType, + extra: RequestHandlerExtra, + ) => GetPromptResult | Promise + : (extra: RequestHandlerExtra) => GetPromptResult | Promise; + +type RegisteredPrompt = { + description?: string; + argsSchema?: ZodObject; + callback: PromptCallback; +}; + +function promptArgumentsFromSchema( + schema: ZodObject, +): PromptArgument[] { + return Object.entries(schema.shape).map( + ([name, field]): PromptArgument => ({ + name, + description: field.description, + required: !field.isOptional(), + }), + ); +} + +function createCompletionResult(suggestions: string[]): CompleteResult { + return { + completion: { + values: suggestions.slice(0, 100), + total: suggestions.length, + hasMore: suggestions.length > 100, + }, + }; +} + +const EMPTY_COMPLETION_RESULT: CompleteResult = { + completion: { + values: [], + hasMore: false, + }, +}; diff --git a/src/shared/protocol.ts b/src/shared/protocol.ts index cf33cd2e..a4f211c6 100644 --- a/src/shared/protocol.ts +++ b/src/shared/protocol.ts @@ -516,9 +516,11 @@ export abstract class Protocol< /** * Asserts that a request handler has not already been set for the given method, in preparation for a new one being automatically installed. */ - protected assertCanSetRequestHandler(method: string): void { + assertCanSetRequestHandler(method: string): void { if (this._requestHandlers.has(method)) { - throw new Error(`A request handler for ${method} already exists, which would be overridden`); + throw new Error( + `A request handler for ${method} already exists, which would be overridden`, + ); } } @@ -550,7 +552,9 @@ export abstract class Protocol< } } -export function mergeCapabilities(base: T, additional: T): T { +export function mergeCapabilities< + T extends ServerCapabilities | ClientCapabilities, +>(base: T, additional: T): T { return Object.entries(additional).reduce( (acc, [key, value]) => { if (value && typeof value === "object") {