Browse Source

Merge pull request #793 from Jiajun0413/fix/cancelled-tool-orphan-child

Alvin 3 weeks ago
parent
commit
0bf7ba3c5f

+ 44 - 0
src/hooks/task-session-manager/index.test.ts

@@ -2281,6 +2281,50 @@ describe('task-session-manager hook', () => {
     });
   });
 
+  test('session.created early-registers board job so after-hook cancellation cannot orphan the child', async () => {
+    // Reproduces #765: parent tool may be cancelled before tool.execute.after,
+    // so the job never lands on the board. Early registration from
+    // session.created keeps runningJobForSession true and lets idle reconcile.
+    const board = new BackgroundJobBoard();
+    const { hook } = createHook({ backgroundJobBoard: board });
+
+    await hook['tool.execute.before'](
+      { tool: 'task', sessionID: 'parent-1', callID: 'call-1' },
+      {
+        args: {
+          subagent_type: 'oracle',
+          description: 'loss design review',
+        },
+      },
+    );
+
+    // Child session is created while the parent tool is still in flight.
+    await hook.event({
+      event: {
+        type: 'session.created',
+        properties: { info: { id: 'child-1', parentID: 'parent-1' } },
+      },
+    });
+
+    expect(board.get('child-1')).toMatchObject({
+      state: 'running',
+      agent: 'oracle',
+      parentSessionID: 'parent-1',
+      description: 'loss design review',
+    });
+
+    // Simulate parent tool never firing tool.execute.after (cancelled).
+    // Child goes idle after finishing — board must still reconcile.
+    await hook.event({
+      event: { type: 'session.idle', properties: { sessionID: 'child-1' } },
+    });
+
+    expect(board.get('child-1')).toMatchObject({
+      state: 'reconciled',
+      terminalState: 'completed',
+    });
+  });
+
   test('cancelled job is not reconciled from idle', async () => {
     const board = new BackgroundJobBoard();
     board.registerLaunch({

+ 29 - 0
src/hooks/task-session-manager/index.ts

@@ -660,6 +660,35 @@ export function createTaskSessionManagerHook(
           options.shouldManageSession(info.parentID)
         ) {
           taskContextTracker.pendingManagedTaskIds.add(info.id);
+          // Early board registration: if the parent tool call is cancelled
+          // before tool.execute.after (e.g. foreground fallback abort), the
+          // after-hook never fires and the job is never tracked — idle then
+          // reports runningJobForSession:false and the orchestrator sees
+          // "Task cancelled" while the child is still working (#765).
+          // Peek (don't take) so tool.execute.after can still re-register.
+          const pending = pendingCallTracker.peekByParent(info.parentID);
+          if (
+            pending &&
+            !pending.resumedTaskId &&
+            !backgroundJobBoard.get(info.id)
+          ) {
+            const record = backgroundJobBoard.registerLaunch({
+              taskID: info.id,
+              parentSessionID: pending.parentSessionId,
+              agent: pending.agentType,
+              description: pending.label,
+              objective: pending.label,
+            });
+            log(
+              '[task-session-manager] early board registration from session.created',
+              {
+                taskID: record.taskID,
+                alias: record.alias,
+                parentSessionID: record.parentSessionID,
+                agent: record.agent,
+              },
+            );
+          }
         }
         return;
       }

+ 8 - 0
src/hooks/task-session-manager/pending-call-tracker.ts

@@ -39,6 +39,14 @@ export function createPendingCallTracker() {
       return pending;
     },
 
+    /** Peek oldest pending call for a parent without removing it. */
+    peekByParent(parentSessionId: string) {
+      for (const call of pendingCalls.values()) {
+        if (call.parentSessionId === parentSessionId) return call;
+      }
+      return undefined;
+    },
+
     clearSession(sessionId: string) {
       for (const [callId, pending] of pendingCalls.entries()) {
         if (pending.parentSessionId === sessionId) {

+ 9 - 0
src/utils/session.test.ts

@@ -55,6 +55,10 @@ describe('session utilities', () => {
     await expect(
       promptWithTimeout(client, { path: { id: 's1' }, body: { parts: [] } }, 5),
     ).rejects.toThrow('Prompt timed out after 5ms');
+
+    // Abort must still be attempted even though it threw; the original
+    // timeout error is preserved.
+    expect(abort).toHaveBeenCalledWith({ path: { id: 's1' } });
   });
 
   test('promptWithTimeout honors abort signal when timeout is disabled', async () => {
@@ -78,6 +82,11 @@ describe('session utilities', () => {
         controller.signal,
       ),
     ).rejects.toThrow('Prompt cancelled');
+
+    // Signal cancel must abort the server-side session, same as timeout.
+    // Without this, the parent tool returns cancelled while the child keeps
+    // running (orphan session).
+    expect(abort).toHaveBeenCalledWith({ path: { id: 's1' } });
   });
 
   test('promptWithTimeout returns when prompt resolves with no timeout', async () => {

+ 14 - 2
src/utils/session.ts

@@ -142,11 +142,17 @@ export async function promptWithTimeout(
 
     await Promise.race(racers);
   } catch (error) {
-    if (error instanceof OperationTimeoutError) {
+    // Abort the server-side session on timeout OR signal cancel. Without
+    // the signal branch, a cancelled parent tool leaves the child running
+    // as an orphan ("Task cancelled" to the orchestrator, child still
+    // working). Match by error identity, not `signal.aborted`, so an
+    // unrelated rejection overlapping a signal abort does not fire an
+    // extra abort round-trip.
+    if (isPromptCancellationError(error)) {
       try {
         await abortSessionWithTimeout(client, sessionId);
       } catch {
-        // Best-effort cleanup: preserve the original prompt timeout error.
+        // Best-effort: preserve the original error.
       }
     }
     throw error;
@@ -156,6 +162,12 @@ export async function promptWithTimeout(
   }
 }
 
+/** OperationTimeoutError or our own "Prompt cancelled" Error. */
+function isPromptCancellationError(error: unknown): boolean {
+  if (error instanceof OperationTimeoutError) return true;
+  return error instanceof Error && error.message === 'Prompt cancelled';
+}
+
 /**
  * Result of extracting session content.
  * `empty` is true when the assistant produced zero text content -