From f1e6acb17d35ccaaef2c156f24c786d0cc123de5 Mon Sep 17 00:00:00 2001 From: ihrpr Date: Tue, 15 Apr 2025 16:37:31 +0100 Subject: [PATCH] resumed tool to return value --- src/client/streamableHttp.ts | 31 +++++++++++++------ src/examples/client/simpleStreamableHttp.ts | 2 +- .../taskResumability.test.ts | 29 ++++++++--------- 3 files changed, 38 insertions(+), 24 deletions(-) diff --git a/src/client/streamableHttp.ts b/src/client/streamableHttp.ts index c1c93ebe..770b5d72 100644 --- a/src/client/streamableHttp.ts +++ b/src/client/streamableHttp.ts @@ -1,5 +1,5 @@ import { Transport } from "../shared/transport.js"; -import { isJSONRPCNotification, JSONRPCMessage, JSONRPCMessageSchema } from "../types.js"; +import { isJSONRPCNotification, isJSONRPCRequest, isJSONRPCResponse, JSONRPCMessage, JSONRPCMessageSchema } from "../types.js"; import { auth, AuthResult, OAuthClientProvider, UnauthorizedError } from "./auth.js"; import { EventSourceParserStream } from "eventsource-parser/stream"; @@ -28,6 +28,14 @@ export interface StartSSEOptions { * The ID of the last received event, used for resuming a disconnected stream */ lastEventId?: string; + /** + * The callback function that is invoked when the last event ID changes + */ + onLastEventIdUpdate?: (event: string) => void + /** + * When reconnecting to a long-running SSE stream, we need to make sure that message id matches + */ + replayMessageId?: string | number; } /** @@ -200,7 +208,7 @@ export class StreamableHTTPClientTransport implements Transport { ); } - this._handleSseStream(response.body); + this._handleSseStream(response.body, options); } catch (error) { this.onerror?.(error as Error); throw error; @@ -231,7 +239,7 @@ export class StreamableHTTPClientTransport implements Transport { * @param lastEventId The ID of the last received event for resumability * @param attemptCount Current reconnection attempt count for this specific stream */ - private _scheduleReconnection(lastEventId: string, attemptCount = 0): void { + private _scheduleReconnection(options: StartSSEOptions, attemptCount = 0): void { // Use provided options or default options const maxRetries = this._reconnectionOptions.maxRetries; @@ -247,18 +255,19 @@ export class StreamableHTTPClientTransport implements Transport { // Schedule the reconnection setTimeout(() => { // Use the last event ID to resume where we left off - this._startOrAuthSse({ lastEventId }).catch(error => { + this._startOrAuthSse(options).catch(error => { this.onerror?.(new Error(`Failed to reconnect SSE stream: ${error instanceof Error ? error.message : String(error)}`)); // Schedule another attempt if this one failed, incrementing the attempt counter - this._scheduleReconnection(lastEventId, attemptCount + 1); + this._scheduleReconnection(options, attemptCount + 1); }); }, delay); } - private _handleSseStream(stream: ReadableStream | null, onLastEventIdUpdate?: (event: string) => void): void { + private _handleSseStream(stream: ReadableStream | null, options: StartSSEOptions): void { if (!stream) { return; } + const { onLastEventIdUpdate, replayMessageId } = options; let lastEventId: string | undefined; const processStream = async () => { @@ -287,6 +296,9 @@ export class StreamableHTTPClientTransport implements Transport { if (!event.event || event.event === "message") { try { const message = JSONRPCMessageSchema.parse(JSON.parse(event.data)); + if (replayMessageId !== undefined && isJSONRPCResponse(message)) { + message.id = replayMessageId; + } this.onmessage?.(message); } catch (error) { this.onerror?.(error as Error); @@ -302,7 +314,7 @@ export class StreamableHTTPClientTransport implements Transport { // Use the exponential backoff reconnection strategy if (lastEventId !== undefined) { try { - this._scheduleReconnection(lastEventId, 0); + this._scheduleReconnection(options, 0); } catch (error) { this.onerror?.(new Error(`Failed to reconnect: ${error instanceof Error ? error.message : String(error)}`)); @@ -352,8 +364,9 @@ export class StreamableHTTPClientTransport implements Transport { const lastEventId = options?.resumptionToken const onLastEventIdUpdate = options?.onresumptiontoken; if (lastEventId) { + // If we have at last event ID, we need to reconnect the SSE stream - this._startOrAuthSse({ lastEventId }).catch(err => this.onerror?.(err)); + this._startOrAuthSse({ lastEventId, replayMessageId: isJSONRPCRequest(message) ? message.id : undefined }).catch(err => this.onerror?.(err)); return; } @@ -418,7 +431,7 @@ export class StreamableHTTPClientTransport implements Transport { // Handle SSE stream responses for requests // We use the same handler as standalone streams, which now supports // reconnection with the last event ID - this._handleSseStream(response.body, onLastEventIdUpdate); + this._handleSseStream(response.body, { onLastEventIdUpdate }); } else if (contentType?.includes("application/json")) { // For non-streaming servers, we might get direct JSON responses const data = await response.json(); diff --git a/src/examples/client/simpleStreamableHttp.ts b/src/examples/client/simpleStreamableHttp.ts index 6cfad987..d0d5408b 100644 --- a/src/examples/client/simpleStreamableHttp.ts +++ b/src/examples/client/simpleStreamableHttp.ts @@ -111,7 +111,7 @@ function commandLoop(): void { case 'start-notifications': { const interval = args[1] ? parseInt(args[1], 10) : 2000; - const count = args[2] ? parseInt(args[2], 10) : 0; + const count = args[2] ? parseInt(args[2], 10) : 10; await startNotifications(interval, count); break; } diff --git a/src/integration-tests/taskResumability.test.ts b/src/integration-tests/taskResumability.test.ts index 7a53ee01..7101925f 100644 --- a/src/integration-tests/taskResumability.test.ts +++ b/src/integration-tests/taskResumability.test.ts @@ -242,7 +242,7 @@ describe('Transport resumability', () => { params: { name: 'run-notifications', arguments: { - count: 5, + count: 3, interval: 10 } } @@ -251,27 +251,23 @@ describe('Transport resumability', () => { onresumptiontoken: onLastEventIdUpdate }); - // Wait for some notifications to arrive (not all) + // Wait for some notifications to arrive (not all) - shorter wait time 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(notifications.length).toBeLessThan(4); 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 => { + // Save the promise so we can catch it after closing + const catchPromise = 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); @@ -279,6 +275,14 @@ describe('Transport resumability', () => { }); + + // Add a short delay to ensure clean disconnect before reconnecting + await new Promise(resolve => setTimeout(resolve, 10)); + + // Wait for the rejection to be handled + await catchPromise; + + // Create second client with same client ID const client2 = new Client({ id: clientId, @@ -307,7 +311,7 @@ describe('Transport resumability', () => { name: 'run-notifications', arguments: { count: 1, - interval: 50 + interval: 5 } } }, CallToolResultSchema, { @@ -316,13 +320,10 @@ describe('Transport resumability', () => { }); // Verify we eventually received at leaset a few motifications - expect(notifications.length).toBeGreaterThan(2); + expect(notifications.length).toBeGreaterThan(1); - // Verify the second client received notifications that the first client didn't - expect(notifications.length).toBeGreaterThan(firstClientNotificationCount); // Clean up - await transport2.close(); });