From c49769b97a8420008d9f53923bd5652f602eda24 Mon Sep 17 00:00:00 2001 From: ihrpr Date: Sun, 13 Apr 2025 23:46:38 +0100 Subject: [PATCH] integration tests --- .../taskResumability.test.ts | 329 ++++++++++++++++++ 1 file changed, 329 insertions(+) create mode 100644 src/integration-tests/taskResumability.test.ts diff --git a/src/integration-tests/taskResumability.test.ts b/src/integration-tests/taskResumability.test.ts new file mode 100644 index 00000000..9397d4a0 --- /dev/null +++ b/src/integration-tests/taskResumability.test.ts @@ -0,0 +1,329 @@ +import { createServer, type Server } from 'node:http'; +import { AddressInfo } from 'node:net'; +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 { CallToolResult, CallToolResultSchema, JSONRPCMessage, LoggingMessageNotificationSchema } from '../types.js'; +import { z } from 'zod'; + +/** + * Simple in-memory event store implementation for resumability + */ +class InMemoryEventStore implements EventStore { + private events: Map = new Map(); + + generateEventId(streamId: string): string { + return `${streamId}_${Date.now()}_${Math.random().toString(36).substring(2, 10)}`; + } + + getStreamIdFromEventId(eventId: string): string { + const parts = eventId.split('_'); + return parts.length > 0 ? parts[0] : ''; + } + + async storeEvent(streamId: string, message: JSONRPCMessage): Promise { + const eventId = this.generateEventId(streamId); + this.events.set(eventId, { streamId, message }); + return eventId; + } + + async getEventsAfter(lastEventId: string): Promise> { + if (!lastEventId || !this.events.has(lastEventId)) { + return []; + } + + // Extract the stream ID from the event ID + const streamId = this.getStreamIdFromEventId(lastEventId); + const result: Array<{ eventId: string, message: JSONRPCMessage }> = []; + 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) { + result.push({ eventId, message }); + } + } + + return result; + } +} + + +describe('Transport resumability', () => { + let server: Server; + let mcpServer: McpServer; + let serverTransport: StreamableHTTPServerTransport; + let baseUrl: URL; + let eventStore: InMemoryEventStore; + + beforeEach(async () => { + // Create event store for resumability + eventStore = new InMemoryEventStore(); + + // Create a simple MCP server + mcpServer = new McpServer( + { name: 'test-server', version: '1.0.0' }, + { capabilities: { logging: {} } } + ); + + // Add a simple notification tool that completes quickly + mcpServer.tool( + 'send-notification', + 'Sends a single notification', + { + message: z.string().describe('Message to send').default('Test notification') + }, + async ({ message }, { sendNotification }) => { + // Send notification immediately + await sendNotification({ + method: "notifications/message", + params: { + level: "info", + data: message + } + }); + + return { + content: [{ type: 'text', text: 'Notification sent' }] + }; + } + ); + + // Add a long-running tool that sends multiple notifications + mcpServer.tool( + 'run-notifications', + 'Sends multiple notifications over time', + { + count: z.number().describe('Number of notifications to send').default(10), + interval: z.number().describe('Interval between notifications in ms').default(50) + }, + async ({ count, interval }, { sendNotification }) => { + // Send notifications at specified intervals + for (let i = 0; i < count; i++) { + await sendNotification({ + method: "notifications/message", + params: { + level: "info", + data: `Notification ${i + 1} of ${count}` + } + }); + + // Wait for the specified interval before sending next notification + if (i < count - 1) { + await new Promise(resolve => setTimeout(resolve, interval)); + } + } + + return { + content: [{ type: 'text', text: `Sent ${count} notifications` }] + }; + } + ); + + // Create a transport with the event store + serverTransport = new StreamableHTTPServerTransport({ + sessionIdGenerator: () => randomUUID(), + eventStore + }); + + // Connect the transport to the MCP server + await mcpServer.connect(serverTransport); + + // Create and start an HTTP server + server = createServer(async (req, res) => { + await serverTransport.handleRequest(req, res); + }); + + // Start the server on a random port + baseUrl = await new Promise((resolve) => { + server.listen(0, '127.0.0.1', () => { + const addr = server.address() as AddressInfo; + resolve(new URL(`http://127.0.0.1:${addr.port}`)); + }); + }); + }); + + afterEach(async () => { + // Clean up resources + await mcpServer.close().catch(() => { }); + await serverTransport.close().catch(() => { }); + server.close(); + }); + + it('should store session ID when client connects', async () => { + // Create and connect a client + const client = new Client({ + name: 'test-client', + version: '1.0.0' + }); + + const transport = new StreamableHTTPClientTransport(baseUrl); + await client.connect(transport); + + // Verify session ID was generated + expect(transport.sessionId).toBeDefined(); + + // Clean up + await transport.close(); + }); + + it('should have session ID functionality', async () => { + // The ability to store a session ID when connecting + const client = new Client({ + name: 'test-client-reconnection', + version: '1.0.0' + }); + + const transport = new StreamableHTTPClientTransport(baseUrl); + + // Make sure the client can connect and get a session ID + await client.connect(transport); + expect(transport.sessionId).toBeDefined(); + + // Clean up + await transport.close(); + }); + + // This test demonstrates the capability to resume long-running tools + // across client disconnection/reconnection + it('should resume long-running notifications with lastEventId', async () => { + // Create unique client ID for this test + const clientId = 'test-client-long-running'; + const notifications: any[] = []; + let sessionId: string | undefined; + let lastEventId: string | undefined; + + // Create first client + let client1 = new Client({ + id: clientId, + name: 'test-client', + version: '1.0.0' + }); + + // Set up notification handler for first client + client1.setNotificationHandler(LoggingMessageNotificationSchema, (notification: any) => { + if (notification.method === 'notifications/message') { + notifications.push(notification.params); + } + }); + + // Connect first client + const transport1 = new StreamableHTTPClientTransport(baseUrl); + await client1.connect(transport1); + sessionId = transport1.sessionId; + expect(sessionId).toBeDefined(); + + // Start a long-running notification stream with tracking of lastEventId + const onLastEventIdUpdate = jest.fn((eventId: string) => { + lastEventId = eventId; + }); + + // Start the notification tool with event tracking using request + const toolPromise = client1.request({ + method: 'tools/call', + params: { + name: 'run-notifications', + arguments: { + count: 5, + interval: 10 + } + } + }, CallToolResultSchema, { + lastEventId, + onLastEventIdUpdate + }); + + // Wait for some notifications to arrive (not all) + await new Promise(resolve => setTimeout(resolve, 20)); + + // Verify we received some notifications and lastEventId was updated + expect(notifications.length).toBeGreaterThan(0); + expect(notifications.length).toBeLessThan(5); + expect(onLastEventIdUpdate).toHaveBeenCalled(); + expect(lastEventId).toBeDefined(); + + // Store original notification count for later comparison + const firstClientNotificationCount = notifications.length; + + // Disconnect first client without waiting for completion + // When we close the connection, it will cause a ConnectionClosed error for + // any in-progress requests, which is expected behavior + // We need to catch the error since closing the transport will + // cause the pending toolPromise to reject with a ConnectionClosed error + await transport1.close(); + + // Try to cancel the promise, but ignore errors since it's already being handled + toolPromise.catch(err => { + // This error is expected - the connection was intentionally closed + if (err?.code !== -32000) { // ConnectionClosed error code + console.error("Unexpected error type during transport close:", err); + } + }); + + + // Create second client with same client ID + const client2 = new Client({ + id: clientId, + name: 'test-client', + version: '1.0.0' + }); + + // Set up notification handler for second client + client2.setNotificationHandler(LoggingMessageNotificationSchema, (notification: any) => { + if (notification.method === 'notifications/message') { + notifications.push(notification.params); + } + }); + + // Connect second client with same session ID + const transport2 = new StreamableHTTPClientTransport(baseUrl, { + sessionId + }); + await client2.connect(transport2); + + // Resume the notification stream using lastEventId + // This is the key part - we're resuming the same long-running tool using lastEventId + const resumedToolPromise = client2.request({ + method: 'tools/call', + params: { + name: 'run-notifications', + arguments: { + count: 5, + interval: 50 + } + } + }, CallToolResultSchema, { + lastEventId, // Pass the lastEventId from the previous session + onLastEventIdUpdate + }); + + // Wait for remaining notifications + await new Promise(resolve => setTimeout(resolve, 200)); + + // Verify we eventually received at leaset a few motifications + expect(notifications.length).toBeGreaterThan(2); + + // Verify the second client received notifications that the first client didn't + expect(notifications.length).toBeGreaterThan(firstClientNotificationCount); + + // Clean up + + await transport2.close(); + + }); +}); \ No newline at end of file