add error handling for onprogress disconnects (stop server crashing)

Author: evalstate

src/examples/server/disconnectTestServer.ts

@@ -6,36 +6,68 @@ import { CallToolResult, isInitializeRequest } from '../../types.js';
 import { randomUUID } from 'node:crypto';
 import * as z from 'zod/v4';
 
+// Usage: npx tsx disconnectTestServer.ts [--abort]
+const useAbort = process.argv.includes('--abort');
+console.log(`Abort on disconnect: ${useAbort ? 'enabled' : 'disabled'}`);
+
 const server = new McpServer(
     { name: 'disconnect-test', version: '1.0.0' },
     { capabilities: { logging: {} } }
 );
 
-server.tool('slow-task', 'Task with progress notifications', { steps: z.number() },
+server.server.onerror = err => console.log('[onerror]', err.message);
+
+server.registerTool(
+    'slow-task',
+    {
+        description: 'Task with progress notifications',
+        inputSchema: { steps: z.number() }
+    },
     async ({ steps }, extra): Promise<CallToolResult> => {
-        for (let i = 1; i <= steps; i++) {
-            console.log(`Sending notification ${i}/${steps}`);
+        // SIMULATING A PROXY: create abort controller for "upstream" request
+        const abortController = new AbortController();
+        if (extra.sessionId) {
+            sessionAbortControllers[extra.sessionId] = abortController;
+        }
 
-            // SIMULATING A PROXY RELAY: onprogress forwards with same progress token
-            const progressToken = extra._meta?.progressToken;
-            if (progressToken !== undefined) {
-                server.server.notification(
-                    {
-                        method: 'notifications/progress',
-                        params: { progressToken, progress: i, total: steps }
-                    },
-                    { relatedRequestId: extra.requestId }
-                );
-            }
+        try {
+            for (let i = 1; i <= steps; i++) {
+                // Check if aborted before each step
+                if (abortController.signal.aborted) {
+                    console.log('Upstream request aborted - stopping work');
+                    break;
+                }
 
-            await new Promise(r => setTimeout(r, 1000));
+                console.log(`Sending notification ${i}/${steps}`);
+
+                // SIMULATING A PROXY RELAY: onprogress forwards with same progress token
+                const progressToken = extra._meta?.progressToken;
+                if (progressToken !== undefined) {
+                    server.server.notification(
+                        {
+                            method: 'notifications/progress',
+                            params: { progressToken, progress: i, total: steps }
+                        },
+                        { relatedRequestId: extra.requestId }
+                    );
+                }
+
+                await new Promise(r => setTimeout(r, 1000));
+            }
+            return { content: [{ type: 'text', text: 'SUCCESS' }] };
+        } finally {
+            // Cleanup abort controller
+            if (extra.sessionId) {
+                delete sessionAbortControllers[extra.sessionId];
+            }
         }
-        return { content: [{ type: 'text', text: 'SUCCESS' }] };
     }
 );
 
 const app = createMcpExpressApp();
 const transports: Record<string, StreamableHTTPServerTransport> = {};
+// SIMULATING A PROXY: track abort controllers for upstream requests per session
+const sessionAbortControllers: Record<string, AbortController> = {};
 
 app.post('/mcp', async (req: Request, res: Response) => {
     const sessionId = req.headers['mcp-session-id'] as string | undefined;
@@ -62,7 +94,16 @@ app.post('/mcp', async (req: Request, res: Response) => {
         res.on('finish', () => { finished = true; });
         res.on('close', () => {
             if (!finished) {
-                console.log('Client disconnected - closing transport to reclaim resources');
+                console.log('Client disconnected unexpectedly');
+                if (useAbort) {
+                    // Abort any in-flight upstream requests for this session
+                    const abortController = sessionAbortControllers[transport!.sessionId!];
+                    if (abortController) {
+                        console.log('Aborting upstream request');
+                        abortController.abort();
+                        delete sessionAbortControllers[transport!.sessionId!];
+                    }
+                }
                 transport!.close();
             }
         });

src/shared/protocol.ts

@@ -854,7 +854,9 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
             }
         }
 
-        handler(params);
+        Promise.resolve()
+            .then(() => handler(params))
+            .catch(error => this._onerror(new Error(`Uncaught error in progress handler: ${error}`)));
     }
 
     private _onresponse(response: JSONRPCResponse | JSONRPCErrorResponse): void {
@@ -1277,7 +1279,8 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
      */
     async notification(notification: SendNotificationT, options?: NotificationOptions): Promise<void> {
         if (!this._transport) {
-            throw new Error('Not connected');
+            this._onerror(new Error('Not connected'));
+            return;
         }
 
         this.assertNotificationCapability(notification.method);