share implementation of in memory store between examples and integration tests

This commit is contained in:
ihrpr 2025-04-16 14:48:47 +01:00
commit c3f08567a2
No known key found for this signature in database

View file

@ -4,68 +4,11 @@ import { randomUUID } from 'node:crypto';
import { Client } from '../client/index.js';
import { StreamableHTTPClientTransport } from '../client/streamableHttp.js';
import { McpServer } from '../server/mcp.js';
import { EventStore, StreamableHTTPServerTransport } from '../server/streamableHttp.js';
import { CallToolResultSchema, JSONRPCMessage, LoggingMessageNotificationSchema } from '../types.js';
import { StreamableHTTPServerTransport } from '../server/streamableHttp.js';
import { CallToolResultSchema, LoggingMessageNotificationSchema } from '../types.js';
import { z } from 'zod';
import { InMemoryEventStore } from '../examples/shared/inMemoryEventStore.js';
/**
* Simple in-memory event store implementation for resumability
*/
class InMemoryEventStore implements EventStore {
private events: Map<string, { streamId: string, message: JSONRPCMessage }> = new Map();
private generateEventId(streamId: string): string {
return `${streamId}_${Date.now()}_${Math.random().toString(36).substring(2, 10)}`;
}
private getStreamIdFromEventId(eventId: string): string {
const parts = eventId.split('_');
return parts.length > 0 ? parts[0] : '';
}
async storeEvent(streamId: string, message: JSONRPCMessage): Promise<string> {
const eventId = this.generateEventId(streamId);
this.events.set(eventId, { streamId, message });
return eventId;
}
async replayEventsAfter(lastEventId: string,
{ send }: { send: (eventId: string, message: JSONRPCMessage) => Promise<void> }
): Promise<string> {
if (!lastEventId || !this.events.has(lastEventId)) {
return '';
}
// Extract the stream ID from the event ID
const streamId = this.getStreamIdFromEventId(lastEventId);
if (!streamId) {
return '';
}
let foundLastEvent = false;
// Sort events by eventId for chronological ordering
const sortedEvents = [...this.events.entries()].sort((a, b) => a[0].localeCompare(b[0]));
for (const [eventId, { streamId: eventStreamId, message }] of sortedEvents) {
// Only include events from the same stream
if (eventStreamId !== streamId) {
continue;
}
// Start collecting events after we find the lastEventId
if (eventId === lastEventId) {
foundLastEvent = true;
continue;
}
if (foundLastEvent) {
await send(eventId, message);
}
}
return streamId;
}
}
describe('Transport resumability', () => {