diff --git a/examples/shared/src/inMemoryEventStore.ts b/examples/shared/src/inMemoryEventStore.ts index 604b84d39c..a417639ddb 100644 --- a/examples/shared/src/inMemoryEventStore.ts +++ b/examples/shared/src/inMemoryEventStore.ts @@ -16,11 +16,10 @@ export class InMemoryEventStore implements EventStore { } /** - * Extracts the stream ID from an event ID + * Looks up the stream ID an event was stored under */ private getStreamIdFromEventId(eventId: string): string { - const parts = eventId.split('_'); - return parts.length > 0 ? parts[0]! : ''; + return this.events.get(eventId)?.streamId ?? ''; } /** @@ -45,7 +44,7 @@ export class InMemoryEventStore implements EventStore { return ''; } - // Extract the stream ID from the event ID + // Look up the stream ID the event was stored under const streamId = this.getStreamIdFromEventId(lastEventId); if (!streamId) { return ''; diff --git a/examples/shared/test/inMemoryEventStore.test.ts b/examples/shared/test/inMemoryEventStore.test.ts new file mode 100644 index 0000000000..72e3e66a2c --- /dev/null +++ b/examples/shared/test/inMemoryEventStore.test.ts @@ -0,0 +1,74 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +import type { JSONRPCMessage } from '@modelcontextprotocol/server'; + +import { InMemoryEventStore } from '../src/inMemoryEventStore'; + +function notification(method: string): JSONRPCMessage { + return { jsonrpc: '2.0', method }; +} + +async function replayAfter(store: InMemoryEventStore, lastEventId: string): Promise<{ streamId: string; replayed: string[] }> { + const replayed: string[] = []; + const streamId = await store.replayEventsAfter(lastEventId, { + send: async eventId => { + replayed.push(eventId); + } + }); + return { streamId, replayed }; +} + +describe('InMemoryEventStore', () => { + beforeEach(() => { + // `generateEventId` stamps each event with `Date.now()`, so two events + // stored in the same millisecond can sort in either order and make the + // replay assertions flaky. Pin timestamps to keep ordering deterministic. + let now = 1000; + vi.spyOn(Date, 'now').mockImplementation(() => now++); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + it('replays events after an event on a standalone GET stream whose ID starts with an underscore', async () => { + const store = new InMemoryEventStore(); + const first = await store.storeEvent('_GET_stream', notification('notifications/message')); + const second = await store.storeEvent('_GET_stream', notification('notifications/resources/updated')); + + const { streamId, replayed } = await replayAfter(store, first); + + expect(streamId).toBe('_GET_stream'); + expect(replayed).toEqual([second]); + }); + + it('replays events after an event on a regular stream ID', async () => { + const store = new InMemoryEventStore(); + const first = await store.storeEvent('stream_1', notification('notifications/message')); + const second = await store.storeEvent('stream_1', notification('notifications/resources/updated')); + + const { streamId, replayed } = await replayAfter(store, first); + + expect(streamId).toBe('stream_1'); + expect(replayed).toEqual([second]); + }); + + it('returns an empty string for an unknown last event ID', async () => { + const store = new InMemoryEventStore(); + await store.storeEvent('_GET_stream', notification('notifications/message')); + + const { streamId, replayed } = await replayAfter(store, 'unknown-event-id'); + + expect(streamId).toBe(''); + expect(replayed).toEqual([]); + }); + + it('returns an empty string for an empty store', async () => { + const store = new InMemoryEventStore(); + + const { streamId, replayed } = await replayAfter(store, 'anything'); + + expect(streamId).toBe(''); + expect(replayed).toEqual([]); + }); +});