resume long running requst from new client
This commit is contained in:
parent
663b68a733
commit
43cb015844
6 changed files with 76 additions and 33 deletions
|
|
@ -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;
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 => {
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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> = {
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Reference in a new issue