improve event store interface
This commit is contained in:
parent
643ef5c5e7
commit
67a4595ae7
4 changed files with 58 additions and 68 deletions
|
|
@ -164,7 +164,7 @@ describe("StreamableHTTPClientTransport", () => {
|
|||
// We expect the 405 error to be caught and handled gracefully
|
||||
// This should not throw an error that breaks the transport
|
||||
await transport.start();
|
||||
await expect(transport["_startOrAuthStandaloneSSE"]()).resolves.not.toThrow("Failed to open SSE stream: Method Not Allowed");
|
||||
await expect(transport["_startOrAuthStandaloneSSE"]({})).resolves.not.toThrow("Failed to open SSE stream: Method Not Allowed");
|
||||
// Check that GET was attempted
|
||||
expect(global.fetch).toHaveBeenCalledWith(
|
||||
expect.anything(),
|
||||
|
|
@ -208,7 +208,7 @@ describe("StreamableHTTPClientTransport", () => {
|
|||
transport.onmessage = messageSpy;
|
||||
|
||||
await transport.start();
|
||||
await transport["_startOrAuthStandaloneSSE"]();
|
||||
await transport["_startOrAuthStandaloneSSE"]({});
|
||||
|
||||
// Give time for the SSE event to be processed
|
||||
await new Promise(resolve => setTimeout(resolve, 50));
|
||||
|
|
@ -289,10 +289,10 @@ describe("StreamableHTTPClientTransport", () => {
|
|||
// Verify options were set correctly (checking implementation details)
|
||||
// Access private properties for testing
|
||||
const transportInstance = transport as unknown as {
|
||||
_defaultReconnectionOptions: StreamableHTTPReconnectionOptions;
|
||||
_reconnectionOptions: StreamableHTTPReconnectionOptions;
|
||||
};
|
||||
expect(transportInstance._defaultReconnectionOptions.initialReconnectionDelay).toBe(1000);
|
||||
expect(transportInstance._defaultReconnectionOptions.maxRetries).toBe(2);
|
||||
expect(transportInstance._reconnectionOptions.initialReconnectionDelay).toBe(500);
|
||||
expect(transportInstance._reconnectionOptions.maxRetries).toBe(5);
|
||||
});
|
||||
|
||||
it("should pass lastEventId when reconnecting", async () => {
|
||||
|
|
@ -313,9 +313,9 @@ describe("StreamableHTTPClientTransport", () => {
|
|||
await transport.start();
|
||||
// Type assertion to access private method
|
||||
const transportWithPrivateMethods = transport as unknown as {
|
||||
_startOrAuthStandaloneSSE: (lastEventId?: string) => Promise<void>
|
||||
_startOrAuthStandaloneSSE: (options: { lastEventId?: string }) => Promise<void>
|
||||
};
|
||||
await transportWithPrivateMethods._startOrAuthStandaloneSSE("test-event-id");
|
||||
await transportWithPrivateMethods._startOrAuthStandaloneSSE({ lastEventId: "test-event-id" });
|
||||
|
||||
// Verify fetch was called with the lastEventId header
|
||||
expect(fetchSpy).toHaveBeenCalled();
|
||||
|
|
@ -382,7 +382,7 @@ describe("StreamableHTTPClientTransport", () => {
|
|||
|
||||
await transport.start();
|
||||
|
||||
await transport["_startOrAuthStandaloneSSE"]();
|
||||
await transport["_startOrAuthStandaloneSSE"]({});
|
||||
expect((actualReqInit.headers as Headers).get("x-custom-header")).toBe("CustomValue");
|
||||
|
||||
requestInit.headers["X-Custom-Header"] = "SecondCustomValue";
|
||||
|
|
|
|||
|
|
@ -12,18 +12,18 @@ class InMemoryEventStore implements EventStore {
|
|||
/**
|
||||
* Generates a unique event ID for a given stream ID
|
||||
*/
|
||||
generateEventId(streamId: string): string {
|
||||
private generateEventId(streamId: string): string {
|
||||
return `${streamId}_${Date.now()}_${Math.random().toString(36).substring(2, 10)}`;
|
||||
}
|
||||
|
||||
getStreamIdFromEventId(eventId: string): string {
|
||||
private getStreamIdFromEventId(eventId: string): string {
|
||||
const parts = eventId.split('_');
|
||||
return parts.length > 0 ? parts[0] : '';
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* Stores an event with a generated event ID
|
||||
* Implements EventStore.storeEvent
|
||||
*/
|
||||
async storeEvent(streamId: string, message: JSONRPCMessage): Promise<string> {
|
||||
const eventId = this.generateEventId(streamId);
|
||||
|
|
@ -33,17 +33,26 @@ class InMemoryEventStore implements EventStore {
|
|||
}
|
||||
|
||||
/**
|
||||
* Retrieves events that occurred after a specific event
|
||||
* Replays events that occurred after a specific event ID
|
||||
* Implements EventStore.replayEventsAfter
|
||||
*/
|
||||
async getEventsAfter(lastEventId: string): Promise<Array<{ eventId: string, message: JSONRPCMessage }>> {
|
||||
async replayEventsAfter(lastEventId: string,
|
||||
{ send }: { send: (eventId: string, message: JSONRPCMessage) => Promise<void> }
|
||||
): Promise<string> {
|
||||
if (!lastEventId || !this.events.has(lastEventId)) {
|
||||
return [];
|
||||
console.log(`No events found for lastEventId: ${lastEventId}`);
|
||||
return '';
|
||||
}
|
||||
|
||||
// Extract the stream ID from the event ID
|
||||
const streamId = lastEventId.split('_')[0];
|
||||
const result: Array<{ eventId: string, message: JSONRPCMessage }> = [];
|
||||
const streamId = this.getStreamIdFromEventId(lastEventId);
|
||||
if (!streamId) {
|
||||
console.log(`Could not extract streamId from lastEventId: ${lastEventId}`);
|
||||
return '';
|
||||
}
|
||||
|
||||
let foundLastEvent = false;
|
||||
let eventCount = 0;
|
||||
|
||||
// Sort events by eventId for chronological ordering
|
||||
const sortedEvents = [...this.events.entries()].sort((a, b) => a[0].localeCompare(b[0]));
|
||||
|
|
@ -54,21 +63,21 @@ class InMemoryEventStore implements EventStore {
|
|||
continue;
|
||||
}
|
||||
|
||||
// Start collecting events after we find the lastEventId
|
||||
// Start sending events after we find the lastEventId
|
||||
if (eventId === lastEventId) {
|
||||
foundLastEvent = true;
|
||||
continue;
|
||||
}
|
||||
|
||||
if (foundLastEvent) {
|
||||
result.push({ eventId, message });
|
||||
await send(eventId, message);
|
||||
eventCount++;
|
||||
}
|
||||
}
|
||||
|
||||
console.log(`Found ${result.length} events after ${lastEventId} for replay`);
|
||||
return result;
|
||||
console.log(`Replayed ${eventCount} events after ${lastEventId} for stream ${streamId}`);
|
||||
return streamId;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
// Create an MCP server with implementation details
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import { createServer, type Server, IncomingMessage, ServerResponse } from "node:http";
|
||||
import { AddressInfo } from "node:net";
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { EventStore, StreamableHTTPServerTransport } from "./streamableHttp.js";
|
||||
import { EventStore, StreamableHTTPServerTransport, EventId, StreamId } from "./streamableHttp.js";
|
||||
import { McpServer } from "./mcp.js";
|
||||
import { CallToolResult, JSONRPCMessage } from "../types.js";
|
||||
import { z } from "zod";
|
||||
|
|
@ -912,31 +912,25 @@ describe("StreamableHTTPServerTransport with resumability", () => {
|
|||
|
||||
// Simple implementation of EventStore
|
||||
const eventStore: EventStore = {
|
||||
generateEventId(streamId: string): string {
|
||||
return `${streamId}_${randomUUID()}`;
|
||||
},
|
||||
getStreamIdFromEventId(eventId: string): string {
|
||||
return eventId.split('_')[0]; // Extract stream ID from the event ID
|
||||
},
|
||||
|
||||
async storeEvent(streamId: string, message: JSONRPCMessage): Promise<string> {
|
||||
const eventId = this.generateEventId(streamId);
|
||||
const eventId = `${streamId}_${randomUUID()}`;
|
||||
storedEvents.set(eventId, { eventId, message });
|
||||
return eventId;
|
||||
},
|
||||
|
||||
async getEventsAfter(lastEventId: string): Promise<Array<{ eventId: string, message: JSONRPCMessage }>> {
|
||||
const streamId = lastEventId.split('_')[0]; // Extract stream ID from the event ID
|
||||
const result: Array<{ eventId: string, message: JSONRPCMessage }> = [];
|
||||
|
||||
async replayEventsAfter(lastEventId: EventId, { send }: {
|
||||
send: (eventId: EventId, message: JSONRPCMessage) => Promise<void>
|
||||
}): Promise<StreamId> {
|
||||
const streamId = lastEventId.split('_')[0];
|
||||
// Extract stream ID from the event ID
|
||||
// For test simplicity, just return all events with matching streamId that aren't the lastEventId
|
||||
// This avoids issues with event ordering in tests
|
||||
for (const [eventId, { message }] of storedEvents.entries()) {
|
||||
if (eventId.startsWith(streamId) && eventId !== lastEventId) {
|
||||
result.push({ eventId, message });
|
||||
await send(eventId, message);
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
return streamId;
|
||||
},
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -1,45 +1,30 @@
|
|||
import { IncomingMessage, ServerResponse } from "node:http";
|
||||
import { Transport } from "../shared/transport.js";
|
||||
import { isJSONRPCNotification, isJSONRPCRequest, isJSONRPCResponse, JSONRPCMessage, JSONRPCMessageSchema, RequestId } from "../types.js";
|
||||
import { isJSONRPCRequest, isJSONRPCResponse, JSONRPCMessage, JSONRPCMessageSchema, RequestId } from "../types.js";
|
||||
import getRawBody from "raw-body";
|
||||
import contentType from "content-type";
|
||||
import { randomUUID } from "node:crypto";
|
||||
|
||||
const MAXIMUM_MESSAGE_SIZE = "4mb";
|
||||
|
||||
export type StreamId = string;
|
||||
export type EventId = string;
|
||||
|
||||
/**
|
||||
* Interface for resumability support via event storage
|
||||
*/
|
||||
export interface EventStore {
|
||||
/**
|
||||
* Generates a unique event ID for a given stream ID
|
||||
* @param streamId The stream ID to include in the event ID
|
||||
* @returns A unique event ID that includes the stream ID
|
||||
*/
|
||||
generateEventId(streamId: string): string;
|
||||
|
||||
/**
|
||||
* Stores an event for later retrieval
|
||||
* @param streamId ID of the stream the event belongs to
|
||||
* @param message The JSON-RPC message to store
|
||||
* @returns The generated event ID for the stored event
|
||||
*/
|
||||
storeEvent(streamId: string, message: JSONRPCMessage): Promise<string>;
|
||||
storeEvent(streamId: StreamId, message: JSONRPCMessage): Promise<EventId>;
|
||||
|
||||
/**
|
||||
* Retrieves events for a stream starting from a given event ID
|
||||
* @param lastEventId The event ID to start from
|
||||
* @returns Array of stored events with their event IDs
|
||||
*/
|
||||
getEventsAfter(lastEventId: string): Promise<Array<{ eventId: string, message: JSONRPCMessage }>>;
|
||||
|
||||
/**
|
||||
* Extracts the stream ID from an event ID
|
||||
* This is necessary for resumability to identify which stream the event belongs to
|
||||
* @param eventId The event ID to extract stream ID from
|
||||
* @returns The stream ID portion of the event ID
|
||||
*/
|
||||
getStreamIdFromEventId(eventId: string): string;
|
||||
replayEventsAfter(lastEventId: EventId, { send }: {
|
||||
send: (eventId: EventId, message: JSONRPCMessage) => Promise<void>
|
||||
}): Promise<StreamId>;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -232,12 +217,7 @@ export class StreamableHTTPServerTransport implements Transport {
|
|||
if (!this._eventStore) {
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
const events = await this._eventStore.getEventsAfter(lastEventId);
|
||||
const streamId = this._eventStore.getStreamIdFromEventId(lastEventId);
|
||||
|
||||
this._streamMapping.set(streamId, res);
|
||||
const headers: Record<string, string> = {
|
||||
"Content-Type": "text/event-stream",
|
||||
"Cache-Control": "no-cache, no-transform",
|
||||
|
|
@ -248,9 +228,16 @@ export class StreamableHTTPServerTransport implements Transport {
|
|||
headers["mcp-session-id"] = this.sessionId;
|
||||
}
|
||||
res.writeHead(200, headers).flushHeaders();
|
||||
for (const { eventId, message } of events) {
|
||||
this.writeSSEEvent(res, message, eventId);
|
||||
}
|
||||
|
||||
const streamId = await this._eventStore?.replayEventsAfter(lastEventId, {
|
||||
send: async (eventId: string, message: JSONRPCMessage) => {
|
||||
if (!this.writeSSEEvent(res, message, eventId)) {
|
||||
this.onerror?.(new Error("Failed replay events"));
|
||||
res.end();
|
||||
}
|
||||
}
|
||||
});
|
||||
this._streamMapping.set(streamId, res);
|
||||
} catch (error) {
|
||||
this.onerror?.(error as Error);
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue