fix: Actor run cancellation handling and add tests for abort scenarios (#662)

* fix: Actor run cancellation handling and add tests for abort scenarios

* fix: Enhance actor cancellation handling and add test for aborted signal scenario

Author: jirispilka

src/tools/core/actor_execution.ts

@@ -59,34 +59,63 @@ export async function callActorGetDataset(options: {
     const CLIENT_ABORT = Symbol('CLIENT_ABORT'); // Just an internal symbol to identify client abort
     const actorClient = apifyClient.actor(actorName);
 
+    const abortActorRun = async (runId: string) => {
+        try {
+            await apifyClient.run(runId).abort({ gracefully: false });
+        } catch (e) {
+            logHttpError(e, 'Error aborting Actor run', { runId });
+        }
+    };
+
+    // The Actor start request itself is not tied to our AbortSignal, so a client can cancel
+    // before we even create the run. In that case we should exit without starting follow-up work.
+    if (abortSignal?.aborted) {
+        log.info('Actor run aborted by client before start', { actorName, mcpSessionId, input: redactSkyfirePayId(input) });
+        return null;
+    }
+
     // Start the actor run
     const actorRun: ActorRun = await actorClient.start(input, callOptions);
 
+    // Cancellation can also arrive while actorClient.start() is still in flight. Once start()
+    // returns we finally have a run ID, so we must immediately check again and abort that run
+    // ourselves; otherwise the abort event is missed and the run continues in the background.
+    if (abortSignal?.aborted) {
+        await abortActorRun(actorRun.id);
+
+        log.info('Actor run aborted by client', { actorName, mcpSessionId, input: redactSkyfirePayId(input) });
+        return null;
+    }
+
     // Start progress tracking if a tracker is provided
     if (progressTracker) {
         progressTracker.startActorRunUpdates(actorRun.id, apifyClient, actorName);
     }
 
-    // Create abort promise that handles both API abort and race rejection
-    const abortPromise = async () => new Promise<typeof CLIENT_ABORT>((resolve) => {
-        abortSignal?.addEventListener('abort', async () => {
-            // Abort the actor run via API
-            try {
-                await apifyClient.run(actorRun.id).abort({ gracefully: false });
-            } catch (e) {
-                logHttpError(e, 'Error aborting Actor run', { runId: actorRun.id });
-            }
-            // Reject to stop waiting
+    // Resolve the race immediately on cancellation and abort the Actor run in the background.
+    // If we waited for the abort API call to finish first, waitForFinish() could win the race
+    // and the run might complete before we treat the request as cancelled.
+    let abortListener: (() => void) | undefined;
+    const abortPromise = new Promise<typeof CLIENT_ABORT>((resolve) => {
+        abortListener = () => {
             resolve(CLIENT_ABORT);
-        }, { once: true });
+            void abortActorRun(actorRun.id);
+        };
+        abortSignal?.addEventListener('abort', abortListener, { once: true });
     });
 
     // Wait for completion or cancellation
     const potentialAbortedRun = await Promise.race([
         apifyClient.run(actorRun.id).waitForFinish(),
-        ...(abortSignal ? [abortPromise()] : []),
+        ...(abortSignal ? [abortPromise] : []),
     ]);
 
+    // Clean up the abort listener if waitForFinish() won the race, so a later signal
+    // doesn't try to abort an already-completed run.
+    if (abortListener) {
+        abortSignal?.removeEventListener('abort', abortListener);
+    }
+
     if (potentialAbortedRun === CLIENT_ABORT) {
         log.info('Actor run aborted by client', { actorName, mcpSessionId, input: redactSkyfirePayId(input) });
         return null;

tests/unit/tools.actor_execution.test.ts

@@ -0,0 +1,113 @@
+import { describe, expect, it, vi } from 'vitest';
+
+import { callActorGetDataset } from '../../src/tools/core/actor_execution.js';
+
+describe('callActorGetDataset', () => {
+    it('should return null without starting a run when the signal is already aborted', async () => {
+        const controller = new AbortController();
+        controller.abort();
+
+        const start = vi.fn();
+
+        const apifyClient = {
+            actor: vi.fn().mockReturnValue({ start }),
+        };
+
+        const result = await callActorGetDataset({
+            actorName: 'apify/rag-web-browser',
+            input: { query: 'https://apify.com' },
+            apifyClient: apifyClient as never,
+            abortSignal: controller.signal,
+        });
+
+        expect(result).toBeNull();
+        expect(start).not.toHaveBeenCalled();
+    });
+
+    it('should abort and return null when the signal is aborted during actor start', async () => {
+        const controller = new AbortController();
+        const actorRun = {
+            id: 'run-123',
+            defaultDatasetId: 'dataset-123',
+            usageTotalUsd: 0,
+            usageUsd: {},
+        };
+
+        const abort = vi.fn().mockResolvedValue(undefined);
+        const waitForFinish = vi.fn().mockResolvedValue(actorRun);
+        const start = vi.fn().mockImplementation(async () => {
+            controller.abort();
+            return actorRun;
+        });
+
+        const apifyClient = {
+            actor: vi.fn().mockReturnValue({
+                start,
+            }),
+            run: vi.fn().mockReturnValue({
+                abort,
+                waitForFinish,
+            }),
+        };
+
+        const result = await callActorGetDataset({
+            actorName: 'apify/rag-web-browser',
+            input: { query: 'https://apify.com' },
+            apifyClient: apifyClient as never,
+            abortSignal: controller.signal,
+        });
+
+        expect(result).toBeNull();
+        expect(abort).toHaveBeenCalledWith({ gracefully: false });
+        expect(waitForFinish).not.toHaveBeenCalled();
+    });
+
+    it('should return null immediately when cancelled after start even if API abort is slow', async () => {
+        const controller = new AbortController();
+        const actorRun = {
+            id: 'run-456',
+            defaultDatasetId: 'dataset-456',
+            usageTotalUsd: 0,
+            usageUsd: {},
+        };
+
+        let resolveAbort: (() => void) | undefined;
+        let resolveWaitForFinish: ((value: typeof actorRun) => void) | undefined;
+        const abort = vi.fn().mockImplementation(async () => await new Promise<void>((resolve) => {
+            resolveAbort = resolve;
+        }));
+        const waitForFinish = vi.fn().mockImplementation(async () => await new Promise<typeof actorRun>((resolve) => {
+            resolveWaitForFinish = resolve;
+        }));
+        const start = vi.fn().mockResolvedValue(actorRun);
+        const listItems = vi.fn().mockResolvedValue({ items: [], total: 0 });
+
+        const apifyClient = {
+            actor: vi.fn().mockReturnValue({ start }),
+            run: vi.fn().mockReturnValue({
+                abort,
+                waitForFinish,
+            }),
+            dataset: vi.fn().mockReturnValue({ listItems }),
+        };
+
+        const resultPromise = callActorGetDataset({
+            actorName: 'apify/rag-web-browser',
+            input: { query: 'https://apify.com' },
+            apifyClient: apifyClient as never,
+            abortSignal: controller.signal,
+        });
+
+        await vi.waitUntil(() => Boolean(resolveWaitForFinish));
+        controller.abort();
+        resolveWaitForFinish!(actorRun);
+
+        const result = await resultPromise;
+
+        expect(result).toBeNull();
+        await vi.waitUntil(() => abort.mock.calls.length > 0);
+        expect(abort).toHaveBeenCalledWith({ gracefully: false });
+        expect(listItems).not.toHaveBeenCalled();
+        resolveAbort!();
+    });
+});