@j0hanz/fetch-url-mcp 1.12.7 → 1.12.8

package/dist/tasks/manager.js

@@ -1,12 +1,459 @@
-import { randomUUID } from 'node:crypto';
-import { createHmac, randomBytes } from 'node:crypto';
+import {} from '@modelcontextprotocol/sdk/server/mcp.js';
+import { CallToolRequestSchema, ErrorCode, McpError, } from '@modelcontextprotocol/sdk/types.js';
+import { AsyncLocalStorage } from 'node:async_hooks';
+import { createHmac, hash, randomBytes, randomUUID } from 'node:crypto';
 import { setInterval } from 'node:timers';
-import { ErrorCode } from '@modelcontextprotocol/sdk/types.js';
-import { config, logInfo, logWarn } from '../lib/core.js';
-import { Loggers } from '../lib/logger-names.js';
-import { createMcpError } from '../lib/mcp-interop.js';
-import { isObject, timingSafeEqualUtf8 } from '../lib/utils.js';
-import { TaskWaiterRegistry, waitForTerminalTask as waitForTerminalTaskWithDeadline, } from './waiters.js';
+import { z } from 'zod';
+import { config, getRequestId, logDebug, logError, Loggers, logInfo, logWarn, resolveMcpSessionOwnerKey, runWithRequestContext, } from '../lib/core.js';
+import { getErrorMessage, handleToolError, stripMcpErrorPrefix, toError, tryReadToolErrorMessage, } from '../lib/error/index.js';
+import { createMcpError, getSdkCallToolHandler, registerServerLifecycleCleanup, } from '../lib/mcp-interop.js';
+import { createUnrefTimeout, formatZodError, isObject, timingSafeEqualUtf8, } from '../lib/utils.js';
+const MIN_TASK_TTL_MS = 1_000;
+const MAX_TASK_TTL_MS = 86_400_000;
+const relatedTaskMetaSchema = z.strictObject({
+    taskId: z.string(),
+});
+const toolCallMetaSchema = z.looseObject({
+    progressToken: z.union([z.string(), z.number()]).optional(),
+    'io.modelcontextprotocol/related-task': relatedTaskMetaSchema.optional(),
+});
+export const extendedCallToolRequestSchema = z.looseObject({
+    method: z.literal('tools/call'),
+    params: z.strictObject({
+        name: z.string().min(1, 'Tool name required'),
+        arguments: z.record(z.string(), z.unknown()).optional(),
+        task: z
+            .strictObject({
+            ttl: z
+                .number()
+                .int()
+                .min(MIN_TASK_TTL_MS, `Minimum ${MIN_TASK_TTL_MS}ms`)
+                .max(MAX_TASK_TTL_MS, `Maximum ${MAX_TASK_TTL_MS}ms`)
+                .optional(),
+        })
+            .optional(),
+        _meta: toolCallMetaSchema.optional(),
+    }),
+});
+export function parseExtendedCallToolRequest(request) {
+    const parsed = extendedCallToolRequestSchema.safeParse(request);
+    if (parsed.success)
+        return parsed.data;
+    throw createMcpError(ErrorCode.InvalidParams, formatZodError(parsed.error));
+}
+export function sanitizeToolCallMeta(meta) {
+    if (!meta)
+        return undefined;
+    const sanitized = { ...meta };
+    delete sanitized['io.modelcontextprotocol/related-task'];
+    return Object.keys(sanitized).length > 0 ? sanitized : undefined;
+}
+export function buildRelatedTaskMeta(taskId, meta) {
+    return {
+        ...(sanitizeToolCallMeta(meta) ?? {}),
+        'io.modelcontextprotocol/related-task': { taskId },
+    };
+}
+export function withRelatedTaskMeta(result, taskId) {
+    return {
+        ...result,
+        _meta: {
+            ...result._meta,
+            'io.modelcontextprotocol/related-task': { taskId },
+        },
+    };
+}
+/* -------------------------------------------------------------------------------------------------
+ * Abort-controller management for in-flight task executions
+ * ------------------------------------------------------------------------------------------------- */
+// Intentionally process-global (not session-scoped): abortAllTaskExecutions() is called
+// during SIGTERM/SIGINT shutdown to cancel every in-flight task across all sessions.
+const taskAbortControllers = new Map();
+function attachAbortController(taskId) {
+    taskAbortControllers.get(taskId)?.abort();
+    if (taskAbortControllers.size >= config.tasks.maxTotal) {
+        logWarn('Abort controller map reached task capacity — possible leak', {
+            size: taskAbortControllers.size,
+            maxTotal: config.tasks.maxTotal,
+        }, Loggers.LOG_TASKS);
+    }
+    const controller = new AbortController();
+    taskAbortControllers.set(taskId, controller);
+    return controller;
+}
+export function abortTaskExecution(taskId) {
+    taskAbortControllers.get(taskId)?.abort();
+    taskAbortControllers.delete(taskId);
+}
+export function cancelTasksForOwner(ownerKey, statusMessage = 'The task was cancelled because its owner session ended.') {
+    if (!ownerKey)
+        return 0;
+    const cancelled = taskManager.cancelTasksByOwner(ownerKey, statusMessage);
+    for (const task of cancelled) {
+        abortTaskExecution(task.taskId);
+    }
+    return cancelled.length;
+}
+export function abortAllTaskExecutions() {
+    for (const taskId of taskAbortControllers.keys())
+        abortTaskExecution(taskId);
+}
+export function toTaskSummary(task) {
+    return {
+        taskId: task.taskId,
+        status: task.status,
+        ...(task.statusMessage ? { statusMessage: task.statusMessage } : {}),
+        ...(task.progress !== undefined ? { progress: task.progress } : {}),
+        ...(task.total !== undefined ? { total: task.total } : {}),
+        createdAt: task.createdAt,
+        lastUpdatedAt: task.lastUpdatedAt,
+        ttl: task.ttl,
+        pollInterval: task.pollInterval,
+    };
+}
+export function emitTaskStatusNotification(server, task) {
+    if (!config.tasks.emitStatusNotifications || !server.isConnected())
+        return;
+    void server.server
+        .notification({
+        method: 'notifications/tasks/status',
+        params: { ...toTaskSummary(task) },
+    })
+        .catch((error) => {
+        logError('Failed to send task status notification', {
+            taskId: task.taskId,
+            status: task.status,
+            error: getErrorMessage(error),
+        }, Loggers.LOG_TASKS);
+    });
+}
+export function throwTaskNotFound() {
+    throw createMcpError(ErrorCode.InvalidParams, 'Task not found');
+}
+/* -------------------------------------------------------------------------------------------------
+ * Execution pipeline
+ * ------------------------------------------------------------------------------------------------- */
+function updateTaskAndEmitStatus(server, taskId, update) {
+    taskManager.updateTask(taskId, update);
+    const task = taskManager.getTask(taskId);
+    if (task)
+        emitTaskStatusNotification(server, task);
+}
+function buildTaskFailureState(error) {
+    const mcpErrorMessage = error instanceof McpError ? stripMcpErrorPrefix(error.message) : undefined;
+    const statusMessage = mcpErrorMessage ?? getErrorMessage(error);
+    if (error instanceof McpError) {
+        return {
+            status: 'failed',
+            statusMessage,
+            error: {
+                code: error.code,
+                ...(error.data !== undefined ? { data: error.data } : {}),
+                message: statusMessage,
+            },
+        };
+    }
+    return {
+        status: 'failed',
+        statusMessage,
+        error: {
+            code: ErrorCode.InternalError,
+            message: statusMessage,
+        },
+    };
+}
+function buildTaskCompletionUpdate(result, tool) {
+    const isError = isObject(result) && 'isError' in result && result.isError === true;
+    return {
+        status: isError ? 'failed' : 'completed',
+        statusMessage: isError
+            ? (tryReadToolErrorMessage(result) ?? 'Execution failed')
+            : (tool.getCompletionStatusMessage?.(result) ??
+                'Task completed successfully.'),
+        result,
+    };
+}
+async function runTaskToolExecution(params) {
+    const { server, taskId, args, tool, meta, sessionId, sendNotification } = params;
+    return runWithRequestContext({
+        requestId: taskId,
+        operationId: taskId,
+        ...(sessionId ? { sessionId } : {}),
+    }, async () => {
+        const controller = attachAbortController(taskId);
+        const progressState = { closed: false };
+        try {
+            logInfo('Task execution started', { taskId, tool: tool.name }, Loggers.LOG_TASKS);
+            const relatedMeta = buildRelatedTaskMeta(taskId, meta);
+            const result = await tool.execute(args, {
+                signal: controller.signal,
+                requestId: taskId,
+                _meta: relatedMeta,
+                progressState,
+                canReportProgress: () => taskManager.getTask(taskId)?.status === 'working',
+                ...compact({ sendNotification }),
+                onProgress: (progress, message, total) => {
+                    const current = taskManager.getTask(taskId);
+                    if (current?.status === 'working' &&
+                        (current.statusMessage !== message ||
+                            current.progress !== progress ||
+                            (total !== undefined && current.total !== total))) {
+                        updateTaskAndEmitStatus(server, taskId, {
+                            statusMessage: message,
+                            progress,
+                            ...(total !== undefined ? { total } : {}),
+                        });
+                    }
+                },
+            });
+            const completionUpdate = buildTaskCompletionUpdate(result, tool);
+            updateTaskAndEmitStatus(server, taskId, completionUpdate);
+            if (completionUpdate.status === 'completed') {
+                logInfo('Task execution completed', { taskId, tool: tool.name }, Loggers.LOG_TASKS);
+            }
+            else {
+                logWarn('Task execution completed with tool error result', { taskId, tool: tool.name }, Loggers.LOG_TASKS);
+            }
+        }
+        catch (error) {
+            logError('Task execution failed', {
+                taskId,
+                tool: tool.name,
+                error: getErrorMessage(error),
+            }, Loggers.LOG_TASKS);
+            updateTaskAndEmitStatus(server, taskId, buildTaskFailureState(error));
+        }
+        finally {
+            progressState.closed = true;
+            taskAbortControllers.delete(taskId);
+        }
+    });
+}
+function extractRawUrl(args) {
+    const url = args?.['url'];
+    return typeof url === 'string' ? url : 'unknown';
+}
+function tryParseArguments(tool, args) {
+    try {
+        return { ok: true, value: tool.parseArguments(args) };
+    }
+    catch (error) {
+        if (error instanceof McpError) {
+            return {
+                ok: false,
+                response: handleToolError(error, extractRawUrl(args)),
+            };
+        }
+        throw error;
+    }
+}
+export async function handleToolCallRequest(server, request, context) {
+    const { params } = request;
+    // Validate the tool name first so an unknown tool always produces MethodNotFound
+    const tool = getTaskCapableTool(server, params.name);
+    if (!tool) {
+        throw createMcpError(ErrorCode.MethodNotFound, `Unknown tool: ${params.name}`);
+    }
+    if (params.task) {
+        if (getTaskCapableToolSupport(server, params.name) === 'forbidden') {
+            throw createMcpError(ErrorCode.MethodNotFound, `Task mode is not supported for tool: ${params.name}`);
+        }
+        const parsed = tryParseArguments(tool, params.arguments);
+        if (!parsed.ok)
+            return parsed.response;
+        const task = taskManager.createTask(params.task.ttl !== undefined ? { ttl: params.task.ttl } : undefined, 'Task started', context.ownerKey);
+        logInfo('Task execution queued', {
+            taskId: task.taskId,
+            tool: params.name,
+            ...(params.task.ttl !== undefined ? { ttl: params.task.ttl } : {}),
+        }, Loggers.LOG_TASKS);
+        void runTaskToolExecution({
+            server,
+            taskId: task.taskId,
+            args: parsed.value,
+            tool,
+            ...compact({
+                meta: params._meta,
+                sessionId: context.sessionId,
+                sendNotification: context.sendNotification,
+            }),
+        });
+        return {
+            task: toTaskSummary(task),
+            ...(tool.immediateResponse
+                ? {
+                    _meta: {
+                        'io.modelcontextprotocol/model-immediate-response': tool.immediateResponse,
+                    },
+                }
+                : {}),
+        };
+    }
+    if (getTaskCapableToolSupport(server, params.name) === 'required') {
+        throw createMcpError(ErrorCode.MethodNotFound, `Task mode is required for tool: ${params.name}`);
+    }
+    const parsed = tryParseArguments(tool, params.arguments);
+    if (!parsed.ok)
+        return parsed.response;
+    const progressState = { closed: false };
+    logDebug('Executing task-capable tool inline', {
+        tool: params.name,
+        hasProgressToken: params._meta?.progressToken !== undefined,
+    }, Loggers.LOG_TASKS);
+    try {
+        return await tool.execute(parsed.value, {
+            ...buildToolHandlerExtra(context, params._meta),
+            progressState,
+        });
+    }
+    finally {
+        progressState.closed = true;
+    }
+}
+/* -------------------------------------------------------------------------------------------------
+ * Task handler schemas and registration
+ * ------------------------------------------------------------------------------------------------- */
+const TaskGetSchema = z.looseObject({
+    method: z.literal('tasks/get', 'Expected "tasks/get"'),
+    params: z.looseObject({ taskId: z.string('Expected string') }, 'Expected object'),
+}, 'Invalid request');
+const TaskListSchema = z.looseObject({
+    method: z.literal('tasks/list', 'Expected "tasks/list"'),
+    params: z
+        .looseObject({
+        cursor: z.string('Expected string').optional(),
+    }, 'Expected object')
+        .optional(),
+}, 'Invalid request');
+const TaskCancelSchema = z.looseObject({
+    method: z.literal('tasks/cancel', 'Expected "tasks/cancel"'),
+    params: z.looseObject({ taskId: z.string('Expected string') }, 'Expected object'),
+}, 'Invalid request');
+const TaskResultSchema = z.looseObject({
+    method: z.literal('tasks/result', 'Expected "tasks/result"'),
+    params: z.looseObject({ taskId: z.string('Expected string') }, 'Expected object'),
+}, 'Invalid request');
+function resolveOwnerScopedExtra(extra) {
+    const parsedExtra = parseHandlerExtra(extra);
+    return {
+        parsedExtra,
+        ownerKey: resolveTaskOwnerKey(parsedExtra),
+    };
+}
+function throwStoredTaskError(task) {
+    if (task.error) {
+        throw createMcpError(task.error.code, task.error.message, task.error.data);
+    }
+    throw createMcpError(ErrorCode.InternalError, task.statusMessage ?? 'Execution failed', { taskId: task.taskId });
+}
+export function registerTaskHandlers(server, options) {
+    const sdkCallToolHandler = getSdkCallToolHandler(server);
+    const taskCapableToolsRegistered = hasRegisteredTaskCapableTools(server);
+    const requireInterception = options?.requireInterception ?? true;
+    if (!sdkCallToolHandler) {
+        if (taskCapableToolsRegistered && requireInterception) {
+            throw Error('Task-capable tools are registered but SDK tools/call interception is unavailable. Upgrade compatibility or disable strict interception with TASKS_REQUIRE_INTERCEPTION=false.');
+        }
+        logWarn('Task call interception disabled: SDK tools/call handler unavailable; task-capable tools require MCP SDK compatibility update', { sdkVersion: 'unknown' }, Loggers.LOG_TASKS);
+    }
+    if (sdkCallToolHandler) {
+        server.server.setRequestHandler(CallToolRequestSchema, async (request, extra) => {
+            const parsedExtra = parseHandlerExtra(extra);
+            const requestId = parsedExtra?.requestId !== undefined
+                ? String(parsedExtra.requestId)
+                : randomUUID();
+            return runWithRequestContext({
+                requestId,
+                operationId: requestId,
+                ...(parsedExtra?.sessionId
+                    ? { sessionId: parsedExtra.sessionId }
+                    : {}),
+            }, () => {
+                const toolName = request.params.name;
+                // Only intercept task-capable tools managed by the local task registry.
+                // Delegate all other tools to the SDK handler to avoid shadowing future tools.
+                if (!hasTaskCapableTool(server, toolName)) {
+                    return sdkCallToolHandler(request, extra);
+                }
+                const parsed = parseExtendedCallToolRequest(request);
+                const context = resolveToolCallContext(parsedExtra, parsed.params._meta);
+                logDebug('Intercepted task-capable tool call', {
+                    tool: toolName,
+                    taskRequested: parsed.params.task !== undefined,
+                    hasProgressToken: parsed.params._meta?.progressToken !== undefined,
+                }, Loggers.LOG_TASKS);
+                return handleToolCallRequest(server, parsed, context);
+            });
+        });
+    }
+    server.server.setRequestHandler(TaskGetSchema, (request, extra) => {
+        const { taskId } = request.params;
+        const { ownerKey } = resolveOwnerScopedExtra(extra);
+        logDebug('tasks/get requested', { taskId }, Loggers.LOG_TASKS);
+        const task = taskManager.getTask(taskId, ownerKey);
+        if (!task)
+            throwTaskNotFound();
+        return toTaskSummary(task);
+    });
+    server.server.setRequestHandler(TaskResultSchema, async (request, extra) => {
+        const { taskId } = request.params;
+        const { parsedExtra, ownerKey } = resolveOwnerScopedExtra(extra);
+        logDebug('tasks/result requested', { taskId }, Loggers.LOG_TASKS);
+        const task = await taskManager.waitForTerminalTask(taskId, ownerKey, parsedExtra?.signal);
+        if (!task)
+            throwTaskNotFound();
+        try {
+            if (task.status === 'cancelled') {
+                throwStoredTaskError(task);
+            }
+            if (task.status === 'failed') {
+                if (task.error) {
+                    throwStoredTaskError(task);
+                }
+                const failedResult = (task.result ?? null);
+                if (failedResult) {
+                    return withRelatedTaskMeta(failedResult, task.taskId);
+                }
+                throwStoredTaskError(task);
+            }
+            const result = isServerResult(task.result)
+                ? task.result
+                : { content: [] };
+            return withRelatedTaskMeta(result, task.taskId);
+        }
+        finally {
+            // Shrink TTL only after the result has been fully constructed and
+            // is about to be delivered — avoids premature expiry if result
+            // construction throws.
+            taskManager.shrinkTtlAfterDelivery(taskId);
+        }
+    });
+    server.server.setRequestHandler(TaskListSchema, (request, extra) => {
+        const { ownerKey } = resolveOwnerScopedExtra(extra);
+        const cursor = request.params?.cursor;
+        logDebug('tasks/list requested', { hasCursor: cursor !== undefined }, Loggers.LOG_TASKS);
+        const { tasks, nextCursor } = taskManager.listTasks(cursor === undefined ? { ownerKey } : { ownerKey, cursor });
+        return {
+            tasks: tasks.map((task) => toTaskSummary(task)),
+            nextCursor,
+        };
+    });
+    server.server.setRequestHandler(TaskCancelSchema, (request, extra) => {
+        const { taskId } = request.params;
+        const { ownerKey } = resolveOwnerScopedExtra(extra);
+        logDebug('tasks/cancel requested', { taskId }, Loggers.LOG_TASKS);
+        const task = taskManager.cancelTask(taskId, ownerKey);
+        if (!task)
+            throwTaskNotFound();
+        abortTaskExecution(taskId);
+        emitTaskStatusNotification(server, task);
+        return toTaskSummary(task);
+    });
+    return {
+        interceptedToolsCall: sdkCallToolHandler !== null,
+        taskCapableToolsRegistered,
+    };
+}
 const DEFAULT_TTL_MS = 60_000;
 const MIN_TTL_MS = 1_000;
 const MAX_TTL_MS = 86_400_000;
@@ -283,7 +730,7 @@ class TaskManager {
         return decoded.anchorTaskId;
     }
     async waitForTerminalTask(taskId, ownerKey, signal) {
-        return waitForTerminalTaskWithDeadline({
+        return waitForTerminalTask({
             taskId,
             ownerKey,
             ...(signal && { signal }),
@@ -306,7 +753,6 @@ class TaskManager {
         }
     }
 }
-export const taskManager = new TaskManager();
 const MAX_CURSOR_LENGTH = 256;
 const MAX_ANCHOR_ID_LENGTH = 128;
 const CURSOR_SECRET = randomBytes(32);
@@ -344,3 +790,292 @@ export function decodeTaskCursor(cursor) {
         return null;
     }
 }
+export function compact(obj) {
+    const result = {};
+    for (const key of Object.keys(obj)) {
+        if (obj[key] !== undefined) {
+            result[key] = obj[key];
+        }
+    }
+    return result;
+}
+function normalizeSendNotification(sendNotification) {
+    if (typeof sendNotification !== 'function')
+        return undefined;
+    const notify = sendNotification;
+    return async (notification) => {
+        await Promise.resolve(notify(notification));
+    };
+}
+function normalizeAuthInfo(authInfo) {
+    if (!isObject(authInfo))
+        return undefined;
+    const { clientId, token } = authInfo;
+    const normalized = {};
+    if (typeof clientId === 'string')
+        normalized.clientId = clientId;
+    if (typeof token === 'string')
+        normalized.token = token;
+    return normalized.clientId || normalized.token ? normalized : undefined;
+}
+export function parseHandlerExtra(extra) {
+    if (!isObject(extra))
+        return undefined;
+    const parsed = {};
+    const { authInfo, signal, requestId, sendNotification } = extra;
+    const sessionId = resolveSessionIdFromExtra(extra);
+    if (sessionId)
+        parsed.sessionId = sessionId;
+    const normalizedAuthInfo = normalizeAuthInfo(authInfo);
+    if (normalizedAuthInfo) {
+        parsed.authInfo = normalizedAuthInfo;
+    }
+    if (signal instanceof AbortSignal)
+        parsed.signal = signal;
+    if (typeof requestId === 'string' || typeof requestId === 'number') {
+        parsed.requestId = requestId;
+    }
+    const normalizedSendNotification = normalizeSendNotification(sendNotification);
+    if (normalizedSendNotification) {
+        parsed.sendNotification = normalizedSendNotification;
+    }
+    return parsed;
+}
+export function buildAuthenticatedOwnerKey(authInfo) {
+    const authClientId = typeof authInfo?.clientId === 'string' ? authInfo.clientId : '';
+    const authToken = typeof authInfo?.token === 'string' ? authInfo.token : '';
+    if (authClientId || authToken) {
+        return `auth:${hash('sha256', `${authClientId}:${authToken}`, 'hex')}`;
+    }
+    return undefined;
+}
+export function resolveTaskOwnerKey(extra) {
+    const authenticatedOwnerKey = buildAuthenticatedOwnerKey(extra?.authInfo);
+    if (authenticatedOwnerKey)
+        return authenticatedOwnerKey;
+    if (extra?.sessionId) {
+        return (resolveMcpSessionOwnerKey(extra.sessionId) ?? `session:${extra.sessionId}`);
+    }
+    return 'default';
+}
+function resolveRequestIdFromExtra(extra) {
+    if (!isObject(extra))
+        return undefined;
+    const { requestId } = extra;
+    if (typeof requestId === 'string')
+        return requestId;
+    if (typeof requestId === 'number')
+        return String(requestId);
+    return undefined;
+}
+function getHeaderString(headers, name) {
+    const value = headers[name];
+    if (typeof value === 'string')
+        return value;
+    if (!Array.isArray(value))
+        return undefined;
+    return value.find((entry) => typeof entry === 'string');
+}
+function resolveSessionIdFromExtra(extra) {
+    if (!isObject(extra))
+        return undefined;
+    const { sessionId } = extra;
+    if (typeof sessionId === 'string')
+        return sessionId;
+    const { requestInfo } = extra;
+    if (!isObject(requestInfo))
+        return undefined;
+    const { headers } = requestInfo;
+    if (!isObject(headers))
+        return undefined;
+    return (getHeaderString(headers, 'mcp-session-id') ??
+        getHeaderString(headers, 'x-mcp-session-id'));
+}
+function resolveToolExecutionContext(extra, requestMeta) {
+    return compact({
+        ownerKey: resolveTaskOwnerKey(extra),
+        sessionId: extra?.sessionId,
+        signal: extra?.signal,
+        requestId: extra?.requestId,
+        sendNotification: extra?.sendNotification,
+        requestMeta: sanitizeToolCallMeta(requestMeta),
+    });
+}
+export function resolveToolCallContext(extra, requestMeta) {
+    return resolveToolExecutionContext(extra, requestMeta);
+}
+export function buildToolHandlerExtra(context, requestMeta) {
+    return compact({
+        signal: context.signal,
+        requestId: context.requestId,
+        sendNotification: context.sendNotification,
+        _meta: sanitizeToolCallMeta(requestMeta ?? context.requestMeta),
+    });
+}
+export function withRequestContextIfMissing(handler) {
+    return async (params, extra) => {
+        const existingRequestId = getRequestId();
+        if (existingRequestId) {
+            return handler(params, extra);
+        }
+        const derivedRequestId = resolveRequestIdFromExtra(extra) ?? randomUUID();
+        const derivedSessionId = resolveSessionIdFromExtra(extra);
+        return runWithRequestContext({
+            requestId: derivedRequestId,
+            operationId: derivedRequestId,
+            ...(derivedSessionId ? { sessionId: derivedSessionId } : {}),
+        }, () => handler(params, extra));
+    };
+}
+export function isServerResult(value) {
+    return isObject(value) && Array.isArray(value['content']);
+}
+const taskCapableToolsByServer = new WeakMap();
+function getServerToolMap(server) {
+    let toolMap = taskCapableToolsByServer.get(server);
+    if (toolMap)
+        return toolMap;
+    toolMap = new Map();
+    taskCapableToolsByServer.set(server, toolMap);
+    registerServerLifecycleCleanup(server, () => {
+        taskCapableToolsByServer.delete(server);
+    });
+    return toolMap;
+}
+export function registerTaskCapableTool(server, descriptor) {
+    getServerToolMap(server).set(descriptor.name, {
+        ...descriptor,
+        taskSupport: descriptor.taskSupport ?? 'optional',
+    });
+}
+export function unregisterTaskCapableTool(server, name) {
+    getServerToolMap(server).delete(name);
+}
+export function getTaskCapableTool(server, name) {
+    return getServerToolMap(server).get(name);
+}
+export function getTaskCapableToolSupport(server, name) {
+    return getServerToolMap(server).get(name)?.taskSupport;
+}
+export function hasTaskCapableTool(server, name) {
+    return getServerToolMap(server).has(name);
+}
+export function hasRegisteredTaskCapableTools(server) {
+    return getServerToolMap(server).size > 0;
+}
+export function setTaskCapableToolSupport(server, name, support) {
+    const descriptor = getServerToolMap(server).get(name);
+    if (!descriptor)
+        return;
+    descriptor.taskSupport = support;
+}
+export class TaskWaiterRegistry {
+    isTerminalStatus;
+    waiters = new Map();
+    constructor(isTerminalStatus) {
+        this.isTerminalStatus = isTerminalStatus;
+    }
+    add(taskId, waiter) {
+        let set = this.waiters.get(taskId);
+        if (!set) {
+            set = new Set();
+            this.waiters.set(taskId, set);
+        }
+        set.add(waiter);
+    }
+    remove(taskId, waiter) {
+        if (!waiter)
+            return;
+        const set = this.waiters.get(taskId);
+        if (!set)
+            return;
+        set.delete(waiter);
+        if (set.size === 0) {
+            this.waiters.delete(taskId);
+        }
+    }
+    notify(task) {
+        if (!this.isTerminalStatus(task.status))
+            return;
+        const waiters = this.waiters.get(task.taskId);
+        if (!waiters)
+            return;
+        this.waiters.delete(task.taskId);
+        for (const waiter of waiters)
+            waiter(task);
+    }
+}
+export async function waitForTerminalTask(options) {
+    const task = options.lookupTask(options.taskId, options.ownerKey);
+    if (!task)
+        return undefined;
+    if (options.isTerminalStatus(task.status))
+        return task;
+    const deadlineMs = task._createdAtMs + task.ttl;
+    const { promise, resolve, reject } = Promise.withResolvers();
+    const resolveInContext = AsyncLocalStorage.bind((value) => {
+        resolve(value);
+    });
+    const rejectInContext = AsyncLocalStorage.bind((error) => {
+        reject(toError(error));
+    });
+    let settled = false;
+    let waiter = null;
+    let deadlineTimeout;
+    const cleanup = () => {
+        if (deadlineTimeout) {
+            deadlineTimeout.cancel();
+            deadlineTimeout = undefined;
+        }
+        if (options.signal) {
+            options.signal.removeEventListener('abort', onAbort);
+        }
+    };
+    const settleOnce = (fn) => {
+        if (settled)
+            return;
+        settled = true;
+        fn();
+    };
+    const onAbort = () => {
+        settleOnce(() => {
+            cleanup();
+            options.registry.remove(options.taskId, waiter);
+            rejectInContext(createMcpError(ErrorCode.ConnectionClosed, 'Request was cancelled'));
+        });
+    };
+    waiter = (updated) => {
+        settleOnce(() => {
+            cleanup();
+            if (updated.ownerKey !== options.ownerKey) {
+                resolveInContext(undefined);
+                return;
+            }
+            resolveInContext(updated);
+        });
+    };
+    if (options.signal?.aborted) {
+        onAbort();
+        return promise;
+    }
+    options.registry.add(options.taskId, waiter);
+    if (options.signal) {
+        options.signal.addEventListener('abort', onAbort, { once: true });
+    }
+    const timeoutMs = Math.max(0, deadlineMs - Date.now());
+    deadlineTimeout = createUnrefTimeout(timeoutMs, { timeout: true });
+    void deadlineTimeout.promise
+        .then(() => {
+        settleOnce(() => {
+            cleanup();
+            options.registry.remove(options.taskId, waiter);
+            options.removeTask(options.taskId);
+            rejectInContext(createMcpError(ErrorCode.InvalidParams, 'Task expired', {
+                taskId: options.taskId,
+            }));
+        });
+    })
+        .catch(rejectInContext);
+    return promise;
+}
+export const taskManager = new TaskManager();