329 lines
No EOL
10 KiB
TypeScript
329 lines
No EOL
10 KiB
TypeScript
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 { 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();
|
|
|
|
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', () => {
|
|
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 = [];
|
|
let lastEventId: string | undefined;
|
|
|
|
// Create first client
|
|
const client1 = new Client({
|
|
id: clientId,
|
|
name: 'test-client',
|
|
version: '1.0.0'
|
|
});
|
|
|
|
// Set up notification handler for first client
|
|
client1.setNotificationHandler(LoggingMessageNotificationSchema, (notification) => {
|
|
if (notification.method === 'notifications/message') {
|
|
notifications.push(notification.params);
|
|
}
|
|
});
|
|
|
|
// Connect first client
|
|
const transport1 = new StreamableHTTPClientTransport(baseUrl);
|
|
await client1.connect(transport1);
|
|
const sessionId = transport1.sessionId;
|
|
expect(sessionId).toBeDefined();
|
|
|
|
// Start a long-running notification stream with tracking of lastEventId
|
|
const onLastEventIdUpdate = jest.fn((eventId: string) => {
|
|
lastEventId = eventId;
|
|
});
|
|
expect(lastEventId).toBeUndefined();
|
|
// 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) => {
|
|
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
|
|
await client2.request({
|
|
method: 'tools/call',
|
|
params: {
|
|
name: 'run-notifications',
|
|
arguments: {
|
|
count: 1,
|
|
interval: 50
|
|
}
|
|
}
|
|
}, CallToolResultSchema, {
|
|
lastEventId, // Pass the lastEventId from the previous session
|
|
onLastEventIdUpdate
|
|
});
|
|
|
|
// 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();
|
|
|
|
});
|
|
}); |