@portel/photon 1.32.1 → 1.32.3

package/dist/auto-ui/streamable-http-transport.js

@@ -24,14 +24,16 @@ import { homedir } from 'os';
 import { PHOTON_VERSION } from '../version.js';
 import { formatToolError } from '../shared/error-handler.js';
 import { SimpleRateLimiter } from '../shared/security.js';
-// Default rate limit: 60 requests/min per source IP. Override via
+// Default rate limit: 600 requests/min per source IP. Beam app UI opens several
+// MCP sessions and can legitimately burst while navigating between photons.
+// Override via
 // PHOTON_MCP_RATE_LIMIT (count) and PHOTON_MCP_RATE_WINDOW_MS (window).
-const MCP_RATE_LIMIT = Math.max(1, parseInt(process.env.PHOTON_MCP_RATE_LIMIT || '60', 10) || 60);
+const MCP_RATE_LIMIT = Math.max(1, parseInt(process.env.PHOTON_MCP_RATE_LIMIT || '600', 10) || 600);
 const MCP_RATE_WINDOW_MS = Math.max(1_000, parseInt(process.env.PHOTON_MCP_RATE_WINDOW_MS || '60000', 10) || 60_000);
 const mcpRateLimiter = new SimpleRateLimiter(MCP_RATE_LIMIT, MCP_RATE_WINDOW_MS);
 import { AGUIEventType } from '../ag-ui/types.js';
 import { proxyExternalAgent, createAGUIOutputHandler } from '../ag-ui/adapter.js';
-import { buildToolMetadataExtensions } from './types.js';
+import { buildResponseUIMetadata, buildToolMCPMeta, buildToolMetadataExtensions } from './types.js';
 import { generateServerCard } from '../server-card.js';
 import { audit } from '../shared/audit.js';
 import { writePhotonEditorDeclaration } from '../photon-editor-declarations.js';
@@ -41,6 +43,53 @@ import { toWireFormat, relatedTaskMeta, TERMINAL_STATES } from '../tasks/types.j
 import { runTaskExecution, resolveTaskInput, waitForTerminalOrInput } from '../tasks/executor.js';
 import { generateAgentCard } from '../a2a/card-generator.js';
 import { isPathInScope } from '../daemon/claims.js';
+const MCP_LIST_PAGE_SIZE = 100;
+class InvalidCursorError extends Error {
+    constructor(cursor) {
+        super(`Invalid pagination cursor: ${String(cursor)}`);
+    }
+}
+function encodeListCursor(offset) {
+    return Buffer.from(JSON.stringify({ offset }), 'utf8').toString('base64url');
+}
+function decodeListCursor(cursor) {
+    if (cursor == null || cursor === '')
+        return 0;
+    if (typeof cursor !== 'string')
+        throw new InvalidCursorError(cursor);
+    // Accept the old task-style numeric cursor shape defensively, but emit opaque
+    // base64url cursors for all new paginated list results.
+    if (/^\d+$/.test(cursor))
+        return Number(cursor);
+    try {
+        const decoded = JSON.parse(Buffer.from(cursor, 'base64url').toString('utf8'));
+        if (!Number.isInteger(decoded.offset) || decoded.offset < 0) {
+            throw new Error('cursor offset must be a non-negative integer');
+        }
+        return decoded.offset;
+    }
+    catch {
+        throw new InvalidCursorError(cursor);
+    }
+}
+function paginateMCPList(items, cursor, pageSize = MCP_LIST_PAGE_SIZE) {
+    const offset = decodeListCursor(cursor);
+    if (offset > items.length)
+        throw new InvalidCursorError(cursor);
+    const page = items.slice(offset, offset + pageSize);
+    const nextOffset = offset + pageSize;
+    return {
+        items: page,
+        ...(nextOffset < items.length ? { nextCursor: encodeListCursor(nextOffset) } : {}),
+    };
+}
+function invalidCursorResponse(id, error) {
+    return {
+        jsonrpc: '2.0',
+        id: id,
+        error: { code: -32602, message: error.message },
+    };
+}
 // ════════════════════════════════════════════════════════════════════════════════
 // JWT HELPERS
 // ════════════════════════════════════════════════════════════════════════════════
