remove the use of enableSessionManagement
This commit is contained in:
ihrpr
parent
b0697195a7
commit
e3bb99cd5c
2 changed files with 104 additions and 127 deletions
|
|
@ -2,7 +2,7 @@ import { IncomingMessage, ServerResponse } from "node:http";
|
|||
import { StreamableHTTPServerTransport } from "./streamable-http.js";
|
||||
import { JSONRPCMessage } from "../types.js";
|
||||
import { Readable } from "node:stream";
|
||||
|
||||
import { randomUUID } from "node:crypto";
|
||||
// Mock IncomingMessage
|
||||
function createMockRequest(options: {
|
||||
method: string;
|
||||
|
|
@ -10,7 +10,7 @@ function createMockRequest(options: {
|
|||
body?: string;
|
||||
}): IncomingMessage {
|
||||
const readable = new Readable();
|
||||
readable._read = () => {};
|
||||
readable._read = () => { };
|
||||
if (options.body) {
|
||||
readable.push(options.body);
|
||||
readable.push(null);
|
||||
|
|
@ -37,12 +37,13 @@ function createMockResponse(): jest.Mocked<ServerResponse> {
|
|||
}
|
||||
|
||||
describe("StreamableHTTPServerTransport", () => {
|
||||
const endpoint = "/mcp";
|
||||
let transport: StreamableHTTPServerTransport;
|
||||
let mockResponse: jest.Mocked<ServerResponse>;
|
||||
|
||||
beforeEach(() => {
|
||||
transport = new StreamableHTTPServerTransport(endpoint);
|
||||
transport = new StreamableHTTPServerTransport({
|
||||
sessionId: randomUUID(),
|
||||
});
|
||||
mockResponse = createMockResponse();
|
||||
});
|
||||
|
||||
|
|
@ -61,7 +62,7 @@ describe("StreamableHTTPServerTransport", () => {
|
|||
const initializeMessage: JSONRPCMessage = {
|
||||
jsonrpc: "2.0",
|
||||
method: "initialize",
|
||||
params: {
|
||||
params: {
|
||||
clientInfo: { name: "test-client", version: "1.0" },
|
||||
protocolVersion: "2025-03-26"
|
||||
},
|
||||
|
|
@ -119,49 +120,13 @@ describe("StreamableHTTPServerTransport", () => {
|
|||
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 always include session ID in initialization response even in stateless mode", async () => {
|
||||
// Create a stateless transport for this test
|
||||
const statelessTransport = new StreamableHTTPServerTransport(endpoint, { enableSessionManagement: false });
|
||||
|
||||
// Create an initialization request
|
||||
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",
|
||||
},
|
||||
body: JSON.stringify(initializeMessage),
|
||||
});
|
||||
|
||||
await statelessTransport.handleRequest(req, mockResponse);
|
||||
|
||||
// In stateless mode, session ID should also be included for initialize responses
|
||||
expect(mockResponse.writeHead).toHaveBeenCalledWith(
|
||||
200,
|
||||
expect.objectContaining({
|
||||
"mcp-session-id": statelessTransport.sessionId,
|
||||
})
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("Stateless Mode", () => {
|
||||
let statelessTransport: StreamableHTTPServerTransport;
|
||||
let mockResponse: jest.Mocked<ServerResponse>;
|
||||
|
||||
beforeEach(() => {
|
||||
statelessTransport = new StreamableHTTPServerTransport(endpoint, { enableSessionManagement: false });
|
||||
statelessTransport = new StreamableHTTPServerTransport({ sessionId: undefined });
|
||||
mockResponse = createMockResponse();
|
||||
});
|
||||
|
||||
|
|
@ -268,7 +233,7 @@ describe("StreamableHTTPServerTransport", () => {
|
|||
});
|
||||
|
||||
await statelessTransport.handleRequest(req2, mockResponse);
|
||||
|
||||
|
||||
// Should still succeed
|
||||
expect(mockResponse.writeHead).toHaveBeenCalledWith(
|
||||
200,
|
||||
|
|
@ -278,12 +243,12 @@ describe("StreamableHTTPServerTransport", () => {
|
|||
);
|
||||
});
|
||||
|
||||
it("should handle initialization requests properly in both modes", async () => {
|
||||
it("should handle initialization requests properly in statefull mode", async () => {
|
||||
// Initialize message that would typically be sent during initialization
|
||||
const initializeMessage: JSONRPCMessage = {
|
||||
jsonrpc: "2.0",
|
||||
method: "initialize",
|
||||
params: {
|
||||
params: {
|
||||
clientInfo: { name: "test-client", version: "1.0" },
|
||||
protocolVersion: "2025-03-26"
|
||||
},
|
||||
|
|
@ -301,7 +266,7 @@ describe("StreamableHTTPServerTransport", () => {
|
|||
});
|
||||
|
||||
await transport.handleRequest(statefulReq, mockResponse);
|
||||
|
||||
|
||||
// In stateful mode, session ID should be included in the response header
|
||||
expect(mockResponse.writeHead).toHaveBeenCalledWith(
|
||||
200,
|
||||
|
|
@ -309,9 +274,19 @@ describe("StreamableHTTPServerTransport", () => {
|
|||
"mcp-session-id": transport.sessionId,
|
||||
})
|
||||
);
|
||||
});
|
||||
|
||||
// Reset mocks for stateless test
|
||||
mockResponse.writeHead.mockClear();
|
||||
it("should handle initialization requests properly in stateless mode", async () => {
|
||||
// Initialize message that would typically be sent during initialization
|
||||
const initializeMessage: JSONRPCMessage = {
|
||||
jsonrpc: "2.0",
|
||||
method: "initialize",
|
||||
params: {
|
||||
clientInfo: { name: "test-client", version: "1.0" },
|
||||
protocolVersion: "2025-03-26"
|
||||
},
|
||||
id: "init-1",
|
||||
};
|
||||
|
||||
// Test stateless transport
|
||||
const statelessReq = createMockRequest({
|
||||
|
|
@ -324,10 +299,11 @@ describe("StreamableHTTPServerTransport", () => {
|
|||
});
|
||||
|
||||
await statelessTransport.handleRequest(statelessReq, mockResponse);
|
||||
|
||||
|
||||
// In stateless mode, session ID should also be included for initialize responses
|
||||
const headers = mockResponse.writeHead.mock.calls[0][1];
|
||||
expect(headers).toHaveProperty("mcp-session-id", statelessTransport.sessionId);
|
||||
expect(headers).not.toHaveProperty("mcp-session-id");
|
||||
|
||||
});
|
||||
});
|
||||
|
||||
|
|
@ -519,14 +495,14 @@ describe("StreamableHTTPServerTransport", () => {
|
|||
|
||||
// Send a message to first connection
|
||||
const message1: JSONRPCMessage = {
|
||||
jsonrpc: "2.0",
|
||||
method: "test1",
|
||||
params: {},
|
||||
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-]+)/);
|
||||
|
|
@ -550,12 +526,12 @@ describe("StreamableHTTPServerTransport", () => {
|
|||
|
||||
// Send a second message
|
||||
const message2: JSONRPCMessage = {
|
||||
jsonrpc: "2.0",
|
||||
method: "test2",
|
||||
params: {},
|
||||
jsonrpc: "2.0",
|
||||
method: "test2",
|
||||
params: {},
|
||||
id: 2
|
||||
};
|
||||
|
||||
|
||||
await transport.send(message2);
|
||||
|
||||
// Verify the second message was received by both connections
|
||||
|
|
@ -596,7 +572,7 @@ describe("StreamableHTTPServerTransport", () => {
|
|||
params: {},
|
||||
id: "test-id",
|
||||
};
|
||||
|
||||
|
||||
const reqPost = createMockRequest({
|
||||
method: "POST",
|
||||
headers: {
|
||||
|
|
@ -605,28 +581,28 @@ describe("StreamableHTTPServerTransport", () => {
|
|||
},
|
||||
body: JSON.stringify(requestMessage),
|
||||
});
|
||||
|
||||
|
||||
await transport.handleRequest(reqPost, mockResponse1);
|
||||
|
||||
|
||||
// Send a response with matching ID
|
||||
const responseMessage: JSONRPCMessage = {
|
||||
jsonrpc: "2.0",
|
||||
result: { success: true },
|
||||
id: "test-id",
|
||||
};
|
||||
|
||||
|
||||
await transport.send(responseMessage);
|
||||
|
||||
|
||||
// Verify response was sent to the right connection
|
||||
expect(mockResponse1.write).toHaveBeenCalledWith(
|
||||
expect.stringContaining(JSON.stringify(responseMessage))
|
||||
);
|
||||
|
||||
|
||||
// Check if write was called with this exact message on the second connection
|
||||
const writeCallsOnSecondConn = mockResponse2.write.mock.calls.filter(call =>
|
||||
const writeCallsOnSecondConn = mockResponse2.write.mock.calls.filter(call =>
|
||||
typeof call[0] === 'string' && call[0].includes(JSON.stringify(responseMessage))
|
||||
);
|
||||
|
||||
|
||||
// Verify the response wasn't broadcast to all connections
|
||||
expect(writeCallsOnSecondConn.length).toBe(0);
|
||||
});
|
||||
|
|
@ -680,7 +656,7 @@ describe("StreamableHTTPServerTransport", () => {
|
|||
const message: JSONRPCMessage = {
|
||||
jsonrpc: "2.0",
|
||||
method: "initialize",
|
||||
params: {
|
||||
params: {
|
||||
clientInfo: { name: "test-client", version: "1.0" },
|
||||
protocolVersion: "2025-03-26"
|
||||
},
|
||||
|
|
@ -715,17 +691,17 @@ describe("StreamableHTTPServerTransport", () => {
|
|||
|
||||
it("should handle pre-parsed batch messages", async () => {
|
||||
const batchMessages: JSONRPCMessage[] = [
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
method: "method1",
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
method: "method1",
|
||||
params: { data: "test1" },
|
||||
id: "batch1"
|
||||
id: "batch1"
|
||||
},
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
method: "method2",
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
method: "method2",
|
||||
params: { data: "test2" },
|
||||
id: "batch2"
|
||||
id: "batch2"
|
||||
},
|
||||
];
|
||||
|
||||
|
|
@ -800,7 +776,7 @@ describe("StreamableHTTPServerTransport", () => {
|
|||
let mockResponse: jest.Mocked<ServerResponse>;
|
||||
|
||||
beforeEach(() => {
|
||||
transportWithHeaders = new StreamableHTTPServerTransport(endpoint, { customHeaders });
|
||||
transportWithHeaders = new StreamableHTTPServerTransport({ sessionId: randomUUID(), customHeaders });
|
||||
mockResponse = createMockResponse();
|
||||
});
|
||||
|
||||
|
|
@ -875,9 +851,10 @@ describe("StreamableHTTPServerTransport", () => {
|
|||
});
|
||||
|
||||
it("should not override essential headers with custom headers", async () => {
|
||||
const transportWithConflictingHeaders = new StreamableHTTPServerTransport(endpoint, {
|
||||
const transportWithConflictingHeaders = new StreamableHTTPServerTransport({
|
||||
sessionId: randomUUID(),
|
||||
customHeaders: {
|
||||
"Content-Type": "text/plain", // 尝试覆盖必要的 Content-Type 头
|
||||
"Content-Type": "text/plain",
|
||||
"X-Custom-Header": "custom-value"
|
||||
}
|
||||
});
|
||||
|
|
@ -902,8 +879,10 @@ describe("StreamableHTTPServerTransport", () => {
|
|||
});
|
||||
|
||||
it("should work with empty custom headers", async () => {
|
||||
const transportWithoutHeaders = new StreamableHTTPServerTransport(endpoint);
|
||||
|
||||
const transportWithoutHeaders = new StreamableHTTPServerTransport({
|
||||
sessionId: randomUUID(),
|
||||
});
|
||||
|
||||
const req = createMockRequest({
|
||||
method: "GET",
|
||||
headers: {
|
||||
|
|
|
|||
|
|
@ -23,11 +23,11 @@ interface StreamConnection {
|
|||
*/
|
||||
export interface StreamableHTTPServerTransportOptions {
|
||||
/**
|
||||
* Whether to enable session management through mcp-session-id headers
|
||||
* When set to false, the transport operates in stateless mode without session validation
|
||||
* @default true
|
||||
* The session ID SHOULD be globally unique and cryptographically secure (e.g., a securely generated UUID, a JWT, or a cryptographic hash)
|
||||
*
|
||||
* When sessionId is not set, the transport will be in stateless mode.
|
||||
*/
|
||||
enableSessionManagement?: boolean;
|
||||
sessionId: string | undefined;
|
||||
|
||||
/**
|
||||
* Custom headers to be included in all responses
|
||||
|
|
@ -38,17 +38,19 @@ export interface StreamableHTTPServerTransportOptions {
|
|||
|
||||
/**
|
||||
* Server transport for Streamable HTTP: this implements the MCP Streamable HTTP transport specification.
|
||||
* It supports both SSE streaming and direct HTTP responses, with session management and message resumability.
|
||||
* It supports both SSE streaming and direct HTTP responses.
|
||||
*
|
||||
* Usage example:
|
||||
*
|
||||
* ```typescript
|
||||
* // Stateful mode (default) - with session management
|
||||
* const statefulTransport = new StreamableHTTPServerTransport("/mcp");
|
||||
* // Stateful mode - server sets the session ID
|
||||
* const statefulTransport = new StreamableHTTPServerTransport({
|
||||
* sessionId: randomUUID(),
|
||||
* });
|
||||
*
|
||||
* // Stateless mode - without session management
|
||||
* const statelessTransport = new StreamableHTTPServerTransport("/mcp", {
|
||||
* enableSessionManagement: false
|
||||
* // Stateless mode - explisitly set session ID to undefined
|
||||
* const statelessTransport = new StreamableHTTPServerTransport({
|
||||
* sessionId: undefined,
|
||||
* });
|
||||
*
|
||||
* // Using with pre-parsed request body
|
||||
|
|
@ -70,23 +72,22 @@ export interface StreamableHTTPServerTransportOptions {
|
|||
*/
|
||||
export class StreamableHTTPServerTransport implements Transport {
|
||||
private _connections: Map<string, StreamConnection> = new Map();
|
||||
private _sessionId: string;
|
||||
// when sessionID is not set, it means the transport is in stateless mode
|
||||
private _sessionId: string | undefined;
|
||||
private _messageHistory: Map<string, {
|
||||
message: JSONRPCMessage;
|
||||
connectionId?: string; // record which connection the message should be sent to
|
||||
}> = new Map();
|
||||
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;
|
||||
onmessage?: (message: JSONRPCMessage) => void;
|
||||
|
||||
constructor(private _endpoint: string, options?: StreamableHTTPServerTransportOptions) {
|
||||
this._sessionId = randomUUID();
|
||||
this._enableSessionManagement = options?.enableSessionManagement !== false;
|
||||
constructor(options: StreamableHTTPServerTransportOptions) {
|
||||
this._sessionId = options?.sessionId;
|
||||
this._customHeaders = options?.customHeaders || {};
|
||||
}
|
||||
|
||||
|
|
@ -106,13 +107,13 @@ export class StreamableHTTPServerTransport implements Transport {
|
|||
*/
|
||||
async handleRequest(req: IncomingMessage, res: ServerResponse, parsedBody?: unknown): Promise<void> {
|
||||
// Only validate session ID for non-initialization requests when session management is enabled
|
||||
if (this._enableSessionManagement) {
|
||||
if (this._sessionId !== undefined) {
|
||||
const sessionId = req.headers["mcp-session-id"];
|
||||
|
||||
|
||||
// Check if this might be an initialization request
|
||||
const isInitializationRequest = req.method === "POST" &&
|
||||
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
|
||||
|
|
@ -189,9 +190,9 @@ export class StreamableHTTPServerTransport implements Transport {
|
|||
"Cache-Control": "no-cache",
|
||||
Connection: "keep-alive",
|
||||
};
|
||||
|
||||
// Only include session ID header if session management is enabled
|
||||
if (this._enableSessionManagement) {
|
||||
|
||||
// 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;
|
||||
}
|
||||
|
||||
|
|
@ -231,8 +232,8 @@ export class StreamableHTTPServerTransport implements Transport {
|
|||
try {
|
||||
// validate the Accept header
|
||||
const acceptHeader = req.headers.accept;
|
||||
if (!acceptHeader ||
|
||||
(!acceptHeader.includes("application/json") && !acceptHeader.includes("text/event-stream"))) {
|
||||
if (!acceptHeader ||
|
||||
(!acceptHeader.includes("application/json") && !acceptHeader.includes("text/event-stream"))) {
|
||||
res.writeHead(406).end(JSON.stringify({
|
||||
jsonrpc: "2.0",
|
||||
error: {
|
||||
|
|
@ -270,7 +271,7 @@ export class StreamableHTTPServerTransport implements Transport {
|
|||
}
|
||||
|
||||
let messages: JSONRPCMessage[];
|
||||
|
||||
|
||||
// handle batch and single messages
|
||||
if (Array.isArray(rawMessage)) {
|
||||
messages = rawMessage.map(msg => JSONRPCMessageSchema.parse(msg));
|
||||
|
|
@ -286,7 +287,7 @@ export class StreamableHTTPServerTransport implements Transport {
|
|||
|
||||
// check if it contains requests
|
||||
const hasRequests = messages.some(msg => 'method' in msg && 'id' in msg);
|
||||
const hasOnlyNotificationsOrResponses = messages.every(msg =>
|
||||
const hasOnlyNotificationsOrResponses = messages.every(msg =>
|
||||
('method' in msg && !('id' in msg)) || ('result' in msg || 'error' in msg));
|
||||
|
||||
if (hasOnlyNotificationsOrResponses) {
|
||||
|
|
@ -300,7 +301,7 @@ export class StreamableHTTPServerTransport implements Transport {
|
|||
} else if (hasRequests) {
|
||||
// if it contains requests, you can choose to return an SSE stream or a JSON response
|
||||
const useSSE = acceptHeader.includes("text/event-stream");
|
||||
|
||||
|
||||
if (useSSE) {
|
||||
const headers: Record<string, string> = {
|
||||
...this._customHeaders,
|
||||
|
|
@ -308,10 +309,8 @@ export class StreamableHTTPServerTransport implements Transport {
|
|||
"Cache-Control": "no-cache",
|
||||
Connection: "keep-alive",
|
||||
};
|
||||
|
||||
// Only include session ID header if session management is enabled
|
||||
// Always include session ID for initialization requests
|
||||
if (this._enableSessionManagement || isInitializationRequest) {
|
||||
|
||||
if (this._sessionId !== undefined) {
|
||||
headers["mcp-session-id"] = this._sessionId;
|
||||
}
|
||||
|
||||
|
|
@ -351,20 +350,19 @@ export class StreamableHTTPServerTransport implements Transport {
|
|||
...this._customHeaders,
|
||||
"Content-Type": "application/json",
|
||||
};
|
||||
|
||||
// Only include session ID header if session management is enabled
|
||||
// Always include session ID for initialization requests
|
||||
if (this._enableSessionManagement || isInitializationRequest) {
|
||||
|
||||
|
||||
if (this._sessionId !== undefined) {
|
||||
headers["mcp-session-id"] = this._sessionId;
|
||||
}
|
||||
|
||||
res.writeHead(200, headers);
|
||||
|
||||
|
||||
// handle each message
|
||||
for (const message of messages) {
|
||||
this.onmessage?.(message);
|
||||
}
|
||||
|
||||
|
||||
res.end();
|
||||
}
|
||||
}
|
||||
|
|
@ -396,11 +394,11 @@ export class StreamableHTTPServerTransport implements Transport {
|
|||
*/
|
||||
private replayMessages(connectionId: string, lastEventId: string): void {
|
||||
if (!lastEventId) return;
|
||||
|
||||
|
||||
// only replay messages that should be sent on this connection
|
||||
const messages = Array.from(this._messageHistory.entries())
|
||||
.filter(([id, { connectionId: msgConnId }]) =>
|
||||
id > lastEventId &&
|
||||
.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));
|
||||
|
||||
|
|
@ -430,11 +428,11 @@ export class StreamableHTTPServerTransport implements Transport {
|
|||
}
|
||||
|
||||
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
|
||||
|
|
@ -458,9 +456,9 @@ export class StreamableHTTPServerTransport implements Transport {
|
|||
}
|
||||
|
||||
const messageId = randomUUID();
|
||||
this._messageHistory.set(messageId, {
|
||||
message,
|
||||
connectionId: targetConnectionId
|
||||
this._messageHistory.set(messageId, {
|
||||
message,
|
||||
connectionId: targetConnectionId
|
||||
});
|
||||
|
||||
// keep the message history in a reasonable range
|
||||
|
|
@ -490,7 +488,7 @@ export class StreamableHTTPServerTransport implements Transport {
|
|||
/**
|
||||
* Returns the session ID for this transport
|
||||
*/
|
||||
get sessionId(): string {
|
||||
get sessionId(): string | undefined {
|
||||
return this._sessionId;
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue