resume long running requst from new client

This commit is contained in:
ihrpr 2025-04-13 21:40:15 +01:00
commit 43cb015844
No known key found for this signature in database
6 changed files with 76 additions and 33 deletions

View file

@ -71,6 +71,11 @@ export type StreamableHTTPClientTransportOptions = {
* Options to configure the reconnection behavior.
*/
reconnectionOptions?: StreamableHTTPReconnectionOptions;
/**
* Session ID for the connection. This is used to identify the session on the server.
* When not provided and connecting to a server that supports session IDs, the server will generate a new session ID.
*/
sessionId?: string;
};
/**
@ -98,6 +103,7 @@ export class StreamableHTTPClientTransport implements Transport {
this._requestInit = opts?.requestInit;
this._authProvider = opts?.authProvider;
this._reconnectionOptions = opts?.reconnectionOptions || this._defaultReconnectionOptions;
this._sessionId = opts?.sessionId;
}
private async _authThenStart(): Promise<void> {
@ -240,7 +246,7 @@ export class StreamableHTTPClientTransport implements Transport {
}, delay);
}
private _handleSseStream(stream: ReadableStream<Uint8Array> | null): void {
private _handleSseStream(stream: ReadableStream<Uint8Array> | null, onLastEventIdUpdate?: (event: string) => void): void {
if (!stream) {
return;
}
@ -266,6 +272,7 @@ export class StreamableHTTPClientTransport implements Transport {
// Update last event ID if provided
if (event.id) {
lastEventId = event.id;
onLastEventIdUpdate?.(lastEventId);
}
if (!event.event || event.event === "message") {
@ -330,8 +337,16 @@ export class StreamableHTTPClientTransport implements Transport {
this.onclose?.();
}
async send(message: JSONRPCMessage | JSONRPCMessage[]): Promise<void> {
async send(message: JSONRPCMessage | JSONRPCMessage[], options?: { lastEventId?: string, onLastEventIdUpdate?: (event: string) => void }): Promise<void> {
try {
// If client passes in a lastEventId in the request options, we need to reconnect the SSE stream
const { lastEventId, onLastEventIdUpdate } = options ?? {};
if (lastEventId) {
// If we have at last event ID, we need to reconnect the SSE stream
this._startOrAuthStandaloneSSE(lastEventId).catch(err => this.onerror?.(err));
return;
}
const headers = await this._commonHeaders();
headers.set("content-type", "application/json");
headers.set("accept", "application/json, text/event-stream");
@ -393,7 +408,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);
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();
@ -416,4 +431,8 @@ export class StreamableHTTPClientTransport implements Transport {
throw error;
}
}
get sessionId(): string | undefined {
return this._sessionId;
}
}

View file

@ -29,6 +29,8 @@ let notificationCount = 0;
let client: Client | null = null;
let transport: StreamableHTTPClientTransport | null = null;
let serverUrl = 'http://localhost:3000/mcp';
let notificationsToolLastEventId: string | undefined = undefined;
let sessionId: string | undefined = undefined;
async function main(): Promise<void> {
console.log('MCP Interactive Client');
@ -186,7 +188,10 @@ async function connect(url?: string): Promise<void> {
}
transport = new StreamableHTTPClientTransport(
new URL(serverUrl)
new URL(serverUrl),
{
sessionId: sessionId
}
);
// Set up notification handlers
@ -218,6 +223,8 @@ async function connect(url?: string): Promise<void> {
// Connect the client
await client.connect(transport);
sessionId = transport.sessionId
console.log('Transport created with session ID:', sessionId);
console.log('Connected to MCP server');
} catch (error) {
console.error('Failed to connect:', error);
@ -291,7 +298,12 @@ async function callTool(name: string, args: Record<string, unknown>): Promise<vo
};
console.log(`Calling tool '${name}' with args:`, args);
const result = await client.request(request, CallToolResultSchema);
const onLastEventIdUpdate = (event: string) => {
notificationsToolLastEventId = event;
};
const result = await client.request(request, CallToolResultSchema, {
lastEventId: notificationsToolLastEventId, onLastEventIdUpdate
});
console.log('Tool result:');
result.content.forEach(item => {

View file

@ -172,14 +172,18 @@ server.tool(
while (count === 0 || counter < count) {
counter++;
await sendNotification({
method: "notifications/message",
params: {
level: "info",
data: `Periodic notification #${counter} at ${new Date().toISOString()}`
}
});
try {
await sendNotification({
method: "notifications/message",
params: {
level: "info",
data: `Periodic notification #${counter} at ${new Date().toISOString()}`
}
});
}
catch (error) {
console.error("Error sending notification:", error);
}
// Wait for the specified interval
await sleep(interval);
}

View file

@ -217,7 +217,6 @@ export class StreamableHTTPServerTransport implements Transport {
// Assign the response to the standalone SSE stream
this._streamMapping.set(this._standaloneSseStreamId, res);
// Set up close handler for client disconnects
res.on("close", () => {
this._streamMapping.delete(this._standaloneSseStreamId);
@ -345,8 +344,10 @@ export class StreamableHTTPServerTransport implements Transport {
const isInitializationRequest = messages.some(
msg => 'method' in msg && msg.method === 'initialize'
);
const mcpSessionId = req.headers["mcp-session-id"] as string | undefined;
if (isInitializationRequest) {
if (this._initialized) {
// if generateSessionId is not set, the server does not support session management
if (this._initialized && this.sessionId !== undefined && mcpSessionId !== this.sessionId) {
res.writeHead(400).end(JSON.stringify({
jsonrpc: "2.0",
error: {
@ -368,7 +369,7 @@ export class StreamableHTTPServerTransport implements Transport {
}));
return;
}
this.sessionId = this.sessionIdGenerator();
this.sessionId = mcpSessionId ?? this.sessionIdGenerator();
this._initialized = true;
}
@ -419,17 +420,8 @@ export class StreamableHTTPServerTransport implements Transport {
this._requestToStreamMapping.set(message.id, streamId);
}
}
// Set up close handler for client disconnects
res.on("close", () => {
// find a stream ID for this response
// Remove all entries that reference this response
for (const [id, stream] of this._requestToStreamMapping.entries()) {
if (streamId === stream) {
this._requestToStreamMapping.delete(id);
this._requestResponseMap.delete(id);
}
}
this._streamMapping.delete(streamId);
});
@ -577,7 +569,7 @@ export class StreamableHTTPServerTransport implements Transport {
// Get the response for this request
const streamId = this._requestToStreamMapping.get(requestId);
const response = this._streamMapping.get(streamId!);
if (!streamId || !response) {
if (!streamId) {
throw new Error(`No connection established for request ID: ${String(requestId)}`);
}
@ -588,9 +580,10 @@ export class StreamableHTTPServerTransport implements Transport {
if (this._eventStore) {
eventId = await this._eventStore.storeEvent(streamId, message);
}
// Write the event to the response stream
this.writeSSEEvent(response, message, eventId);
if (response) {
// Write the event to the response stream
this.writeSSEEvent(response, message, eventId);
}
}
if (isJSONRPCResponse(message)) {
@ -603,6 +596,9 @@ export class StreamableHTTPServerTransport implements Transport {
const allResponsesReady = relatedIds.every(id => this._requestResponseMap.has(id));
if (allResponsesReady) {
if (!response) {
throw new Error(`No connection established for request ID: ${String(requestId)}`);
}
if (this._enableJsonResponse) {
// All responses ready, send as JSON
const headers: Record<string, string> = {

View file

@ -87,6 +87,18 @@ export type RequestOptions = {
* May be used to indicate to the transport which incoming request to associate this outgoing request with.
*/
relatedRequestId?: RequestId;
/**
* May be used to indicate to the transport which last event ID to associate this outgoing request with.
* This is used to resume a long-running requests that may have been interrupted and a new instance of a client is being created.
*/
lastEventId?: string;
/**
* A callback that is invoked when the last event ID is updated.
* This is used to notidy the client that the last event ID has changed, so that client can update its state accordingly.
*/
onLastEventIdUpdate?: (event: string) => void;
};
/**
@ -501,7 +513,7 @@ export abstract class Protocol<
resultSchema: T,
options?: RequestOptions,
): Promise<z.infer<T>> {
const { relatedRequestId } = options ?? {};
const { relatedRequestId, lastEventId, onLastEventIdUpdate } = options ?? {};
return new Promise((resolve, reject) => {
if (!this._transport) {
@ -543,7 +555,7 @@ export abstract class Protocol<
requestId: messageId,
reason: String(reason),
},
}, { relatedRequestId })
}, { relatedRequestId, lastEventId, onLastEventIdUpdate })
.catch((error) =>
this._onerror(new Error(`Failed to send cancellation: ${error}`)),
);
@ -581,7 +593,7 @@ export abstract class Protocol<
this._setupTimeout(messageId, timeout, options?.maxTotalTimeout, timeoutHandler, options?.resetTimeoutOnProgress ?? false);
this._transport.send(jsonrpcRequest, { relatedRequestId }).catch((error) => {
this._transport.send(jsonrpcRequest, { relatedRequestId, lastEventId, onLastEventIdUpdate }).catch((error) => {
this._cleanupTimeout(messageId);
reject(error);
});

View file

@ -18,7 +18,7 @@ export interface Transport {
*
* If present, `relatedRequestId` is used to indicate to the transport which incoming request to associate this outgoing message with.
*/
send(message: JSONRPCMessage, options?: { relatedRequestId?: RequestId }): Promise<void>;
send(message: JSONRPCMessage, options?: { relatedRequestId?: RequestId, lastEventId?: string, onLastEventIdUpdate?: (event: string) => void }): Promise<void>;
/**
* Closes the connection.