test: Pin sync/task tool-call contracts before dispatch dedup

tests/unit/helpers/mcp_server.ts

@@ -0,0 +1,106 @@
+import { InMemoryTaskStore } from '@modelcontextprotocol/sdk/experimental/tasks/stores/in-memory.js';
+
+import type { ALLOWED_TASK_TOOL_EXECUTION_MODES } from '../../../src/const.js';
+import type { ActorsMcpServer } from '../../../src/mcp/server.js';
+import type { ActorsMcpServerOptions, InternalToolArgs, ToolEntry, ToolInputSchema } from '../../../src/types.js';
+import { TOOL_TYPE } from '../../../src/types.js';
+import { compileSchema } from '../../../src/utils/ajv.js';
+
+/**
+ * Signature of an SDK request handler reached via the private `_requestHandlers` map. The
+ * `mcp.server.*` tests drive these handlers directly (no transport, no `server.request()`).
+ */
+export type HandlerFn = (
+    req: Record<string, unknown>,
+    extra: Record<string, unknown>,
+) => Promise<Record<string, unknown>>;
+
+/**
+ * Returns the real request handler the SDK registered for `method` (e.g. 'tools/call',
+ * 'tasks/result'), reached through the server's private `_requestHandlers` map so a test can invoke
+ * it directly. Throws if the handler is not registered. This reach into an SDK-internal seam is
+ * centralized here so an SDK upgrade only needs one fix.
+ */
+export function getRequestHandler(server: unknown, method: string): HandlerFn {
+    // eslint-disable-next-line no-underscore-dangle
+    const handler = (server as { server: { _requestHandlers: Map<string, HandlerFn> } }).server._requestHandlers.get(
+        method,
+    );
+    if (!handler) throw new Error(`Handler "${method}" not registered`);
+    return handler;
+}
+
+/**
+ * Constructs a real `ActorsMcpServer` backed by an `InMemoryTaskStore`, runs `run` against it, and
+ * always closes it. Defaults match the existing `mcp.server.*` tests (telemetry off, placeholder
+ * token); pass `options` to override (e.g. telemetry on with no token for the shape tests).
+ */
+export async function withServer<T>(
+    run: (server: ActorsMcpServer) => Promise<T>,
+    options?: Partial<ActorsMcpServerOptions>,
+): Promise<T> {
+    const { ActorsMcpServer } = await import('../../../src/mcp/server.js');
+    const server = new ActorsMcpServer({
+        taskStore: new InMemoryTaskStore(),
+        setupSigintHandler: false,
+        telemetry: { enabled: false },
+        token: 'fake-token',
+        ...options,
+    });
+    try {
+        return await run(server);
+    } finally {
+        await server.close();
+    }
+}
+
+/**
+ * A synthetic internal tool whose `call` throws `error` (default: a plain `Error('boom')`), so
+ * dispatch falls through to the outer catch. An empty input schema validates against `{}`. Set
+ * `taskSupport` to make the tool eligible for the task path (it otherwise fails the pre-dispatch gate).
+ */
+export function makeThrowingTool(
+    options: { name?: string; error?: unknown; taskSupport?: (typeof ALLOWED_TASK_TOOL_EXECUTION_MODES)[number] } = {},
+): ToolEntry {
+    const { name = 'test-throwing-tool', error = new Error('boom'), taskSupport } = options;
+    return {
+        type: TOOL_TYPE.INTERNAL,
+        name,
+        description: 'throws',
+        inputSchema: { type: 'object', properties: {} } as ToolInputSchema,
+        ajvValidate: compileSchema({ type: 'object', properties: {} }),
+        ...(taskSupport ? { execution: { taskSupport } } : {}),
+        call: async (_toolArgs: InternalToolArgs) => {
+            throw error;
+        },
+    };
+}
+
+/**
+ * A synthetic internal tool that records what the server passed into `call` (whether it ran, and the
+ * `progressTracker` it received). Generalizes to any "did the server pass X to the tool?" assertion.
+ */
+export function makeRecorderTool(name: string): {
+    tool: ToolEntry;
+    received: { called: boolean; progressTracker: InternalToolArgs['progressTracker'] | undefined };
+} {
+    const received: { called: boolean; progressTracker: InternalToolArgs['progressTracker'] | undefined } = {
+        called: false,
+        progressTracker: undefined,
+    };
+    const tool: ToolEntry = {
+        type: TOOL_TYPE.INTERNAL,
+        name,
+        description: 'recorder tool for progress wiring tests',
+        inputSchema: { type: 'object', properties: {}, additionalProperties: true },
+        ajvValidate: Object.assign(() => true, { errors: null }) as unknown as ToolEntry['ajvValidate'],
+        paymentRequired: false,
+        annotations: {},
+        call: async (toolArgs: InternalToolArgs) => {
+            received.called = true;
+            received.progressTracker = toolArgs.progressTracker;
+            return { content: [{ type: 'text', text: 'ok' }] };
+        },
+    } as ToolEntry;
+    return { tool, received };
+}

