✨ feat: StreamableHTTPServerTransport implement
This commit is contained in:
Geng Yan
ihrpr
parent
e2ab858308
commit
cc9ae5b4ac
2 changed files with 815 additions and 0 deletions
427
src/server/streamable-http.test.ts
Normal file
427
src/server/streamable-http.test.ts
Normal file
|
|
@ -0,0 +1,427 @@
|
|||
import { IncomingMessage, ServerResponse } from "node:http";
|
||||
import { StreamableHTTPServerTransport } from "./streamable-http.js";
|
||||
import { JSONRPCMessage } from "../types.js";
|
||||
import { Readable } from "node:stream";
|
||||
|
||||
// Mock IncomingMessage
|
||||
function createMockRequest(options: {
|
||||
method: string;
|
||||
headers: Record<string, string | string[] | undefined>;
|
||||
body?: string;
|
||||
}): IncomingMessage {
|
||||
const readable = new Readable();
|
||||
readable._read = () => {};
|
||||
if (options.body) {
|
||||
readable.push(options.body);
|
||||
readable.push(null);
|
||||
}
|
||||
|
||||
return Object.assign(readable, {
|
||||
method: options.method,
|
||||
headers: options.headers,
|
||||
}) as IncomingMessage;
|
||||
}
|
||||
|
||||
// Mock ServerResponse
|
||||
function createMockResponse(): jest.Mocked<ServerResponse> {
|
||||
const response = {
|
||||
writeHead: jest.fn().mockReturnThis(),
|
||||
write: jest.fn().mockReturnThis(),
|
||||
end: jest.fn().mockReturnThis(),
|
||||
on: jest.fn().mockReturnThis(),
|
||||
emit: jest.fn().mockReturnThis(),
|
||||
getHeader: jest.fn(),
|
||||
setHeader: jest.fn(),
|
||||
} as unknown as jest.Mocked<ServerResponse>;
|
||||
return response;
|
||||
}
|
||||
|
||||
describe("StreamableHTTPServerTransport", () => {
|
||||
const endpoint = "/mcp";
|
||||
let transport: StreamableHTTPServerTransport;
|
||||
let mockResponse: jest.Mocked<ServerResponse>;
|
||||
|
||||
beforeEach(() => {
|
||||
transport = new StreamableHTTPServerTransport(endpoint);
|
||||
mockResponse = createMockResponse();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
describe("Session Management", () => {
|
||||
it("should generate a valid session ID", () => {
|
||||
expect(transport.sessionId).toBeTruthy();
|
||||
expect(typeof transport.sessionId).toBe("string");
|
||||
});
|
||||
|
||||
it("should include session ID in response headers", async () => {
|
||||
const req = createMockRequest({
|
||||
method: "GET",
|
||||
headers: {
|
||||
accept: "text/event-stream"
|
||||
},
|
||||
});
|
||||
|
||||
await transport.handleRequest(req, mockResponse);
|
||||
|
||||
expect(mockResponse.writeHead).toHaveBeenCalledWith(
|
||||
200,
|
||||
expect.objectContaining({
|
||||
"Mcp-Session-Id": transport.sessionId,
|
||||
})
|
||||
);
|
||||
});
|
||||
|
||||
it("should reject invalid session ID", async () => {
|
||||
const req = createMockRequest({
|
||||
method: "GET",
|
||||
headers: {
|
||||
"mcp-session-id": "invalid-session-id",
|
||||
},
|
||||
});
|
||||
|
||||
await transport.handleRequest(req, mockResponse);
|
||||
|
||||
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"'));
|
||||
});
|
||||
});
|
||||
|
||||
describe("Request Handling", () => {
|
||||
it("should reject GET requests without Accept: text/event-stream header", async () => {
|
||||
const req = createMockRequest({
|
||||
method: "GET",
|
||||
headers: {},
|
||||
});
|
||||
|
||||
await transport.handleRequest(req, mockResponse);
|
||||
|
||||
expect(mockResponse.writeHead).toHaveBeenCalledWith(406);
|
||||
expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"jsonrpc":"2.0"'));
|
||||
});
|
||||
|
||||
it("should properly handle GET requests with Accept header and establish SSE connection", async () => {
|
||||
const req = createMockRequest({
|
||||
method: "GET",
|
||||
headers: {
|
||||
accept: "text/event-stream",
|
||||
},
|
||||
});
|
||||
|
||||
await transport.handleRequest(req, mockResponse);
|
||||
|
||||
expect(mockResponse.writeHead).toHaveBeenCalledWith(
|
||||
200,
|
||||
expect.objectContaining({
|
||||
"Content-Type": "text/event-stream",
|
||||
"Cache-Control": "no-cache",
|
||||
Connection: "keep-alive",
|
||||
})
|
||||
);
|
||||
});
|
||||
|
||||
it("should reject POST requests without proper Accept header", async () => {
|
||||
const message: JSONRPCMessage = {
|
||||
jsonrpc: "2.0",
|
||||
method: "test",
|
||||
params: {},
|
||||
id: 1,
|
||||
};
|
||||
|
||||
const req = createMockRequest({
|
||||
method: "POST",
|
||||
headers: {
|
||||
"content-type": "application/json",
|
||||
},
|
||||
body: JSON.stringify(message),
|
||||
});
|
||||
|
||||
await transport.handleRequest(req, mockResponse);
|
||||
|
||||
expect(mockResponse.writeHead).toHaveBeenCalledWith(406);
|
||||
});
|
||||
|
||||
it("should properly handle JSON-RPC request messages in POST requests", async () => {
|
||||
const message: JSONRPCMessage = {
|
||||
jsonrpc: "2.0",
|
||||
method: "test",
|
||||
params: {},
|
||||
id: 1,
|
||||
};
|
||||
|
||||
const req = createMockRequest({
|
||||
method: "POST",
|
||||
headers: {
|
||||
"content-type": "application/json",
|
||||
"accept": "text/event-stream",
|
||||
},
|
||||
body: JSON.stringify(message),
|
||||
});
|
||||
|
||||
const onMessageMock = jest.fn();
|
||||
transport.onmessage = onMessageMock;
|
||||
|
||||
await transport.handleRequest(req, mockResponse);
|
||||
|
||||
expect(onMessageMock).toHaveBeenCalledWith(message);
|
||||
expect(mockResponse.writeHead).toHaveBeenCalledWith(
|
||||
200,
|
||||
expect.objectContaining({
|
||||
"Content-Type": "text/event-stream",
|
||||
})
|
||||
);
|
||||
});
|
||||
|
||||
it("should properly handle JSON-RPC notification or response messages in POST requests", async () => {
|
||||
const notification: JSONRPCMessage = {
|
||||
jsonrpc: "2.0",
|
||||
method: "test",
|
||||
params: {},
|
||||
};
|
||||
|
||||
const req = createMockRequest({
|
||||
method: "POST",
|
||||
headers: {
|
||||
"content-type": "application/json",
|
||||
"accept": "application/json, text/event-stream",
|
||||
},
|
||||
body: JSON.stringify(notification),
|
||||
});
|
||||
|
||||
const onMessageMock = jest.fn();
|
||||
transport.onmessage = onMessageMock;
|
||||
|
||||
await transport.handleRequest(req, mockResponse);
|
||||
|
||||
expect(onMessageMock).toHaveBeenCalledWith(notification);
|
||||
expect(mockResponse.writeHead).toHaveBeenCalledWith(202);
|
||||
});
|
||||
|
||||
it("should handle batch messages properly", async () => {
|
||||
const batchMessages: JSONRPCMessage[] = [
|
||||
{ jsonrpc: "2.0", method: "test1", params: {} },
|
||||
{ jsonrpc: "2.0", method: "test2", params: {} },
|
||||
];
|
||||
|
||||
const req = createMockRequest({
|
||||
method: "POST",
|
||||
headers: {
|
||||
"content-type": "application/json",
|
||||
"accept": "application/json",
|
||||
},
|
||||
body: JSON.stringify(batchMessages),
|
||||
});
|
||||
|
||||
const onMessageMock = jest.fn();
|
||||
transport.onmessage = onMessageMock;
|
||||
|
||||
await transport.handleRequest(req, mockResponse);
|
||||
|
||||
expect(onMessageMock).toHaveBeenCalledTimes(2);
|
||||
expect(mockResponse.writeHead).toHaveBeenCalledWith(202);
|
||||
});
|
||||
|
||||
it("should reject unsupported Content-Type", async () => {
|
||||
const req = createMockRequest({
|
||||
method: "POST",
|
||||
headers: {
|
||||
"content-type": "text/plain",
|
||||
"accept": "application/json",
|
||||
},
|
||||
body: "test",
|
||||
});
|
||||
|
||||
await transport.handleRequest(req, mockResponse);
|
||||
|
||||
expect(mockResponse.writeHead).toHaveBeenCalledWith(415);
|
||||
expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"jsonrpc":"2.0"'));
|
||||
});
|
||||
|
||||
it("should properly handle DELETE requests and close session", async () => {
|
||||
const req = createMockRequest({
|
||||
method: "DELETE",
|
||||
headers: {},
|
||||
});
|
||||
|
||||
const onCloseMock = jest.fn();
|
||||
transport.onclose = onCloseMock;
|
||||
|
||||
await transport.handleRequest(req, mockResponse);
|
||||
|
||||
expect(mockResponse.writeHead).toHaveBeenCalledWith(200);
|
||||
expect(onCloseMock).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe("Message Replay", () => {
|
||||
it("should replay messages after specified Last-Event-ID", async () => {
|
||||
// Establish first connection with Accept header
|
||||
const req1 = createMockRequest({
|
||||
method: "GET",
|
||||
headers: {
|
||||
"accept": "text/event-stream"
|
||||
},
|
||||
});
|
||||
await transport.handleRequest(req1, mockResponse);
|
||||
|
||||
// Send a message to first connection
|
||||
const message1: JSONRPCMessage = {
|
||||
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-]+)/);
|
||||
if (!idMatch) {
|
||||
throw new Error("Message ID not found in write call");
|
||||
}
|
||||
const lastEventId = idMatch[1];
|
||||
|
||||
// Create a second connection with last-event-id
|
||||
const mockResponse2 = createMockResponse();
|
||||
const req2 = createMockRequest({
|
||||
method: "GET",
|
||||
headers: {
|
||||
"accept": "text/event-stream",
|
||||
"last-event-id": lastEventId,
|
||||
},
|
||||
});
|
||||
|
||||
await transport.handleRequest(req2, mockResponse2);
|
||||
|
||||
// Send a second message
|
||||
const message2: JSONRPCMessage = {
|
||||
jsonrpc: "2.0",
|
||||
method: "test2",
|
||||
params: {},
|
||||
id: 2
|
||||
};
|
||||
|
||||
await transport.send(message2);
|
||||
|
||||
// Verify the second message was received by both connections
|
||||
expect(mockResponse.write).toHaveBeenCalledWith(
|
||||
expect.stringContaining(JSON.stringify(message1))
|
||||
);
|
||||
expect(mockResponse2.write).toHaveBeenCalledWith(
|
||||
expect.stringContaining(JSON.stringify(message2))
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("Message Targeting", () => {
|
||||
it("should send response messages to the connection that sent the request", async () => {
|
||||
// Create two connections
|
||||
const mockResponse1 = createMockResponse();
|
||||
const req1 = createMockRequest({
|
||||
method: "GET",
|
||||
headers: {
|
||||
"accept": "text/event-stream",
|
||||
},
|
||||
});
|
||||
await transport.handleRequest(req1, mockResponse1);
|
||||
|
||||
const mockResponse2 = createMockResponse();
|
||||
const req2 = createMockRequest({
|
||||
method: "GET",
|
||||
headers: {
|
||||
"accept": "text/event-stream",
|
||||
},
|
||||
});
|
||||
await transport.handleRequest(req2, mockResponse2);
|
||||
|
||||
// Send a request through the first connection
|
||||
const requestMessage: JSONRPCMessage = {
|
||||
jsonrpc: "2.0",
|
||||
method: "test",
|
||||
params: {},
|
||||
id: "test-id",
|
||||
};
|
||||
|
||||
const reqPost = createMockRequest({
|
||||
method: "POST",
|
||||
headers: {
|
||||
"content-type": "application/json",
|
||||
"accept": "text/event-stream",
|
||||
},
|
||||
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 =>
|
||||
typeof call[0] === 'string' && call[0].includes(JSON.stringify(responseMessage))
|
||||
);
|
||||
|
||||
// Verify the response wasn't broadcast to all connections
|
||||
expect(writeCallsOnSecondConn.length).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe("Error Handling", () => {
|
||||
it("should handle invalid JSON data", async () => {
|
||||
const req = createMockRequest({
|
||||
method: "POST",
|
||||
headers: {
|
||||
"content-type": "application/json",
|
||||
"accept": "application/json",
|
||||
},
|
||||
body: "invalid json",
|
||||
});
|
||||
|
||||
const onErrorMock = jest.fn();
|
||||
transport.onerror = onErrorMock;
|
||||
|
||||
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('"code":-32700'));
|
||||
expect(onErrorMock).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("should handle invalid JSON-RPC messages", async () => {
|
||||
const req = createMockRequest({
|
||||
method: "POST",
|
||||
headers: {
|
||||
"content-type": "application/json",
|
||||
"accept": "application/json",
|
||||
},
|
||||
body: JSON.stringify({ invalid: "message" }),
|
||||
});
|
||||
|
||||
const onErrorMock = jest.fn();
|
||||
transport.onerror = onErrorMock;
|
||||
|
||||
await transport.handleRequest(req, mockResponse);
|
||||
|
||||
expect(mockResponse.writeHead).toHaveBeenCalledWith(400);
|
||||
expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"jsonrpc":"2.0"'));
|
||||
expect(onErrorMock).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
});
|
||||
388
src/server/streamable-http.ts
Normal file
388
src/server/streamable-http.ts
Normal file
|
|
@ -0,0 +1,388 @@
|
|||
import { randomUUID } from "node:crypto";
|
||||
import { IncomingMessage, ServerResponse } from "node:http";
|
||||
import { Transport } from "../shared/transport.js";
|
||||
import { JSONRPCMessage, JSONRPCMessageSchema } from "../types.js";
|
||||
import getRawBody from "raw-body";
|
||||
import contentType from "content-type";
|
||||
|
||||
const MAXIMUM_MESSAGE_SIZE = "4mb";
|
||||
|
||||
interface StreamConnection {
|
||||
response: ServerResponse;
|
||||
lastEventId?: string;
|
||||
messages: Array<{
|
||||
id: string;
|
||||
message: JSONRPCMessage;
|
||||
}>;
|
||||
// mark this connection as a response to a specific request
|
||||
requestId?: string | null;
|
||||
}
|
||||
|
||||
/**
|
||||
* 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.
|
||||
*/
|
||||
export class StreamableHTTPServerTransport implements Transport {
|
||||
private _connections: Map<string, StreamConnection> = new Map();
|
||||
private _sessionId: string;
|
||||
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
|
||||
|
||||
onclose?: () => void;
|
||||
onerror?: (error: Error) => void;
|
||||
onmessage?: (message: JSONRPCMessage) => void;
|
||||
|
||||
constructor(private _endpoint: string) {
|
||||
this._sessionId = randomUUID();
|
||||
}
|
||||
|
||||
/**
|
||||
* Starts the transport. This is required by the Transport interface but is a no-op
|
||||
* for the Streamable HTTP transport as connections are managed per-request.
|
||||
*/
|
||||
async start(): Promise<void> {
|
||||
if (this._started) {
|
||||
throw new Error("Transport already started");
|
||||
}
|
||||
this._started = true;
|
||||
}
|
||||
|
||||
/**
|
||||
* 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;
|
||||
}
|
||||
|
||||
if (req.method === "GET") {
|
||||
await this.handleGetRequest(req, res);
|
||||
} else if (req.method === "POST") {
|
||||
await this.handlePostRequest(req, res);
|
||||
} else if (req.method === "DELETE") {
|
||||
await this.handleDeleteRequest(req, res);
|
||||
} else {
|
||||
res.writeHead(405).end(JSON.stringify({
|
||||
jsonrpc: "2.0",
|
||||
error: {
|
||||
code: -32000,
|
||||
message: "Method not allowed"
|
||||
},
|
||||
id: null
|
||||
}));
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Handles GET requests to establish SSE connections
|
||||
*/
|
||||
private async handleGetRequest(req: IncomingMessage, res: ServerResponse): Promise<void> {
|
||||
// validate the Accept header
|
||||
const acceptHeader = req.headers.accept;
|
||||
if (!acceptHeader || !acceptHeader.includes("text/event-stream")) {
|
||||
res.writeHead(406).end(JSON.stringify({
|
||||
jsonrpc: "2.0",
|
||||
error: {
|
||||
code: -32000,
|
||||
message: "Not Acceptable: Client must accept text/event-stream"
|
||||
},
|
||||
id: null
|
||||
}));
|
||||
return;
|
||||
}
|
||||
|
||||
const connectionId = randomUUID();
|
||||
const lastEventId = req.headers["last-event-id"];
|
||||
const lastEventIdStr = Array.isArray(lastEventId) ? lastEventId[0] : lastEventId;
|
||||
|
||||
res.writeHead(200, {
|
||||
"Content-Type": "text/event-stream",
|
||||
"Cache-Control": "no-cache",
|
||||
Connection: "keep-alive",
|
||||
"Mcp-Session-Id": this._sessionId,
|
||||
});
|
||||
|
||||
const connection: StreamConnection = {
|
||||
response: res,
|
||||
lastEventId: lastEventIdStr,
|
||||
messages: [],
|
||||
};
|
||||
|
||||
this._connections.set(connectionId, connection);
|
||||
|
||||
// if there is a Last-Event-ID, replay messages on this connection
|
||||
if (lastEventIdStr) {
|
||||
this.replayMessages(connectionId, lastEventIdStr);
|
||||
}
|
||||
|
||||
res.on("close", () => {
|
||||
this._connections.delete(connectionId);
|
||||
// remove all request mappings associated with this connection
|
||||
for (const [reqId, connId] of this._requestConnections.entries()) {
|
||||
if (connId === connectionId) {
|
||||
this._requestConnections.delete(reqId);
|
||||
}
|
||||
}
|
||||
if (this._connections.size === 0) {
|
||||
this.onclose?.();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Handles POST requests containing JSON-RPC messages
|
||||
*/
|
||||
private async handlePostRequest(req: IncomingMessage, res: ServerResponse): Promise<void> {
|
||||
try {
|
||||
// validate the Accept header
|
||||
const acceptHeader = req.headers.accept;
|
||||
if (!acceptHeader ||
|
||||
(!acceptHeader.includes("application/json") && !acceptHeader.includes("text/event-stream"))) {
|
||||
res.writeHead(406).end(JSON.stringify({
|
||||
jsonrpc: "2.0",
|
||||
error: {
|
||||
code: -32000,
|
||||
message: "Not Acceptable: Client must accept application/json and/or text/event-stream"
|
||||
},
|
||||
id: null
|
||||
}));
|
||||
return;
|
||||
}
|
||||
|
||||
const ct = req.headers["content-type"];
|
||||
if (!ct || !ct.includes("application/json")) {
|
||||
res.writeHead(415).end(JSON.stringify({
|
||||
jsonrpc: "2.0",
|
||||
error: {
|
||||
code: -32000,
|
||||
message: "Unsupported Media Type: Content-Type must be application/json"
|
||||
},
|
||||
id: null
|
||||
}));
|
||||
return;
|
||||
}
|
||||
|
||||
const parsedCt = contentType.parse(ct);
|
||||
const body = await getRawBody(req, {
|
||||
limit: MAXIMUM_MESSAGE_SIZE,
|
||||
encoding: parsedCt.parameters.charset ?? "utf-8",
|
||||
});
|
||||
|
||||
const rawMessage = JSON.parse(body.toString());
|
||||
let messages: JSONRPCMessage[];
|
||||
|
||||
// handle batch and single messages
|
||||
if (Array.isArray(rawMessage)) {
|
||||
messages = rawMessage.map(msg => JSONRPCMessageSchema.parse(msg));
|
||||
} else {
|
||||
messages = [JSONRPCMessageSchema.parse(rawMessage)];
|
||||
}
|
||||
|
||||
// check if it contains requests
|
||||
const hasRequests = messages.some(msg => 'method' in msg && 'id' in msg);
|
||||
const hasOnlyNotificationsOrResponses = messages.every(msg =>
|
||||
('method' in msg && !('id' in msg)) || ('result' in msg || 'error' in msg));
|
||||
|
||||
if (hasOnlyNotificationsOrResponses) {
|
||||
// if it only contains notifications or responses, return 202
|
||||
res.writeHead(202).end();
|
||||
|
||||
// handle each message
|
||||
for (const message of messages) {
|
||||
this.onmessage?.(message);
|
||||
}
|
||||
} 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) {
|
||||
res.writeHead(200, {
|
||||
"Content-Type": "text/event-stream",
|
||||
"Cache-Control": "no-cache",
|
||||
Connection: "keep-alive",
|
||||
"Mcp-Session-Id": this._sessionId,
|
||||
});
|
||||
|
||||
const connectionId = randomUUID();
|
||||
const connection: StreamConnection = {
|
||||
response: res,
|
||||
messages: [],
|
||||
};
|
||||
|
||||
this._connections.set(connectionId, connection);
|
||||
|
||||
// map each request to a connection ID
|
||||
for (const message of messages) {
|
||||
if ('method' in message && 'id' in message) {
|
||||
this._requestConnections.set(String(message.id), connectionId);
|
||||
}
|
||||
this.onmessage?.(message);
|
||||
}
|
||||
|
||||
res.on("close", () => {
|
||||
this._connections.delete(connectionId);
|
||||
// remove all request mappings associated with this connection
|
||||
for (const [reqId, connId] of this._requestConnections.entries()) {
|
||||
if (connId === connectionId) {
|
||||
this._requestConnections.delete(reqId);
|
||||
}
|
||||
}
|
||||
if (this._connections.size === 0) {
|
||||
this.onclose?.();
|
||||
}
|
||||
});
|
||||
} else {
|
||||
// use direct JSON response
|
||||
res.writeHead(200, {
|
||||
"Content-Type": "application/json",
|
||||
"Mcp-Session-Id": this._sessionId,
|
||||
});
|
||||
|
||||
// handle each message
|
||||
for (const message of messages) {
|
||||
this.onmessage?.(message);
|
||||
}
|
||||
|
||||
res.end();
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
// return JSON-RPC formatted error
|
||||
res.writeHead(400).end(JSON.stringify({
|
||||
jsonrpc: "2.0",
|
||||
error: {
|
||||
code: -32700,
|
||||
message: "Parse error",
|
||||
data: String(error)
|
||||
},
|
||||
id: null
|
||||
}));
|
||||
this.onerror?.(error as Error);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Handles DELETE requests to terminate sessions
|
||||
*/
|
||||
private async handleDeleteRequest(req: IncomingMessage, res: ServerResponse): Promise<void> {
|
||||
await this.close();
|
||||
res.writeHead(200).end();
|
||||
}
|
||||
|
||||
/**
|
||||
* Replays messages after the specified event ID for a specific connection
|
||||
*/
|
||||
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 &&
|
||||
(!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));
|
||||
|
||||
const connection = this._connections.get(connectionId);
|
||||
if (!connection) return;
|
||||
|
||||
for (const [id, { message }] of messages) {
|
||||
connection.response.write(
|
||||
`id: ${id}\nevent: message\ndata: ${JSON.stringify(message)}\n\n`
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async close(): Promise<void> {
|
||||
for (const connection of this._connections.values()) {
|
||||
connection.response.end();
|
||||
}
|
||||
this._connections.clear();
|
||||
this._messageHistory.clear();
|
||||
this._requestConnections.clear();
|
||||
this.onclose?.();
|
||||
}
|
||||
|
||||
async send(message: JSONRPCMessage): Promise<void> {
|
||||
if (this._connections.size === 0) {
|
||||
throw new Error("No active connections");
|
||||
}
|
||||
|
||||
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
|
||||
const firstConnId = this._connections.keys().next().value;
|
||||
if (firstConnId) {
|
||||
targetConnectionId = firstConnId;
|
||||
} else {
|
||||
throw new Error("No available connections");
|
||||
}
|
||||
} else {
|
||||
targetConnectionId = connId;
|
||||
}
|
||||
} else {
|
||||
// for other messages, select an available connection
|
||||
const firstConnId = this._connections.keys().next().value;
|
||||
if (firstConnId) {
|
||||
targetConnectionId = firstConnId;
|
||||
} else {
|
||||
throw new Error("No available connections");
|
||||
}
|
||||
}
|
||||
|
||||
const messageId = randomUUID();
|
||||
this._messageHistory.set(messageId, {
|
||||
message,
|
||||
connectionId: targetConnectionId
|
||||
});
|
||||
|
||||
// keep the message history in a reasonable range
|
||||
if (this._messageHistory.size > 1000) {
|
||||
const oldestKey = Array.from(this._messageHistory.keys())[0];
|
||||
this._messageHistory.delete(oldestKey);
|
||||
}
|
||||
|
||||
// send the message to all active connections
|
||||
for (const [connId, connection] of this._connections.entries()) {
|
||||
// if it is a response message, only send to the target connection
|
||||
if ('id' in message && ('result' in message || 'error' in message)) {
|
||||
if (connId === targetConnectionId) {
|
||||
connection.response.write(
|
||||
`id: ${messageId}\nevent: message\ndata: ${JSON.stringify(message)}\n\n`
|
||||
);
|
||||
}
|
||||
} else {
|
||||
// for other messages, send to all connections
|
||||
connection.response.write(
|
||||
`id: ${messageId}\nevent: message\ndata: ${JSON.stringify(message)}\n\n`
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns the session ID for this transport
|
||||
*/
|
||||
get sessionId(): string {
|
||||
return this._sessionId;
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue