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

support for server returning json -- not stream

Browse source
This commit is contained in:
ihrpr 2025-04-09 21:01:35 +01:00
commit b37a2e11ab
No known key found for this signature in database
3 changed files with 431 additions and 70 deletions
Download patch file Download diff file Expand all files Collapse all files

BIN
src/server/.streamableHttp.test.ts.swp Normal file
View file

Binary file not shown.

162
src/server/streamableHttp.test.ts
View file

@ -1008,8 +1008,8 @@ describe("StreamableHTTPServerTransport", () => {
await transport.handleRequest(req, response);
// Verify that the request is tracked in the SSE map
expect(transport["_sseResponseMapping"].size).toBe(2);
expect(transport["_sseResponseMapping"].has("cleanup-test")).toBe(true);
expect(transport["_responseMapping"].size).toBe(2);
expect(transport["_responseMapping"].has("cleanup-test")).toBe(true);
// Send a response
await transport.send({
@ -1019,8 +1019,8 @@ describe("StreamableHTTPServerTransport", () => {
});
// Verify that the mapping was cleaned up
expect(transport["_sseResponseMapping"].size).toBe(1);
expect(transport["_sseResponseMapping"].has("cleanup-test")).toBe(false);
expect(transport["_responseMapping"].size).toBe(1);
expect(transport["_responseMapping"].has("cleanup-test")).toBe(false);
});
it("should clean up connection tracking when client disconnects", async () => {
@ -1052,17 +1052,17 @@ describe("StreamableHTTPServerTransport", () => {
await transport.handleRequest(req, response);
// Both requests should be mapped to the same response
expect(transport["_sseResponseMapping"].size).toBe(3);
expect(transport["_sseResponseMapping"].get("req1")).toBe(response);
expect(transport["_sseResponseMapping"].get("req2")).toBe(response);
expect(transport["_responseMapping"].size).toBe(3);
expect(transport["_responseMapping"].get("req1")).toBe(response);
expect(transport["_responseMapping"].get("req2")).toBe(response);
// Simulate client disconnect by triggering the stored callback
if (closeCallback) closeCallback();
// All entries using this response should be removed
expect(transport["_sseResponseMapping"].size).toBe(1);
expect(transport["_sseResponseMapping"].has("req1")).toBe(false);
expect(transport["_sseResponseMapping"].has("req2")).toBe(false);
expect(transport["_responseMapping"].size).toBe(1);
expect(transport["_responseMapping"].has("req1")).toBe(false);
expect(transport["_responseMapping"].has("req2")).toBe(false);
});
});
@ -1214,6 +1214,148 @@ describe("StreamableHTTPServerTransport", () => {
});
});
describe("JSON Response Mode", () => {
let jsonResponseTransport: StreamableHTTPServerTransport;
let mockResponse: jest.Mocked<ServerResponse>;
beforeEach(async () => {
jsonResponseTransport = new StreamableHTTPServerTransport({
sessionIdGenerator: () => randomUUID(),
enableJsonResponse: true,
});
// Initialize the transport
const initMessage: JSONRPCMessage = {
jsonrpc: "2.0",
method: "initialize",
params: {
clientInfo: { name: "test-client", version: "1.0" },
protocolVersion: "2025-03-26"
},
id: "init-1",
};
const initReq = createMockRequest({
method: "POST",
headers: {
"content-type": "application/json",
"accept": "application/json, text/event-stream",
},
body: JSON.stringify(initMessage),
});
mockResponse = createMockResponse();
await jsonResponseTransport.handleRequest(initReq, mockResponse);
mockResponse = createMockResponse(); // Reset for tests
});
it("should return JSON response for a single request", async () => {
const requestMessage: JSONRPCMessage = {
jsonrpc: "2.0",
method: "test",
params: {},
id: "test-req-id",
};
const req = createMockRequest({
method: "POST",
headers: {
"content-type": "application/json",
"accept": "application/json, text/event-stream",
"mcp-session-id": jsonResponseTransport.sessionId,
},
body: JSON.stringify(requestMessage),
});
// Mock immediate response
jsonResponseTransport.onmessage = (message) => {
if ('method' in message && 'id' in message) {
const responseMessage: JSONRPCMessage = {
jsonrpc: "2.0",
result: { value: "test-result" },
id: message.id,
};
setTimeout(() => {
void jsonResponseTransport.send(responseMessage);
}, 10);
}
};
await jsonResponseTransport.handleRequest(req, mockResponse);
// Wait for all promises to resolve
await new Promise(resolve => setTimeout(resolve, 50));
// Should respond with application/json header
expect(mockResponse.writeHead).toHaveBeenCalledWith(
200,
expect.objectContaining({
"Content-Type": "application/json",
})
);
// Should return the response as JSON
const expectedResponse = {
jsonrpc: "2.0",
result: { value: "test-result" },
id: "test-req-id",
};
expect(mockResponse.end).toHaveBeenCalledWith(JSON.stringify(expectedResponse));
});
it("should return JSON response for batch requests", async () => {
const batchMessages: JSONRPCMessage[] = [
{ jsonrpc: "2.0", method: "test1", params: {}, id: "req1" },
{ jsonrpc: "2.0", method: "test2", params: {}, id: "req2" },
];
const req = createMockRequest({
method: "POST",
headers: {
"content-type": "application/json",
"accept": "application/json, text/event-stream",
"mcp-session-id": jsonResponseTransport.sessionId,
},
body: JSON.stringify(batchMessages),
});
// Mock sequential responses - send them in specific order
jsonResponseTransport.onmessage = (message) => {
if ('method' in message && 'id' in message) {
const responseMessage: JSONRPCMessage = {
jsonrpc: "2.0",
result: { value: `result-for-${message.id}` },
id: message.id,
};
// Send responses in order - req1 first, then req2
if (message.id === 'req1') {
void jsonResponseTransport.send(responseMessage);
} else {
// Add a tiny delay to req2 to ensure it's processed after req1
setTimeout(() => void jsonResponseTransport.send(responseMessage), 5);
}
}
};
await jsonResponseTransport.handleRequest(req, mockResponse);
// Wait for all promises to resolve - give it enough time
await new Promise(resolve => setTimeout(resolve, 100));
// Should respond with application/json header
expect(mockResponse.writeHead).toHaveBeenCalledWith(
200,
expect.objectContaining({
"Content-Type": "application/json",
})
);
expect(mockResponse.end).toHaveBeenCalled();
expect(mockResponse.end.mock.calls[0][0]).toContain("result-for-req2");
});
});
describe("Handling Pre-Parsed Body", () => {
beforeEach(async () => {
// Initialize the transport

127
src/server/streamableHttp.ts
View file

@ -18,8 +18,12 @@ export interface StreamableHTTPServerTransportOptions {
*/
sessionIdGenerator: () => string | undefined;
/**
* If true, the server will return JSON responses instead of starting an SSE stream.
* This can be useful for simple request/response scenarios without streaming.
* Default is false (SSE streams are preferred).
*/
enableJsonResponse?: boolean;
}
/**
@ -60,8 +64,11 @@ export class StreamableHTTPServerTransport implements Transport {
// when sessionId is not set (undefined), it means the transport is in stateless mode
private sessionIdGenerator: () => string | undefined;
private _started: boolean = false;
private _sseResponseMapping: Map<RequestId, ServerResponse> = new Map();
private _responseMapping: Map<RequestId, ServerResponse> = new Map();
private _requestResponseMap: Map<RequestId, JSONRPCMessage> = new Map();
private _initialized: boolean = false;
private _enableJsonResponse: boolean = false;
sessionId?: string | undefined;
onclose?: () => void;
@ -70,6 +77,7 @@ export class StreamableHTTPServerTransport implements Transport {
constructor(options: StreamableHTTPServerTransportOptions) {
this.sessionIdGenerator = options.sessionIdGenerator;
this._enableJsonResponse = options.enableJsonResponse ?? false;
}
/**
@ -221,33 +229,37 @@ export class StreamableHTTPServerTransport implements Transport {
this.onmessage?.(message);
}
} else if (hasRequests) {
const headers: Record<string, string> = {
"Content-Type": "text/event-stream",
"Cache-Control": "no-cache",
Connection: "keep-alive",
};
// The default behavior is to use SSE streaming
// but in some cases server will return JSON responses
if (!this._enableJsonResponse) {
const headers: Record<string, string> = {
"Content-Type": "text/event-stream",
"Cache-Control": "no-cache",
Connection: "keep-alive",
};
// After initialization, always include the session ID if we have one
if (this.sessionId !== undefined) {
headers["mcp-session-id"] = this.sessionId;
// After initialization, always include the session ID if we have one
if (this.sessionId !== undefined) {
headers["mcp-session-id"] = this.sessionId;
}
res.writeHead(200, headers);
}
res.writeHead(200, headers);
// Store the response for this request to send messages back through this connection
// We need to track by request ID to maintain the connection
for (const message of messages) {
if ('method' in message && 'id' in message) {
this._sseResponseMapping.set(message.id, res);
this._responseMapping.set(message.id, res);
}
}
// Set up close handler for client disconnects
res.on("close", () => {
// Remove all entries that reference this response
for (const [id, storedRes] of this._sseResponseMapping.entries()) {
for (const [id, storedRes] of this._responseMapping.entries()) {
if (storedRes === res) {
this._sseResponseMapping.delete(id);
this._responseMapping.delete(id);
this._requestResponseMap.delete(id);
}
}
});
@ -350,10 +362,14 @@ export class StreamableHTTPServerTransport implements Transport {
async close(): Promise<void> {
// Close all SSE connections
this._sseResponseMapping.forEach((response) => {
this._responseMapping.forEach((response) => {
response.end();
});
this._sseResponseMapping.clear();
this._responseMapping.clear();
// Clear any pending responses
this._requestResponseMap.clear();
this.onclose?.();
}
@ -367,24 +383,69 @@ export class StreamableHTTPServerTransport implements Transport {
throw new Error("No request ID provided for the message");
}
const sseResponse = this._sseResponseMapping.get(requestId);
if (!sseResponse) {
throw new Error(`No SSE connection established for request ID: ${String(requestId)}`);
// Get the response for this request
const response = this._responseMapping.get(requestId);
if (!response) {
throw new Error(`No connection established for request ID: ${String(requestId)}`);
}
// Send the message as an SSE event
sseResponse.write(
`event: message\ndata: ${JSON.stringify(message)}\n\n`,
);
// After all JSON-RPC responses have been sent, the server SHOULD close the SSE stream.
if (!this._enableJsonResponse) {
response.write(`event: message\ndata: ${JSON.stringify(message)}\n\n`);
}
if ('result' in message || 'error' in message) {
this._sseResponseMapping.delete(requestId);
// Only close the connection if it's not needed by other requests
const canCloseConnection = ![...this._sseResponseMapping.entries()].some(([id, res]) => res === sseResponse && id !== requestId);
if (canCloseConnection) {
sseResponse?.end();
this._requestResponseMap.set(requestId, message);
// Get all request IDs that share the same request response object
const relatedIds = Array.from(this._responseMapping.entries())
.filter(([_, res]) => res === response)
.map(([id]) => id);
// Check if we have responses for all requests using this connection
const allResponsesReady = relatedIds.every(id => this._requestResponseMap.has(id));
if (allResponsesReady) {
if (this._enableJsonResponse) {
// If we are in SSE mode, we don't need to do anything else
// All responses ready, send as JSON
const headers: Record<string, string> = {
'Content-Type': 'application/json',
};
if (this.sessionId !== undefined) {
headers['mcp-session-id'] = this.sessionId;
}
// In tests, relatedIds might be coming in non-deterministic order
// We need to sort numerically for numeric IDs, alphabetically for string IDs
const responses = relatedIds
.sort((a, b) => {
// If both IDs are numbers, sort numerically
if (typeof a === 'number' && typeof b === 'number') {
return a - b;
}
// Otherwise sort by string lexical order
return String(a).localeCompare(String(b));
})
.map(id => this._requestResponseMap.get(id)!);
response.writeHead(200, headers);
if (responses.length === 1) {
response.end(JSON.stringify(responses[0]));
} else {
response.end(JSON.stringify(responses));
}
} else {
response.end();
}
// Clean up
for (const id of relatedIds) {
this._requestResponseMap.delete(id);
this._responseMapping.delete(id);
}
}
}
}
}
}