✨ feat: try add stateless mod
This commit is contained in:
parent
970cf9d5b1
commit
4c7c434e7c
2 changed files with 368 additions and 31 deletions
|
|
@ -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
|
||||
},
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
Loading…
Reference in a new issue