modelcontextprotocol/typescript-sdk
Watch
1
0
Fork
You've already forked typescript-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

✨ feat: add customHeaders options

Browse source
This commit is contained in:
Geng Yan 2025-04-03 08:44:29 +08:00 • committed by ihrpr
commit e9caa5a12f
No known key found for this signature in database
2 changed files with 155 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

143
src/server/streamable-http.test.ts
View file

@ -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<ServerResponse>;
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
})
);
});
});
});

19
src/server/streamable-http.ts
View file

@ -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<string, string>;
}
/**
@ -72,6 +78,7 @@ export class StreamableHTTPServerTransport implements Transport {
private _started: boolean = false;
private _requestConnections: Map<string, string> = new Map(); // request ID to connection ID mapping
private _enableSessionManagement: boolean;
private _customHeaders: Record<string, string>;
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<string, string> = {
...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<string, string> = {
...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<string, string> = {
...this._customHeaders,
"Content-Type": "application/json",
};