From 098266ae22d10737a3c319927dce82c447bea13b Mon Sep 17 00:00:00 2001 From: Justin Spahr-Summers Date: Tue, 7 Jan 2025 19:12:03 +0000 Subject: [PATCH] Add simplified API for registering resources and resource templates --- src/server/index.test.ts | 275 +++++++++++++++++++++++++++++++++++++++ src/server/index.ts | 242 +++++++++++++++++++++++++++++++++- 2 files changed, 512 insertions(+), 5 deletions(-) diff --git a/src/server/index.test.ts b/src/server/index.test.ts index ed614a2b..8a5e2971 100644 --- a/src/server/index.test.ts +++ b/src/server/index.test.ts @@ -17,10 +17,14 @@ import { ErrorCode, ListToolsResultSchema, CallToolResultSchema, + ListResourcesResultSchema, + ListResourceTemplatesResultSchema, + ReadResourceResultSchema, } 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"; test("should accept latest protocol version", async () => { let sendPromiseResolve: (value: unknown) => void; @@ -478,6 +482,7 @@ test("should handle server cancelling a request", async () => { // Request should be rejected await expect(createMessagePromise).rejects.toBe("Cancelled by test"); }); + test("should handle request timeout", async () => { const server = new Server( { @@ -927,3 +932,273 @@ describe("Server.tool", () => { ).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 UriTemplate("test://resource/{id}"), + 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 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 UriTemplate("test://resource/{id}"), + async () => ({ + contents: [ + { + uri: "test://resource/123", + text: "Test content", + }, + ], + }), + ); + + expect(() => { + server.resource( + "test", + new UriTemplate("test://resource/{id}"), + 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/); + }); +}); diff --git a/src/server/index.ts b/src/server/index.ts index 469a3a77..8bb900ef 100644 --- a/src/server/index.ts +++ b/src/server/index.ts @@ -21,6 +21,9 @@ import { InitializeRequestSchema, InitializeResult, LATEST_PROTOCOL_VERSION, + ListResourcesRequestSchema, + ListResourcesResult, + ListResourceTemplatesRequestSchema, ListRootsRequest, ListRootsResultSchema, ListToolsRequestSchema, @@ -28,7 +31,10 @@ import { LoggingMessageNotification, McpError, Notification, + ReadResourceRequestSchema, + ReadResourceResult, Request, + Resource, ResourceUpdatedNotification, Result, ServerCapabilities, @@ -38,6 +44,7 @@ import { SUPPORTED_PROTOCOL_VERSIONS, Tool, } from "../types.js"; +import { UriTemplate, Variables } from "../shared/uriTemplate.js"; export type ServerOptions = ProtocolOptions & { /** @@ -53,11 +60,11 @@ export type ServerOptions = ProtocolOptions & { */ export type ToolCallback = Args extends ZodRawShape - ? ( - args: z.objectOutputType, - extra: RequestHandlerExtra, - ) => CallToolResult | Promise - : (extra: RequestHandlerExtra) => CallToolResult | Promise; + ? ( + args: z.objectOutputType, + extra: RequestHandlerExtra, + ) => CallToolResult | Promise + : (extra: RequestHandlerExtra) => CallToolResult | Promise; type RegisteredTool = { description?: string; @@ -69,6 +76,30 @@ const EMPTY_OBJECT_JSON_SCHEMA = { type: "object" as const, }; +export type ResourceMetadata = Omit; + +export type ReadResourceCallback = ( + uri: URL, + variables?: Variables, +) => ReadResourceResult | Promise; + +export type ListResourcesCallback = () => + | ListResourcesResult + | Promise; + +type RegisteredResource = { + name: string; + metadata?: ResourceMetadata; + readCallback: ReadResourceCallback; +}; + +type RegisteredResourceTemplate = { + uriTemplate: UriTemplate; + metadata?: ResourceMetadata; + listCallback?: ListResourcesCallback; + readCallback: ReadResourceCallback; +}; + /** * An MCP server on top of a pluggable transport. * @@ -107,6 +138,10 @@ export class Server< private _clientVersion?: Implementation; private _capabilities: ServerCapabilities; private _registeredTools: { [name: string]: RegisteredTool } = {}; + private _registeredResources: { [uri: string]: RegisteredResource } = {}; + private _registeredResourceTemplates: { + [name: string]: RegisteredResourceTemplate; + } = {}; /** * Callback for when initialization has fully completed (i.e., the client has sent an `initialized` notification). @@ -492,4 +527,201 @@ export class Server< this.setToolRequestHandlers(); } + + 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 () => { + 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.listCallback) { + continue; + } + + const result = await template.listCallback(); + 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.uriTemplate.toString(), + ...template.metadata, + })); + + return { resourceTemplates }; + }); + + this.setRequestHandler(ReadResourceRequestSchema, async (request) => { + 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); + } + + // Then check templates + for (const template of Object.values(this._registeredResourceTemplates)) { + const variables = template.uriTemplate.match(uri.toString()); + if (variables) { + return template.readCallback(uri, variables); + } + } + + throw new McpError(ErrorCode.InvalidParams, `Resource ${uri} not found`); + }); + } + + resource(name: string, uri: string, readCallback: ReadResourceCallback): void; + + resource( + name: string, + uri: string, + metadata: ResourceMetadata, + readCallback: ReadResourceCallback, + ): void; + + resource( + name: string, + uriTemplate: UriTemplate, + readCallback: ReadResourceCallback, + ): void; + + resource( + name: string, + uriTemplate: UriTemplate, + metadata: ResourceMetadata, + readCallback: ReadResourceCallback, + ): void; + + resource( + name: string, + uriTemplate: UriTemplate, + listCallback: ListResourcesCallback, + readCallback: ReadResourceCallback, + ): void; + + resource( + name: string, + uriTemplate: UriTemplate, + metadata: ResourceMetadata, + listCallback: ListResourcesCallback, + readCallback: ReadResourceCallback, + ): void; + + resource( + name: string, + uriOrTemplate: string | UriTemplate, + ...rest: unknown[] + ): void { + let metadata: ResourceMetadata | undefined; + if (typeof rest[0] === "object") { + metadata = rest.shift() as ResourceMetadata; + } + + let listCallback: ListResourcesCallback | undefined; + if (rest.length > 1) { + listCallback = rest.shift() as ListResourcesCallback; + } + + const readCallback = rest[0] as ReadResourceCallback; + if (typeof uriOrTemplate === "string") { + this.registerResource({ + name, + uri: uriOrTemplate, + metadata, + readCallback, + }); + } else { + this.registerResourceTemplate({ + name, + uriTemplate: uriOrTemplate, + metadata, + listCallback, + readCallback, + }); + } + } + + 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, + uriTemplate, + metadata, + listCallback, + readCallback, + }: { + name: string; + uriTemplate: UriTemplate; + metadata?: ResourceMetadata; + listCallback?: ListResourcesCallback; + readCallback: ReadResourceCallback; + }): void { + if (this._registeredResourceTemplates[name]) { + throw new Error(`Resource template ${name} is already registered`); + } + + this._registeredResourceTemplates[name] = { + uriTemplate, + metadata, + listCallback, + readCallback, + }; + + this.setResourceRequestHandlers(); + } }