integration tests

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<string, { streamId: string, message: JSONRPCMessage }> = 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<string> {
+    const eventId = this.generateEventId(streamId);
+    this.events.set(eventId, { streamId, message });
+    return eventId;
+  }
+
+  async getEventsAfter(lastEventId: string): Promise<Array<{ eventId: string, message: JSONRPCMessage }>> {
+    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<URL>((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();
+
+  });
+});