@@ -146,6 +195,119 @@ function requestSession(sessionId, method, params, timeoutMs = 30 * 60_000) {
         }
     });
 }
+function sessionSupportsFormElicitation(session) {
+    const elicitation = session.clientCapabilities?.elicitation;
+    if (!elicitation || typeof elicitation !== 'object' || Array.isArray(elicitation)) {
+        return false;
+    }
+    return Object.keys(elicitation).length === 0 || 'form' in elicitation;
+}
+function buildMcpElicitParamsFromAsk(ask) {
+    const message = ask.message || 'Please provide input';
+    switch (ask.ask) {
+        case 'confirm':
+            return {
+                mode: 'form',
+                message,
+                requestedSchema: {
+                    type: 'object',
+                    properties: {
+                        confirmed: {
+                            type: 'boolean',
+                            title: ask.label || 'Confirm',
+                            description: ask.hint || ask.message,
+                            default: ask.default ?? false,
+                        },
+                    },
+                    required: ['confirmed'],
+                },
+            };
+        case 'number':
+            return {
+                mode: 'form',
+                message,
+                requestedSchema: {
+                    type: 'object',
+                    properties: {
+                        value: {
+                            type: 'number',
+                            title: ask.label || 'Number',
+                            description: ask.hint || ask.message,
+                            default: ask.default,
+                            minimum: ask.min,
+                            maximum: ask.max,
+                        },
+                    },
+                    required: ask.required !== false ? ['value'] : [],
+                },
+            };
+        case 'select': {
+            const options = ask.options || [];
+            const optionItems = options.map((option) => typeof option === 'string'
+                ? { const: option, title: option }
+                : {
+                    const: option.value,
+                    title: option.label || String(option.value),
+                    ...(option.description ? { description: option.description } : {}),
+                });
+            return {
+                mode: 'form',
+                message,
+                requestedSchema: {
+                    type: 'object',
+                    properties: {
+                        selection: ask.multi
+                            ? {
+                                type: 'array',
+                                title: ask.label || 'Selection',
+                                description: ask.hint || ask.message,
+                                items: { anyOf: optionItems },
+                                default: ask.default,
+                            }
+                            : {
+                                type: 'string',
+                                title: ask.label || 'Selection',
+                                description: ask.hint || ask.message,
+                                anyOf: optionItems,
+                                default: ask.default,
+                            },
+                    },
+                    required: ask.required !== false ? ['selection'] : [],
+                },
+            };
+        }
+        case 'text':
+        case 'password':
+        default:
+            return {
+                mode: 'form',
+                message,
+                requestedSchema: {
+                    type: 'object',
+                    properties: {
+                        value: {
+                            type: 'string',
+                            title: ask.label || 'Input',
+                            description: ask.hint || ask.message,
+                            default: ask.default,
+                        },
+                    },
+                    required: ask.required !== false ? ['value'] : [],
+                },
+            };
+    }
+}
+function extractMcpElicitValue(ask, result) {
+    if (result?.action !== 'accept') {
+        return ask.multi ? [] : ask.ask === 'confirm' ? false : null;
+    }
+    const content = result.content || {};
+    if (ask.ask === 'confirm')
+        return content.confirmed ?? false;
+    if (ask.ask === 'select')
+        return content.selection;
+    return content.value;
+}
 /**
  * Build a sampling provider for a specific Beam session — the
  * person at the browser plays the role of the LLM. The provider
@@ -583,6 +745,8 @@ const handlers = {
             session.clientInfo = clientInfo;
             session.isBeam = clientInfo.name === 'beam';
         }
+        session.clientCapabilities =
+            req.params?.capabilities || {};
         // Generate configuration schema for unconfigured photons
         const configurationSchema = generateConfigurationSchema(ctx.photons);
         return {
@@ -916,6 +1080,10 @@ const handlers = {
             const webRootRoute = httpRoutes?.find((r) => r.method === 'GET' && r.path === '/');
             const webUrl = webRootRoute ? `/web/${photon.name}/` : undefined;
             for (const method of photon.methods) {
+                const uiResourceUri = method.linkedUi
+                    ? `ui://${photon.name}/${method.linkedUi}`
+                    : undefined;
+                const meta = buildToolMCPMeta(method, { uiResourceUri });
                 tools.push({
                     name: `${photon.name}/${method.name}`,
                     description: method.description || `Execute ${method.name}`,
@@ -944,19 +1112,7 @@ const handlers = {
                     'x-photon-resource-count': photon.resourceCount ?? 0,
                     ...(webUrl ? { 'x-web-url': webUrl, 'x-web-description': photon.description } : {}),
                     ...buildToolMetadataExtensions(method),
-                    // MCP Apps standard: _meta.ui for linked UI resources and visibility
-                    ...(method.linkedUi || method.visibility
-                        ? {
-                            _meta: {
-                                ui: {
-                                    ...(method.linkedUi
-                                        ? { resourceUri: `ui://${photon.name}/${method.linkedUi}` }
-                                        : {}),
-                                    ...(method.visibility ? { visibility: method.visibility } : {}),
-                                },
-                            },
-                        }
-                        : {}),
+                    ...(Object.keys(meta).length > 0 ? { _meta: meta } : {}),
                 });
             }
         }
@@ -1007,6 +1163,7 @@ const handlers = {
                 if (!mcp.connected || !mcp.methods)
                     continue;
                 for (const method of mcp.methods) {
+                    const meta = buildToolMCPMeta(method, { uiResourceUri: method.linkedUi });
                     tools.push({
                         name: `${mcp.name}/${method.name}`,
                         description: method.description || `Execute ${method.name}`,
@@ -1021,17 +1178,7 @@ const handlers = {
                         'x-mcp-app-uri': mcp.appResourceUri, // MCP App resource URI (default/first)
                         'x-mcp-app-uris': mcp.appResourceUris || [], // All MCP App resource URIs
                         ...buildToolMetadataExtensions(method),
-                        // MCP Apps standard: _meta.ui for linked UI resources and visibility
-                        ...(method.linkedUi || method.visibility
-                            ? {
-                                _meta: {
-                                    ui: {
-                                        ...(method.linkedUi ? { resourceUri: method.linkedUi } : {}),
-                                        ...(method.visibility ? { visibility: method.visibility } : {}),
-                                    },
-                                },
-                            }
-                            : {}),
+                        ...(Object.keys(meta).length > 0 ? { _meta: meta } : {}),
                     });
                 }
             }
@@ -1267,7 +1414,22 @@ const handlers = {
                 }
                 return true;
             });
-        return { jsonrpc: '2.0', id: req.id, result: { tools: visibleTools } };
+        try {
+            const page = paginateMCPList(visibleTools, req.params?.cursor);
+            return {
+                jsonrpc: '2.0',
+                id: req.id,
+                result: {
+                    tools: page.items,
+                    ...(page.nextCursor ? { nextCursor: page.nextCursor } : {}),
+                },
+            };
+        }
+        catch (error) {
+            if (error instanceof InvalidCursorError)
+                return invalidCursorResponse(req.id, error);
+            throw error;
+        }
     },
     'tools/call': async (req, session, ctx) => {
         const { name, arguments: args } = req.params;
@@ -1437,11 +1599,7 @@ const handlers = {
         const methodInfo = photonInfo?.configured
             ? photonInfo.methods?.find((m) => m.name === methodName)
             : undefined;
-        // Build UI metadata
-        const uiMetadata = {};
-        if (methodInfo?.outputFormat) {
-            uiMetadata['x-output-format'] = methodInfo.outputFormat;
-        }
+        const uiMetadata = buildResponseUIMetadata(photonName, methodInfo);
         // Auto-confirm @destructive operations before execution (any transport path)
         if (methodInfo?.destructiveHint) {
             const elicitResult = await requestBeamElicitation({
@@ -1867,8 +2025,19 @@ const handlers = {
             // Create inputProvider to handle ask yields (elicitation)
             // Supports persistent: true for durable approvals that survive navigation/restart
             const inputProvider = async (ask) => {
+                if (!session.isBeam) {
+                    if (!sessionSupportsFormElicitation(session)) {
+                        throw new Error(`Tool ${photonName}/${methodName} requires MCP elicitation, but this client did not advertise the elicitation capability. ` +
+                            'Call it from an MCP client that supports elicitation/create, or use the Beam UI.');
+                    }
+                    if (!session.sseResponse || session.sseResponse.writableEnded) {
+                        throw new Error(`Tool ${photonName}/${methodName} requires MCP elicitation, but this session has no live SSE stream for server-initiated requests.`);
+                    }
+                    const result = await requestSession(session.id, 'elicitation/create', buildMcpElicitParamsFromAsk(ask), 300000);
+                    return extractMcpElicitValue(ask, result);
+                }
                 if (!ctx.broadcast) {
-                    throw new Error('No broadcast connection for elicitation');
+                    throw new Error('No broadcast connection for Beam elicitation');
                 }
                 // Generate unique elicitation ID
                 const elicitationId = randomUUID();
@@ -2184,7 +2353,22 @@ const handlers = {
                 description: `Approval request from ${approval.photon}.${approval.method}`,
             });
         }
-        return { jsonrpc: '2.0', id: req.id, result: { resources } };
+        try {
+            const page = paginateMCPList(resources, req.params?.cursor);
+            return {
+                jsonrpc: '2.0',
+                id: req.id,
+                result: {
+                    resources: page.items,
+                    ...(page.nextCursor ? { nextCursor: page.nextCursor } : {}),
+                },
+            };
+        }
+        catch (error) {
+            if (error instanceof InvalidCursorError)
+                return invalidCursorResponse(req.id, error);
+            throw error;
+        }
     },
     'resources/templates/list': async (req, _session, ctx) => {
         const resourceTemplates = [];
@@ -2205,7 +2389,22 @@ const handlers = {
                 });
             }
         }
-        return { jsonrpc: '2.0', id: req.id, result: { resourceTemplates } };
+        try {
+            const page = paginateMCPList(resourceTemplates, req.params?.cursor);
+            return {
+                jsonrpc: '2.0',
+                id: req.id,
+                result: {
+                    resourceTemplates: page.items,
+                    ...(page.nextCursor ? { nextCursor: page.nextCursor } : {}),
+                },
+            };
+        }
+        catch (error) {
+            if (error instanceof InvalidCursorError)
+                return invalidCursorResponse(req.id, error);
+            throw error;
+        }
     },
     'resources/read': async (req, session, ctx) => {
         const { uri } = req.params;
@@ -2376,7 +2575,22 @@ const handlers = {
                 });
             }
         }
-        return { jsonrpc: '2.0', id: req.id, result: { prompts } };
+        try {
+            const page = paginateMCPList(prompts, req.params?.cursor);
+            return {
+                jsonrpc: '2.0',
+                id: req.id,
+                result: {
+                    prompts: page.items,
+                    ...(page.nextCursor ? { nextCursor: page.nextCursor } : {}),
+                },
+            };
+        }
+        catch (error) {
+            if (error instanceof InvalidCursorError)
+                return invalidCursorResponse(req.id, error);
+            throw error;
+        }
     },
     'prompts/get': async (req, _session, ctx) => {
         const { name } = req.params;
@@ -3415,11 +3629,14 @@ async function handleBeamStudioWrite(req, ctx, args) {
         catch {
             // Parse is best-effort — don't fail the write
         }
-        // Trigger hot-reload if available
-        if (ctx.reloadPhoton) {
+        // Trigger hot-reload if available. Studio writes go through the debounced
+        // reload path so streamed saves coalesce with file watcher events.
+        const useScheduledReload = !!ctx.schedulePhotonReload;
+        const reloadPhoton = ctx.schedulePhotonReload || ctx.reloadPhoton;
+        if (reloadPhoton) {
             try {
-                const reloadResult = await ctx.reloadPhoton(photonName);
-                if (reloadResult.success) {
+                const reloadResult = await reloadPhoton(photonName);
+                if (!useScheduledReload && reloadResult.success && reloadResult.photon) {
                     broadcastToBeam('beam/hot-reload', { photon: reloadResult.photon });
                 }
             }
@@ -3543,10 +3760,12 @@ async function handleBeamStudioApplyFiles(req, ctx, args) {
                 await writePhotonEditorDeclaration(file.path, file.source, ctx.workingDir).catch(() => null);
             }
         }
-        if (ctx.reloadPhoton) {
+        const useScheduledReload = !!ctx.schedulePhotonReload;
+        const reloadPhoton = ctx.schedulePhotonReload || ctx.reloadPhoton;
+        if (reloadPhoton) {
             try {
-                const reloadResult = await ctx.reloadPhoton(photonName);
-                if (reloadResult.success) {
+                const reloadResult = await reloadPhoton(photonName);
+                if (!useScheduledReload && reloadResult.success && reloadResult.photon) {
                     broadcastToBeam('beam/hot-reload', { photon: reloadResult.photon });
                 }
             }
@@ -3844,148 +4063,160 @@ export async function handleStreamableHTTP(req, res, options) {
         req.on('aborted', abortRequest);
         res.on('close', abortRequest);
         res.on('error', abortRequest);
-        req.socket?.on('close', abortRequest);
-        res.socket?.on('close', abortRequest);
-        let responseStreamStarted = false;
-        const ensureResponseStream = () => {
-            if (responseStreamStarted)
-                return;
-            res.writeHead(200, {
-                'Content-Type': 'text/event-stream',
-                'Cache-Control': 'no-cache',
-                Connection: 'keep-alive',
-                'X-Accel-Buffering': 'no',
-                'Mcp-Session-Id': session.id,
-            });
-            res.socket?.setNoDelay(true);
-            responseStreamStarted = true;
-        };
-        const sendResponseStreamMessage = (message) => {
-            if (!wantsSSE || res.writableEnded || res.destroyed)
-                return;
-            ensureResponseStream();
-            res.write(`data: ${JSON.stringify(message)}\n\n`);
-        };
-        // Read body
-        let body = '';
-        for await (const chunk of req) {
-            body += chunk;
+        const sockets = new Set();
+        if (req.socket)
+            sockets.add(req.socket);
+        if (res.socket)
+            sockets.add(res.socket);
+        for (const socket of sockets) {
+            socket.on('close', abortRequest);
         }
-        let requests;
         try {
-            const parsed = JSON.parse(body);
-            requests = Array.isArray(parsed) ? parsed : [parsed];
-        }
-        catch {
-            res.writeHead(400);
-            res.end(JSON.stringify({ error: 'Invalid JSON' }));
-            return true;
-        }
-        // Extract caller identity from Authorization header or query token (MCP OAuth)
-        const caller = decodeJWTCaller(authHeader);
-        const context = {
-            photons: options.photons,
-            photonMCPs: options.photonMCPs,
-            externalMCPs: options.externalMCPs,
-            externalMCPClients: options.externalMCPClients,
-            externalMCPSDKClients: options.externalMCPSDKClients,
-            reconnectExternalMCP: options.reconnectExternalMCP,
-            loadUIAsset: options.loadUIAsset,
-            configurePhoton: options.configurePhoton,
-            reloadPhoton: options.reloadPhoton,
-            removePhoton: options.removePhoton,
-            updateMetadata: options.updateMetadata,
-            generatePhotonHelp: options.generatePhotonHelp,
-            loader: options.loader,
-            broadcast: options.broadcast,
-            responseStream: wantsSSE ? { send: sendResponseStreamMessage } : undefined,
-            signal: requestAbort.signal,
-            subscriptionManager: options.subscriptionManager,
-            workingDir: options.workingDir,
-            caller,
-        };
-        // Process requests
-        const responses = [];
-        for (const request of requests) {
-            // Response to a server→client request: no method, has id, has
-            // either result or error. Route to the pending-request map so
-            // the samplingProvider / future server-initiated primitives see
-            // the browser's reply. These never produce an outgoing response.
-            //
-            // SECURITY: the reply's session MUST match the session that
-            // originated the server→client request. Ids are globally
-            // monotonic (`srv-1`, `srv-2`, ...), so without a session cross
-            // check, session A could POST a reply carrying session B's id
-            // and inject a fabricated sampling result into B's photon.
-            // Drop mismatched replies silently — the original timeout on
-            // the pending entry stays the only way to fail legitimately.
-            if (!request.method && request.id !== undefined) {
-                const msg = request;
-                const pending = pendingServerRequests.get(msg.id);
-                if (pending && pending.sessionId === session.id) {
-                    if (pending.timer)
-                        clearTimeout(pending.timer);
-                    pendingServerRequests.delete(msg.id);
-                    if (msg.error) {
-                        pending.reject(new Error(msg.error.message || 'server→client request failed'));
+            let responseStreamStarted = false;
+            const ensureResponseStream = () => {
+                if (responseStreamStarted)
+                    return;
+                res.writeHead(200, {
+                    'Content-Type': 'text/event-stream',
+                    'Cache-Control': 'no-cache',
+                    Connection: 'keep-alive',
+                    'X-Accel-Buffering': 'no',
+                    'Mcp-Session-Id': session.id,
+                });
+                res.socket?.setNoDelay(true);
+                responseStreamStarted = true;
+            };
+            const sendResponseStreamMessage = (message) => {
+                if (!wantsSSE || res.writableEnded || res.destroyed)
+                    return;
+                ensureResponseStream();
+                res.write(`data: ${JSON.stringify(message)}\n\n`);
+            };
+            // Read body
+            let body = '';
+            for await (const chunk of req) {
+                body += chunk;
+            }
+            let requests;
+            try {
+                const parsed = JSON.parse(body);
+                requests = Array.isArray(parsed) ? parsed : [parsed];
+            }
+            catch {
+                res.writeHead(400);
+                res.end(JSON.stringify({ error: 'Invalid JSON' }));
+                return true;
+            }
+            // Extract caller identity from Authorization header or query token (MCP OAuth)
+            const caller = decodeJWTCaller(authHeader);
+            const context = {
+                photons: options.photons,
+                photonMCPs: options.photonMCPs,
+                externalMCPs: options.externalMCPs,
+                externalMCPClients: options.externalMCPClients,
+                externalMCPSDKClients: options.externalMCPSDKClients,
+                reconnectExternalMCP: options.reconnectExternalMCP,
+                loadUIAsset: options.loadUIAsset,
+                configurePhoton: options.configurePhoton,
+                reloadPhoton: options.reloadPhoton,
+                schedulePhotonReload: options.schedulePhotonReload,
+                removePhoton: options.removePhoton,
+                updateMetadata: options.updateMetadata,
+                generatePhotonHelp: options.generatePhotonHelp,
+                loader: options.loader,
+                broadcast: options.broadcast,
+                responseStream: wantsSSE ? { send: sendResponseStreamMessage } : undefined,
+                signal: requestAbort.signal,
+                subscriptionManager: options.subscriptionManager,
+                workingDir: options.workingDir,
+                caller,
+            };
+            // Process requests
+            const responses = [];
+            for (const request of requests) {
+                // Response to a server→client request: no method, has id, has
+                // either result or error. Route to the pending-request map so
+                // the samplingProvider / future server-initiated primitives see
+                // the browser's reply. These never produce an outgoing response.
+                //
+                // SECURITY: the reply's session MUST match the session that
+                // originated the server→client request. Ids are globally
+                // monotonic (`srv-1`, `srv-2`, ...), so without a session cross
+                // check, session A could POST a reply carrying session B's id
+                // and inject a fabricated sampling result into B's photon.
+                // Drop mismatched replies silently — the original timeout on
+                // the pending entry stays the only way to fail legitimately.
+                if (!request.method && request.id !== undefined) {
+                    const msg = request;
+                    const pending = pendingServerRequests.get(msg.id);
+                    if (pending && pending.sessionId === session.id) {
+                        if (pending.timer)
+                            clearTimeout(pending.timer);
+                        pendingServerRequests.delete(msg.id);
+                        if (msg.error) {
+                            pending.reject(new Error(msg.error.message || 'server→client request failed'));
+                        }
+                        else {
+                            pending.resolve(msg.result);
+                        }
                     }
-                    else {
-                        pending.resolve(msg.result);
+                    continue;
+                }
+                const handler = handlers[request.method];
+                if (!handler) {
+                    if (request.id !== undefined) {
+                        responses.push({
+                            jsonrpc: '2.0',
+                            id: request.id,
+                            error: { code: -32601, message: `Method not found: ${request.method}` },
+                        });
                     }
+                    continue;
                 }
-                continue;
-            }
-            const handler = handlers[request.method];
-            if (!handler) {
-                if (request.id !== undefined) {
-                    responses.push({
-                        jsonrpc: '2.0',
-                        id: request.id,
-                        error: { code: -32601, message: `Method not found: ${request.method}` },
-                    });
+                const response = await handler(request, session, context);
+                // Only include responses for requests (not notifications)
+                if (request.id !== undefined && response.id !== undefined) {
+                    responses.push(response);
                 }
-                continue;
             }
-            const response = await handler(request, session, context);
-            // Only include responses for requests (not notifications)
-            if (request.id !== undefined && response.id !== undefined) {
-                responses.push(response);
+            // Send response
+            if (responses.length === 0) {
+                // All were notifications
+                if (responseStreamStarted) {
+                    res.end();
+                }
+                else {
+                    res.writeHead(202);
+                    res.end();
+                }
             }
-        }
-        // Send response
-        if (responses.length === 0) {
-            // All were notifications
-            if (responseStreamStarted) {
+            else if (wantsSSE) {
+                // SSE response
+                ensureResponseStream();
+                for (const response of responses) {
+                    sendResponseStreamMessage(response);
+                }
                 res.end();
             }
             else {
-                res.writeHead(202);
-                res.end();
+                // JSON response
+                res.writeHead(200, {
+                    'Content-Type': 'application/json',
+                    'Mcp-Session-Id': session.id,
+                });
+                const result = responses.length === 1 ? responses[0] : responses;
+                res.end(JSON.stringify(result));
             }
+            return true;
         }
-        else if (wantsSSE) {
-            // SSE response
-            ensureResponseStream();
-            for (const response of responses) {
-                sendResponseStreamMessage(response);
+        finally {
+            res.off('close', abortRequest);
+            res.off('error', abortRequest);
+            req.off('aborted', abortRequest);
+            for (const socket of sockets) {
+                socket.off('close', abortRequest);
             }
-            res.end();
         }
-        else {
-            // JSON response
-            res.writeHead(200, {
-                'Content-Type': 'application/json',
-                'Mcp-Session-Id': session.id,
-            });
-            const result = responses.length === 1 ? responses[0] : responses;
-            res.end(JSON.stringify(result));
-        }
-        res.off('close', abortRequest);
-        res.off('error', abortRequest);
-        req.off('aborted', abortRequest);
-        req.socket?.off('close', abortRequest);
-        res.socket?.off('close', abortRequest);
-        return true;
     }
     // Method not allowed
     res.writeHead(405);