diff --git a/README.md b/README.md index a6ec645c..4fcdff17 100644 --- a/README.md +++ b/README.md @@ -380,6 +380,68 @@ server.tool( ## Advanced Usage +### Dynamic Servers + +If you want to offer an initial set of tools/prompts/resources, but later add additional ones based on user action or external state change, you can add/update/remove them _after_ the Server is connected. This will automatically emit the corresponding `listChanged` notificaions: + +```ts +import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; +import { z } from "zod"; + +const server = new McpServer({ + name: "Dynamic Example", + version: "1.0.0" +}); + +const listMessageTool = server.tool( + "listMessages", + { channel: z.string() }, + async ({ channel }) => ({ + content: [{ type: "text", text: await listMessages(channel) }] + }) +); + +const putMessageTool = server.tool( + "putMessage", + { channel: z.string(), message: z.string() }, + async ({ channel, message }) => ({ + content: [{ type: "text", text: await putMessage(channel, string) }] + }) +); +// Until we upgrade auth, `putMessage` is disabled (won't show up in listTools) +putMessageTool.disable() + +const upgradeAuthTool = server.tool( + "upgradeAuth", + { permission: z.enum(["write', vadmin"])}, + // Any mutations here will automatically emit `listChanged` notifications + async ({ permission }) => { + const { ok, err, previous } = await upgradeAuthAndStoreToken(permission) + if (!ok) return {content: [{ type: "text", text: `Error: ${err}` }]} + + // If we previously had read-only access, 'putMessage' is now available + if (previous === "read") { + putMessageTool.enable() + } + + if (permission === 'write') { + // If we've just upgraded to 'write' permissions, we can still call 'upgradeAuth' + // but can only upgrade to 'admin'. + upgradeAuthTool.update({ + paramSchema: { permission: z.enum(["admin"]) }, // change validation rules + }) + } else { + // If we're now an admin, we no longer have anywhere to upgrade to, so fully remove that tool + upgradeAuthTool.remove() + } + } +) + +// Connect as normal +const transport = new StdioServerTransport(); +await server.connect(transport); +``` + ### Low-Level Server For more control, you can use the low-level Server class directly: diff --git a/package.json b/package.json index 1b9fd44d..c7963ec6 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@modelcontextprotocol/sdk", - "version": "1.9.0", + "version": "1.10.0", "description": "Model Context Protocol implementation for TypeScript", "license": "MIT", "author": "Anthropic, PBC (https://anthropic.com)", diff --git a/src/client/index.test.ts b/src/client/index.test.ts index a1b43b2b..36dd6518 100644 --- a/src/client/index.test.ts +++ b/src/client/index.test.ts @@ -66,6 +66,9 @@ test("should initialize with matching protocol version", async () => { protocolVersion: LATEST_PROTOCOL_VERSION, }), }), + expect.objectContaining({ + relatedRequestId: undefined, + }), ); // Should have the instructions returned diff --git a/src/client/streamableHttp.test.ts b/src/client/streamableHttp.test.ts new file mode 100644 index 00000000..1572e912 --- /dev/null +++ b/src/client/streamableHttp.test.ts @@ -0,0 +1,378 @@ +import { StreamableHTTPClientTransport } from "./streamableHttp.js"; +import { JSONRPCMessage } from "../types.js"; + + +describe("StreamableHTTPClientTransport", () => { + let transport: StreamableHTTPClientTransport; + + beforeEach(() => { + transport = new StreamableHTTPClientTransport(new URL("http://localhost:1234/mcp")); + jest.spyOn(global, "fetch"); + }); + + afterEach(async () => { + await transport.close().catch(() => { }); + jest.clearAllMocks(); + }); + + it("should send JSON-RPC messages via POST", async () => { + const message: JSONRPCMessage = { + jsonrpc: "2.0", + method: "test", + params: {}, + id: "test-id" + }; + + (global.fetch as jest.Mock).mockResolvedValueOnce({ + ok: true, + status: 202, + headers: new Headers(), + }); + + await transport.send(message); + + expect(global.fetch).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ + method: "POST", + headers: expect.any(Headers), + body: JSON.stringify(message) + }) + ); + }); + + it("should send batch messages", async () => { + const messages: JSONRPCMessage[] = [ + { jsonrpc: "2.0", method: "test1", params: {}, id: "id1" }, + { jsonrpc: "2.0", method: "test2", params: {}, id: "id2" } + ]; + + (global.fetch as jest.Mock).mockResolvedValueOnce({ + ok: true, + status: 200, + headers: new Headers({ "content-type": "text/event-stream" }), + body: null + }); + + await transport.send(messages); + + expect(global.fetch).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ + method: "POST", + headers: expect.any(Headers), + body: JSON.stringify(messages) + }) + ); + }); + + it("should store session ID received during initialization", async () => { + const message: JSONRPCMessage = { + jsonrpc: "2.0", + method: "initialize", + params: { + clientInfo: { name: "test-client", version: "1.0" }, + protocolVersion: "2025-03-26" + }, + id: "init-id" + }; + + (global.fetch as jest.Mock).mockResolvedValueOnce({ + ok: true, + status: 200, + headers: new Headers({ "content-type": "text/event-stream", "mcp-session-id": "test-session-id" }), + }); + + await transport.send(message); + + // Send a second message that should include the session ID + (global.fetch as jest.Mock).mockResolvedValueOnce({ + ok: true, + status: 202, + headers: new Headers() + }); + + await transport.send({ jsonrpc: "2.0", method: "test", params: {} } as JSONRPCMessage); + + // Check that second request included session ID header + const calls = (global.fetch as jest.Mock).mock.calls; + const lastCall = calls[calls.length - 1]; + expect(lastCall[1].headers).toBeDefined(); + expect(lastCall[1].headers.get("mcp-session-id")).toBe("test-session-id"); + }); + + it("should handle 404 response when session expires", async () => { + const message: JSONRPCMessage = { + jsonrpc: "2.0", + method: "test", + params: {}, + id: "test-id" + }; + + (global.fetch as jest.Mock).mockResolvedValueOnce({ + ok: false, + status: 404, + statusText: "Not Found", + text: () => Promise.resolve("Session not found"), + headers: new Headers() + }); + + const errorSpy = jest.fn(); + transport.onerror = errorSpy; + + await expect(transport.send(message)).rejects.toThrow("Error POSTing to endpoint (HTTP 404)"); + expect(errorSpy).toHaveBeenCalled(); + }); + + it("should handle non-streaming JSON response", async () => { + const message: JSONRPCMessage = { + jsonrpc: "2.0", + method: "test", + params: {}, + id: "test-id" + }; + + const responseMessage: JSONRPCMessage = { + jsonrpc: "2.0", + result: { success: true }, + id: "test-id" + }; + + (global.fetch as jest.Mock).mockResolvedValueOnce({ + ok: true, + status: 200, + headers: new Headers({ "content-type": "application/json" }), + json: () => Promise.resolve(responseMessage) + }); + + const messageSpy = jest.fn(); + transport.onmessage = messageSpy; + + await transport.send(message); + + expect(messageSpy).toHaveBeenCalledWith(responseMessage); + }); + + it("should attempt initial GET connection and handle 405 gracefully", async () => { + // Mock the server not supporting GET for SSE (returning 405) + (global.fetch as jest.Mock).mockResolvedValueOnce({ + ok: false, + status: 405, + statusText: "Method Not Allowed" + }); + + // We expect the 405 error to be caught and handled gracefully + // This should not throw an error that breaks the transport + await transport.start(); + await expect(transport["_startOrAuthStandaloneSSE"]()).resolves.not.toThrow("Failed to open SSE stream: Method Not Allowed"); + // Check that GET was attempted + expect(global.fetch).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ + method: "GET", + headers: expect.any(Headers) + }) + ); + + // Verify transport still works after 405 + (global.fetch as jest.Mock).mockResolvedValueOnce({ + ok: true, + status: 202, + headers: new Headers() + }); + + await transport.send({ jsonrpc: "2.0", method: "test", params: {} } as JSONRPCMessage); + expect(global.fetch).toHaveBeenCalledTimes(2); + }); + + it("should handle successful initial GET connection for SSE", async () => { + // Set up readable stream for SSE events + const encoder = new TextEncoder(); + const stream = new ReadableStream({ + start(controller) { + // Send a server notification via SSE + const event = "event: message\ndata: {\"jsonrpc\": \"2.0\", \"method\": \"serverNotification\", \"params\": {}}\n\n"; + controller.enqueue(encoder.encode(event)); + } + }); + + // Mock successful GET connection + (global.fetch as jest.Mock).mockResolvedValueOnce({ + ok: true, + status: 200, + headers: new Headers({ "content-type": "text/event-stream" }), + body: stream + }); + + const messageSpy = jest.fn(); + transport.onmessage = messageSpy; + + await transport.start(); + await transport["_startOrAuthStandaloneSSE"](); + + // Give time for the SSE event to be processed + await new Promise(resolve => setTimeout(resolve, 50)); + + expect(messageSpy).toHaveBeenCalledWith( + expect.objectContaining({ + jsonrpc: "2.0", + method: "serverNotification", + params: {} + }) + ); + }); + + it("should handle multiple concurrent SSE streams", async () => { + // Mock two POST requests that return SSE streams + const makeStream = (id: string) => { + const encoder = new TextEncoder(); + return new ReadableStream({ + start(controller) { + const event = `event: message\ndata: {"jsonrpc": "2.0", "result": {"id": "${id}"}, "id": "${id}"}\n\n`; + controller.enqueue(encoder.encode(event)); + } + }); + }; + + (global.fetch as jest.Mock) + .mockResolvedValueOnce({ + ok: true, + status: 200, + headers: new Headers({ "content-type": "text/event-stream" }), + body: makeStream("request1") + }) + .mockResolvedValueOnce({ + ok: true, + status: 200, + headers: new Headers({ "content-type": "text/event-stream" }), + body: makeStream("request2") + }); + + const messageSpy = jest.fn(); + transport.onmessage = messageSpy; + + // Send two concurrent requests + await Promise.all([ + transport.send({ jsonrpc: "2.0", method: "test1", params: {}, id: "request1" }), + transport.send({ jsonrpc: "2.0", method: "test2", params: {}, id: "request2" }) + ]); + + // Give time for SSE processing + await new Promise(resolve => setTimeout(resolve, 100)); + + // Both streams should have delivered their messages + expect(messageSpy).toHaveBeenCalledTimes(2); + + // Verify received messages without assuming specific order + expect(messageSpy.mock.calls.some(call => { + const msg = call[0]; + return msg.id === "request1" && msg.result?.id === "request1"; + })).toBe(true); + + expect(messageSpy.mock.calls.some(call => { + const msg = call[0]; + return msg.id === "request2" && msg.result?.id === "request2"; + })).toBe(true); + }); + + it("should include last-event-id header when resuming a broken connection", async () => { + // First make a successful connection that provides an event ID + const encoder = new TextEncoder(); + const stream = new ReadableStream({ + start(controller) { + const event = "id: event-123\nevent: message\ndata: {\"jsonrpc\": \"2.0\", \"method\": \"serverNotification\", \"params\": {}}\n\n"; + controller.enqueue(encoder.encode(event)); + controller.close(); + } + }); + + (global.fetch as jest.Mock).mockResolvedValueOnce({ + ok: true, + status: 200, + headers: new Headers({ "content-type": "text/event-stream" }), + body: stream + }); + + await transport.start(); + await transport["_startOrAuthStandaloneSSE"](); + await new Promise(resolve => setTimeout(resolve, 50)); + + // Now simulate attempting to reconnect + (global.fetch as jest.Mock).mockResolvedValueOnce({ + ok: true, + status: 200, + headers: new Headers({ "content-type": "text/event-stream" }), + body: null + }); + + await transport["_startOrAuthStandaloneSSE"](); + + // Check that Last-Event-ID was included + const calls = (global.fetch as jest.Mock).mock.calls; + const lastCall = calls[calls.length - 1]; + expect(lastCall[1].headers.get("last-event-id")).toBe("event-123"); + }); + + it("should throw error when invalid content-type is received", async () => { + const message: JSONRPCMessage = { + jsonrpc: "2.0", + method: "test", + params: {}, + id: "test-id" + }; + + const stream = new ReadableStream({ + start(controller) { + controller.enqueue("invalid text response"); + controller.close(); + } + }); + + const errorSpy = jest.fn(); + transport.onerror = errorSpy; + + (global.fetch as jest.Mock).mockResolvedValueOnce({ + ok: true, + status: 200, + headers: new Headers({ "content-type": "text/plain" }), + body: stream + }); + + await transport.start(); + await expect(transport.send(message)).rejects.toThrow("Unexpected content type: text/plain"); + expect(errorSpy).toHaveBeenCalled(); + }); + + + it("should always send specified custom headers", async () => { + const requestInit = { + headers: { + "X-Custom-Header": "CustomValue" + } + }; + transport = new StreamableHTTPClientTransport(new URL("http://localhost:1234/mcp"), { + requestInit: requestInit + }); + + let actualReqInit: RequestInit = {}; + + ((global.fetch as jest.Mock)).mockImplementation( + async (_url, reqInit) => { + actualReqInit = reqInit; + return new Response(null, { status: 200, headers: { "content-type": "text/event-stream" } }); + } + ); + + await transport.start(); + + await transport["_startOrAuthStandaloneSSE"](); + expect((actualReqInit.headers as Headers).get("x-custom-header")).toBe("CustomValue"); + + requestInit.headers["X-Custom-Header"] = "SecondCustomValue"; + + await transport.send({ jsonrpc: "2.0", method: "test", params: {} } as JSONRPCMessage); + expect((actualReqInit.headers as Headers).get("x-custom-header")).toBe("SecondCustomValue"); + + expect(global.fetch).toHaveBeenCalledTimes(2); + }); +}); diff --git a/src/client/streamableHttp.ts b/src/client/streamableHttp.ts new file mode 100644 index 00000000..ea69ee77 --- /dev/null +++ b/src/client/streamableHttp.ts @@ -0,0 +1,294 @@ +import { Transport } from "../shared/transport.js"; +import { isJSONRPCNotification, JSONRPCMessage, JSONRPCMessageSchema } from "../types.js"; +import { auth, AuthResult, OAuthClientProvider, UnauthorizedError } from "./auth.js"; +import { EventSourceParserStream } from "eventsource-parser/stream"; + +export class StreamableHTTPError extends Error { + constructor( + public readonly code: number | undefined, + message: string | undefined, + ) { + super(`Streamable HTTP error: ${message}`); + } +} + +/** + * Configuration options for the `StreamableHTTPClientTransport`. + */ +export type StreamableHTTPClientTransportOptions = { + /** + * An OAuth client provider to use for authentication. + * + * When an `authProvider` is specified and the connection is started: + * 1. The connection is attempted with any existing access token from the `authProvider`. + * 2. If the access token has expired, the `authProvider` is used to refresh the token. + * 3. If token refresh fails or no access token exists, and auth is required, `OAuthClientProvider.redirectToAuthorization` is called, and an `UnauthorizedError` will be thrown from `connect`/`start`. + * + * After the user has finished authorizing via their user agent, and is redirected back to the MCP client application, call `StreamableHTTPClientTransport.finishAuth` with the authorization code before retrying the connection. + * + * If an `authProvider` is not provided, and auth is required, an `UnauthorizedError` will be thrown. + * + * `UnauthorizedError` might also be thrown when sending any message over the transport, indicating that the session has expired, and needs to be re-authed and reconnected. + */ + authProvider?: OAuthClientProvider; + + /** + * Customizes HTTP requests to the server. + */ + requestInit?: RequestInit; +}; + +/** + * Client transport for Streamable HTTP: this implements the MCP Streamable HTTP transport specification. + * It will connect to a server using HTTP POST for sending messages and HTTP GET with Server-Sent Events + * for receiving messages. + */ +export class StreamableHTTPClientTransport implements Transport { + private _abortController?: AbortController; + private _url: URL; + private _requestInit?: RequestInit; + private _authProvider?: OAuthClientProvider; + private _sessionId?: string; + private _lastEventId?: string; + + onclose?: () => void; + onerror?: (error: Error) => void; + onmessage?: (message: JSONRPCMessage) => void; + + constructor( + url: URL, + opts?: StreamableHTTPClientTransportOptions, + ) { + this._url = url; + this._requestInit = opts?.requestInit; + this._authProvider = opts?.authProvider; + } + + private async _authThenStart(): Promise { + if (!this._authProvider) { + throw new UnauthorizedError("No auth provider"); + } + + let result: AuthResult; + try { + result = await auth(this._authProvider, { serverUrl: this._url }); + } catch (error) { + this.onerror?.(error as Error); + throw error; + } + + if (result !== "AUTHORIZED") { + throw new UnauthorizedError(); + } + + return await this._startOrAuthStandaloneSSE(); + } + + private async _commonHeaders(): Promise { + const headers: HeadersInit = {}; + if (this._authProvider) { + const tokens = await this._authProvider.tokens(); + if (tokens) { + headers["Authorization"] = `Bearer ${tokens.access_token}`; + } + } + + if (this._sessionId) { + headers["mcp-session-id"] = this._sessionId; + } + + return new Headers( + { ...headers, ...this._requestInit?.headers } + ); + } + + private async _startOrAuthStandaloneSSE(): Promise { + try { + // Try to open an initial SSE stream with GET to listen for server messages + // This is optional according to the spec - server may not support it + const headers = await this._commonHeaders(); + headers.set("Accept", "text/event-stream"); + + // Include Last-Event-ID header for resumable streams + if (this._lastEventId) { + headers.set("last-event-id", this._lastEventId); + } + + const response = await fetch(this._url, { + method: "GET", + headers, + signal: this._abortController?.signal, + }); + + if (!response.ok) { + if (response.status === 401 && this._authProvider) { + // Need to authenticate + return await this._authThenStart(); + } + + // 405 indicates that the server does not offer an SSE stream at GET endpoint + // This is an expected case that should not trigger an error + if (response.status === 405) { + return; + } + + throw new StreamableHTTPError( + response.status, + `Failed to open SSE stream: ${response.statusText}`, + ); + } + // Successful connection, handle the SSE stream as a standalone listener + this._handleSseStream(response.body); + } catch (error) { + this.onerror?.(error as Error); + throw error; + } + } + + private _handleSseStream(stream: ReadableStream | null): void { + if (!stream) { + return; + } + + const processStream = async () => { + // Create a pipeline: binary stream -> text decoder -> SSE parser + const eventStream = stream + .pipeThrough(new TextDecoderStream()) + .pipeThrough(new EventSourceParserStream()); + + for await (const event of eventStream) { + // Update last event ID if provided + if (event.id) { + this._lastEventId = event.id; + } + // Handle message events (default event type is undefined per docs) + // or explicit 'message' event type + if (!event.event || event.event === "message") { + try { + const message = JSONRPCMessageSchema.parse(JSON.parse(event.data)); + this.onmessage?.(message); + } catch (error) { + this.onerror?.(error as Error); + } + } + } + }; + + processStream().catch(err => this.onerror?.(err)); + } + + async start() { + if (this._abortController) { + throw new Error( + "StreamableHTTPClientTransport already started! If using Client class, note that connect() calls start() automatically.", + ); + } + + this._abortController = new AbortController(); + } + + /** + * Call this method after the user has finished authorizing via their user agent and is redirected back to the MCP client application. This will exchange the authorization code for an access token, enabling the next connection attempt to successfully auth. + */ + async finishAuth(authorizationCode: string): Promise { + if (!this._authProvider) { + throw new UnauthorizedError("No auth provider"); + } + + const result = await auth(this._authProvider, { serverUrl: this._url, authorizationCode }); + if (result !== "AUTHORIZED") { + throw new UnauthorizedError("Failed to authorize"); + } + } + + async close(): Promise { + // Abort any pending requests + this._abortController?.abort(); + + this.onclose?.(); + } + + async send(message: JSONRPCMessage | JSONRPCMessage[]): Promise { + try { + const headers = await this._commonHeaders(); + headers.set("content-type", "application/json"); + headers.set("accept", "application/json, text/event-stream"); + + const init = { + ...this._requestInit, + method: "POST", + headers, + body: JSON.stringify(message), + signal: this._abortController?.signal, + }; + + const response = await fetch(this._url, init); + + // Handle session ID received during initialization + const sessionId = response.headers.get("mcp-session-id"); + if (sessionId) { + this._sessionId = sessionId; + } + + if (!response.ok) { + if (response.status === 401 && this._authProvider) { + const result = await auth(this._authProvider, { serverUrl: this._url }); + if (result !== "AUTHORIZED") { + throw new UnauthorizedError(); + } + + // Purposely _not_ awaited, so we don't call onerror twice + return this.send(message); + } + + const text = await response.text().catch(() => null); + throw new Error( + `Error POSTing to endpoint (HTTP ${response.status}): ${text}`, + ); + } + + // If the response is 202 Accepted, there's no body to process + if (response.status === 202) { + // if the accepted notification is initialized, we start the SSE stream + // if it's supported by the server + if (isJSONRPCNotification(message) && message.method === "notifications/initialized") { + // We don't need to handle 405 here anymore as it's handled in _startOrAuthStandaloneSSE + this._startOrAuthStandaloneSSE().catch(err => this.onerror?.(err)); + } + return; + } + + // Get original message(s) for detecting request IDs + const messages = Array.isArray(message) ? message : [message]; + + const hasRequests = messages.filter(msg => "method" in msg && "id" in msg && msg.id !== undefined).length > 0; + + // Check the response type + const contentType = response.headers.get("content-type"); + + if (hasRequests) { + if (contentType?.includes("text/event-stream")) { + this._handleSseStream(response.body); + } else if (contentType?.includes("application/json")) { + // For non-streaming servers, we might get direct JSON responses + const data = await response.json(); + const responseMessages = Array.isArray(data) + ? data.map(msg => JSONRPCMessageSchema.parse(msg)) + : [JSONRPCMessageSchema.parse(data)]; + + for (const msg of responseMessages) { + this.onmessage?.(msg); + } + } else { + throw new StreamableHTTPError( + -1, + `Unexpected content type: ${contentType}`, + ); + } + } + } catch (error) { + this.onerror?.(error as Error); + throw error; + } + } +} diff --git a/src/examples/README.md b/src/examples/README.md new file mode 100644 index 00000000..2ea1a596 --- /dev/null +++ b/src/examples/README.md @@ -0,0 +1,137 @@ +# MCP TypeScript SDK Examples + +This directory contains example implementations of MCP clients and servers using the TypeScript SDK. + +## Table of Contents + +- [Streamable HTTP Servers - Single Node Deployment](#streamable-http---single-node-deployment-with-basic-session-state-management) + - [Simple Server with Streamable HTTP](#simple-server-with-streamable-http-transport-serversimplestreamablehttpts) + - [Server Supporting SSE via GET](#server-supporting-with-sse-via-get-serverstandalonessewithgetstreamablehttpts) + - [Server with JSON Response Mode](#server-with-json-response-mode-serverjsonresponsestreamablehttpts) +- [Client Example - Streamable HTTP](#client-clientsimplestreamablehttpts) +- [Useful bash commands for testing](#useful-commands-for-testing) + +## Streamable HTTP - single node deployment with basic session state management + +Multi node with state management example will be added soon after we add support. + + +### Simple server with Streamable HTTP transport (`server/simpleStreamableHttp.ts`) + +A simple MCP server that uses the Streamable HTTP transport, implemented with Express. The server provides: + +- A simple `greet` tool that returns a greeting for a name +- A `greeting-template` prompt that generates a greeting template +- A static `greeting-resource` resource + +#### Running the server + +```bash +npx tsx src/examples/server/simpleStreamableHttp.ts +``` + +The server will start on port 3000. You can test the initialization and tool listing: + +### Server supporting SSE via GET (`server/standaloneSseWithGetStreamableHttp.ts`) + +An MCP server that demonstrates how to support SSE notifications via GET requests using the Streamable HTTP transport with Express. The server dynamically adds resources at regular intervals and supports notifications for resource list changes (server notifications are available through the standalone SSE connection established by GET request). + +#### Running the server + +```bash +npx tsx src/examples/server/standaloneSseWithGetStreamableHttp.ts +``` + +The server will start on port 3000 and automatically create new resources every 5 seconds. + +### Server with JSON response mode (`server/jsonResponseStreamableHttp.ts`) + +A simple MCP server that uses the Streamable HTTP transport with JSON response mode enabled, implemented with Express. The server provides a simple `greet` tool that returns a greeting for a name. + +_NOTE: This demonstrates a server that does not use SSE at all. Note that this limits its support for MCP features; for example, it cannot provide logging and progress notifications for tool execution._ + +#### Running the server + +```bash +npx tsx src/examples/server/jsonResponseStreamableHttp.ts +``` + + +### Client (`client/simpleStreamableHttp.ts`) + +A client that connects to the server, initializes it, and demonstrates how to: + +- List available tools and call the `greet` tool +- List available prompts and get the `greeting-template` prompt +- List available resources + +#### Running the client + +```bash +npx tsx src/examples/client/simpleStreamableHttp.ts +``` + +Make sure the server is running before starting the client. + + +### Useful commands for testing + +#### Initialize +Streamable HTTP transport requires to do the initialization first. + +```bash +# First initialize the server and save the session ID to a variable +SESSION_ID=$(curl -X POST \ + -H "Content-Type: application/json" \ + -H "Accept: application/json" \ + -H "Accept: text/event-stream" \ + -d '{ + "jsonrpc": "2.0", + "method": "initialize", + "params": { + "capabilities": {}, + "protocolVersion": "2025-03-26", + "clientInfo": { + "name": "test", + "version": "1.0.0" + } + }, + "id": "1" + }' \ + -i http://localhost:3000/mcp 2>&1 | grep -i "mcp-session-id" | cut -d' ' -f2 | tr -d '\r') +echo "Session ID: $SESSION_ID + +``` +Once a session is established, we can send POST requests: + +#### List tools +```bash +# Then list tools using the saved session ID +curl -X POST -H "Content-Type: application/json" -H "Accept: application/json, text/event-stream" \ + -H "mcp-session-id: $SESSION_ID" \ + -d '{"jsonrpc":"2.0","method":"tools/list","params":{},"id":"2"}' \ + http://localhost:3000/mcp +``` + +#### Call tool + +```bash +# Call the greet tool using the saved session ID +curl -X POST \ + -H "Content-Type: application/json" \ + -H "Accept: application/json" \ + -H "Accept: text/event-stream" \ + -H "mcp-session-id: $SESSION_ID" \ + -d '{ + "jsonrpc": "2.0", + "method": "tools/call", + "params": { + "name": "greet", + "arguments": { + "name": "World" + } + }, + "id": "2" + }' \ + http://localhost:3000/mcp +``` diff --git a/src/examples/client/simpleStreamableHttp.ts b/src/examples/client/simpleStreamableHttp.ts new file mode 100644 index 00000000..739e1164 --- /dev/null +++ b/src/examples/client/simpleStreamableHttp.ts @@ -0,0 +1,154 @@ +import { Client } from '../../client/index.js'; +import { StreamableHTTPClientTransport } from '../../client/streamableHttp.js'; +import { + ListToolsRequest, + ListToolsResultSchema, + CallToolRequest, + CallToolResultSchema, + ListPromptsRequest, + ListPromptsResultSchema, + GetPromptRequest, + GetPromptResultSchema, + ListResourcesRequest, + ListResourcesResultSchema, + LoggingMessageNotificationSchema, + ResourceListChangedNotificationSchema, +} from '../../types.js'; + +async function main(): Promise { + // Create a new client with streamable HTTP transport + const client = new Client({ + name: 'example-client', + version: '1.0.0' + }); + + const transport = new StreamableHTTPClientTransport( + new URL('http://localhost:3000/mcp') + ); + + // Connect the client using the transport and initialize the server + await client.connect(transport); + console.log('Connected to MCP server'); + + // Set up notification handlers for server-initiated messages + client.setNotificationHandler(LoggingMessageNotificationSchema, (notification) => { + console.log(`Notification received: ${notification.params.level} - ${notification.params.data}`); + }); + client.setNotificationHandler(ResourceListChangedNotificationSchema, async (_) => { + console.log(`Resource list changed notification received!`); + const resourcesRequest: ListResourcesRequest = { + method: 'resources/list', + params: {} + }; + const resourcesResult = await client.request(resourcesRequest, ListResourcesResultSchema); + console.log('Available resources count:', resourcesResult.resources.length); + }); + + // List and call tools + await listTools(client); + + await callGreetTool(client); + await callMultiGreetTool(client); + + + // List available prompts + try { + const promptsRequest: ListPromptsRequest = { + method: 'prompts/list', + params: {} + }; + const promptsResult = await client.request(promptsRequest, ListPromptsResultSchema); + console.log('Available prompts:', promptsResult.prompts); + } catch (error) { + console.log(`Prompts not supported by this server (${error})`); + } + + // Get a prompt + try { + const promptRequest: GetPromptRequest = { + method: 'prompts/get', + params: { + name: 'greeting-template', + arguments: { name: 'MCP User' } + } + }; + const promptResult = await client.request(promptRequest, GetPromptResultSchema); + console.log('Prompt template:', promptResult.messages[0].content.text); + } catch (error) { + console.log(`Prompt retrieval not supported by this server (${error})`); + } + + // List available resources + try { + const resourcesRequest: ListResourcesRequest = { + method: 'resources/list', + params: {} + }; + const resourcesResult = await client.request(resourcesRequest, ListResourcesResultSchema); + console.log('Available resources:', resourcesResult.resources); + } catch (error) { + console.log(`Resources not supported by this server (${error})`); + } + // Keep the connection open to receive notifications + console.log('\nKeeping connection open to receive notifications. Press Ctrl+C to exit.'); +} + +async function listTools(client: Client): Promise { + try { + const toolsRequest: ListToolsRequest = { + method: 'tools/list', + params: {} + }; + const toolsResult = await client.request(toolsRequest, ListToolsResultSchema); + console.log('Available tools:', toolsResult.tools); + if (toolsResult.tools.length === 0) { + console.log('No tools available from the server'); + } + } catch (error) { + console.log(`Tools not supported by this server (${error})`); + return + } +} + +async function callGreetTool(client: Client): Promise { + try { + const greetRequest: CallToolRequest = { + method: 'tools/call', + params: { + name: 'greet', + arguments: { name: 'MCP User' } + } + }; + const greetResult = await client.request(greetRequest, CallToolResultSchema); + console.log('Greeting result:', greetResult.content[0].text); + } catch (error) { + console.log(`Error calling greet tool: ${error}`); + } +} + +async function callMultiGreetTool(client: Client): Promise { + try { + console.log('\nCalling multi-greet tool (with notifications)...'); + const multiGreetRequest: CallToolRequest = { + method: 'tools/call', + params: { + name: 'multi-greet', + arguments: { name: 'MCP User' } + } + }; + const multiGreetResult = await client.request(multiGreetRequest, CallToolResultSchema); + console.log('Multi-greet results:'); + multiGreetResult.content.forEach(item => { + if (item.type === 'text') { + console.log(`- ${item.text}`); + } + }); + } catch (error) { + console.log(`Error calling multi-greet tool: ${error}`); + } +} + +main().catch((error: unknown) => { + console.error('Error running MCP client:', error); + process.exit(1); +}); \ No newline at end of file diff --git a/src/examples/server/jsonResponseStreamableHttp.ts b/src/examples/server/jsonResponseStreamableHttp.ts new file mode 100644 index 00000000..34ab65d1 --- /dev/null +++ b/src/examples/server/jsonResponseStreamableHttp.ts @@ -0,0 +1,189 @@ +import express, { Request, Response } from 'express'; +import { randomUUID } from 'node:crypto'; +import { McpServer } from '../../server/mcp.js'; +import { StreamableHTTPServerTransport } from '../../server/streamableHttp.js'; +import { z } from 'zod'; +import { CallToolResult } from '../../types.js'; + +// Create an MCP server with implementation details +const server = new McpServer({ + name: 'json-response-streamable-http-server', + version: '1.0.0', +}, { + capabilities: { + logging: {}, + } +}); + +// Register a simple tool that returns a greeting +server.tool( + 'greet', + 'A simple greeting tool', + { + name: z.string().describe('Name to greet'), + }, + async ({ name }): Promise => { + return { + content: [ + { + type: 'text', + text: `Hello, ${name}!`, + }, + ], + }; + } +); + +// Register a tool that sends multiple greetings with notifications +server.tool( + 'multi-greet', + 'A tool that sends different greetings with delays between them', + { + name: z.string().describe('Name to greet'), + }, + async ({ name }, { sendNotification }): Promise => { + const sleep = (ms: number) => new Promise(resolve => setTimeout(resolve, ms)); + + await sendNotification({ + method: "notifications/message", + params: { level: "debug", data: `Starting multi-greet for ${name}` } + }); + + await sleep(1000); // Wait 1 second before first greeting + + await sendNotification({ + method: "notifications/message", + params: { level: "info", data: `Sending first greeting to ${name}` } + }); + + await sleep(1000); // Wait another second before second greeting + + await sendNotification({ + method: "notifications/message", + params: { level: "info", data: `Sending second greeting to ${name}` } + }); + + return { + content: [ + { + type: 'text', + text: `Good morning, ${name}!`, + } + ], + }; + } +); + +const app = express(); +app.use(express.json()); + +// Map to store transports by session ID +const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {}; + +app.post('/mcp', async (req: Request, res: Response) => { + console.log('Received MCP request:', req.body); + try { + // Check for existing session ID + const sessionId = req.headers['mcp-session-id'] as string | undefined; + let transport: StreamableHTTPServerTransport; + + if (sessionId && transports[sessionId]) { + // Reuse existing transport + transport = transports[sessionId]; + } else if (!sessionId && isInitializeRequest(req.body)) { + // New initialization request - use JSON response mode + transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: () => randomUUID(), + enableJsonResponse: true, // Enable JSON response mode + }); + + // Connect the transport to the MCP server BEFORE handling the request + await server.connect(transport); + + // After handling the request, if we get a session ID back, store the transport + await transport.handleRequest(req, res, req.body); + + // Store the transport by session ID for future requests + if (transport.sessionId) { + transports[transport.sessionId] = transport; + } + return; // Already handled + } else { + // Invalid request - no session ID or not initialization request + res.status(400).json({ + jsonrpc: '2.0', + error: { + code: -32000, + message: 'Bad Request: No valid session ID provided', + }, + id: null, + }); + return; + } + + // Handle the request with existing transport - no need to reconnect + await transport.handleRequest(req, res, req.body); + } catch (error) { + console.error('Error handling MCP request:', error); + if (!res.headersSent) { + res.status(500).json({ + jsonrpc: '2.0', + error: { + code: -32603, + message: 'Internal server error', + }, + id: null, + }); + } + } +}); + +// Handle GET requests for SSE streams according to spec +app.get('/mcp', async (req: Request, res: Response) => { + // Since this is a very simple example, we don't support GET requests for this server + // The spec requires returning 405 Method Not Allowed in this case + res.status(405).set('Allow', 'POST').send('Method Not Allowed'); +}); + +// Helper function to detect initialize requests +function isInitializeRequest(body: unknown): boolean { + if (Array.isArray(body)) { + return body.some(msg => typeof msg === 'object' && msg !== null && 'method' in msg && msg.method === 'initialize'); + } + return typeof body === 'object' && body !== null && 'method' in body && body.method === 'initialize'; +} + +// Start the server +const PORT = 3000; +app.listen(PORT, () => { + console.log(`MCP Streamable HTTP Server listening on port ${PORT}`); + console.log(`Initialize session with the command below id you are using curl for testing: + ----------------------------- + SESSION_ID=$(curl -X POST \ + -H "Content-Type: application/json" \ + -H "Accept: application/json" \ + -H "Accept: text/event-stream" \ + -d '{ + "jsonrpc": "2.0", + "method": "initialize", + "params": { + "capabilities": {}, + "protocolVersion": "2025-03-26", + "clientInfo": { + "name": "test", + "version": "1.0.0" + } + }, + "id": "1" + }' \ + -i http://localhost:3000/mcp 2>&1 | grep -i "mcp-session-id" | cut -d' ' -f2 | tr -d '\\r') + echo "Session ID: $SESSION_ID" + -----------------------------`); +}); + +// Handle server shutdown +process.on('SIGINT', async () => { + console.log('Shutting down server...'); + await server.close(); + process.exit(0); +}); \ No newline at end of file diff --git a/src/examples/server/simpleStreamableHttp.ts b/src/examples/server/simpleStreamableHttp.ts new file mode 100644 index 00000000..f0f74439 --- /dev/null +++ b/src/examples/server/simpleStreamableHttp.ts @@ -0,0 +1,231 @@ +import express, { Request, Response } from 'express'; +import { randomUUID } from 'node:crypto'; +import { McpServer } from '../../server/mcp.js'; +import { StreamableHTTPServerTransport } from '../../server/streamableHttp.js'; +import { z } from 'zod'; +import { CallToolResult, GetPromptResult, ReadResourceResult } from '../../types.js'; + +// Create an MCP server with implementation details +const server = new McpServer({ + name: 'simple-streamable-http-server', + version: '1.0.0', +}, { capabilities: { logging: {} } }); + +// Register a simple tool that returns a greeting +server.tool( + 'greet', + 'A simple greeting tool', + { + name: z.string().describe('Name to greet'), + }, + async ({ name }): Promise => { + return { + content: [ + { + type: 'text', + text: `Hello, ${name}!`, + }, + ], + }; + } +); + +// Register a tool that sends multiple greetings with notifications +server.tool( + 'multi-greet', + 'A tool that sends different greetings with delays between them', + { + name: z.string().describe('Name to greet'), + }, + async ({ name }, { sendNotification }): Promise => { + const sleep = (ms: number) => new Promise(resolve => setTimeout(resolve, ms)); + + await sendNotification({ + method: "notifications/message", + params: { level: "debug", data: `Starting multi-greet for ${name}` } + }); + + await sleep(1000); // Wait 1 second before first greeting + + await sendNotification({ + method: "notifications/message", + params: { level: "info", data: `Sending first greeting to ${name}` } + }); + + await sleep(1000); // Wait another second before second greeting + + await sendNotification({ + method: "notifications/message", + params: { level: "info", data: `Sending second greeting to ${name}` } + }); + + return { + content: [ + { + type: 'text', + text: `Good morning, ${name}!`, + } + ], + }; + } +); + +// Register a simple prompt +server.prompt( + 'greeting-template', + 'A simple greeting prompt template', + { + name: z.string().describe('Name to include in greeting'), + }, + async ({ name }): Promise => { + return { + messages: [ + { + role: 'user', + content: { + type: 'text', + text: `Please greet ${name} in a friendly manner.`, + }, + }, + ], + }; + } +); + +// Create a simple resource at a fixed URI +server.resource( + 'greeting-resource', + 'https://example.com/greetings/default', + { mimeType: 'text/plain' }, + async (): Promise => { + return { + contents: [ + { + uri: 'https://example.com/greetings/default', + text: 'Hello, world!', + }, + ], + }; + } +); + +const app = express(); +app.use(express.json()); + +// Map to store transports by session ID +const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {}; + +app.post('/mcp', async (req: Request, res: Response) => { + console.log('Received MCP request:', req.body); + try { + // Check for existing session ID + const sessionId = req.headers['mcp-session-id'] as string | undefined; + let transport: StreamableHTTPServerTransport; + + if (sessionId && transports[sessionId]) { + // Reuse existing transport + transport = transports[sessionId]; + } else if (!sessionId && isInitializeRequest(req.body)) { + // New initialization request + transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: () => randomUUID(), + }); + + // Connect the transport to the MCP server BEFORE handling the request + // so responses can flow back through the same transport + await server.connect(transport); + + // After handling the request, if we get a session ID back, store the transport + await transport.handleRequest(req, res, req.body); + + // Store the transport by session ID for future requests + if (transport.sessionId) { + transports[transport.sessionId] = transport; + } + return; // Already handled + } else { + // Invalid request - no session ID or not initialization request + res.status(400).json({ + jsonrpc: '2.0', + error: { + code: -32000, + message: 'Bad Request: No valid session ID provided', + }, + id: null, + }); + return; + } + + // Handle the request with existing transport - no need to reconnect + // The existing transport is already connected to the server + await transport.handleRequest(req, res, req.body); + } catch (error) { + console.error('Error handling MCP request:', error); + if (!res.headersSent) { + res.status(500).json({ + jsonrpc: '2.0', + error: { + code: -32603, + message: 'Internal server error', + }, + id: null, + }); + } + } +}); + +// Handle GET requests for SSE streams (using built-in support from StreamableHTTP) +app.get('/mcp', async (req: Request, res: Response) => { + const sessionId = req.headers['mcp-session-id'] as string | undefined; + if (!sessionId || !transports[sessionId]) { + res.status(400).send('Invalid or missing session ID'); + return; + } + + console.log(`Establishing SSE stream for session ${sessionId}`); + const transport = transports[sessionId]; + await transport.handleRequest(req, res); +}); + +// Helper function to detect initialize requests +function isInitializeRequest(body: unknown): boolean { + if (Array.isArray(body)) { + return body.some(msg => typeof msg === 'object' && msg !== null && 'method' in msg && msg.method === 'initialize'); + } + return typeof body === 'object' && body !== null && 'method' in body && body.method === 'initialize'; +} + +// Start the server +const PORT = 3000; +app.listen(PORT, () => { + console.log(`MCP Streamable HTTP Server listening on port ${PORT}`); + console.log(`Initialize session with the command below id you are using curl for testing: + ----------------------------- + SESSION_ID=$(curl -X POST \ + -H "Content-Type: application/json" \ + -H "Accept: application/json" \ + -H "Accept: text/event-stream" \ + -d '{ + "jsonrpc": "2.0", + "method": "initialize", + "params": { + "capabilities": {}, + "protocolVersion": "2025-03-26", + "clientInfo": { + "name": "test", + "version": "1.0.0" + } + }, + "id": "1" + }' \ + -i http://localhost:3000/mcp 2>&1 | grep -i "mcp-session-id" | cut -d' ' -f2 | tr -d '\\r') + echo "Session ID: $SESSION_ID" + -----------------------------`); +}); + +// Handle server shutdown +process.on('SIGINT', async () => { + console.log('Shutting down server...'); + await server.close(); + process.exit(0); +}); \ No newline at end of file diff --git a/src/examples/server/standaloneSseWithGetStreamableHttp.ts b/src/examples/server/standaloneSseWithGetStreamableHttp.ts new file mode 100644 index 00000000..f9d3696b --- /dev/null +++ b/src/examples/server/standaloneSseWithGetStreamableHttp.ts @@ -0,0 +1,130 @@ +import express, { Request, Response } from 'express'; +import { randomUUID } from 'node:crypto'; +import { McpServer } from '../../server/mcp.js'; +import { StreamableHTTPServerTransport } from '../../server/streamableHttp.js'; +import { ReadResourceResult } from '../../types.js'; + +// Create an MCP server with implementation details +const server = new McpServer({ + name: 'resource-list-changed-notification-server', + version: '1.0.0', +}); + +// Store transports by session ID to send notifications +const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {}; + +const addResource = (name: string, content: string) => { + const uri = `https://mcp-example.com/dynamic/${encodeURIComponent(name)}`; + server.resource( + name, + uri, + { mimeType: 'text/plain', description: `Dynamic resource: ${name}` }, + async (): Promise => { + return { + contents: [{ uri, text: content }], + }; + } + ); + +}; + +addResource('example-resource', 'Initial content for example-resource'); + +const resourceChangeInterval = setInterval(() => { + const name = randomUUID(); + addResource(name, `Content for ${name}`); +}, 5000); // Change resources every 5 seconds for testing + +const app = express(); +app.use(express.json()); + +app.post('/mcp', async (req: Request, res: Response) => { + console.log('Received MCP request:', req.body); + try { + // Check for existing session ID + const sessionId = req.headers['mcp-session-id'] as string | undefined; + let transport: StreamableHTTPServerTransport; + + if (sessionId && transports[sessionId]) { + // Reuse existing transport + transport = transports[sessionId]; + } else if (!sessionId && isInitializeRequest(req.body)) { + // New initialization request + transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: () => randomUUID(), + }); + + // Connect the transport to the MCP server + await server.connect(transport); + + await transport.handleRequest(req, res, req.body); + + // Store the transport by session ID for future requests + if (transport.sessionId) { + transports[transport.sessionId] = transport; + } + return; // Already handled + } else { + // Invalid request - no session ID or not initialization request + res.status(400).json({ + jsonrpc: '2.0', + error: { + code: -32000, + message: 'Bad Request: No valid session ID provided', + }, + id: null, + }); + return; + } + + // Handle the request with existing transport + await transport.handleRequest(req, res, req.body); + } catch (error) { + console.error('Error handling MCP request:', error); + if (!res.headersSent) { + res.status(500).json({ + jsonrpc: '2.0', + error: { + code: -32603, + message: 'Internal server error', + }, + id: null, + }); + } + } +}); + +// Handle GET requests for SSE streams (now using built-in support from StreamableHTTP) +app.get('/mcp', async (req: Request, res: Response) => { + const sessionId = req.headers['mcp-session-id'] as string | undefined; + if (!sessionId || !transports[sessionId]) { + res.status(400).send('Invalid or missing session ID'); + return; + } + + console.log(`Establishing SSE stream for session ${sessionId}`); + const transport = transports[sessionId]; + await transport.handleRequest(req, res); +}); + +// Helper function to detect initialize requests +function isInitializeRequest(body: unknown): boolean { + if (Array.isArray(body)) { + return body.some(msg => typeof msg === 'object' && msg !== null && 'method' in msg && msg.method === 'initialize'); + } + return typeof body === 'object' && body !== null && 'method' in body && body.method === 'initialize'; +} + +// Start the server +const PORT = 3000; +app.listen(PORT, () => { + console.log(`Server listening on port ${PORT}`); +}); + +// Handle server shutdown +process.on('SIGINT', async () => { + console.log('Shutting down server...'); + clearInterval(resourceChangeInterval); + await server.close(); + process.exit(0); +}); \ No newline at end of file diff --git a/src/server/mcp.test.ts b/src/server/mcp.test.ts index 2e91a568..eaac5c71 100644 --- a/src/server/mcp.test.ts +++ b/src/server/mcp.test.ts @@ -11,6 +11,8 @@ import { ListPromptsResultSchema, GetPromptResultSchema, CompleteResultSchema, + LoggingMessageNotificationSchema, + Notification, } from "../types.js"; import { ResourceTemplate } from "./mcp.js"; import { completable } from "./completable.js"; @@ -35,10 +37,14 @@ describe("McpServer", () => { { capabilities: { logging: {} } }, ); + const notifications: Notification[] = [] const client = new Client({ name: "test client", version: "1.0", }); + client.fallbackNotificationHandler = async (notification) => { + notifications.push(notification) + } const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); @@ -55,6 +61,16 @@ describe("McpServer", () => { data: "Test log message", }), ).resolves.not.toThrow(); + + expect(notifications).toMatchObject([ + { + "method": "notifications/message", + params: { + level: "info", + data: "Test log message", + } + } + ]) }); }); @@ -85,6 +101,8 @@ describe("ResourceTemplate", () => { const abortController = new AbortController(); const result = await template.listCallback?.({ signal: abortController.signal, + sendRequest: () => { throw new Error("Not implemented") }, + sendNotification: () => { throw new Error("Not implemented") } }); expect(result?.resources).toHaveLength(1); expect(list).toHaveBeenCalled(); @@ -97,10 +115,14 @@ describe("tool()", () => { name: "test server", version: "1.0", }); + const notifications: Notification[] = [] const client = new Client({ name: "test client", version: "1.0", }); + client.fallbackNotificationHandler = async (notification) => { + notifications.push(notification) + } mcpServer.tool("test", async () => ({ content: [ @@ -116,7 +138,7 @@ describe("tool()", () => { await Promise.all([ client.connect(clientTransport), - mcpServer.server.connect(serverTransport), + mcpServer.connect(serverTransport), ]); const result = await client.request( @@ -131,6 +153,254 @@ describe("tool()", () => { expect(result.tools[0].inputSchema).toEqual({ type: "object", }); + + // Adding the tool before the connection was established means no notification was sent + expect(notifications).toHaveLength(0) + + // Adding another tool triggers the update notification + mcpServer.tool("test2", async () => ({ + content: [ + { + type: "text", + text: "Test response", + }, + ], + })); + + // Yield event loop to let the notification fly + await new Promise(process.nextTick) + + expect(notifications).toMatchObject([ + { + method: "notifications/tools/list_changed", + } + ]) + }); + + test("should update existing tool", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const notifications: Notification[] = [] + const client = new Client({ + name: "test client", + version: "1.0", + }); + client.fallbackNotificationHandler = async (notification) => { + notifications.push(notification) + } + + // Register initial tool + const tool = mcpServer.tool("test", async () => ({ + content: [ + { + type: "text", + text: "Initial response", + }, + ], + })); + + // Update the tool + tool.update({ + callback: async () => ({ + content: [ + { + type: "text", + text: "Updated response", + }, + ], + }) + }); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.connect(serverTransport), + ]); + + // Call the tool and verify we get the updated response + const result = await client.request( + { + method: "tools/call", + params: { + name: "test", + }, + }, + CallToolResultSchema, + ); + + expect(result.content).toEqual([ + { + type: "text", + text: "Updated response", + }, + ]); + + // Update happened before transport was connected, so no notifications should be expected + expect(notifications).toHaveLength(0) + }); + + test("should update tool with schema", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const notifications: Notification[] = [] + const client = new Client({ + name: "test client", + version: "1.0", + }); + client.fallbackNotificationHandler = async (notification) => { + notifications.push(notification) + } + + // Register initial tool + const tool = mcpServer.tool( + "test", + { + name: z.string(), + }, + async ({ name }) => ({ + content: [ + { + type: "text", + text: `Initial: ${name}`, + }, + ], + }), + ); + + // Update the tool with a different schema + tool.update({ + paramsSchema: { + name: z.string(), + value: z.number(), + }, + callback: async ({name, value}) => ({ + content: [ + { + type: "text", + text: `Updated: ${name}, ${value}`, + }, + ], + }) + }); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.connect(serverTransport), + ]); + + // Verify the schema was updated + const listResult = await client.request( + { + method: "tools/list", + }, + ListToolsResultSchema, + ); + + expect(listResult.tools[0].inputSchema).toMatchObject({ + properties: { + name: { type: "string" }, + value: { type: "number" }, + }, + }); + + // Call the tool with the new schema + const callResult = await client.request( + { + method: "tools/call", + params: { + name: "test", + arguments: { + name: "test", + value: 42, + }, + }, + }, + CallToolResultSchema, + ); + + expect(callResult.content).toEqual([ + { + type: "text", + text: "Updated: test, 42", + }, + ]); + + // Update happened before transport was connected, so no notifications should be expected + expect(notifications).toHaveLength(0) + }); + + test("should send tool list changed notifications when connected", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const notifications: Notification[] = [] + const client = new Client({ + name: "test client", + version: "1.0", + }); + client.fallbackNotificationHandler = async (notification) => { + notifications.push(notification) + } + + // Register initial tool + const tool = mcpServer.tool("test", async () => ({ + content: [ + { + type: "text", + text: "Test response", + }, + ], + })); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.connect(serverTransport), + ]); + + expect(notifications).toHaveLength(0) + + // Now update the tool + tool.update({ + callback: async () => ({ + content: [ + { + type: "text", + text: "Updated response", + }, + ], + }) + }); + + // Yield event loop to let the notification fly + await new Promise(process.nextTick) + + expect(notifications).toMatchObject([ + { method: "notifications/tools/list_changed" } + ]) + + // Now delete the tool + tool.remove(); + + // Yield event loop to let the notification fly + await new Promise(process.nextTick) + + expect(notifications).toMatchObject([ + { method: "notifications/tools/list_changed" }, + { method: "notifications/tools/list_changed" }, + ]) }); test("should register tool with args schema", async () => { @@ -318,7 +588,7 @@ describe("tool()", () => { // This should succeed mcpServer.tool("tool1", () => ({ content: [] })); - + // This should also succeed and not throw about request handlers mcpServer.tool("tool2", () => ({ content: [] })); }); @@ -376,6 +646,63 @@ describe("tool()", () => { expect(receivedSessionId).toBe("test-session-123"); }); + test("should provide sendNotification within tool call", async () => { + const mcpServer = new McpServer( + { + name: "test server", + version: "1.0", + }, + { capabilities: { logging: {} } }, + ); + + const client = new Client( + { + name: "test client", + version: "1.0", + }, + { + capabilities: { + tools: {}, + }, + }, + ); + + let receivedLogMessage: string | undefined; + const loggingMessage = "hello here is log message 1"; + + client.setNotificationHandler(LoggingMessageNotificationSchema, (notification) => { + receivedLogMessage = notification.params.data as string; + }); + + mcpServer.tool("test-tool", async ({ sendNotification }) => { + await sendNotification({ method: "notifications/message", params: { level: "debug", data: loggingMessage } }); + return { + content: [ + { + type: "text", + text: "Test response", + }, + ], + }; + }); + + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + await client.request( + { + method: "tools/call", + params: { + name: "test-tool", + }, + }, + CallToolResultSchema, + ); + expect(receivedLogMessage).toBe(loggingMessage); + }); + test("should allow client to call server tools", async () => { const mcpServer = new McpServer({ name: "test server", @@ -577,6 +904,312 @@ describe("resource()", () => { expect(result.resources[0].uri).toBe("test://resource"); }); + test("should update resource with uri", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const notifications: Notification[] = []; + const client = new Client({ + name: "test client", + version: "1.0", + }); + client.fallbackNotificationHandler = async (notification) => { + notifications.push(notification); + }; + + // Register initial resource + const resource = mcpServer.resource("test", "test://resource", async () => ({ + contents: [ + { + uri: "test://resource", + text: "Initial content", + }, + ], + })); + + // Update the resource + resource.update({ + callback: async () => ({ + contents: [ + { + uri: "test://resource", + text: "Updated content", + }, + ], + }) + }); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.connect(serverTransport), + ]); + + // Read the resource and verify we get the updated content + const result = await client.request( + { + method: "resources/read", + params: { + uri: "test://resource", + }, + }, + ReadResourceResultSchema, + ); + + expect(result.contents).toHaveLength(1); + expect(result.contents[0].text).toBe("Updated content"); + + // Update happened before transport was connected, so no notifications should be expected + expect(notifications).toHaveLength(0); + }); + + test("should update resource template", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const notifications: Notification[] = []; + const client = new Client({ + name: "test client", + version: "1.0", + }); + client.fallbackNotificationHandler = async (notification) => { + notifications.push(notification); + }; + + // Register initial resource template + const resourceTemplate = mcpServer.resource( + "test", + new ResourceTemplate("test://resource/{id}", { list: undefined }), + async (uri) => ({ + contents: [ + { + uri: uri.href, + text: "Initial content", + }, + ], + }), + ); + + // Update the resource template + resourceTemplate.update({ + callback: async (uri) => ({ + contents: [ + { + uri: uri.href, + text: "Updated content", + }, + ], + }) + }); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.connect(serverTransport), + ]); + + // Read the resource and verify we get the updated content + const result = await client.request( + { + method: "resources/read", + params: { + uri: "test://resource/123", + }, + }, + ReadResourceResultSchema, + ); + + expect(result.contents).toHaveLength(1); + expect(result.contents[0].text).toBe("Updated content"); + + // Update happened before transport was connected, so no notifications should be expected + expect(notifications).toHaveLength(0); + }); + + test("should send resource list changed notification when connected", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const notifications: Notification[] = []; + const client = new Client({ + name: "test client", + version: "1.0", + }); + client.fallbackNotificationHandler = async (notification) => { + notifications.push(notification); + }; + + // Register initial resource + const resource = 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.connect(serverTransport), + ]); + + expect(notifications).toHaveLength(0); + + // Now update the resource while connected + resource.update({ + callback: async () => ({ + contents: [ + { + uri: "test://resource", + text: "Updated content", + }, + ], + }) + }); + + // Yield event loop to let the notification fly + await new Promise(process.nextTick); + + expect(notifications).toMatchObject([ + { method: "notifications/resources/list_changed" } + ]); + }); + + test("should remove resource and send notification when connected", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const notifications: Notification[] = []; + const client = new Client({ + name: "test client", + version: "1.0", + }); + client.fallbackNotificationHandler = async (notification) => { + notifications.push(notification); + }; + + // Register initial resources + const resource1 = mcpServer.resource("resource1", "test://resource1", async () => ({ + contents: [{ uri: "test://resource1", text: "Resource 1 content" }], + })); + + mcpServer.resource("resource2", "test://resource2", async () => ({ + contents: [{ uri: "test://resource2", text: "Resource 2 content" }], + })); + + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.connect(serverTransport), + ]); + + // Verify both resources are registered + let result = await client.request( + { method: "resources/list" }, + ListResourcesResultSchema, + ); + + expect(result.resources).toHaveLength(2); + + expect(notifications).toHaveLength(0); + + // Remove a resource + resource1.remove() + + // Yield event loop to let the notification fly + await new Promise(process.nextTick); + + // Should have sent notification + expect(notifications).toMatchObject([ + { method: "notifications/resources/list_changed" } + ]); + + // Verify the resource was removed + result = await client.request( + { method: "resources/list" }, + ListResourcesResultSchema, + ); + + expect(result.resources).toHaveLength(1); + expect(result.resources[0].uri).toBe("test://resource2"); + }); + + test("should remove resource template and send notification when connected", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const notifications: Notification[] = []; + const client = new Client({ + name: "test client", + version: "1.0", + }); + client.fallbackNotificationHandler = async (notification) => { + notifications.push(notification); + }; + + // Register resource template + const resourceTemplate = mcpServer.resource( + "template", + new ResourceTemplate("test://resource/{id}", { list: undefined }), + async (uri) => ({ + contents: [ + { + uri: uri.href, + text: "Template content", + }, + ], + }), + ); + + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.connect(serverTransport), + ]); + + // Verify template is registered + const result = await client.request( + { method: "resources/templates/list" }, + ListResourceTemplatesResultSchema, + ); + + expect(result.resourceTemplates).toHaveLength(1); + expect(notifications).toHaveLength(0); + + // Remove the template + resourceTemplate.remove() + + // Yield event loop to let the notification fly + await new Promise(process.nextTick); + + // Should have sent notification + expect(notifications).toMatchObject([ + { method: "notifications/resources/list_changed" } + ]); + + // Verify the template was removed + const result2 = await client.request( + { method: "resources/templates/list" }, + ListResourceTemplatesResultSchema, + ); + + expect(result2.resourceTemplates).toHaveLength(0); + }); + test("should register resource with metadata", async () => { const mcpServer = new McpServer({ name: "test server", @@ -815,7 +1448,7 @@ describe("resource()", () => { }, ], })); - + // This should also succeed and not throw about request handlers mcpServer.resource("resource2", "test://resource2", async () => ({ contents: [ @@ -1113,6 +1746,303 @@ describe("prompt()", () => { expect(result.prompts[0].name).toBe("test"); expect(result.prompts[0].arguments).toBeUndefined(); }); + test("should update existing prompt", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const notifications: Notification[] = []; + const client = new Client({ + name: "test client", + version: "1.0", + }); + client.fallbackNotificationHandler = async (notification) => { + notifications.push(notification); + }; + + // Register initial prompt + const prompt = mcpServer.prompt("test", async () => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: "Initial response", + }, + }, + ], + })); + + // Update the prompt + prompt.update({ + callback: async () => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: "Updated response", + }, + }, + ], + }) + }); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.connect(serverTransport), + ]); + + // Call the prompt and verify we get the updated response + const result = await client.request( + { + method: "prompts/get", + params: { + name: "test", + }, + }, + GetPromptResultSchema, + ); + + expect(result.messages).toHaveLength(1); + expect(result.messages[0].content.text).toBe("Updated response"); + + // Update happened before transport was connected, so no notifications should be expected + expect(notifications).toHaveLength(0); + }); + + test("should update prompt with schema", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const notifications: Notification[] = []; + const client = new Client({ + name: "test client", + version: "1.0", + }); + client.fallbackNotificationHandler = async (notification) => { + notifications.push(notification); + }; + + // Register initial prompt + const prompt = mcpServer.prompt( + "test", + { + name: z.string(), + }, + async ({ name }) => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: `Initial: ${name}`, + }, + }, + ], + }), + ); + + // Update the prompt with a different schema + prompt.update({ + argsSchema: { + name: z.string(), + value: z.string(), + }, + callback: async ({name, value}) => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: `Updated: ${name}, ${value}`, + }, + }, + ], + }) + }); + + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.connect(serverTransport), + ]); + + // Verify the schema was updated + const listResult = await client.request( + { + method: "prompts/list", + }, + ListPromptsResultSchema, + ); + + expect(listResult.prompts[0].arguments).toHaveLength(2); + expect(listResult.prompts[0].arguments?.map(a => a.name).sort()).toEqual(["name", "value"]); + + // Call the prompt with the new schema + const getResult = await client.request( + { + method: "prompts/get", + params: { + name: "test", + arguments: { + name: "test", + value: "value", + }, + }, + }, + GetPromptResultSchema, + ); + + expect(getResult.messages).toHaveLength(1); + expect(getResult.messages[0].content.text).toBe("Updated: test, value"); + + // Update happened before transport was connected, so no notifications should be expected + expect(notifications).toHaveLength(0); + }); + + test("should send prompt list changed notification when connected", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const notifications: Notification[] = []; + const client = new Client({ + name: "test client", + version: "1.0", + }); + client.fallbackNotificationHandler = async (notification) => { + notifications.push(notification); + }; + + // Register initial prompt + const prompt = 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.connect(serverTransport), + ]); + + expect(notifications).toHaveLength(0); + + // Now update the prompt while connected + prompt.update({ + callback: async () => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: "Updated response", + }, + }, + ], + }) + }); + + // Yield event loop to let the notification fly + await new Promise(process.nextTick); + + expect(notifications).toMatchObject([ + { method: "notifications/prompts/list_changed" } + ]); + }); + + test("should remove prompt and send notification when connected", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + const notifications: Notification[] = []; + const client = new Client({ + name: "test client", + version: "1.0", + }); + client.fallbackNotificationHandler = async (notification) => { + notifications.push(notification); + }; + + // Register initial prompts + const prompt1 = mcpServer.prompt("prompt1", async () => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: "Prompt 1 response", + }, + }, + ], + })); + + mcpServer.prompt("prompt2", async () => ({ + messages: [ + { + role: "assistant", + content: { + type: "text", + text: "Prompt 2 response", + }, + }, + ], + })); + + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.connect(serverTransport), + ]); + + // Verify both prompts are registered + let result = await client.request( + { method: "prompts/list" }, + ListPromptsResultSchema, + ); + + expect(result.prompts).toHaveLength(2); + expect(result.prompts.map(p => p.name).sort()).toEqual(["prompt1", "prompt2"]); + + expect(notifications).toHaveLength(0); + + // Remove a prompt + prompt1.remove() + + // Yield event loop to let the notification fly + await new Promise(process.nextTick); + + // Should have sent notification + expect(notifications).toMatchObject([ + { method: "notifications/prompts/list_changed" } + ]); + + // Verify the prompt was removed + result = await client.request( + { method: "prompts/list" }, + ListPromptsResultSchema, + ); + + expect(result.prompts).toHaveLength(1); + expect(result.prompts[0].name).toBe("prompt2"); + }); test("should register prompt with args schema", async () => { const mcpServer = new McpServer({ @@ -1321,7 +2251,7 @@ describe("prompt()", () => { }, ], })); - + // This should also succeed and not throw about request handlers mcpServer.prompt("prompt2", async () => ({ messages: [ diff --git a/src/server/mcp.ts b/src/server/mcp.ts index 8f4a909c..652f9774 100644 --- a/src/server/mcp.ts +++ b/src/server/mcp.ts @@ -37,6 +37,8 @@ import { PromptArgument, GetPromptResult, ReadResourceResult, + ServerRequest, + ServerNotification, } from "../types.js"; import { Completable, CompletableDef } from "./completable.js"; import { UriTemplate, Variables } from "../shared/uriTemplate.js"; @@ -96,13 +98,17 @@ export class McpServer { ); this.server.registerCapabilities({ - tools: {}, - }); + tools: { + listChanged: true + } + }) this.server.setRequestHandler( ListToolsRequestSchema, (): ListToolsResult => ({ - tools: Object.entries(this._registeredTools).map( + tools: Object.entries(this._registeredTools).filter( + ([, tool]) => tool.enabled, + ).map( ([name, tool]): Tool => { return { name, @@ -129,6 +135,13 @@ export class McpServer { ); } + if (!tool.enabled) { + throw new McpError( + ErrorCode.InvalidParams, + `Tool ${request.params.name} disabled`, + ); + } + if (tool.inputSchema) { const parseResult = await tool.inputSchema.safeParseAsync( request.params.arguments, @@ -218,7 +231,14 @@ export class McpServer { if (!prompt) { throw new McpError( ErrorCode.InvalidParams, - `Prompt ${request.params.ref.name} not found`, + `Prompt ${ref.name} not found`, + ); + } + + if (!prompt.enabled) { + throw new McpError( + ErrorCode.InvalidParams, + `Prompt ${ref.name} disabled`, ); } @@ -285,13 +305,17 @@ export class McpServer { ); this.server.registerCapabilities({ - resources: {}, - }); + resources: { + listChanged: true + } + }) this.server.setRequestHandler( ListResourcesRequestSchema, async (request, extra) => { - const resources = Object.entries(this._registeredResources).map( + const resources = Object.entries(this._registeredResources).filter( + ([_, resource]) => resource.enabled, + ).map( ([uri, resource]) => ({ uri, name: resource.name, @@ -343,6 +367,12 @@ export class McpServer { // First check for exact resource match const resource = this._registeredResources[uri.toString()]; if (resource) { + if (!resource.enabled) { + throw new McpError( + ErrorCode.InvalidParams, + `Resource ${uri} disabled`, + ); + } return resource.readCallback(uri, extra); } @@ -385,13 +415,17 @@ export class McpServer { ); this.server.registerCapabilities({ - prompts: {}, - }); + prompts: { + listChanged: true + } + }) this.server.setRequestHandler( ListPromptsRequestSchema, (): ListPromptsResult => ({ - prompts: Object.entries(this._registeredPrompts).map( + prompts: Object.entries(this._registeredPrompts).filter( + ([, prompt]) => prompt.enabled, + ).map( ([name, prompt]): Prompt => { return { name, @@ -416,6 +450,13 @@ export class McpServer { ); } + if (!prompt.enabled) { + throw new McpError( + ErrorCode.InvalidParams, + `Prompt ${request.params.name} disabled`, + ); + } + if (prompt.argsSchema) { const parseResult = await prompt.argsSchema.safeParseAsync( request.params.arguments, @@ -445,7 +486,7 @@ export class McpServer { /** * 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; + resource(name: string, uri: string, readCallback: ReadResourceCallback): RegisteredResource; /** * Registers a resource `name` at a fixed URI with metadata, which will use the given callback to respond to read requests. @@ -455,7 +496,7 @@ export class McpServer { uri: string, metadata: ResourceMetadata, readCallback: ReadResourceCallback, - ): void; + ): RegisteredResource; /** * Registers a resource `name` with a template pattern, which will use the given callback to respond to read requests. @@ -464,7 +505,7 @@ export class McpServer { name: string, template: ResourceTemplate, readCallback: ReadResourceTemplateCallback, - ): void; + ): RegisteredResourceTemplate; /** * Registers a resource `name` with a template pattern and metadata, which will use the given callback to respond to read requests. @@ -474,13 +515,13 @@ export class McpServer { template: ResourceTemplate, metadata: ResourceMetadata, readCallback: ReadResourceTemplateCallback, - ): void; + ): RegisteredResourceTemplate; resource( name: string, uriOrTemplate: string | ResourceTemplate, ...rest: unknown[] - ): void { + ): RegisteredResource | RegisteredResourceTemplate { let metadata: ResourceMetadata | undefined; if (typeof rest[0] === "object") { metadata = rest.shift() as ResourceMetadata; @@ -495,35 +536,73 @@ export class McpServer { throw new Error(`Resource ${uriOrTemplate} is already registered`); } - this._registeredResources[uriOrTemplate] = { + const registeredResource: RegisteredResource = { name, metadata, readCallback: readCallback as ReadResourceCallback, + enabled: true, + disable: () => registeredResource.update({ enabled: false }), + enable: () => registeredResource.update({ enabled: true }), + remove: () => registeredResource.update({ uri: null }), + update: (updates) => { + if (typeof updates.uri !== "undefined" && updates.uri !== uriOrTemplate) { + delete this._registeredResources[uriOrTemplate] + if (updates.uri) this._registeredResources[updates.uri] = registeredResource + } + if (typeof updates.name !== "undefined") registeredResource.name = updates.name + if (typeof updates.metadata !== "undefined") registeredResource.metadata = updates.metadata + if (typeof updates.callback !== "undefined") registeredResource.readCallback = updates.callback + if (typeof updates.enabled !== "undefined") registeredResource.enabled = updates.enabled + this.sendResourceListChanged() + }, }; + this._registeredResources[uriOrTemplate] = registeredResource; + + this.setResourceRequestHandlers(); + this.sendResourceListChanged(); + return registeredResource; } else { if (this._registeredResourceTemplates[name]) { throw new Error(`Resource template ${name} is already registered`); } - this._registeredResourceTemplates[name] = { + const registeredResourceTemplate: RegisteredResourceTemplate = { resourceTemplate: uriOrTemplate, metadata, readCallback: readCallback as ReadResourceTemplateCallback, + enabled: true, + disable: () => registeredResourceTemplate.update({ enabled: false }), + enable: () => registeredResourceTemplate.update({ enabled: true }), + remove: () => registeredResourceTemplate.update({ name: null }), + update: (updates) => { + if (typeof updates.name !== "undefined" && updates.name !== name) { + delete this._registeredResourceTemplates[name] + if (updates.name) this._registeredResourceTemplates[updates.name] = registeredResourceTemplate + } + if (typeof updates.template !== "undefined") registeredResourceTemplate.resourceTemplate = updates.template + if (typeof updates.metadata !== "undefined") registeredResourceTemplate.metadata = updates.metadata + if (typeof updates.callback !== "undefined") registeredResourceTemplate.readCallback = updates.callback + if (typeof updates.enabled !== "undefined") registeredResourceTemplate.enabled = updates.enabled + this.sendResourceListChanged() + }, }; - } + this._registeredResourceTemplates[name] = registeredResourceTemplate; - this.setResourceRequestHandlers(); + this.setResourceRequestHandlers(); + this.sendResourceListChanged(); + return registeredResourceTemplate; + } } /** * Registers a zero-argument tool `name`, which will run the given function when the client calls it. */ - tool(name: string, cb: ToolCallback): void; + tool(name: string, cb: ToolCallback): RegisteredTool; /** * 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; + tool(name: string, description: string, cb: ToolCallback): RegisteredTool; /** * 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. @@ -532,7 +611,7 @@ export class McpServer { name: string, paramsSchema: Args, cb: ToolCallback, - ): void; + ): RegisteredTool; /** * 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. @@ -542,9 +621,9 @@ export class McpServer { description: string, paramsSchema: Args, cb: ToolCallback, - ): void; + ): RegisteredTool; - tool(name: string, ...rest: unknown[]): void { + tool(name: string, ...rest: unknown[]): RegisteredTool { if (this._registeredTools[name]) { throw new Error(`Tool ${name} is already registered`); } @@ -560,25 +639,44 @@ export class McpServer { } const cb = rest[0] as ToolCallback; - this._registeredTools[name] = { + const registeredTool: RegisteredTool = { description, inputSchema: paramsSchema === undefined ? undefined : z.object(paramsSchema), callback: cb, + enabled: true, + disable: () => registeredTool.update({ enabled: false }), + enable: () => registeredTool.update({ enabled: true }), + remove: () => registeredTool.update({ name: null }), + update: (updates) => { + if (typeof updates.name !== "undefined" && updates.name !== name) { + delete this._registeredTools[name] + if (updates.name) this._registeredTools[updates.name] = registeredTool + } + if (typeof updates.description !== "undefined") registeredTool.description = updates.description + if (typeof updates.paramsSchema !== "undefined") registeredTool.inputSchema = z.object(updates.paramsSchema) + if (typeof updates.callback !== "undefined") registeredTool.callback = updates.callback + if (typeof updates.enabled !== "undefined") registeredTool.enabled = updates.enabled + this.sendToolListChanged() + }, }; + this._registeredTools[name] = registeredTool; this.setToolRequestHandlers(); + this.sendToolListChanged() + + return registeredTool } /** * Registers a zero-argument prompt `name`, which will run the given function when the client calls it. */ - prompt(name: string, cb: PromptCallback): void; + prompt(name: string, cb: PromptCallback): RegisteredPrompt; /** * 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; + prompt(name: string, description: string, cb: PromptCallback): RegisteredPrompt; /** * 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. @@ -587,7 +685,7 @@ export class McpServer { name: string, argsSchema: Args, cb: PromptCallback, - ): void; + ): RegisteredPrompt; /** * 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. @@ -597,9 +695,9 @@ export class McpServer { description: string, argsSchema: Args, cb: PromptCallback, - ): void; + ): RegisteredPrompt; - prompt(name: string, ...rest: unknown[]): void { + prompt(name: string, ...rest: unknown[]): RegisteredPrompt { if (this._registeredPrompts[name]) { throw new Error(`Prompt ${name} is already registered`); } @@ -615,13 +713,67 @@ export class McpServer { } const cb = rest[0] as PromptCallback; - this._registeredPrompts[name] = { + const registeredPrompt: RegisteredPrompt = { description, argsSchema: argsSchema === undefined ? undefined : z.object(argsSchema), callback: cb, + enabled: true, + disable: () => registeredPrompt.update({ enabled: false }), + enable: () => registeredPrompt.update({ enabled: true }), + remove: () => registeredPrompt.update({ name: null }), + update: (updates) => { + if (typeof updates.name !== "undefined" && updates.name !== name) { + delete this._registeredPrompts[name] + if (updates.name) this._registeredPrompts[updates.name] = registeredPrompt + } + if (typeof updates.description !== "undefined") registeredPrompt.description = updates.description + if (typeof updates.argsSchema !== "undefined") registeredPrompt.argsSchema = z.object(updates.argsSchema) + if (typeof updates.callback !== "undefined") registeredPrompt.callback = updates.callback + if (typeof updates.enabled !== "undefined") registeredPrompt.enabled = updates.enabled + this.sendPromptListChanged() + }, }; + this._registeredPrompts[name] = registeredPrompt; this.setPromptRequestHandlers(); + this.sendPromptListChanged() + + return registeredPrompt + } + + /** + * Checks if the server is connected to a transport. + * @returns True if the server is connected + */ + isConnected() { + return this.server.transport !== undefined + } + + /** + * Sends a resource list changed event to the client, if connected. + */ + sendResourceListChanged() { + if (this.isConnected()) { + this.server.sendResourceListChanged(); + } + } + + /** + * Sends a tool list changed event to the client, if connected. + */ + sendToolListChanged() { + if (this.isConnected()) { + this.server.sendToolListChanged(); + } + } + + /** + * Sends a prompt list changed event to the client, if connected. + */ + sendPromptListChanged() { + if (this.isConnected()) { + this.server.sendPromptListChanged(); + } } } @@ -694,14 +846,19 @@ export type ToolCallback = Args extends ZodRawShape ? ( args: z.objectOutputType, - extra: RequestHandlerExtra, + extra: RequestHandlerExtra, ) => CallToolResult | Promise - : (extra: RequestHandlerExtra) => CallToolResult | Promise; + : (extra: RequestHandlerExtra) => CallToolResult | Promise; -type RegisteredTool = { +export type RegisteredTool = { description?: string; inputSchema?: AnyZodObject; callback: ToolCallback; + enabled: boolean; + enable(): void; + disable(): void; + update(updates: { name?: string | null, description?: string, paramsSchema?: Args, callback?: ToolCallback, enabled?: boolean }): void + remove(): void }; const EMPTY_OBJECT_JSON_SCHEMA = { @@ -717,7 +874,7 @@ export type ResourceMetadata = Omit; * Callback to list all resources matching a given template. */ export type ListResourcesCallback = ( - extra: RequestHandlerExtra, + extra: RequestHandlerExtra, ) => ListResourcesResult | Promise; /** @@ -725,13 +882,18 @@ export type ListResourcesCallback = ( */ export type ReadResourceCallback = ( uri: URL, - extra: RequestHandlerExtra, + extra: RequestHandlerExtra, ) => ReadResourceResult | Promise; -type RegisteredResource = { +export type RegisteredResource = { name: string; metadata?: ResourceMetadata; readCallback: ReadResourceCallback; + enabled: boolean; + enable(): void; + disable(): void; + update(updates: { name?: string, uri?: string | null, metadata?: ResourceMetadata, callback?: ReadResourceCallback, enabled?: boolean }): void + remove(): void }; /** @@ -740,13 +902,18 @@ type RegisteredResource = { export type ReadResourceTemplateCallback = ( uri: URL, variables: Variables, - extra: RequestHandlerExtra, + extra: RequestHandlerExtra, ) => ReadResourceResult | Promise; -type RegisteredResourceTemplate = { +export type RegisteredResourceTemplate = { resourceTemplate: ResourceTemplate; metadata?: ResourceMetadata; readCallback: ReadResourceTemplateCallback; + enabled: boolean; + enable(): void; + disable(): void; + update(updates: { name?: string | null, template?: ResourceTemplate, metadata?: ResourceMetadata, callback?: ReadResourceTemplateCallback, enabled?: boolean }): void + remove(): void }; type PromptArgsRawShape = { @@ -760,14 +927,19 @@ export type PromptCallback< > = Args extends PromptArgsRawShape ? ( args: z.objectOutputType, - extra: RequestHandlerExtra, + extra: RequestHandlerExtra, ) => GetPromptResult | Promise - : (extra: RequestHandlerExtra) => GetPromptResult | Promise; + : (extra: RequestHandlerExtra) => GetPromptResult | Promise; -type RegisteredPrompt = { +export type RegisteredPrompt = { description?: string; argsSchema?: ZodObject; callback: PromptCallback; + enabled: boolean; + enable(): void; + disable(): void; + update(updates: { name?: string | null, description?: string, argsSchema?: Args, callback?: PromptCallback, enabled?: boolean }): void + remove(): void }; function promptArgumentsFromSchema( diff --git a/src/server/sse.test.ts b/src/server/sse.test.ts new file mode 100644 index 00000000..2fd2c042 --- /dev/null +++ b/src/server/sse.test.ts @@ -0,0 +1,109 @@ +import http from 'http'; +import { jest } from '@jest/globals'; +import { SSEServerTransport } from './sse.js'; + +const createMockResponse = () => { + const res = { + writeHead: jest.fn(), + write: jest.fn().mockReturnValue(true), + on: jest.fn(), + }; + res.writeHead.mockReturnThis(); + res.on.mockReturnThis(); + + return res as unknown as http.ServerResponse; +}; + +describe('SSEServerTransport', () => { + describe('start method', () => { + it('should correctly append sessionId to a simple relative endpoint', async () => { + const mockRes = createMockResponse(); + const endpoint = '/messages'; + const transport = new SSEServerTransport(endpoint, mockRes); + const expectedSessionId = transport.sessionId; + + await transport.start(); + + expect(mockRes.writeHead).toHaveBeenCalledWith(200, expect.any(Object)); + expect(mockRes.write).toHaveBeenCalledTimes(1); + expect(mockRes.write).toHaveBeenCalledWith( + `event: endpoint\ndata: /messages?sessionId=${expectedSessionId}\n\n` + ); + }); + + it('should correctly append sessionId to an endpoint with existing query parameters', async () => { + const mockRes = createMockResponse(); + const endpoint = '/messages?foo=bar&baz=qux'; + const transport = new SSEServerTransport(endpoint, mockRes); + const expectedSessionId = transport.sessionId; + + await transport.start(); + + expect(mockRes.writeHead).toHaveBeenCalledWith(200, expect.any(Object)); + expect(mockRes.write).toHaveBeenCalledTimes(1); + expect(mockRes.write).toHaveBeenCalledWith( + `event: endpoint\ndata: /messages?foo=bar&baz=qux&sessionId=${expectedSessionId}\n\n` + ); + }); + + it('should correctly append sessionId to an endpoint with a hash fragment', async () => { + const mockRes = createMockResponse(); + const endpoint = '/messages#section1'; + const transport = new SSEServerTransport(endpoint, mockRes); + const expectedSessionId = transport.sessionId; + + await transport.start(); + + expect(mockRes.writeHead).toHaveBeenCalledWith(200, expect.any(Object)); + expect(mockRes.write).toHaveBeenCalledTimes(1); + expect(mockRes.write).toHaveBeenCalledWith( + `event: endpoint\ndata: /messages?sessionId=${expectedSessionId}#section1\n\n` + ); + }); + + it('should correctly append sessionId to an endpoint with query parameters and a hash fragment', async () => { + const mockRes = createMockResponse(); + const endpoint = '/messages?key=value#section2'; + const transport = new SSEServerTransport(endpoint, mockRes); + const expectedSessionId = transport.sessionId; + + await transport.start(); + + expect(mockRes.writeHead).toHaveBeenCalledWith(200, expect.any(Object)); + expect(mockRes.write).toHaveBeenCalledTimes(1); + expect(mockRes.write).toHaveBeenCalledWith( + `event: endpoint\ndata: /messages?key=value&sessionId=${expectedSessionId}#section2\n\n` + ); + }); + + it('should correctly handle the root path endpoint "/"', async () => { + const mockRes = createMockResponse(); + const endpoint = '/'; + const transport = new SSEServerTransport(endpoint, mockRes); + const expectedSessionId = transport.sessionId; + + await transport.start(); + + expect(mockRes.writeHead).toHaveBeenCalledWith(200, expect.any(Object)); + expect(mockRes.write).toHaveBeenCalledTimes(1); + expect(mockRes.write).toHaveBeenCalledWith( + `event: endpoint\ndata: /?sessionId=${expectedSessionId}\n\n` + ); + }); + + it('should correctly handle an empty string endpoint ""', async () => { + const mockRes = createMockResponse(); + const endpoint = ''; + const transport = new SSEServerTransport(endpoint, mockRes); + const expectedSessionId = transport.sessionId; + + await transport.start(); + + expect(mockRes.writeHead).toHaveBeenCalledWith(200, expect.any(Object)); + expect(mockRes.write).toHaveBeenCalledTimes(1); + expect(mockRes.write).toHaveBeenCalledWith( + `event: endpoint\ndata: /?sessionId=${expectedSessionId}\n\n` + ); + }); + }); +}); diff --git a/src/server/sse.ts b/src/server/sse.ts index ae993dd7..03f6fefc 100644 --- a/src/server/sse.ts +++ b/src/server/sse.ts @@ -5,6 +5,7 @@ import { JSONRPCMessage, JSONRPCMessageSchema } from "../types.js"; import getRawBody from "raw-body"; import contentType from "content-type"; import { AuthInfo } from "./auth/types.js"; +import { URL } from 'url'; const MAXIMUM_MESSAGE_SIZE = "4mb"; @@ -50,8 +51,17 @@ export class SSEServerTransport implements Transport { }); // Send the endpoint event + // Use a dummy base URL because this._endpoint is relative. + // This allows using URL/URLSearchParams for robust parameter handling. + const dummyBase = 'http://localhost'; // Any valid base works + const endpointUrl = new URL(this._endpoint, dummyBase); + endpointUrl.searchParams.set('sessionId', this._sessionId); + + // Reconstruct the relative URL string (pathname + search + hash) + const relativeUrlWithSession = endpointUrl.pathname + endpointUrl.search + endpointUrl.hash; + this.res.write( - `event: endpoint\ndata: ${encodeURI(this._endpoint)}?sessionId=${this._sessionId}\n\n`, + `event: endpoint\ndata: ${relativeUrlWithSession}\n\n`, ); this._sseResponse = this.res; diff --git a/src/server/streamableHttp.test.ts b/src/server/streamableHttp.test.ts new file mode 100644 index 00000000..85fcae2f --- /dev/null +++ b/src/server/streamableHttp.test.ts @@ -0,0 +1,1590 @@ +import { IncomingMessage, ServerResponse } from "node:http"; +import { StreamableHTTPServerTransport } from "./streamableHttp.js"; +import { JSONRPCMessage } from "../types.js"; +import { Readable } from "node:stream"; +import { randomUUID } from "node:crypto"; +// Mock IncomingMessage +function createMockRequest(options: { + method: string; + headers: Record; + body?: string; +}): IncomingMessage { + const readable = new Readable(); + readable._read = () => { }; + if (options.body) { + readable.push(options.body); + readable.push(null); + } + + return Object.assign(readable, { + method: options.method, + headers: options.headers, + }) as IncomingMessage; +} + +// Mock ServerResponse +function createMockResponse(): jest.Mocked { + const response = { + writeHead: jest.fn().mockReturnThis(), + write: jest.fn().mockReturnThis(), + end: jest.fn().mockReturnThis(), + on: jest.fn().mockReturnThis(), + emit: jest.fn().mockReturnThis(), + getHeader: jest.fn(), + setHeader: jest.fn(), + flushHeaders: jest.fn(), + } as unknown as jest.Mocked; + return response; +} + +describe("StreamableHTTPServerTransport", () => { + let transport: StreamableHTTPServerTransport; + let mockResponse: jest.Mocked; + let mockRequest: string; + + beforeEach(() => { + transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: () => randomUUID(), + }); + mockResponse = createMockResponse(); + mockRequest = JSON.stringify({ + jsonrpc: "2.0", + method: "test", + params: {}, + id: 1, + }); + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + describe("Session Management", () => { + it("should generate session ID during initialization", async () => { + const initializeMessage: JSONRPCMessage = { + jsonrpc: "2.0", + method: "initialize", + params: { + clientInfo: { name: "test-client", version: "1.0" }, + protocolVersion: "2025-03-26" + }, + id: "init-1", + }; + + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + }, + body: JSON.stringify(initializeMessage), + }); + + expect(transport.sessionId).toBeUndefined(); + expect(transport["_initialized"]).toBe(false); + + await transport.handleRequest(req, mockResponse); + + expect(transport.sessionId).toBeDefined(); + expect(transport["_initialized"]).toBe(true); + expect(mockResponse.writeHead).toHaveBeenCalledWith( + 200, + expect.objectContaining({ + "mcp-session-id": transport.sessionId, + }) + ); + }); + + it("should reject second initialization request", async () => { + // First initialize + const initMessage1: JSONRPCMessage = { + jsonrpc: "2.0", + method: "initialize", + params: { + clientInfo: { name: "test-client", version: "1.0" }, + protocolVersion: "2025-03-26" + }, + id: "init-1", + }; + + const req1 = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + }, + body: JSON.stringify(initMessage1), + }); + + await transport.handleRequest(req1, mockResponse); + expect(transport["_initialized"]).toBe(true); + + // Reset mock for second request + mockResponse.writeHead.mockClear(); + mockResponse.end.mockClear(); + + // Try second initialize + const initMessage2: JSONRPCMessage = { + jsonrpc: "2.0", + method: "initialize", + params: { + clientInfo: { name: "test-client", version: "1.0" }, + protocolVersion: "2025-03-26" + }, + id: "init-2", + }; + + const req2 = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + }, + body: JSON.stringify(initMessage2), + }); + + await transport.handleRequest(req2, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalledWith(400); + expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"message":"Invalid Request: Server already initialized"')); + }); + + it("should reject batch initialize request", async () => { + const batchInitialize: JSONRPCMessage[] = [ + { + jsonrpc: "2.0", + method: "initialize", + params: { + clientInfo: { name: "test-client", version: "1.0" }, + protocolVersion: "2025-03-26" + }, + id: "init-1", + }, + { + jsonrpc: "2.0", + method: "initialize", + params: { + clientInfo: { name: "test-client-2", version: "1.0" }, + protocolVersion: "2025-03-26" + }, + id: "init-2", + } + ]; + + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + }, + body: JSON.stringify(batchInitialize), + }); + + await transport.handleRequest(req, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalledWith(400); + expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"message":"Invalid Request: Only one initialization request is allowed"')); + }); + + it("should reject invalid session ID", async () => { + // First initialize the transport + const initMessage: JSONRPCMessage = { + jsonrpc: "2.0", + method: "initialize", + params: { + clientInfo: { name: "test-client", version: "1.0" }, + protocolVersion: "2025-03-26" + }, + id: "init-1", + }; + + const initReq = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + }, + body: JSON.stringify(initMessage), + }); + + await transport.handleRequest(initReq, mockResponse); + mockResponse.writeHead.mockClear(); + + // Now try with an invalid session ID + const req = createMockRequest({ + method: "POST", + headers: { + "mcp-session-id": "invalid-session-id", + "accept": "application/json, text/event-stream", + "content-type": "application/json", + }, + body: mockRequest, + }); + + await transport.handleRequest(req, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalledWith(404); + expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"jsonrpc":"2.0"')); + expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"message":"Session not found"')); + }); + + it("should reject non-initialization requests without session ID with 400 Bad Request", async () => { + // First initialize the transport + const initMessage: JSONRPCMessage = { + jsonrpc: "2.0", + method: "initialize", + params: { + clientInfo: { name: "test-client", version: "1.0" }, + protocolVersion: "2025-03-26" + }, + id: "init-1", + }; + + const initReq = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + }, + body: JSON.stringify(initMessage), + }); + + await transport.handleRequest(initReq, mockResponse); + mockResponse.writeHead.mockClear(); + + // Now try without session ID + const req = createMockRequest({ + method: "POST", + headers: { + "accept": "application/json, text/event-stream", + "content-type": "application/json", + // No mcp-session-id header + }, + body: mockRequest + }); + + await transport.handleRequest(req, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalledWith(400); + expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"jsonrpc":"2.0"')); + expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"message":"Bad Request: Mcp-Session-Id header is required"')); + }); + + it("should reject requests to uninitialized server", async () => { + // Create a new transport that hasn't been initialized + const uninitializedTransport = new StreamableHTTPServerTransport({ + sessionIdGenerator: () => randomUUID(), + }); + + const req = createMockRequest({ + method: "POST", + headers: { + "accept": "application/json, text/event-stream", + "content-type": "application/json", + "mcp-session-id": "any-session-id", + }, + body: mockRequest + }); + + await uninitializedTransport.handleRequest(req, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalledWith(400); + expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"message":"Bad Request: Server not initialized"')); + }); + + it("should reject session ID as array", async () => { + // First initialize the transport + const initMessage: JSONRPCMessage = { + jsonrpc: "2.0", + method: "initialize", + params: { + clientInfo: { name: "test-client", version: "1.0" }, + protocolVersion: "2025-03-26" + }, + id: "init-1", + }; + + const initReq = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + }, + body: JSON.stringify(initMessage), + }); + + await transport.handleRequest(initReq, mockResponse); + mockResponse.writeHead.mockClear(); + + // Now try with an array session ID + const req = createMockRequest({ + method: "POST", + headers: { + "mcp-session-id": ["session1", "session2"], + "accept": "application/json, text/event-stream", + "content-type": "application/json", + }, + body: mockRequest, + }); + + await transport.handleRequest(req, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalledWith(400); + expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"message":"Bad Request: Mcp-Session-Id header must be a single value"')); + }); + }); + describe("Mode without state management", () => { + let transportWithoutState: StreamableHTTPServerTransport; + let mockResponse: jest.Mocked; + + beforeEach(async () => { + transportWithoutState = new StreamableHTTPServerTransport({ sessionIdGenerator: () => undefined }); + mockResponse = createMockResponse(); + + // Initialize the transport for each test + const initMessage: JSONRPCMessage = { + jsonrpc: "2.0", + method: "initialize", + params: { + clientInfo: { name: "test-client", version: "1.0" }, + protocolVersion: "2025-03-26" + }, + id: "init-1", + }; + + const initReq = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + }, + body: JSON.stringify(initMessage), + }); + + await transportWithoutState.handleRequest(initReq, mockResponse); + mockResponse.writeHead.mockClear(); + }); + + it("should not include session ID in response headers when in mode without state management", async () => { + // Use a non-initialization request + const message: JSONRPCMessage = { + jsonrpc: "2.0", + method: "test", + params: {}, + id: 1, + }; + + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + }, + body: JSON.stringify(message), + }); + + await transportWithoutState.handleRequest(req, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalled(); + // Extract the headers from writeHead call + const headers = mockResponse.writeHead.mock.calls[0][1]; + expect(headers).not.toHaveProperty("mcp-session-id"); + }); + + it("should not validate session ID in mode without state management", async () => { + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + "mcp-session-id": "invalid-session-id", // This would cause a 404 in mode with state management + }, + body: JSON.stringify({ + jsonrpc: "2.0", + method: "test", + params: {}, + id: 1 + }), + }); + + await transportWithoutState.handleRequest(req, mockResponse); + + // Should still get 200 OK, not 404 Not Found + expect(mockResponse.writeHead).toHaveBeenCalledWith( + 200, + expect.not.objectContaining({ + "mcp-session-id": expect.anything(), + }) + ); + }); + + it("should handle POST requests without session validation in mode without state management", async () => { + const message: JSONRPCMessage = { + jsonrpc: "2.0", + method: "test", + params: {}, + id: 1, + }; + + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + "mcp-session-id": "non-existent-session-id", // This would be rejected in mode with state management + }, + body: JSON.stringify(message), + }); + + const onMessageMock = jest.fn(); + transportWithoutState.onmessage = onMessageMock; + + await transportWithoutState.handleRequest(req, mockResponse); + + // Message should be processed despite invalid session ID + expect(onMessageMock).toHaveBeenCalledWith(message); + }); + + it("should work with a mix of requests with and without session IDs in mode without state management", async () => { + // First request without session ID + const req1 = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + accept: "application/json, text/event-stream", + }, + body: JSON.stringify({ + jsonrpc: "2.0", + method: "test", + params: {}, + id: "test-id" + }) + }); + + await transportWithoutState.handleRequest(req1, mockResponse); + expect(mockResponse.writeHead).toHaveBeenCalledWith( + 200, + expect.objectContaining({ + "Content-Type": "text/event-stream", + }) + ); + + // Reset mock for second request + mockResponse.writeHead.mockClear(); + + // Second request with a session ID (which would be invalid in mode with state management) + const req2 = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + accept: "application/json, text/event-stream", + "mcp-session-id": "some-random-session-id", + }, + body: JSON.stringify({ + jsonrpc: "2.0", + method: "test2", + params: {}, + id: "test-id-2" + }) + }); + + await transportWithoutState.handleRequest(req2, mockResponse); + + // Should still succeed + expect(mockResponse.writeHead).toHaveBeenCalledWith( + 200, + expect.objectContaining({ + "Content-Type": "text/event-stream", + }) + ); + }); + + it("should handle initialization in mode without state management", async () => { + const transportWithoutState = new StreamableHTTPServerTransport({ sessionIdGenerator: () => undefined }); + + // Initialize message + const initializeMessage: JSONRPCMessage = { + jsonrpc: "2.0", + method: "initialize", + params: { + clientInfo: { name: "test-client", version: "1.0" }, + protocolVersion: "2025-03-26" + }, + id: "init-1", + }; + + expect(transportWithoutState.sessionId).toBeUndefined(); + expect(transportWithoutState["_initialized"]).toBe(false); + + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + }, + body: JSON.stringify(initializeMessage), + }); + + const newResponse = createMockResponse(); + await transportWithoutState.handleRequest(req, newResponse); + + // After initialization, the sessionId should still be undefined + expect(transportWithoutState.sessionId).toBeUndefined(); + expect(transportWithoutState["_initialized"]).toBe(true); + + // Headers should NOT include session ID in mode without state management + const headers = newResponse.writeHead.mock.calls[0][1]; + expect(headers).not.toHaveProperty("mcp-session-id"); + }); + }); + + describe("Request Handling", () => { + // Initialize the transport before tests that need initialization + beforeEach(async () => { + // For tests that need initialization, initialize here + const initMessage: JSONRPCMessage = { + jsonrpc: "2.0", + method: "initialize", + params: { + clientInfo: { name: "test-client", version: "1.0" }, + protocolVersion: "2025-03-26" + }, + id: "init-1", + }; + + const initReq = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + }, + body: JSON.stringify(initMessage), + }); + + await transport.handleRequest(initReq, mockResponse); + mockResponse.writeHead.mockClear(); + }); + + it("should accept GET requests for SSE with proper Accept header", async () => { + const req = createMockRequest({ + method: "GET", + headers: { + "accept": "text/event-stream", + "mcp-session-id": transport.sessionId, + }, + }); + + await transport.handleRequest(req, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalledWith(200, expect.objectContaining({ + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache, no-transform", + "Connection": "keep-alive", + "mcp-session-id": transport.sessionId, + })); + }); + + it("should reject GET requests without Accept: text/event-stream header", async () => { + const req = createMockRequest({ + method: "GET", + headers: { + "accept": "application/json", + "mcp-session-id": transport.sessionId, + }, + }); + + await transport.handleRequest(req, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalledWith(406); + expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"message":"Not Acceptable: Client must accept text/event-stream"')); + }); + + it("should send server-initiated requests to GET SSE stream", async () => { + // Open a standalone SSE stream with GET + const req = createMockRequest({ + method: "GET", + headers: { + "accept": "text/event-stream", + "mcp-session-id": transport.sessionId, + }, + }); + + const sseResponse = createMockResponse(); + await transport.handleRequest(req, sseResponse); + + // Send a notification without a related request ID + const notification: JSONRPCMessage = { + jsonrpc: "2.0", + method: "notifications/resources/updated", + params: { uri: "someuri" } + }; + + await transport.send(notification); + + // Verify notification was sent on SSE stream + expect(sseResponse.write).toHaveBeenCalledWith( + expect.stringContaining(`event: message\ndata: ${JSON.stringify(notification)}\n\n`) + ); + }); + + it("should not close GET SSE stream after sending server requests or notifications", async () => { + // Open a standalone SSE stream + const req = createMockRequest({ + method: "GET", + headers: { + "accept": "text/event-stream", + "mcp-session-id": transport.sessionId, + }, + }); + + const sseResponse = createMockResponse(); + await transport.handleRequest(req, sseResponse); + + // Send multiple notifications + const notification1: JSONRPCMessage = { jsonrpc: "2.0", method: "event1", params: {} }; + const notification2: JSONRPCMessage = { jsonrpc: "2.0", method: "event2", params: {} }; + + await transport.send(notification1); + await transport.send(notification2); + + // Stream should remain open + expect(sseResponse.end).not.toHaveBeenCalled(); + }); + + it("should reject second GET SSE stream for the same session", async () => { + // Open first SSE stream - should succeed + const req1 = createMockRequest({ + method: "GET", + headers: { + "accept": "text/event-stream", + "mcp-session-id": transport.sessionId, + }, + }); + + const sseResponse1 = createMockResponse(); + await transport.handleRequest(req1, sseResponse1); + + // Try to open a second SSE stream - should be rejected + const req2 = createMockRequest({ + method: "GET", + headers: { + "accept": "text/event-stream", + "mcp-session-id": transport.sessionId, + }, + }); + + const sseResponse2 = createMockResponse(); + await transport.handleRequest(req2, sseResponse2); + + // First stream should be good + expect(sseResponse1.writeHead).toHaveBeenCalledWith(200, expect.anything()); + + // Second stream should get 409 Conflict + expect(sseResponse2.writeHead).toHaveBeenCalledWith(409); + expect(sseResponse2.end).toHaveBeenCalledWith(expect.stringContaining('"message":"Conflict: Only one SSE stream is allowed per session"')); + }); + + it("should reject POST requests without proper Accept header", async () => { + const message: JSONRPCMessage = { + jsonrpc: "2.0", + method: "test", + params: {}, + id: 1, + }; + + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "mcp-session-id": transport.sessionId, + }, + body: JSON.stringify(message), + }); + + await transport.handleRequest(req, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalledWith(406); + }); + + it("should properly handle JSON-RPC request messages in POST requests", async () => { + const message: JSONRPCMessage = { + jsonrpc: "2.0", + method: "test", + params: {}, + id: 1, + }; + + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + "mcp-session-id": transport.sessionId, + }, + body: JSON.stringify(message), + }); + + const onMessageMock = jest.fn(); + transport.onmessage = onMessageMock; + + await transport.handleRequest(req, mockResponse); + + expect(onMessageMock).toHaveBeenCalledWith(message); + expect(mockResponse.writeHead).toHaveBeenCalledWith( + 200, + expect.objectContaining({ + "Content-Type": "text/event-stream", + }) + ); + }); + + it("should properly handle JSON-RPC notification or response messages in POST requests", async () => { + const notification: JSONRPCMessage = { + jsonrpc: "2.0", + method: "test", + params: {}, + }; + + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + "mcp-session-id": transport.sessionId, + }, + body: JSON.stringify(notification), + }); + + const onMessageMock = jest.fn(); + transport.onmessage = onMessageMock; + + await transport.handleRequest(req, mockResponse); + + expect(onMessageMock).toHaveBeenCalledWith(notification); + expect(mockResponse.writeHead).toHaveBeenCalledWith(202); + }); + + it("should handle batch notification messages properly with 202 response", async () => { + const batchMessages: JSONRPCMessage[] = [ + { jsonrpc: "2.0", method: "test1", params: {} }, + { jsonrpc: "2.0", method: "test2", params: {} }, + ]; + + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + "mcp-session-id": transport.sessionId, + }, + body: JSON.stringify(batchMessages), + }); + + const onMessageMock = jest.fn(); + transport.onmessage = onMessageMock; + + await transport.handleRequest(req, mockResponse); + + expect(onMessageMock).toHaveBeenCalledTimes(2); + expect(mockResponse.writeHead).toHaveBeenCalledWith(202); + }); + + it("should handle batch request messages with SSE when Accept header includes text/event-stream", async () => { + const batchMessages: JSONRPCMessage[] = [ + { jsonrpc: "2.0", method: "test1", params: {}, id: "req1" }, + { jsonrpc: "2.0", method: "test2", params: {}, id: "req2" }, + ]; + + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "text/event-stream, application/json", + "mcp-session-id": transport.sessionId, + }, + body: JSON.stringify(batchMessages), + }); + + const onMessageMock = jest.fn(); + transport.onmessage = onMessageMock; + + mockResponse = createMockResponse(); // Create fresh mock + await transport.handleRequest(req, mockResponse); + + // Should establish SSE connection + expect(mockResponse.writeHead).toHaveBeenCalledWith( + 200, + expect.objectContaining({ + "Content-Type": "text/event-stream" + }) + ); + expect(onMessageMock).toHaveBeenCalledTimes(2); + // Stream should remain open until responses are sent + expect(mockResponse.end).not.toHaveBeenCalled(); + }); + + it("should reject unsupported Content-Type", async () => { + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "text/plain", + "accept": "application/json, text/event-stream", + "mcp-session-id": transport.sessionId, + }, + body: "test", + }); + + await transport.handleRequest(req, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalledWith(415); + expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"jsonrpc":"2.0"')); + }); + + it("should properly handle DELETE requests and close session", async () => { + // First initialize the transport + const initMessage: JSONRPCMessage = { + jsonrpc: "2.0", + method: "initialize", + params: { + clientInfo: { name: "test-client", version: "1.0" }, + protocolVersion: "2025-03-26" + }, + id: "init-1", + }; + + const initReq = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + }, + body: JSON.stringify(initMessage), + }); + + await transport.handleRequest(initReq, mockResponse); + mockResponse.writeHead.mockClear(); + + // Now try DELETE with proper session ID + const req = createMockRequest({ + method: "DELETE", + headers: { + "mcp-session-id": transport.sessionId, + }, + }); + + const onCloseMock = jest.fn(); + transport.onclose = onCloseMock; + + await transport.handleRequest(req, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalledWith(200); + expect(onCloseMock).toHaveBeenCalled(); + }); + + it("should reject DELETE requests with invalid session ID", async () => { + // First initialize the transport + const initMessage: JSONRPCMessage = { + jsonrpc: "2.0", + method: "initialize", + params: { + clientInfo: { name: "test-client", version: "1.0" }, + protocolVersion: "2025-03-26" + }, + id: "init-1", + }; + + const initReq = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + }, + body: JSON.stringify(initMessage), + }); + + await transport.handleRequest(initReq, mockResponse); + mockResponse.writeHead.mockClear(); + + // Now try DELETE with invalid session ID + const req = createMockRequest({ + method: "DELETE", + headers: { + "mcp-session-id": "invalid-session-id", + }, + }); + + const onCloseMock = jest.fn(); + transport.onclose = onCloseMock; + + await transport.handleRequest(req, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalledWith(404); + expect(onCloseMock).not.toHaveBeenCalled(); + }); + }); + + describe("SSE Response Handling", () => { + beforeEach(async () => { + // Initialize the transport + const initMessage: JSONRPCMessage = { + jsonrpc: "2.0", + method: "initialize", + params: { + clientInfo: { name: "test-client", version: "1.0" }, + protocolVersion: "2025-03-26" + }, + id: "init-1", + }; + + const initReq = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + }, + body: JSON.stringify(initMessage), + }); + const initResponse = createMockResponse(); + await transport.handleRequest(initReq, initResponse); + mockResponse.writeHead.mockClear(); + }); + + it("should send response messages as SSE events", async () => { + // Setup a POST request with JSON-RPC request that accepts SSE + const requestMessage: JSONRPCMessage = { + jsonrpc: "2.0", + method: "test", + params: {}, + id: "test-req-id" + }; + + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + "mcp-session-id": transport.sessionId + }, + body: JSON.stringify(requestMessage) + }); + + await transport.handleRequest(req, mockResponse); + + // Send a response to the request + const responseMessage: JSONRPCMessage = { + jsonrpc: "2.0", + result: { value: "test-result" }, + id: "test-req-id" + }; + + await transport.send(responseMessage, { relatedRequestId: "test-req-id" }); + + // Verify response was sent as SSE event + expect(mockResponse.write).toHaveBeenCalledWith( + expect.stringContaining(`event: message\ndata: ${JSON.stringify(responseMessage)}\n\n`) + ); + + // Stream should be closed after sending response + expect(mockResponse.end).toHaveBeenCalled(); + }); + + it("should keep stream open when sending intermediate notifications and requests", async () => { + // Setup a POST request with JSON-RPC request that accepts SSE + const requestMessage: JSONRPCMessage = { + jsonrpc: "2.0", + method: "test", + params: {}, + id: "test-req-id" + }; + + // Create fresh response for this test + mockResponse = createMockResponse(); + + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + "mcp-session-id": transport.sessionId + }, + body: JSON.stringify(requestMessage) + }); + + await transport.handleRequest(req, mockResponse); + + // Send an intermediate notification + const notification: JSONRPCMessage = { + jsonrpc: "2.0", + method: "progress", + params: { progress: "50%" } + }; + + await transport.send(notification, { relatedRequestId: "test-req-id" }); + + // Stream should remain open + expect(mockResponse.end).not.toHaveBeenCalled(); + + // Send the final response + const responseMessage: JSONRPCMessage = { + jsonrpc: "2.0", + result: { value: "test-result" }, + id: "test-req-id" + }; + + await transport.send(responseMessage, { relatedRequestId: "test-req-id" }); + + // Now stream should be closed + expect(mockResponse.end).toHaveBeenCalled(); + }); + + it("should keep stream open when multiple requests share the same connection", async () => { + // Create a fresh response for this test + const sharedResponse = createMockResponse(); + + // Send two requests in a batch that will share the same connection + const batchRequests: JSONRPCMessage[] = [ + { jsonrpc: "2.0", method: "method1", params: {}, id: "req1" }, + { jsonrpc: "2.0", method: "method2", params: {}, id: "req2" } + ]; + + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + "mcp-session-id": transport.sessionId + }, + body: JSON.stringify(batchRequests) + }); + + await transport.handleRequest(req, sharedResponse); + + // Respond to first request + const response1: JSONRPCMessage = { + jsonrpc: "2.0", + result: { value: "result1" }, + id: "req1" + }; + + await transport.send(response1); + + // Connection should remain open because req2 is still pending + expect(sharedResponse.write).toHaveBeenCalledWith( + expect.stringContaining(`event: message\ndata: ${JSON.stringify(response1)}\n\n`) + ); + expect(sharedResponse.end).not.toHaveBeenCalled(); + + // Respond to second request + const response2: JSONRPCMessage = { + jsonrpc: "2.0", + result: { value: "result2" }, + id: "req2" + }; + + await transport.send(response2); + + // Now connection should close as all requests are complete + expect(sharedResponse.write).toHaveBeenCalledWith( + expect.stringContaining(`event: message\ndata: ${JSON.stringify(response2)}\n\n`) + ); + expect(sharedResponse.end).toHaveBeenCalled(); + }); + + it("should clean up connection tracking when a response is sent", async () => { + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + "mcp-session-id": transport.sessionId + }, + body: JSON.stringify({ + jsonrpc: "2.0", + method: "test", + params: {}, + id: "cleanup-test" + }) + }); + + const response = createMockResponse(); + await transport.handleRequest(req, response); + + // Verify that the request is tracked in the SSE map + expect(transport["_responseMapping"].size).toBe(2); + expect(transport["_responseMapping"].has("cleanup-test")).toBe(true); + + // Send a response + await transport.send({ + jsonrpc: "2.0", + result: {}, + id: "cleanup-test" + }); + + // Verify that the mapping was cleaned up + expect(transport["_responseMapping"].size).toBe(1); + expect(transport["_responseMapping"].has("cleanup-test")).toBe(false); + }); + + it("should clean up connection tracking when client disconnects", async () => { + // Setup two requests that share a connection + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + "mcp-session-id": transport.sessionId + }, + body: JSON.stringify([ + { jsonrpc: "2.0", method: "longRunning1", params: {}, id: "req1" }, + { jsonrpc: "2.0", method: "longRunning2", params: {}, id: "req2" } + ]) + }); + + const response = createMockResponse(); + + // We need to manually store the callback to trigger it later + let closeCallback: (() => void) | undefined; + response.on.mockImplementation((event, callback: () => void) => { + if (typeof event === "string" && event === "close") { + closeCallback = callback; + } + return response; + }); + + await transport.handleRequest(req, response); + + // Both requests should be mapped to the same response + expect(transport["_responseMapping"].size).toBe(3); + expect(transport["_responseMapping"].get("req1")).toBe(response); + expect(transport["_responseMapping"].get("req2")).toBe(response); + + // Simulate client disconnect by triggering the stored callback + if (closeCallback) closeCallback(); + + // All entries using this response should be removed + expect(transport["_responseMapping"].size).toBe(1); + expect(transport["_responseMapping"].has("req1")).toBe(false); + expect(transport["_responseMapping"].has("req2")).toBe(false); + }); + }); + + describe("Message Targeting", () => { + beforeEach(async () => { + // Initialize the transport + const initMessage: JSONRPCMessage = { + jsonrpc: "2.0", + method: "initialize", + params: { + clientInfo: { name: "test-client", version: "1.0" }, + protocolVersion: "2025-03-26" + }, + id: "init-1", + }; + + const initReq = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + }, + body: JSON.stringify(initMessage), + }); + + await transport.handleRequest(initReq, mockResponse); + mockResponse.writeHead.mockClear(); + }); + + it("should send response messages to the connection that sent the request", async () => { + // Create request with two separate connections + const requestMessage1: JSONRPCMessage = { + jsonrpc: "2.0", + method: "test1", + params: {}, + id: "req-id-1", + }; + + const mockResponse1 = createMockResponse(); + const req1 = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + "mcp-session-id": transport.sessionId + }, + body: JSON.stringify(requestMessage1), + }); + await transport.handleRequest(req1, mockResponse1); + + const requestMessage2: JSONRPCMessage = { + jsonrpc: "2.0", + method: "test2", + params: {}, + id: "req-id-2", + }; + + const mockResponse2 = createMockResponse(); + const req2 = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + "mcp-session-id": transport.sessionId + }, + body: JSON.stringify(requestMessage2), + }); + await transport.handleRequest(req2, mockResponse2); + + // Send responses with matching IDs + const responseMessage1: JSONRPCMessage = { + jsonrpc: "2.0", + result: { success: true }, + id: "req-id-1", + }; + + await transport.send(responseMessage1, { relatedRequestId: "req-id-1" }); + + const responseMessage2: JSONRPCMessage = { + jsonrpc: "2.0", + result: { success: true }, + id: "req-id-2", + }; + + await transport.send(responseMessage2, { relatedRequestId: "req-id-2" }); + + // Verify responses were sent to the right connections + expect(mockResponse1.write).toHaveBeenCalledWith( + expect.stringContaining(JSON.stringify(responseMessage1)) + ); + + expect(mockResponse2.write).toHaveBeenCalledWith( + expect.stringContaining(JSON.stringify(responseMessage2)) + ); + + // Verify responses were not sent to the wrong connections + const resp1HasResp2 = mockResponse1.write.mock.calls.some(call => + typeof call[0] === 'string' && call[0].includes(JSON.stringify(responseMessage2)) + ); + expect(resp1HasResp2).toBe(false); + + const resp2HasResp1 = mockResponse2.write.mock.calls.some(call => + typeof call[0] === 'string' && call[0].includes(JSON.stringify(responseMessage1)) + ); + expect(resp2HasResp1).toBe(false); + }); + }); + + describe("Error Handling", () => { + it("should return 400 error for invalid JSON data", async () => { + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + }, + body: "invalid json", + }); + + const onErrorMock = jest.fn(); + transport.onerror = onErrorMock; + + await transport.handleRequest(req, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalledWith(400); + expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"jsonrpc":"2.0"')); + expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"code":-32700')); + expect(onErrorMock).toHaveBeenCalled(); + }); + + it("should return 400 error for invalid JSON-RPC messages", async () => { + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + }, + body: JSON.stringify({ invalid: "message" }), + }); + + const onErrorMock = jest.fn(); + transport.onerror = onErrorMock; + + await transport.handleRequest(req, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalledWith(400); + expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"jsonrpc":"2.0"')); + expect(onErrorMock).toHaveBeenCalled(); + }); + }); + + describe("JSON Response Mode", () => { + let jsonResponseTransport: StreamableHTTPServerTransport; + let mockResponse: jest.Mocked; + + beforeEach(async () => { + jsonResponseTransport = new StreamableHTTPServerTransport({ + sessionIdGenerator: () => randomUUID(), + enableJsonResponse: true, + }); + + // Initialize the transport + const initMessage: JSONRPCMessage = { + jsonrpc: "2.0", + method: "initialize", + params: { + clientInfo: { name: "test-client", version: "1.0" }, + protocolVersion: "2025-03-26" + }, + id: "init-1", + }; + + const initReq = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + }, + body: JSON.stringify(initMessage), + }); + + mockResponse = createMockResponse(); + await jsonResponseTransport.handleRequest(initReq, mockResponse); + mockResponse = createMockResponse(); // Reset for tests + }); + + it("should return JSON response for a single request", async () => { + const requestMessage: JSONRPCMessage = { + jsonrpc: "2.0", + method: "tools/list", + params: {}, + id: "test-req-id", + }; + + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + "mcp-session-id": jsonResponseTransport.sessionId, + }, + body: JSON.stringify(requestMessage), + }); + + // Mock immediate response + jsonResponseTransport.onmessage = (message) => { + if ('method' in message && 'id' in message) { + const responseMessage: JSONRPCMessage = { + jsonrpc: "2.0", + result: { value: `test-result` }, + id: message.id, + }; + void jsonResponseTransport.send(responseMessage); + } + }; + + await jsonResponseTransport.handleRequest(req, mockResponse); + // Should respond with application/json header + expect(mockResponse.writeHead).toHaveBeenCalledWith( + 200, + expect.objectContaining({ + "Content-Type": "application/json", + }) + ); + + // Should return the response as JSON + const expectedResponse = { + jsonrpc: "2.0", + result: { value: "test-result" }, + id: "test-req-id", + }; + + expect(mockResponse.end).toHaveBeenCalledWith(JSON.stringify(expectedResponse)); + }); + + it("should return JSON response for batch requests", async () => { + const batchMessages: JSONRPCMessage[] = [ + { jsonrpc: "2.0", method: "tools/list", params: {}, id: "req1" }, + { jsonrpc: "2.0", method: "tools/call", params: {}, id: "req2" }, + ]; + + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + "mcp-session-id": jsonResponseTransport.sessionId, + }, + body: JSON.stringify(batchMessages), + }); + + // Mock responses without enforcing specific order + jsonResponseTransport.onmessage = (message) => { + if ('method' in message && 'id' in message) { + const responseMessage: JSONRPCMessage = { + jsonrpc: "2.0", + result: { value: `result-for-${message.id}` }, + id: message.id, + }; + void jsonResponseTransport.send(responseMessage); + } + }; + + await jsonResponseTransport.handleRequest(req, mockResponse); + + // Should respond with application/json header + expect(mockResponse.writeHead).toHaveBeenCalledWith( + 200, + expect.objectContaining({ + "Content-Type": "application/json", + }) + ); + + // Verify response was sent but don't assume specific order + expect(mockResponse.end).toHaveBeenCalled(); + const responseJson = JSON.parse(mockResponse.end.mock.calls[0][0] as string); + expect(Array.isArray(responseJson)).toBe(true); + expect(responseJson).toHaveLength(2); + + // Check each response exists separately without assuming order + expect(responseJson).toContainEqual(expect.objectContaining({ id: "req1", result: { value: "result-for-req1" } })); + expect(responseJson).toContainEqual(expect.objectContaining({ id: "req2", result: { value: "result-for-req2" } })); + }); + }); + + describe("Handling Pre-Parsed Body", () => { + beforeEach(async () => { + // Initialize the transport + const initMessage: JSONRPCMessage = { + jsonrpc: "2.0", + method: "initialize", + params: { + clientInfo: { name: "test-client", version: "1.0" }, + protocolVersion: "2025-03-26" + }, + id: "init-1", + }; + + const initReq = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + }, + body: JSON.stringify(initMessage), + }); + + await transport.handleRequest(initReq, mockResponse); + mockResponse.writeHead.mockClear(); + }); + + it("should accept pre-parsed request body", async () => { + const message: JSONRPCMessage = { + jsonrpc: "2.0", + method: "test", + params: {}, + id: "pre-parsed-test", + }; + + // Create a request without actual body content + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + "mcp-session-id": transport.sessionId, + }, + // No body provided here - it will be passed as parsedBody + }); + + const onMessageMock = jest.fn(); + transport.onmessage = onMessageMock; + + // Pass the pre-parsed body directly + await transport.handleRequest(req, mockResponse, message); + + // Verify the message was processed correctly + expect(onMessageMock).toHaveBeenCalledWith(message); + expect(mockResponse.writeHead).toHaveBeenCalledWith( + 200, + expect.objectContaining({ + "Content-Type": "text/event-stream", + }) + ); + }); + + it("should handle pre-parsed batch messages", async () => { + const batchMessages: JSONRPCMessage[] = [ + { + jsonrpc: "2.0", + method: "method1", + params: { data: "test1" }, + id: "batch1" + }, + { + jsonrpc: "2.0", + method: "method2", + params: { data: "test2" }, + id: "batch2" + }, + ]; + + // Create a request without actual body content + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + "mcp-session-id": transport.sessionId, + }, + // No body provided here - it will be passed as parsedBody + }); + + const onMessageMock = jest.fn(); + transport.onmessage = onMessageMock; + + // Pass the pre-parsed body directly + await transport.handleRequest(req, mockResponse, batchMessages); + + // Should be called for each message in the batch + expect(onMessageMock).toHaveBeenCalledTimes(2); + expect(onMessageMock).toHaveBeenCalledWith(batchMessages[0]); + expect(onMessageMock).toHaveBeenCalledWith(batchMessages[1]); + }); + + it("should prefer pre-parsed body over request body", async () => { + const requestBodyMessage: JSONRPCMessage = { + jsonrpc: "2.0", + method: "fromRequestBody", + params: {}, + id: "request-body", + }; + + const parsedBodyMessage: JSONRPCMessage = { + jsonrpc: "2.0", + method: "fromParsedBody", + params: {}, + id: "parsed-body", + }; + + // Create a request with actual body content + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + "mcp-session-id": transport.sessionId, + }, + body: JSON.stringify(requestBodyMessage), + }); + + const onMessageMock = jest.fn(); + transport.onmessage = onMessageMock; + + // Pass the pre-parsed body directly + await transport.handleRequest(req, mockResponse, parsedBodyMessage); + + // Should use the parsed body instead of the request body + expect(onMessageMock).toHaveBeenCalledWith(parsedBodyMessage); + expect(onMessageMock).not.toHaveBeenCalledWith(requestBodyMessage); + }); + }); +}); \ No newline at end of file diff --git a/src/server/streamableHttp.ts b/src/server/streamableHttp.ts new file mode 100644 index 00000000..ec8d2aa7 --- /dev/null +++ b/src/server/streamableHttp.ts @@ -0,0 +1,523 @@ +import { IncomingMessage, ServerResponse } from "node:http"; +import { Transport } from "../shared/transport.js"; +import { JSONRPCMessage, JSONRPCMessageSchema, RequestId } from "../types.js"; +import getRawBody from "raw-body"; +import contentType from "content-type"; + +const MAXIMUM_MESSAGE_SIZE = "4mb"; + +/** + * Configuration options for StreamableHTTPServerTransport + */ +export interface StreamableHTTPServerTransportOptions { + /** + * Function that generates a session ID for the transport. + * The session ID SHOULD be globally unique and cryptographically secure (e.g., a securely generated UUID, a JWT, or a cryptographic hash) + * + * Return undefined to disable session management. + */ + sessionIdGenerator: () => string | undefined; + + /** + * If true, the server will return JSON responses instead of starting an SSE stream. + * This can be useful for simple request/response scenarios without streaming. + * Default is false (SSE streams are preferred). + */ + enableJsonResponse?: boolean; +} + +/** + * Server transport for Streamable HTTP: this implements the MCP Streamable HTTP transport specification. + * It supports both SSE streaming and direct HTTP responses. + * + * Usage example: + * + * ```typescript + * // Stateful mode - server sets the session ID + * const statefulTransport = new StreamableHTTPServerTransport({ + * sessionId: randomUUID(), + * }); + * + * // Stateless mode - explicitly set session ID to undefined + * const statelessTransport = new StreamableHTTPServerTransport({ + * sessionId: undefined, + * }); + * + * // Using with pre-parsed request body + * app.post('/mcp', (req, res) => { + * transport.handleRequest(req, res, req.body); + * }); + * ``` + * + * In stateful mode: + * - Session ID is generated and included in response headers + * - Session ID is always included in initialization responses + * - Requests with invalid session IDs are rejected with 404 Not Found + * - Non-initialization requests without a session ID are rejected with 400 Bad Request + * - State is maintained in-memory (connections, message history) + * + * In stateless mode: + * - Session ID is only included in initialization responses + * - No session validation is performed + */ +export class StreamableHTTPServerTransport implements Transport { + // when sessionId is not set (undefined), it means the transport is in stateless mode + private sessionIdGenerator: () => string | undefined; + private _started: boolean = false; + private _responseMapping: Map = new Map(); + private _requestResponseMap: Map = new Map(); + private _initialized: boolean = false; + private _enableJsonResponse: boolean = false; + private _standaloneSSE: ServerResponse | undefined; + + + sessionId?: string | undefined; + onclose?: () => void; + onerror?: (error: Error) => void; + onmessage?: (message: JSONRPCMessage) => void; + + constructor(options: StreamableHTTPServerTransportOptions) { + this.sessionIdGenerator = options.sessionIdGenerator; + this._enableJsonResponse = options.enableJsonResponse ?? false; + } + + /** + * Starts the transport. This is required by the Transport interface but is a no-op + * for the Streamable HTTP transport as connections are managed per-request. + */ + async start(): Promise { + if (this._started) { + throw new Error("Transport already started"); + } + this._started = true; + } + + /** + * Handles an incoming HTTP request, whether GET or POST + */ + async handleRequest(req: IncomingMessage, res: ServerResponse, parsedBody?: unknown): Promise { + if (req.method === "POST") { + await this.handlePostRequest(req, res, parsedBody); + } else if (req.method === "GET") { + await this.handleGetRequest(req, res); + } else if (req.method === "DELETE") { + await this.handleDeleteRequest(req, res); + } else { + await this.handleUnsupportedRequest(res); + } + } + + /** + * Handles GET requests for SSE stream + */ + private async handleGetRequest(req: IncomingMessage, res: ServerResponse): Promise { + // The client MUST include an Accept header, listing text/event-stream as a supported content type. + const acceptHeader = req.headers.accept; + if (!acceptHeader?.includes("text/event-stream")) { + res.writeHead(406).end(JSON.stringify({ + jsonrpc: "2.0", + error: { + code: -32000, + message: "Not Acceptable: Client must accept text/event-stream" + }, + id: null + })); + return; + } + + // If an Mcp-Session-Id is returned by the server during initialization, + // clients using the Streamable HTTP transport MUST include it + // in the Mcp-Session-Id header on all of their subsequent HTTP requests. + if (!this.validateSession(req, res)) { + return; + } + + // The server MUST either return Content-Type: text/event-stream in response to this HTTP GET, + // or else return HTTP 405 Method Not Allowed + const headers: Record = { + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache, no-transform", + Connection: "keep-alive", + }; + + // After initialization, always include the session ID if we have one + if (this.sessionId !== undefined) { + headers["mcp-session-id"] = this.sessionId; + } + // The server MAY include a Last-Event-ID header in the response to this HTTP GET. + // Resumability will be supported in the future + + // Check if there's already an active standalone SSE stream for this session + + if (this._standaloneSSE !== undefined) { + // Only one GET SSE stream is allowed per session + res.writeHead(409).end(JSON.stringify({ + jsonrpc: "2.0", + error: { + code: -32000, + message: "Conflict: Only one SSE stream is allowed per session" + }, + id: null + })); + return; + } + // We need to send headers immediately as message will arrive much later, + // otherwise the client will just wait for the first message + res.writeHead(200, headers).flushHeaders(); + + // Assing the response to the standalone SSE stream + this._standaloneSSE = res; + + // Set up close handler for client disconnects + res.on("close", () => { + this._standaloneSSE = undefined; + }); + } + + /** + * Handles unsupported requests (PUT, PATCH, etc.) + */ + private async handleUnsupportedRequest(res: ServerResponse): Promise { + res.writeHead(405, { + "Allow": "GET, POST, DELETE" + }).end(JSON.stringify({ + jsonrpc: "2.0", + error: { + code: -32000, + message: "Method not allowed." + }, + id: null + })); + } + + /** + * Handles POST requests containing JSON-RPC messages + */ + private async handlePostRequest(req: IncomingMessage, res: ServerResponse, parsedBody?: unknown): Promise { + try { + // Validate the Accept header + const acceptHeader = req.headers.accept; + // The client MUST include an Accept header, listing both application/json and text/event-stream as supported content types. + if (!acceptHeader?.includes("application/json") || !acceptHeader.includes("text/event-stream")) { + res.writeHead(406).end(JSON.stringify({ + jsonrpc: "2.0", + error: { + code: -32000, + message: "Not Acceptable: Client must accept both application/json and text/event-stream" + }, + id: null + })); + return; + } + + const ct = req.headers["content-type"]; + if (!ct || !ct.includes("application/json")) { + res.writeHead(415).end(JSON.stringify({ + jsonrpc: "2.0", + error: { + code: -32000, + message: "Unsupported Media Type: Content-Type must be application/json" + }, + id: null + })); + return; + } + + let rawMessage; + if (parsedBody !== undefined) { + rawMessage = parsedBody; + } else { + const parsedCt = contentType.parse(ct); + const body = await getRawBody(req, { + limit: MAXIMUM_MESSAGE_SIZE, + encoding: parsedCt.parameters.charset ?? "utf-8", + }); + rawMessage = JSON.parse(body.toString()); + } + + let messages: JSONRPCMessage[]; + + // handle batch and single messages + if (Array.isArray(rawMessage)) { + messages = rawMessage.map(msg => JSONRPCMessageSchema.parse(msg)); + } else { + messages = [JSONRPCMessageSchema.parse(rawMessage)]; + } + + // Check if this is an initialization request + // https://spec.modelcontextprotocol.io/specification/2025-03-26/basic/lifecycle/ + const isInitializationRequest = messages.some( + msg => 'method' in msg && msg.method === 'initialize' + ); + if (isInitializationRequest) { + if (this._initialized) { + res.writeHead(400).end(JSON.stringify({ + jsonrpc: "2.0", + error: { + code: -32600, + message: "Invalid Request: Server already initialized" + }, + id: null + })); + return; + } + if (messages.length > 1) { + res.writeHead(400).end(JSON.stringify({ + jsonrpc: "2.0", + error: { + code: -32600, + message: "Invalid Request: Only one initialization request is allowed" + }, + id: null + })); + return; + } + this.sessionId = this.sessionIdGenerator(); + this._initialized = true; + + } + // If an Mcp-Session-Id is returned by the server during initialization, + // clients using the Streamable HTTP transport MUST include it + // in the Mcp-Session-Id header on all of their subsequent HTTP requests. + if (!isInitializationRequest && !this.validateSession(req, res)) { + return; + } + + + // check if it contains requests + const hasRequests = messages.some(msg => 'method' in msg && 'id' in msg); + const hasOnlyNotificationsOrResponses = messages.every(msg => + ('method' in msg && !('id' in msg)) || ('result' in msg || 'error' in msg)); + + if (hasOnlyNotificationsOrResponses) { + // if it only contains notifications or responses, return 202 + res.writeHead(202).end(); + + // handle each message + for (const message of messages) { + this.onmessage?.(message); + } + } else if (hasRequests) { + // The default behavior is to use SSE streaming + // but in some cases server will return JSON responses + if (!this._enableJsonResponse) { + const headers: Record = { + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + Connection: "keep-alive", + }; + + // After initialization, always include the session ID if we have one + if (this.sessionId !== undefined) { + headers["mcp-session-id"] = this.sessionId; + } + + res.writeHead(200, headers); + } + // Store the response for this request to send messages back through this connection + // We need to track by request ID to maintain the connection + for (const message of messages) { + if ('method' in message && 'id' in message) { + this._responseMapping.set(message.id, res); + } + } + + // Set up close handler for client disconnects + res.on("close", () => { + // Remove all entries that reference this response + for (const [id, storedRes] of this._responseMapping.entries()) { + if (storedRes === res) { + this._responseMapping.delete(id); + this._requestResponseMap.delete(id); + } + } + }); + + // handle each message + for (const message of messages) { + this.onmessage?.(message); + } + // The server SHOULD NOT close the SSE stream before sending all JSON-RPC responses + // This will be handled by the send() method when responses are ready + } + } catch (error) { + // return JSON-RPC formatted error + res.writeHead(400).end(JSON.stringify({ + jsonrpc: "2.0", + error: { + code: -32700, + message: "Parse error", + data: String(error) + }, + id: null + })); + this.onerror?.(error as Error); + } + } + + /** + * Handles DELETE requests to terminate sessions + */ + private async handleDeleteRequest(req: IncomingMessage, res: ServerResponse): Promise { + if (!this.validateSession(req, res)) { + return; + } + await this.close(); + res.writeHead(200).end(); + } + + /** + * Validates session ID for non-initialization requests + * Returns true if the session is valid, false otherwise + */ + private validateSession(req: IncomingMessage, res: ServerResponse): boolean { + if (!this._initialized) { + // If the server has not been initialized yet, reject all requests + res.writeHead(400).end(JSON.stringify({ + jsonrpc: "2.0", + error: { + code: -32000, + message: "Bad Request: Server not initialized" + }, + id: null + })); + return false; + } + if (this.sessionId === undefined) { + // If the session ID is not set, the session management is disabled + // and we don't need to validate the session ID + return true; + } + const sessionId = req.headers["mcp-session-id"]; + + if (!sessionId) { + // Non-initialization requests without a session ID should return 400 Bad Request + res.writeHead(400).end(JSON.stringify({ + jsonrpc: "2.0", + error: { + code: -32000, + message: "Bad Request: Mcp-Session-Id header is required" + }, + id: null + })); + return false; + } else if (Array.isArray(sessionId)) { + res.writeHead(400).end(JSON.stringify({ + jsonrpc: "2.0", + error: { + code: -32000, + message: "Bad Request: Mcp-Session-Id header must be a single value" + }, + id: null + })); + return false; + } + else if (sessionId !== this.sessionId) { + // Reject requests with invalid session ID with 404 Not Found + res.writeHead(404).end(JSON.stringify({ + jsonrpc: "2.0", + error: { + code: -32001, + message: "Session not found" + }, + id: null + })); + return false; + } + + return true; + } + + + async close(): Promise { + // Close all SSE connections + this._responseMapping.forEach((response) => { + response.end(); + }); + this._responseMapping.clear(); + + // Clear any pending responses + this._requestResponseMap.clear(); + + this.onclose?.(); + } + + async send(message: JSONRPCMessage, options?: { relatedRequestId?: RequestId }): Promise { + let requestId = options?.relatedRequestId; + if ('result' in message || 'error' in message) { + // If the message is a response, use the request ID from the message + requestId = message.id; + } + + // Check if this message should be sent on the standalone SSE stream (no request ID) + // Ignore notifications from tools (which have relatedRequestId set) + // Those will be sent via dedicated response SSE streams + if (requestId === undefined) { + // For standalone SSE streams, we can only send requests and notifications + if ('result' in message || 'error' in message) { + throw new Error("Cannot send a response on a standalone SSE stream unless resuming a previous client request"); + } + + if (this._standaloneSSE === undefined) { + // The spec says the server MAY send messages on the stream, so it's ok to discard if no stream + return; + } + + // Send the message to the standalone SSE stream + this._standaloneSSE.write(`event: message\ndata: ${JSON.stringify(message)}\n\n`); + return; + } + + // Get the response for this request + const response = this._responseMapping.get(requestId); + if (!response) { + throw new Error(`No connection established for request ID: ${String(requestId)}`); + } + + if (!this._enableJsonResponse) { + response.write(`event: message\ndata: ${JSON.stringify(message)}\n\n`); + } + if ('result' in message || 'error' in message) { + this._requestResponseMap.set(requestId, message); + + // Get all request IDs that share the same request response object + const relatedIds = Array.from(this._responseMapping.entries()) + .filter(([_, res]) => res === response) + .map(([id]) => id); + + // Check if we have responses for all requests using this connection + const allResponsesReady = relatedIds.every(id => this._requestResponseMap.has(id)); + + if (allResponsesReady) { + if (this._enableJsonResponse) { + // All responses ready, send as JSON + const headers: Record = { + 'Content-Type': 'application/json', + }; + if (this.sessionId !== undefined) { + headers['mcp-session-id'] = this.sessionId; + } + + const responses = relatedIds + .map(id => this._requestResponseMap.get(id)!); + + response.writeHead(200, headers); + if (responses.length === 1) { + response.end(JSON.stringify(responses[0])); + } else { + response.end(JSON.stringify(responses)); + } + } else { + // End the SSE stream + response.end(); + } + // Clean up + for (const id of relatedIds) { + this._requestResponseMap.delete(id); + this._responseMapping.delete(id); + } + } + } + } +} + diff --git a/src/shared/protocol.ts b/src/shared/protocol.ts index a39eba68..39483b77 100644 --- a/src/shared/protocol.ts +++ b/src/shared/protocol.ts @@ -3,6 +3,10 @@ import { CancelledNotificationSchema, ClientCapabilities, ErrorCode, + isJSONRPCError, + isJSONRPCRequest, + isJSONRPCResponse, + isJSONRPCNotification, JSONRPCError, JSONRPCNotification, JSONRPCRequest, @@ -79,27 +83,57 @@ export type RequestOptions = { * If not specified, there is no maximum total timeout. */ maxTotalTimeout?: number; + + /** + * May be used to indicate to the transport which incoming request to associate this outgoing request with. + */ + relatedRequestId?: RequestId; }; +/** + * Options that can be given per notification. + */ +export type NotificationOptions = { + /** + * May be used to indicate to the transport which incoming request to associate this outgoing notification with. + */ + relatedRequestId?: RequestId; +} + /** * Extra data given to request handlers. */ -export type RequestHandlerExtra = { - /** - * An abort signal used to communicate if the request was cancelled from the sender's side. - */ - signal: AbortSignal; +export type RequestHandlerExtra = { + /** + * An abort signal used to communicate if the request was cancelled from the sender's side. + */ + signal: AbortSignal; - /** - * Information about a validated access token, provided to request handlers. - */ - authInfo?: AuthInfo; + /** + * Information about a validated access token, provided to request handlers. + */ + authInfo?: AuthInfo; - /** - * The session ID from the transport, if available. - */ - sessionId?: string; -}; + /** + * The session ID from the transport, if available. + */ + sessionId?: string; + + /** + * Sends a notification that relates to the current request being handled. + * + * This is used by certain transports to correctly associate related messages. + */ + sendNotification: (notification: SendNotificationT) => Promise; + + /** + * Sends a request that relates to the current request being handled. + * + * This is used by certain transports to correctly associate related messages. + */ + sendRequest: >(request: SendRequestT, resultSchema: U, options?: RequestOptions) => Promise>; + }; /** * Information about a request's timeout state @@ -128,7 +162,7 @@ export abstract class Protocol< string, ( request: JSONRPCRequest, - extra: RequestHandlerExtra, + extra: RequestHandlerExtra, ) => Promise > = new Map(); private _requestHandlerAbortControllers: Map = @@ -247,12 +281,14 @@ export abstract class Protocol< }; this._transport.onmessage = (message, extra) => { - if (!("method" in message)) { + if (isJSONRPCResponse(message) || isJSONRPCError(message)) { this._onresponse(message); - } else if ("id" in message) { + } else if (isJSONRPCRequest(message)) { this._onrequest(message, extra); - } else { + } else if (isJSONRPCNotification(message)) { this._onnotification(message); + } else { + this._onerror(new Error(`Unknown message type: ${JSON.stringify(message)}`)); } }; @@ -321,11 +357,16 @@ export abstract class Protocol< const abortController = new AbortController(); this._requestHandlerAbortControllers.set(request.id, abortController); - // Create fullExtra object with both abort signal and sessionId from transport + // Create `fullExtra` object with both abort signal and sessionId from transport // in addition to any authInfo provided - const fullExtra: RequestHandlerExtra = { + const fullExtra: RequestHandlerExtra = { signal: abortController.signal, sessionId: this._transport?.sessionId, + sendNotification: + (notification) => + this.notification(notification, { relatedRequestId: request.id }), + sendRequest: (r, resultSchema, options?) => + this.request(r, resultSchema, { ...options, relatedRequestId: request.id }), authInfo: extra?.authInfo, }; @@ -372,7 +413,7 @@ export abstract class Protocol< private _onprogress(notification: ProgressNotification): void { const { progressToken, ...params } = notification.params; const messageId = Number(progressToken); - + const handler = this._progressHandlers.get(messageId); if (!handler) { this._onerror(new Error(`Received a progress notification for an unknown token: ${JSON.stringify(notification)}`)); @@ -381,7 +422,7 @@ export abstract class Protocol< const responseHandler = this._responseHandlers.get(messageId); const timeoutInfo = this._timeoutInfo.get(messageId); - + if (timeoutInfo && responseHandler && timeoutInfo.resetTimeoutOnProgress) { try { this._resetTimeout(messageId); @@ -410,7 +451,7 @@ export abstract class Protocol< this._progressHandlers.delete(messageId); this._cleanupTimeout(messageId); - if ("result" in response) { + if (isJSONRPCResponse(response)) { handler(response); } else { const error = new McpError( @@ -468,6 +509,8 @@ export abstract class Protocol< resultSchema: T, options?: RequestOptions, ): Promise> { + const { relatedRequestId } = options ?? {}; + return new Promise((resolve, reject) => { if (!this._transport) { reject(new Error("Not connected")); @@ -508,7 +551,7 @@ export abstract class Protocol< requestId: messageId, reason: String(reason), }, - }) + }, { relatedRequestId }) .catch((error) => this._onerror(new Error(`Failed to send cancellation: ${error}`)), ); @@ -546,7 +589,7 @@ export abstract class Protocol< this._setupTimeout(messageId, timeout, options?.maxTotalTimeout, timeoutHandler, options?.resetTimeoutOnProgress ?? false); - this._transport.send(jsonrpcRequest).catch((error) => { + this._transport.send(jsonrpcRequest, { relatedRequestId }).catch((error) => { this._cleanupTimeout(messageId); reject(error); }); @@ -556,7 +599,7 @@ export abstract class Protocol< /** * Emits a notification, which is a one-way message that does not expect a response. */ - async notification(notification: SendNotificationT): Promise { + async notification(notification: SendNotificationT, options?: NotificationOptions): Promise { if (!this._transport) { throw new Error("Not connected"); } @@ -568,7 +611,7 @@ export abstract class Protocol< jsonrpc: "2.0", }; - await this._transport.send(jsonrpcNotification); + await this._transport.send(jsonrpcNotification, options); } /** @@ -584,14 +627,15 @@ export abstract class Protocol< requestSchema: T, handler: ( request: z.infer, - extra: RequestHandlerExtra, + extra: RequestHandlerExtra, ) => SendResultT | Promise, ): void { const method = requestSchema.shape.method.value; this.assertRequestHandlerCapability(method); - this._requestHandlers.set(method, (request, extra) => - Promise.resolve(handler(requestSchema.parse(request), extra)), - ); + + this._requestHandlers.set(method, (request, extra) => { + return Promise.resolve(handler(requestSchema.parse(request), extra)); + }); } /** diff --git a/src/shared/transport.ts b/src/shared/transport.ts index ffc77d0b..c2732391 100644 --- a/src/shared/transport.ts +++ b/src/shared/transport.ts @@ -1,5 +1,5 @@ import { AuthInfo } from "../server/auth/types.js"; -import { JSONRPCMessage } from "../types.js"; +import { JSONRPCMessage, RequestId } from "../types.js"; /** * Describes the minimal contract for a MCP transport that a client or server can communicate over. @@ -16,8 +16,10 @@ export interface Transport { /** * Sends a JSON-RPC message (request or response). + * + * If present, `relatedRequestId` is used to indicate to the transport which incoming request to associate this outgoing message with. */ - send(message: JSONRPCMessage): Promise; + send(message: JSONRPCMessage, options?: { relatedRequestId?: RequestId }): Promise; /** * Closes the connection. diff --git a/src/types.ts b/src/types.ts index da4afd2a..db6bf125 100644 --- a/src/types.ts +++ b/src/types.ts @@ -78,6 +78,9 @@ export const JSONRPCRequestSchema = z .merge(RequestSchema) .strict(); +export const isJSONRPCRequest = (value: unknown): value is JSONRPCRequest => + JSONRPCRequestSchema.safeParse(value).success; + /** * A notification which does not expect a response. */ @@ -88,6 +91,11 @@ export const JSONRPCNotificationSchema = z .merge(NotificationSchema) .strict(); +export const isJSONRPCNotification = ( + value: unknown +): value is JSONRPCNotification => + JSONRPCNotificationSchema.safeParse(value).success; + /** * A successful (non-error) response to a request. */ @@ -99,6 +107,9 @@ export const JSONRPCResponseSchema = z }) .strict(); +export const isJSONRPCResponse = (value: unknown): value is JSONRPCResponse => + JSONRPCResponseSchema.safeParse(value).success; + /** * Error codes defined by the JSON-RPC specification. */ @@ -139,6 +150,9 @@ export const JSONRPCErrorSchema = z }) .strict(); +export const isJSONRPCError = (value: unknown): value is JSONRPCError => + JSONRPCErrorSchema.safeParse(value).success; + export const JSONRPCMessageSchema = z.union([ JSONRPCRequestSchema, JSONRPCNotificationSchema,