From e9caa5a12f8cba52b734102cd3cd28836f6cc113 Mon Sep 17 00:00:00 2001 From: Geng Yan Date: Thu, 3 Apr 2025 08:44:29 +0800 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20feat:=20add=20customHeaders=20optio?= =?UTF-8?q?ns?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/server/streamable-http.test.ts | 143 ++++++++++++++++++++++++++++- src/server/streamable-http.ts | 19 +++- 2 files changed, 155 insertions(+), 7 deletions(-) diff --git a/src/server/streamable-http.test.ts b/src/server/streamable-http.test.ts index 13f70b59..74b92c1a 100644 --- a/src/server/streamable-http.test.ts +++ b/src/server/streamable-http.test.ts @@ -98,7 +98,7 @@ describe("StreamableHTTPServerTransport", () => { await transport.handleRequest(req, mockResponse); - expect(mockResponse.writeHead).toHaveBeenCalledWith(404); + expect(mockResponse.writeHead).toHaveBeenCalledWith(404, {}); // check if the error response is a valid JSON-RPC error format expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"jsonrpc":"2.0"')); expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"error"')); @@ -115,7 +115,7 @@ describe("StreamableHTTPServerTransport", () => { await transport.handleRequest(req, mockResponse); - expect(mockResponse.writeHead).toHaveBeenCalledWith(400); + 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"')); }); @@ -342,7 +342,7 @@ describe("StreamableHTTPServerTransport", () => { await transport.handleRequest(req, mockResponse); - expect(mockResponse.writeHead).toHaveBeenCalledWith(406); + expect(mockResponse.writeHead).toHaveBeenCalledWith(406, {}); expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"jsonrpc":"2.0"')); }); @@ -788,4 +788,141 @@ describe("StreamableHTTPServerTransport", () => { expect(onMessageMock).not.toHaveBeenCalledWith(requestBodyMessage); }); }); + + describe("Custom Headers", () => { + const customHeaders = { + "X-Custom-Header": "custom-value", + "X-API-Version": "1.0", + "Access-Control-Allow-Origin": "*" + }; + + let transportWithHeaders: StreamableHTTPServerTransport; + let mockResponse: jest.Mocked; + + beforeEach(() => { + transportWithHeaders = new StreamableHTTPServerTransport(endpoint, { customHeaders }); + mockResponse = createMockResponse(); + }); + + it("should include custom headers in SSE response", async () => { + const req = createMockRequest({ + method: "GET", + headers: { + accept: "text/event-stream", + "mcp-session-id": transportWithHeaders.sessionId + }, + }); + + await transportWithHeaders.handleRequest(req, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalledWith( + 200, + expect.objectContaining({ + ...customHeaders, + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "mcp-session-id": transportWithHeaders.sessionId + }) + ); + }); + + it("should include custom headers in JSON response", 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", + "mcp-session-id": transportWithHeaders.sessionId + }, + body: JSON.stringify(message), + }); + + await transportWithHeaders.handleRequest(req, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalledWith( + 200, + expect.objectContaining({ + ...customHeaders, + "Content-Type": "application/json", + "mcp-session-id": transportWithHeaders.sessionId + }) + ); + }); + + it("should include custom headers in error responses", async () => { + const req = createMockRequest({ + method: "GET", + headers: { + accept: "text/event-stream", + "mcp-session-id": "invalid-session-id" + }, + }); + + await transportWithHeaders.handleRequest(req, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalledWith( + 404, + expect.objectContaining(customHeaders) + ); + }); + + it("should not override essential headers with custom headers", async () => { + const transportWithConflictingHeaders = new StreamableHTTPServerTransport(endpoint, { + customHeaders: { + "Content-Type": "text/plain", // 尝试覆盖必要的 Content-Type 头 + "X-Custom-Header": "custom-value" + } + }); + + const req = createMockRequest({ + method: "GET", + headers: { + accept: "text/event-stream", + "mcp-session-id": transportWithConflictingHeaders.sessionId + }, + }); + + await transportWithConflictingHeaders.handleRequest(req, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalledWith( + 200, + expect.objectContaining({ + "Content-Type": "text/event-stream", // 应该保持原有的 Content-Type + "X-Custom-Header": "custom-value" + }) + ); + }); + + it("should work with empty custom headers", async () => { + const transportWithoutHeaders = new StreamableHTTPServerTransport(endpoint); + + const req = createMockRequest({ + method: "GET", + headers: { + accept: "text/event-stream", + "mcp-session-id": transportWithoutHeaders.sessionId + }, + }); + + await transportWithoutHeaders.handleRequest(req, mockResponse); + + expect(mockResponse.writeHead).toHaveBeenCalledWith( + 200, + expect.objectContaining({ + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "mcp-session-id": transportWithoutHeaders.sessionId + }) + ); + }); + }); }); \ No newline at end of file diff --git a/src/server/streamable-http.ts b/src/server/streamable-http.ts index bec0a1ad..13d28968 100644 --- a/src/server/streamable-http.ts +++ b/src/server/streamable-http.ts @@ -28,6 +28,12 @@ export interface StreamableHTTPServerTransportOptions { * @default true */ enableSessionManagement?: boolean; + + /** + * Custom headers to be included in all responses + * These headers will be added to both SSE and regular HTTP responses + */ + customHeaders?: Record; } /** @@ -72,6 +78,7 @@ export class StreamableHTTPServerTransport implements Transport { private _started: boolean = false; private _requestConnections: Map = new Map(); // request ID to connection ID mapping private _enableSessionManagement: boolean; + private _customHeaders: Record; onclose?: () => void; onerror?: (error: Error) => void; @@ -80,6 +87,7 @@ export class StreamableHTTPServerTransport implements Transport { constructor(private _endpoint: string, options?: StreamableHTTPServerTransportOptions) { this._sessionId = randomUUID(); this._enableSessionManagement = options?.enableSessionManagement !== false; + this._customHeaders = options?.customHeaders || {}; } /** @@ -111,7 +119,7 @@ export class StreamableHTTPServerTransport implements Transport { // Continue processing normally } else if (!sessionId) { // Non-initialization requests without a session ID should return 400 Bad Request - res.writeHead(400).end(JSON.stringify({ + res.writeHead(400, this._customHeaders).end(JSON.stringify({ jsonrpc: "2.0", error: { code: -32000, @@ -122,7 +130,7 @@ export class StreamableHTTPServerTransport implements Transport { return; } else if ((Array.isArray(sessionId) ? sessionId[0] : sessionId) !== this._sessionId) { // Reject requests with invalid session ID with 404 Not Found - res.writeHead(404).end(JSON.stringify({ + res.writeHead(404, this._customHeaders).end(JSON.stringify({ jsonrpc: "2.0", error: { code: -32001, @@ -141,7 +149,7 @@ export class StreamableHTTPServerTransport implements Transport { } else if (req.method === "DELETE") { await this.handleDeleteRequest(req, res); } else { - res.writeHead(405).end(JSON.stringify({ + res.writeHead(405, this._customHeaders).end(JSON.stringify({ jsonrpc: "2.0", error: { code: -32000, @@ -159,7 +167,7 @@ export class StreamableHTTPServerTransport implements Transport { // validate the Accept header const acceptHeader = req.headers.accept; if (!acceptHeader || !acceptHeader.includes("text/event-stream")) { - res.writeHead(406).end(JSON.stringify({ + res.writeHead(406, this._customHeaders).end(JSON.stringify({ jsonrpc: "2.0", error: { code: -32000, @@ -176,6 +184,7 @@ export class StreamableHTTPServerTransport implements Transport { // Prepare response headers const headers: Record = { + ...this._customHeaders, "Content-Type": "text/event-stream", "Cache-Control": "no-cache", Connection: "keep-alive", @@ -294,6 +303,7 @@ export class StreamableHTTPServerTransport implements Transport { if (useSSE) { const headers: Record = { + ...this._customHeaders, "Content-Type": "text/event-stream", "Cache-Control": "no-cache", Connection: "keep-alive", @@ -338,6 +348,7 @@ export class StreamableHTTPServerTransport implements Transport { } else { // use direct JSON response const headers: Record = { + ...this._customHeaders, "Content-Type": "application/json", };