tests/unit/mcp.server.error_handling.test.ts

@@ -1,14 +1,10 @@
-import { InMemoryTaskStore } from '@modelcontextprotocol/sdk/experimental/tasks/stores/in-memory.js';
 import { afterEach, describe, expect, it, vi } from 'vitest';
 
 import log from '@apify/log';
 
 import { TOOL_STATUS } from '../../src/const.js';
-import { ActorsMcpServer } from '../../src/mcp/server.js';
-import type { InternalToolArgs, ToolEntry, ToolInputSchema } from '../../src/types.js';
-import { TOOL_TYPE } from '../../src/types.js';
-import { compileSchema } from '../../src/utils/ajv.js';
 import { getToolCallErrorUserText } from '../../src/utils/mcp.js';
+import { getRequestHandler, makeThrowingTool, withServer } from './helpers/mcp_server.js';
 
 /**
  * Covers the `server.onerror` wiring in `setupErrorHandling()`: client faults softFail with a
@@ -19,13 +15,7 @@ describe('ActorsMcpServer onerror', () => {
     afterEach(() => vi.restoreAllMocks());
 
     it('soft-fails client faults with a sanitized message and error-logs the rest', async () => {
-        const server = new ActorsMcpServer({
-            taskStore: new InMemoryTaskStore(),
-            setupSigintHandler: false,
-            telemetry: { enabled: false },
-            token: 'fake-token',
-        });
-        try {
+        await withServer(async (server) => {
             const softFail = vi.spyOn(log, 'softFail').mockImplementation(() => log);
             const errorLog = vi.spyOn(log, 'error').mockImplementation(() => log);
 
@@ -37,9 +27,7 @@ describe('ActorsMcpServer onerror', () => {
 
             server.server.onerror?.(new Error('Unexpected internal failure'));
             expect(errorLog).toHaveBeenCalledTimes(1);
-        } finally {
-            await server.close();
-        }
+        });
     });
 });
 
@@ -50,54 +38,20 @@ describe('ActorsMcpServer onerror', () => {
  * in `toolTelemetry` — no `failureCategory`/`failureHttpStatus` (which would leak onto the wire).
  */
 describe('CallToolRequestSchema handler outer catch', () => {
-    type HandlerFn = (req: Record<string, unknown>, extra: Record<string, unknown>) => Promise<Record<string, unknown>>;
-
-    // A synthetic internal tool whose call throws a plain (non-McpError) error, so dispatch falls
-    // through to the outer catch. An empty input schema validates against `{}`.
-    function makeThrowingTool(): ToolEntry {
-        return {
-            type: TOOL_TYPE.INTERNAL,
-            name: 'test-throwing-tool',
-            description: 'throws',
-            inputSchema: { type: 'object', properties: {} } as ToolInputSchema,
-            ajvValidate: compileSchema({ type: 'object', properties: {} }),
-            call: async (_toolArgs: InternalToolArgs) => {
-                throw new Error('boom');
-            },
-        };
-    }
-
-    function getToolCallHandler(server: ActorsMcpServer): HandlerFn {
-        // eslint-disable-next-line no-underscore-dangle
-        const handler = (
-            server as unknown as { server: { _requestHandlers: Map<string, HandlerFn> } }
-        ).server._requestHandlers.get('tools/call');
-        if (!handler) throw new Error('tools/call handler not registered');
-        return handler;
-    }
-
     async function dispatchThrow(aborted: boolean) {
-        const server = new ActorsMcpServer({
-            taskStore: new InMemoryTaskStore(),
-            setupSigintHandler: false,
-            telemetry: { enabled: false },
-            token: 'fake-token',
-        });
-        try {
+        return withServer(async (server) => {
             vi.spyOn(log, 'error').mockImplementation(() => log);
             vi.spyOn(log, 'exception').mockImplementation(() => log);
             server.upsertTools([makeThrowingTool()]);
-            const handler = getToolCallHandler(server);
-            return await handler(
+            const handler = getRequestHandler(server, 'tools/call');
+            return handler(
                 {
                     method: 'tools/call',
                     params: { name: 'test-throwing-tool', arguments: {}, _meta: { mcpSessionId: 's1' } },
                 },
                 { signal: { aborted }, sendNotification: vi.fn() },
             );
-        } finally {
-            await server.close();
-        }
+        });
     }
 
     it('preserves ABORTED toolStatus and emits only { toolStatus } when the request was aborted', async () => {

tests/unit/mcp.server.progress_wiring.test.ts

@@ -1,79 +1,20 @@
-import { InMemoryTaskStore } from '@modelcontextprotocol/sdk/experimental/tasks/stores/in-memory.js';
 import { describe, expect, it, vi } from 'vitest';
 
 import { HELPER_TOOLS } from '../../src/const.js';
-import type { ActorsMcpServer } from '../../src/mcp/server.js';
-import type { InternalToolArgs, ToolEntry } from '../../src/types.js';
-import { TOOL_TYPE } from '../../src/types.js';
 import { ProgressTracker } from '../../src/utils/progress.js';
+import { getRequestHandler, makeRecorderTool, withServer } from './helpers/mcp_server.js';
 
 /**
  * Covers the request-metadata → `createProgressTracker` wiring in `tools/call`. The unit tests
  * for `get-actor-run` itself inject a fake tracker directly into the tool, so they would not
  * catch a regression in the server-level opt-in.
- *
- * TODO(followup): lift `withServer`, `getCallToolHandler`, and `makeRecorderTool` to
- * `tests/unit/_helpers/` and reuse from `mcp.task_notifications.test.ts` and
- * `mcp.server.capability_gating.test.ts` (both already duplicate the same server bootstrap and
- * `_requestHandlers.get(...)` getter). The recorder pattern generalizes to any "did the server
- * pass X to the tool?" assertion (apifyClient, extra.signal, future opt-ins).
  */
 
-type HandlerFn = (req: Record<string, unknown>, extra: Record<string, unknown>) => Promise<Record<string, unknown>>;
-
-function getCallToolHandler(server: unknown): HandlerFn {
-    // eslint-disable-next-line no-underscore-dangle
-    const handler = (server as { server: { _requestHandlers: Map<string, HandlerFn> } }).server._requestHandlers.get(
-        'tools/call',
-    );
-    if (!handler) throw new Error('Handler "tools/call" not registered');
-    return handler;
-}
-
-function makeRecorderTool(name: string): {
-    tool: ToolEntry;
-    received: { progressTracker: InternalToolArgs['progressTracker'] | undefined };
-} {
-    const received: { progressTracker: InternalToolArgs['progressTracker'] | undefined } = {
-        progressTracker: undefined,
-    };
-    const tool: ToolEntry = {
-        type: TOOL_TYPE.INTERNAL,
-        name,
-        description: 'recorder tool for progress wiring tests',
-        inputSchema: { type: 'object', properties: {}, additionalProperties: true },
-        ajvValidate: Object.assign(() => true, { errors: null }) as unknown as ToolEntry['ajvValidate'],
-        paymentRequired: false,
-        annotations: {},
-        call: async (toolArgs: InternalToolArgs) => {
-            received.progressTracker = toolArgs.progressTracker;
-            return { content: [{ type: 'text', text: 'ok' }] };
-        },
-    } as ToolEntry;
-    return { tool, received };
-}
-
-async function withServer<T>(run: (server: ActorsMcpServer) => Promise<T>): Promise<T> {
-    const { ActorsMcpServer: ActorsMcpServerClass } = await import('../../src/mcp/server.js');
-    const taskStore = new InMemoryTaskStore();
-    const server = new ActorsMcpServerClass({
-        taskStore,
-        setupSigintHandler: false,
-        telemetry: { enabled: false },
-        token: 'fake-token',
-    });
-    try {
-        return await run(server);
-    } finally {
-        await server.close();
-    }
-}
-
 async function runRecorder(toolName: string, meta: Record<string, unknown>) {
     return withServer(async (server) => {
         const { tool, received } = makeRecorderTool(toolName);
         server.upsertTools([tool]);
-        const handler = getCallToolHandler(server as never);
+        const handler = getRequestHandler(server, 'tools/call');
         await handler(
             { method: 'tools/call', params: { name: toolName, arguments: {}, _meta: meta } },
             { sendNotification: vi.fn() },