✨ feat: try add stateless mod

This commit is contained in:
Geng Yan 2025-04-02 08:47:35 +08:00 • committed by ihrpr
commit 4c7c434e7c
No known key found for this signature in database
2 changed files with 368 additions and 31 deletions

View file

@ -57,11 +57,24 @@ describe("StreamableHTTPServerTransport", () => {
});
it("should include session ID in response headers", async () => {
const req = createMockRequest({
method: "GET",
headers: {
accept: "text/event-stream"
// Use POST with initialize method to avoid session ID requirement
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 transport.handleRequest(req, mockResponse);
@ -79,6 +92,7 @@ describe("StreamableHTTPServerTransport", () => {
method: "GET",
headers: {
"mcp-session-id": "invalid-session-id",
"accept": "text/event-stream"
},
});
@ -89,13 +103,241 @@ describe("StreamableHTTPServerTransport", () => {
expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"jsonrpc":"2.0"'));
expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"error"'));
});
it("should reject non-initialization requests without session ID with 400 Bad Request", async () => {
const req = createMockRequest({
method: "GET",
headers: {
accept: "text/event-stream",
// No mcp-session-id header
},
});
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 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 });
mockResponse = createMockResponse();
});
it("should not include session ID in response headers when in stateless mode", 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",
},
body: JSON.stringify(message),
});
await statelessTransport.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 stateless mode", async () => {
const req = createMockRequest({
method: "GET",
headers: {
accept: "text/event-stream",
"mcp-session-id": "invalid-session-id", // This would cause a 404 in stateful mode
},
});
await statelessTransport.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 stateless mode", 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": "non-existent-session-id", // This would be rejected in stateful mode
},
body: JSON.stringify(message),
});
const onMessageMock = jest.fn();
statelessTransport.onmessage = onMessageMock;
await statelessTransport.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 stateless mode", async () => {
// First request without session ID
const req1 = createMockRequest({
method: "GET",
headers: {
accept: "text/event-stream",
},
});
await statelessTransport.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 stateful mode)
const req2 = createMockRequest({
method: "GET",
headers: {
accept: "text/event-stream",
"mcp-session-id": "some-random-session-id",
},
});
await statelessTransport.handleRequest(req2, mockResponse);
// Should still succeed
expect(mockResponse.writeHead).toHaveBeenCalledWith(
200,
expect.objectContaining({
"Content-Type": "text/event-stream",
})
);
});
it("should handle initialization requests properly in both modes", 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 stateful transport (default)
const statefulReq = createMockRequest({
method: "POST",
headers: {
"content-type": "application/json",
"accept": "application/json",
},
body: JSON.stringify(initializeMessage),
});
await transport.handleRequest(statefulReq, mockResponse);
// In stateful mode, session ID should be included in the response header
expect(mockResponse.writeHead).toHaveBeenCalledWith(
200,
expect.objectContaining({
"mcp-session-id": transport.sessionId,
})
);
// Reset mocks for stateless test
mockResponse.writeHead.mockClear();
// Test stateless transport
const statelessReq = createMockRequest({
method: "POST",
headers: {
"content-type": "application/json",
"accept": "application/json",
},
body: JSON.stringify(initializeMessage),
});
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);
});
});
describe("Request Handling", () => {
it("should reject GET requests without Accept: text/event-stream header", async () => {
const req = createMockRequest({
method: "GET",
headers: {},
headers: {
"mcp-session-id": transport.sessionId,
},
});
await transport.handleRequest(req, mockResponse);
@ -108,7 +350,8 @@ describe("StreamableHTTPServerTransport", () => {
const req = createMockRequest({
method: "GET",
headers: {
accept: "text/event-stream",
"accept": "text/event-stream",
"mcp-session-id": transport.sessionId,
},
});
@ -127,7 +370,7 @@ describe("StreamableHTTPServerTransport", () => {
it("should reject POST requests without proper Accept header", async () => {
const message: JSONRPCMessage = {
jsonrpc: "2.0",
method: "test",
method: "initialize", // Use initialize to bypass session ID check
params: {},
id: 1,
};
@ -148,7 +391,7 @@ describe("StreamableHTTPServerTransport", () => {
it("should properly handle JSON-RPC request messages in POST requests", async () => {
const message: JSONRPCMessage = {
jsonrpc: "2.0",
method: "test",
method: "initialize", // Use initialize to bypass session ID check
params: {},
id: 1,
};
@ -188,6 +431,7 @@ describe("StreamableHTTPServerTransport", () => {
headers: {
"content-type": "application/json",
"accept": "application/json, text/event-stream",
"mcp-session-id": transport.sessionId,
},
body: JSON.stringify(notification),
});
@ -212,6 +456,7 @@ describe("StreamableHTTPServerTransport", () => {
headers: {
"content-type": "application/json",
"accept": "application/json",
"mcp-session-id": transport.sessionId,
},
body: JSON.stringify(batchMessages),
});
@ -231,6 +476,7 @@ describe("StreamableHTTPServerTransport", () => {
headers: {
"content-type": "text/plain",
"accept": "application/json",
"mcp-session-id": transport.sessionId,
},
body: "test",
});
@ -244,7 +490,9 @@ describe("StreamableHTTPServerTransport", () => {
it("should properly handle DELETE requests and close session", async () => {
const req = createMockRequest({
method: "DELETE",
headers: {},
headers: {
"mcp-session-id": transport.sessionId,
},
});
const onCloseMock = jest.fn();
@ -259,11 +507,12 @@ describe("StreamableHTTPServerTransport", () => {
describe("Message Replay", () => {
it("should replay messages after specified Last-Event-ID", async () => {
// Establish first connection with Accept header
// Establish first connection with Accept header and session ID
const req1 = createMockRequest({
method: "GET",
headers: {
"accept": "text/event-stream"
"accept": "text/event-stream",
"mcp-session-id": transport.sessionId
},
});
await transport.handleRequest(req1, mockResponse);
@ -293,6 +542,7 @@ describe("StreamableHTTPServerTransport", () => {
headers: {
"accept": "text/event-stream",
"last-event-id": lastEventId,
"mcp-session-id": transport.sessionId
},
});

View file

@ -18,9 +18,44 @@ interface StreamConnection {
requestId?: string | null;
}
/**
* Configuration options for StreamableHTTPServerTransport
*/
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
*/
enableSessionManagement?: boolean;
}
/**
* 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.
*
* Usage example:
*
* ```typescript
* // Stateful mode (default) - with session management
* const statefulTransport = new StreamableHTTPServerTransport("/mcp");
*
* // Stateless mode - without session management
* const statelessTransport = new StreamableHTTPServerTransport("/mcp", {
* enableSessionManagement: false
* });
* ```
*
* 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 {
private _connections: Map<string, StreamConnection> = new Map();
@ -31,13 +66,15 @@ export class StreamableHTTPServerTransport implements Transport {
}> = new Map();
private _started: boolean = false;
private _requestConnections: Map<string, string> = new Map(); // request ID to connection ID mapping
private _enableSessionManagement: boolean;
onclose?: () => void;
onerror?: (error: Error) => void;
onmessage?: (message: JSONRPCMessage) => void;
constructor(private _endpoint: string) {
constructor(private _endpoint: string, options?: StreamableHTTPServerTransportOptions) {
this._sessionId = randomUUID();
this._enableSessionManagement = options?.enableSessionManagement !== false;
}
/**
@ -55,18 +92,41 @@ export class StreamableHTTPServerTransport implements Transport {
* Handles an incoming HTTP request, whether GET or POST
*/
async handleRequest(req: IncomingMessage, res: ServerResponse): Promise<void> {
// validate the session ID
const sessionId = req.headers["mcp-session-id"];
if (sessionId && (Array.isArray(sessionId) ? sessionId[0] : sessionId) !== this._sessionId) {
res.writeHead(404).end(JSON.stringify({
jsonrpc: "2.0",
error: {
code: -32001,
message: "Session not found"
},
id: null
}));
return;
// Only validate session ID for non-initialization requests when session management is enabled
if (this._enableSessionManagement) {
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).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).end(JSON.stringify({
jsonrpc: "2.0",
error: {
code: -32001,
message: "Session not found"
},
id: null
}));
return;
}
}
if (req.method === "GET") {
@ -109,12 +169,19 @@ export class StreamableHTTPServerTransport implements Transport {
const lastEventId = req.headers["last-event-id"];
const lastEventIdStr = Array.isArray(lastEventId) ? lastEventId[0] : lastEventId;
res.writeHead(200, {
// Prepare response headers
const headers: Record<string, string> = {
"Content-Type": "text/event-stream",
"Cache-Control": "no-cache",
Connection: "keep-alive",
"mcp-session-id": this._sessionId,
});
};
// Only include session ID header if session management is enabled
if (this._enableSessionManagement) {
headers["mcp-session-id"] = this._sessionId;
}
res.writeHead(200, headers);
const connection: StreamConnection = {
response: res,
@ -192,6 +259,12 @@ 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
);
// check if it contains requests
const hasRequests = messages.some(msg => 'method' in msg && 'id' in msg);
const hasOnlyNotificationsOrResponses = messages.every(msg =>
@ -210,12 +283,19 @@ export class StreamableHTTPServerTransport implements Transport {
const useSSE = acceptHeader.includes("text/event-stream");
if (useSSE) {
res.writeHead(200, {
const headers: Record<string, string> = {
"Content-Type": "text/event-stream",
"Cache-Control": "no-cache",
Connection: "keep-alive",
"mcp-session-id": this._sessionId,
});
};
// Only include session ID header if session management is enabled
// Always include session ID for initialization requests
if (this._enableSessionManagement || isInitializationRequest) {
headers["mcp-session-id"] = this._sessionId;
}
res.writeHead(200, headers);
const connectionId = randomUUID();
const connection: StreamConnection = {
@ -247,10 +327,17 @@ export class StreamableHTTPServerTransport implements Transport {
});
} else {
// use direct JSON response
res.writeHead(200, {
const headers: Record<string, string> = {
"Content-Type": "application/json",
"mcp-session-id": this._sessionId,
});
};
// Only include session ID header if session management is enabled
// Always include session ID for initialization requests
if (this._enableSessionManagement || isInitializationRequest) {
headers["mcp-session-id"] = this._sessionId;
}
res.writeHead(200, headers);
// handle each message
for (const message of messages) {