inMemoryEventStore.ts

import type { EventStore, JSONRPCMessage } from '@modelcontextprotocol/server';

/**
 * Simple in-memory implementation of the EventStore interface for resumability
 * This is primarily intended for examples and testing, not for production use
 * where a persistent storage solution would be more appropriate.
 */
export class InMemoryEventStore implements EventStore {
    private events: Map<string, { streamId: string; message: JSONRPCMessage }> = new Map();

    /**
     * Generates a unique event ID for a given stream ID
     */
    private generateEventId(streamId: string): string {
        return `${streamId}_${Date.now()}_${Math.random().toString(36).slice(2, 10)}`;
    }

    /**
     * Extracts the stream ID from an event ID
     */
    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);
        this.events.set(eventId, { streamId, message });
        return eventId;
    }

    /**
     * Replays events that occurred after a specific event ID
     * Implements EventStore.replayEventsAfter
     */
    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()].toSorted((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 sending events after we find the lastEventId
            if (eventId === lastEventId) {
                foundLastEvent = true;
                continue;
            }

            if (foundLastEvent) {
                await send(eventId, message);
            }
        }
        return streamId;
    }
}