Browse Source

fix: abort child session on signal cancel + early-register board job from session.created

promptWithTimeout aborted the child session on timeout but not on AbortSignal
cancel, leaving the server-side session running after the parent tool call
returned 'Task cancelled' (orphan). Also, when the parent tool is cancelled
before tool.execute.after fires (e.g. foreground fallback abort), the job is
never registered on the board, so idle reconciliation reports
runningJobForSession:false and the orchestrator sees 'Task cancelled' while
the child is still working (#765).

Changes:
- src/utils/session.ts: also call abortSessionWithTimeout when signal.aborted.
- src/utils/session.test.ts: assert abort is called on signal cancel.
- src/hooks/task-session-manager/pending-call-tracker.ts: add peekByParent.
- src/hooks/task-session-manager/index.ts: early-register the board job in
  the session.created handler by peeking the pending call (does not consume
  it; tool.execute.after can still re-register). Guards the #765 case where
  the after-hook never fires.
- src/hooks/task-session-manager/index.test.ts: regression test that
  session.created keeps the job tracked through idle even when
  tool.execute.after is skipped.

Refs #765
Jiajun0413 3 weeks ago
parent
commit
bd74e53796

+ 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 -