From 6c1b9ba06a403f4498efa50b29a0cfa9db1f0f3f Mon Sep 17 00:00:00 2001 From: ihrpr Date: Mon, 7 Apr 2025 12:12:43 +0100 Subject: [PATCH] remove local storage for replaying and remove support for sse connection thought get --- src/server/streamableHttp.test.ts | 337 +++++++++++++++++++----------- src/server/streamableHttp.ts | 304 +++++++++------------------ src/shared/transport.ts | 2 +- 3 files changed, 314 insertions(+), 329 deletions(-) diff --git a/src/server/streamableHttp.test.ts b/src/server/streamableHttp.test.ts index baf4d80a..8fb55e79 100644 --- a/src/server/streamableHttp.test.ts +++ b/src/server/streamableHttp.test.ts @@ -52,7 +52,7 @@ describe("StreamableHTTPServerTransport", () => { }); describe("Session Management", () => { - it("should generate a valid session ID", () => { + it("should store a valid session ID", () => { expect(transport.sessionId).toBeTruthy(); expect(typeof transport.sessionId).toBe("string"); }); @@ -158,11 +158,18 @@ describe("StreamableHTTPServerTransport", () => { it("should not validate session ID in stateless mode", async () => { const req = createMockRequest({ - method: "GET", + method: "POST", headers: { - accept: "text/event-stream", + "content-type": "application/json", + "accept": "application/json", "mcp-session-id": "invalid-session-id", // This would cause a 404 in stateful mode }, + body: JSON.stringify({ + jsonrpc: "2.0", + method: "test", + params: {}, + id: 1 + }), }); await statelessTransport.handleRequest(req, mockResponse); @@ -206,10 +213,17 @@ describe("StreamableHTTPServerTransport", () => { it("should work with a mix of requests with and without session IDs in stateless mode", async () => { // First request without session ID const req1 = createMockRequest({ - method: "GET", + method: "POST", headers: { + "content-type": "application/json", accept: "text/event-stream", }, + body: JSON.stringify({ + jsonrpc: "2.0", + method: "test", + params: {}, + id: "test-id" + }) }); await statelessTransport.handleRequest(req1, mockResponse); @@ -225,11 +239,18 @@ describe("StreamableHTTPServerTransport", () => { // Second request with a session ID (which would be invalid in stateful mode) const req2 = createMockRequest({ - method: "GET", + method: "POST", headers: { + "content-type": "application/json", accept: "text/event-stream", "mcp-session-id": "some-random-session-id", }, + body: JSON.stringify({ + jsonrpc: "2.0", + method: "test2", + params: {}, + id: "test-id-2" + }) }); await statelessTransport.handleRequest(req2, mockResponse); @@ -308,21 +329,7 @@ describe("StreamableHTTPServerTransport", () => { }); describe("Request Handling", () => { - it("should reject GET requests without Accept: text/event-stream header", async () => { - const req = createMockRequest({ - method: "GET", - headers: { - "mcp-session-id": transport.sessionId, - }, - }); - - await transport.handleRequest(req, mockResponse); - - expect(mockResponse.writeHead).toHaveBeenCalledWith(406, {}); - expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"jsonrpc":"2.0"')); - }); - - it("should properly handle GET requests with Accept header and establish SSE connection", async () => { + it("should reject GET requests for SSE with 405 Method Not Allowed", async () => { const req = createMockRequest({ method: "GET", headers: { @@ -333,14 +340,11 @@ describe("StreamableHTTPServerTransport", () => { await transport.handleRequest(req, mockResponse); - expect(mockResponse.writeHead).toHaveBeenCalledWith( - 200, - expect.objectContaining({ - "Content-Type": "text/event-stream", - "Cache-Control": "no-cache", - Connection: "keep-alive", - }) - ); + expect(mockResponse.writeHead).toHaveBeenCalledWith(405, expect.objectContaining({ + "Allow": "POST, DELETE" + })); + expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"jsonrpc":"2.0"')); + expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('Server does not offer an SSE stream at this endpoint')); }); it("should reject POST requests without proper Accept header", async () => { @@ -421,7 +425,7 @@ describe("StreamableHTTPServerTransport", () => { expect(mockResponse.writeHead).toHaveBeenCalledWith(202); }); - it("should handle batch messages properly", async () => { + 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: {} }, @@ -446,6 +450,39 @@ describe("StreamableHTTPServerTransport", () => { 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; + + 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", @@ -481,130 +518,170 @@ describe("StreamableHTTPServerTransport", () => { }); }); - describe("Message Replay", () => { - it("should replay messages after specified Last-Event-ID", async () => { - // Establish first connection with Accept header and session ID - const req1 = createMockRequest({ - method: "GET", + describe("SSE Response Handling", () => { + 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": "text/event-stream", "mcp-session-id": transport.sessionId }, - }); - await transport.handleRequest(req1, mockResponse); - - // Send a message to first connection - const message1: JSONRPCMessage = { - jsonrpc: "2.0", - method: "test1", - params: {}, - id: 1 - }; - - await transport.send(message1); - - // Get message ID (captured from write call) - const writeCall = mockResponse.write.mock.calls[0][0] as string; - const idMatch = writeCall.match(/id: ([a-f0-9-]+)/); - if (!idMatch) { - throw new Error("Message ID not found in write call"); - } - const lastEventId = idMatch[1]; - - // Create a second connection with last-event-id - const mockResponse2 = createMockResponse(); - const req2 = createMockRequest({ - method: "GET", - headers: { - "accept": "text/event-stream", - "last-event-id": lastEventId, - "mcp-session-id": transport.sessionId - }, + body: JSON.stringify(requestMessage) }); - await transport.handleRequest(req2, mockResponse2); + await transport.handleRequest(req, mockResponse); - // Send a second message - const message2: JSONRPCMessage = { + // Send a response to the request + const responseMessage: JSONRPCMessage = { jsonrpc: "2.0", - method: "test2", - params: {}, - id: 2 + result: { value: "test-result" }, + id: "test-req-id" }; - await transport.send(message2); + await transport.send(responseMessage, { relatedRequestId: "test-req-id" }); - // Verify the second message was received by both connections + // Verify response was sent as SSE event expect(mockResponse.write).toHaveBeenCalledWith( - expect.stringContaining(JSON.stringify(message1)) - ); - expect(mockResponse2.write).toHaveBeenCalledWith( - expect.stringContaining(JSON.stringify(message2)) + 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" + }; + + const req = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "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(); }); }); describe("Message Targeting", () => { it("should send response messages to the connection that sent the request", async () => { - // Create two connections - const mockResponse1 = createMockResponse(); - const req1 = createMockRequest({ - method: "GET", - headers: { - "accept": "text/event-stream", - }, - }); - await transport.handleRequest(req1, mockResponse1); - - const mockResponse2 = createMockResponse(); - const req2 = createMockRequest({ - method: "GET", - headers: { - "accept": "text/event-stream", - }, - }); - await transport.handleRequest(req2, mockResponse2); - - // Send a request through the first connection - const requestMessage: JSONRPCMessage = { + // Create request with two separate connections + const requestMessage1: JSONRPCMessage = { jsonrpc: "2.0", - method: "test", + method: "test1", params: {}, - id: "test-id", + id: "req-id-1", }; - const reqPost = createMockRequest({ + const mockResponse1 = createMockResponse(); + const req1 = createMockRequest({ method: "POST", headers: { "content-type": "application/json", "accept": "text/event-stream", + "mcp-session-id": transport.sessionId }, - body: JSON.stringify(requestMessage), + body: JSON.stringify(requestMessage1), }); + await transport.handleRequest(req1, mockResponse1); - await transport.handleRequest(reqPost, mockResponse1); - - // Send a response with matching ID - const responseMessage: JSONRPCMessage = { + const requestMessage2: JSONRPCMessage = { jsonrpc: "2.0", - result: { success: true }, - id: "test-id", + method: "test2", + params: {}, + id: "req-id-2", }; - await transport.send(responseMessage); + const mockResponse2 = createMockResponse(); + const req2 = createMockRequest({ + method: "POST", + headers: { + "content-type": "application/json", + "accept": "text/event-stream", + "mcp-session-id": transport.sessionId + }, + body: JSON.stringify(requestMessage2), + }); + await transport.handleRequest(req2, mockResponse2); - // Verify response was sent to the right connection + // 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(responseMessage)) + expect.stringContaining(JSON.stringify(responseMessage1)) ); - // Check if write was called with this exact message on the second connection - const writeCallsOnSecondConn = mockResponse2.write.mock.calls.filter(call => - typeof call[0] === 'string' && call[0].includes(JSON.stringify(responseMessage)) + expect(mockResponse2.write).toHaveBeenCalledWith( + expect.stringContaining(JSON.stringify(responseMessage2)) ); - // Verify the response wasn't broadcast to all connections - expect(writeCallsOnSecondConn.length).toBe(0); + // 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); }); }); @@ -749,6 +826,7 @@ describe("StreamableHTTPServerTransport", () => { headers: { "content-type": "application/json", "accept": "application/json", + "mcp-session-id": transport.sessionId, }, body: JSON.stringify(requestBodyMessage), }); @@ -782,11 +860,18 @@ describe("StreamableHTTPServerTransport", () => { it("should include custom headers in SSE response", async () => { const req = createMockRequest({ - method: "GET", + method: "POST", headers: { + "content-type": "application/json", accept: "text/event-stream", "mcp-session-id": transportWithHeaders.sessionId }, + body: JSON.stringify({ + jsonrpc: "2.0", + method: "test", + params: {}, + id: "test-headers-id" + }) }); await transportWithHeaders.handleRequest(req, mockResponse); @@ -860,11 +945,18 @@ describe("StreamableHTTPServerTransport", () => { }); const req = createMockRequest({ - method: "GET", + method: "POST", headers: { + "content-type": "application/json", accept: "text/event-stream", "mcp-session-id": transportWithConflictingHeaders.sessionId }, + body: JSON.stringify({ + jsonrpc: "2.0", + method: "test", + params: {}, + id: "test-conflict-id" + }) }); await transportWithConflictingHeaders.handleRequest(req, mockResponse); @@ -872,7 +964,7 @@ describe("StreamableHTTPServerTransport", () => { expect(mockResponse.writeHead).toHaveBeenCalledWith( 200, expect.objectContaining({ - "Content-Type": "text/event-stream", // 应该保持原有的 Content-Type + "Content-Type": "text/event-stream", "X-Custom-Header": "custom-value" }) ); @@ -884,11 +976,18 @@ describe("StreamableHTTPServerTransport", () => { }); const req = createMockRequest({ - method: "GET", + method: "POST", headers: { + "content-type": "application/json", accept: "text/event-stream", "mcp-session-id": transportWithoutHeaders.sessionId }, + body: JSON.stringify({ + jsonrpc: "2.0", + method: "test", + params: {}, + id: "test-empty-headers-id" + }) }); await transportWithoutHeaders.handleRequest(req, mockResponse); diff --git a/src/server/streamableHttp.ts b/src/server/streamableHttp.ts index ef654f85..6afe826a 100644 --- a/src/server/streamableHttp.ts +++ b/src/server/streamableHttp.ts @@ -1,23 +1,11 @@ -import { randomUUID } from "node:crypto"; import { IncomingMessage, ServerResponse } from "node:http"; import { Transport } from "../shared/transport.js"; -import { JSONRPCMessage, JSONRPCMessageSchema } from "../types.js"; +import { JSONRPCMessage, JSONRPCMessageSchema, RequestId } from "../types.js"; import getRawBody from "raw-body"; import contentType from "content-type"; const MAXIMUM_MESSAGE_SIZE = "4mb"; -interface StreamConnection { - response: ServerResponse; - lastEventId?: string; - messages: Array<{ - id: string; - message: JSONRPCMessage; - }>; - // mark this connection as a response to a specific request - requestId?: string | null; -} - /** * Configuration options for StreamableHTTPServerTransport */ @@ -34,6 +22,7 @@ export interface StreamableHTTPServerTransportOptions { * These headers will be added to both SSE and regular HTTP responses */ customHeaders?: Record; + } /** @@ -71,16 +60,11 @@ export interface StreamableHTTPServerTransportOptions { * - No session validation is performed */ export class StreamableHTTPServerTransport implements Transport { - private _connections: Map = new Map(); // when sessionID is not set, it means the transport is in stateless mode private _sessionId: string | undefined; - private _messageHistory: Map = new Map(); private _started: boolean = false; - private _requestConnections: Map = new Map(); // request ID to connection ID mapping private _customHeaders: Record; + private _sseResponseMapping: Map = new Map(); onclose?: () => void; onerror?: (error: Error) => void; @@ -108,37 +92,10 @@ export class StreamableHTTPServerTransport implements Transport { async handleRequest(req: IncomingMessage, res: ServerResponse, parsedBody?: unknown): Promise { // Only validate session ID for non-initialization requests when session management is enabled if (this._sessionId !== undefined) { - const sessionId = req.headers["mcp-session-id"]; - - // Check if this might be an initialization request const isInitializationRequest = req.method === "POST" && req.headers["content-type"]?.includes("application/json"); - if (isInitializationRequest) { - // For POST requests with JSON content, we need to check if it's an initialization request - // This will be done in handlePostRequest, as we need to parse the body - // Continue processing normally - } else if (!sessionId) { - // Non-initialization requests without a session ID should return 400 Bad Request - res.writeHead(400, this._customHeaders).end(JSON.stringify({ - jsonrpc: "2.0", - error: { - code: -32000, - message: "Bad Request: Mcp-Session-Id header is required" - }, - id: null - })); - return; - } else if ((Array.isArray(sessionId) ? sessionId[0] : sessionId) !== this._sessionId) { - // Reject requests with invalid session ID with 404 Not Found - res.writeHead(404, this._customHeaders).end(JSON.stringify({ - jsonrpc: "2.0", - error: { - code: -32001, - message: "Session not found" - }, - id: null - })); + if (!isInitializationRequest && !this.validateSession(req, res)) { return; } } @@ -163,66 +120,22 @@ export class StreamableHTTPServerTransport implements Transport { /** * Handles GET requests to establish SSE connections + * According to the MCP Streamable HTTP transport spec, the server MUST either return SSE or 405. + * We choose to return 405 Method Not Allowed as we don't support GET SSE connections yet. */ private async handleGetRequest(req: IncomingMessage, res: ServerResponse): Promise { - // validate the Accept header - const acceptHeader = req.headers.accept; - if (!acceptHeader || !acceptHeader.includes("text/event-stream")) { - res.writeHead(406, this._customHeaders).end(JSON.stringify({ - jsonrpc: "2.0", - error: { - code: -32000, - message: "Not Acceptable: Client must accept text/event-stream" - }, - id: null - })); - return; - } - - const connectionId = randomUUID(); - const lastEventId = req.headers["last-event-id"]; - const lastEventIdStr = Array.isArray(lastEventId) ? lastEventId[0] : lastEventId; - - // Prepare response headers - const headers: Record = { + // Return 405 Method Not Allowed + res.writeHead(405, { ...this._customHeaders, - "Content-Type": "text/event-stream", - "Cache-Control": "no-cache", - Connection: "keep-alive", - }; - - // Only include session ID header if session management is enabled by assigning a value to _sessionId - if (this._sessionId !== undefined) { - headers["mcp-session-id"] = this._sessionId; - } - - res.writeHead(200, headers); - - const connection: StreamConnection = { - response: res, - lastEventId: lastEventIdStr, - messages: [], - }; - - this._connections.set(connectionId, connection); - - // if there is a Last-Event-ID, replay messages on this connection - if (lastEventIdStr) { - this.replayMessages(connectionId, lastEventIdStr); - } - - res.on("close", () => { - this._connections.delete(connectionId); - // remove all request mappings associated with this connection - for (const [reqId, connId] of this._requestConnections.entries()) { - if (connId === connectionId) { - this._requestConnections.delete(reqId); - } - } - if (this._connections.size === 0) { - this.onclose?.(); - } - }); + "Allow": "POST, DELETE" + }).end(JSON.stringify({ + jsonrpc: "2.0", + error: { + code: -32000, + message: "Method not allowed: Server does not offer an SSE stream at this endpoint" + }, + id: null + })); } /** @@ -279,11 +192,18 @@ export class StreamableHTTPServerTransport implements Transport { 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' && 'id' in msg - ); + if (this._sessionId !== undefined) { + // 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' && 'id' in msg + ); + + if (!isInitializationRequest && !this.validateSession(req, res)) { + return; + } + } + // check if it contains requests const hasRequests = messages.some(msg => 'method' in msg && 'id' in msg); @@ -310,40 +230,29 @@ export class StreamableHTTPServerTransport implements Transport { Connection: "keep-alive", }; + // For initialization requests, always include the session ID if we have one + // even if we're in stateless mode if (this._sessionId !== undefined) { headers["mcp-session-id"] = this._sessionId; } res.writeHead(200, headers); - const connectionId = randomUUID(); - const connection: StreamConnection = { - response: res, - messages: [], - }; - - this._connections.set(connectionId, connection); - - // map each request to a connection ID + // 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._requestConnections.set(String(message.id), connectionId); + this._sseResponseMapping.set(message.id, res); } + } + + // handle each message + for (const message of messages) { this.onmessage?.(message); } - res.on("close", () => { - this._connections.delete(connectionId); - // remove all request mappings associated with this connection - for (const [reqId, connId] of this._requestConnections.entries()) { - if (connId === connectionId) { - this._requestConnections.delete(reqId); - } - } - if (this._connections.size === 0) { - this.onclose?.(); - } - }); + // 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 } else { // use direct JSON response const headers: Record = { @@ -351,7 +260,8 @@ export class StreamableHTTPServerTransport implements Transport { "Content-Type": "application/json", }; - + // For initialization requests, always include the session ID if we have one + // even if we're in stateless mode if (this._sessionId !== undefined) { headers["mcp-session-id"] = this._sessionId; } @@ -390,97 +300,73 @@ export class StreamableHTTPServerTransport implements Transport { } /** - * Replays messages after the specified event ID for a specific connection + * Validates session ID for non-initialization requests when session management is enabled + * Returns true if the session is valid, false otherwise */ - private replayMessages(connectionId: string, lastEventId: string): void { - if (!lastEventId) return; + private validateSession(req: IncomingMessage, res: ServerResponse): boolean { + const sessionId = req.headers["mcp-session-id"]; - // only replay messages that should be sent on this connection - const messages = Array.from(this._messageHistory.entries()) - .filter(([id, { connectionId: msgConnId }]) => - id > lastEventId && - (!msgConnId || msgConnId === connectionId)) // only replay messages that are not specified to a connection or specified to the current connection - .sort(([a], [b]) => a.localeCompare(b)); - - const connection = this._connections.get(connectionId); - if (!connection) return; - - for (const [id, { message }] of messages) { - connection.response.write( - `id: ${id}\nevent: message\ndata: ${JSON.stringify(message)}\n\n` - ); + if (!sessionId) { + // Non-initialization requests without a session ID should return 400 Bad Request + res.writeHead(400, this._customHeaders).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) ? sessionId[0] : sessionId) !== this._sessionId) { + // Reject requests with invalid session ID with 404 Not Found + res.writeHead(404, this._customHeaders).end(JSON.stringify({ + jsonrpc: "2.0", + error: { + code: -32001, + message: "Session not found" + }, + id: null + })); + return false; } + + return true; } + async close(): Promise { - for (const connection of this._connections.values()) { - connection.response.end(); - } - this._connections.clear(); - this._messageHistory.clear(); - this._requestConnections.clear(); + // Close all SSE connections + this._sseResponseMapping.forEach((response) => { + response.end(); + }); + this._sseResponseMapping.clear(); this.onclose?.(); } - async send(message: JSONRPCMessage): Promise { - if (this._connections.size === 0) { - throw new Error("No active connections"); + async send(message: JSONRPCMessage, options?: { relatedRequestId?: RequestId }): Promise { + const relatedRequestId = options?.relatedRequestId; + if (relatedRequestId === undefined) { + throw new Error("relatedRequestId is required"); } - let targetConnectionId = ""; - - // if it is a response, find the corresponding request connection - if ('id' in message && ('result' in message || 'error' in message)) { - const connId = this._requestConnections.get(String(message.id)); - - // if the corresponding connection is not found, the connection may be disconnected - if (!connId || !this._connections.has(connId)) { - // select an available connection - const firstConnId = this._connections.keys().next().value; - if (firstConnId) { - targetConnectionId = firstConnId; - } else { - throw new Error("No available connections"); - } - } else { - targetConnectionId = connId; - } - } else { - // for other messages, select an available connection - const firstConnId = this._connections.keys().next().value; - if (firstConnId) { - targetConnectionId = firstConnId; - } else { - throw new Error("No available connections"); - } + const sseResponse = this._sseResponseMapping.get(relatedRequestId); + if (!sseResponse) { + throw new Error("No SSE connection established"); } - const messageId = randomUUID(); - this._messageHistory.set(messageId, { - message, - connectionId: targetConnectionId - }); + // Send the message as an SSE event + sseResponse.write( + `event: message\ndata: ${JSON.stringify(message)}\n\n`, + ); - // keep the message history in a reasonable range - if (this._messageHistory.size > 1000) { - const oldestKey = Array.from(this._messageHistory.keys())[0]; - this._messageHistory.delete(oldestKey); - } - - // send the message to all active connections - for (const [connId, connection] of this._connections.entries()) { - // if it is a response message, only send to the target connection - if ('id' in message && ('result' in message || 'error' in message)) { - if (connId === targetConnectionId) { - connection.response.write( - `id: ${messageId}\nevent: message\ndata: ${JSON.stringify(message)}\n\n` - ); - } - } else { - // for other messages, send to all connections - connection.response.write( - `id: ${messageId}\nevent: message\ndata: ${JSON.stringify(message)}\n\n` - ); + // If this is a response message with the same ID as the request, we can check + // if we need to close the stream after sending the response + if ('result' in message || 'error' in message) { + if (message.id === relatedRequestId) { + // This is a response to the original request, we can close the stream + // after sending all related responses + this._sseResponseMapping.delete(relatedRequestId); + sseResponse.end(); } } } diff --git a/src/shared/transport.ts b/src/shared/transport.ts index 72db1da8..d855ceab 100644 --- a/src/shared/transport.ts +++ b/src/shared/transport.ts @@ -16,7 +16,7 @@ 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. + * If present, `relatedRequestId` is used to indicate to the transport which incoming request to associate this outgoing message with. */ send(message: JSONRPCMessage, options?: { relatedRequestId?: RequestId }): Promise;