Browse Source

fix: recover timed-out subagent sessions

Zerdeşt Taifour 1 month ago
parent
commit
d1c196b4ae

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

@@ -207,6 +207,108 @@ describe('task-session-manager hook', () => {
     );
   });
 
+  test('reuses timed-out running aliases for safe recovery', async () => {
+    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: 'explorer',
+          description: 'map timed out session',
+        },
+      },
+    );
+    await hook['tool.execute.after'](
+      { tool: 'task', sessionID: 'parent-1', callID: 'call-1' },
+      {
+        output: [
+          'task_id: child-1',
+          'state: running',
+          '',
+          '<task_result>',
+          'Timed out after 120000ms while waiting for task completion.',
+          '</task_result>',
+        ].join('\n'),
+      },
+    );
+
+    expect(
+      board.resolveRecoverable('parent-1', 'exp-1', 'explorer')?.taskID,
+    ).toBe('child-1');
+
+    const resume = {
+      args: { subagent_type: 'explorer', task_id: 'exp-1' },
+    };
+    await hook['tool.execute.before'](
+      { tool: 'task', sessionID: 'parent-1', callID: 'resume-1' },
+      resume,
+    );
+
+    expect(resume.args.task_id).toBe('child-1');
+    expect(board.get('child-1')).toMatchObject({
+      state: 'running',
+      timedOut: true,
+    });
+  });
+
+  test('busy timeout recovery clears timeout overlay from prompt', async () => {
+    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: 'explorer',
+          description: 'recover timed out child',
+        },
+      },
+    );
+    await hook['tool.execute.after'](
+      { tool: 'task', sessionID: 'parent-1', callID: 'call-1' },
+      {
+        output: [
+          'task_id: child-1',
+          'state: running',
+          '',
+          '<task_result>',
+          'Timed out after 120000ms while waiting for task completion.',
+          '</task_result>',
+        ].join('\n'),
+      },
+    );
+
+    const beforeMessages = createMessages('parent-1', 'before busy');
+    await hook['experimental.chat.messages.transform']({}, beforeMessages);
+    expect(beforeMessages.messages[0].parts[0].text).toContain(
+      'running, timed out',
+    );
+
+    await hook.event({
+      event: {
+        type: 'session.status',
+        properties: {
+          sessionID: 'child-1',
+          status: { type: 'busy' },
+        },
+      },
+    });
+
+    expect(board.get('child-1')).toMatchObject({
+      state: 'running',
+      timedOut: false,
+      statusUncertain: false,
+    });
+
+    const afterMessages = createMessages('parent-1', 'after busy');
+    await hook['experimental.chat.messages.transform']({}, afterMessages);
+    expect(afterMessages.messages[0].parts[0].text).not.toContain(
+      'running, timed out',
+    );
+  });
+
   test('updates background job board from injected completion messages', async () => {
     const board = new BackgroundJobBoard();
     const { hook } = createHook({ backgroundJobBoard: board });

+ 11 - 5
src/hooks/task-session-manager/index.ts

@@ -511,11 +511,17 @@ export function createTaskSessionManagerHook(
       }
 
       const requested = args.task_id.trim();
-      const remembered = backgroundJobBoard.resolveReusable(
-        input.sessionID,
-        requested,
-        agentType,
-      );
+      const remembered =
+        backgroundJobBoard.resolveReusable(
+          input.sessionID,
+          requested,
+          agentType,
+        ) ??
+        backgroundJobBoard.resolveRecoverable(
+          input.sessionID,
+          requested,
+          agentType,
+        );
 
       if (!remembered) {
         if (RAW_SESSION_ID_PATTERN.test(requested)) {

+ 63 - 0
src/multiplexer/session-manager.test.ts

@@ -410,6 +410,69 @@ describe('MultiplexerSessionManager', () => {
       expect(mockMultiplexer.closePane).not.toHaveBeenCalled();
     });
 
+    test('timed out running jobs still close after safe recovery and completion', async () => {
+      const ctx = createMockContext();
+      const board = new BackgroundJobBoard();
+      board.registerLaunch({
+        taskID: 'timedout-child',
+        parentSessionID: 'parent-1',
+        agent: 'explorer',
+      });
+      board.updateStatus({
+        taskID: 'timedout-child',
+        state: 'running',
+        timedOut: true,
+        now: 100,
+      });
+      mockMultiplexer.spawnPane.mockResolvedValue({
+        success: true,
+        paneId: 'p-timedout-child',
+      });
+      const manager = new MultiplexerSessionManager(
+        ctx,
+        defaultMultiplexerConfig,
+        board,
+      );
+      board.setTerminalStateListener((taskID) => {
+        void manager.retryDeferredIdleClose(taskID);
+      });
+
+      await manager.onSessionCreated({
+        type: 'session.created',
+        properties: {
+          info: { id: 'timedout-child', parentID: 'parent-1' },
+        },
+      });
+
+      await manager.onSessionStatus({
+        type: 'session.status',
+        properties: {
+          sessionID: 'timedout-child',
+          status: { type: 'idle' },
+        },
+      });
+
+      expect(mockMultiplexer.closePane).not.toHaveBeenCalled();
+
+      board.markRunningFromLiveSession('timedout-child', 200);
+      expect(board.get('timedout-child')).toMatchObject({
+        state: 'running',
+        timedOut: false,
+        lastLiveBusyAt: 200,
+      });
+
+      board.updateStatus({
+        taskID: 'timedout-child',
+        state: 'completed',
+        resultSummary: 'done',
+      });
+      await Promise.resolve();
+
+      expect(mockMultiplexer.closePane).toHaveBeenCalledWith(
+        'p-timedout-child',
+      );
+    });
+
     test('deferred idle closes retry on terminal status updates', async () => {
       for (const state of ['completed', 'error', 'cancelled'] as const) {
         resetMultiplexerSessionManagerState();

+ 73 - 3
src/utils/background-job-board.test.ts

@@ -164,7 +164,7 @@ describe('BackgroundJobBoard', () => {
     expect(board.resolveReusable('parent-1', 'ses_error')).toBeUndefined();
   });
 
-  test('prompt tells orchestrator to reuse completed sessions only', () => {
+  test('prompt distinguishes reusable and recoverable sessions', () => {
     const board = new BackgroundJobBoard();
     board.registerLaunch({
       taskID: 'ses_1',
@@ -174,10 +174,24 @@ describe('BackgroundJobBoard', () => {
     });
     board.updateStatus({ taskID: 'ses_1', state: 'completed' });
     board.markReconciled('ses_1');
+    board.registerLaunch({
+      taskID: 'ses_2',
+      parentSessionID: 'parent-1',
+      agent: 'explorer',
+      description: 'recover timed out task',
+    });
+    board.updateStatus({ taskID: 'ses_2', state: 'running', timedOut: true });
 
-    expect(board.formatForPrompt('parent-1')).toContain(
-      'Reuse only completed sessions',
+    const prompt = board.formatForPrompt('parent-1');
+
+    expect(prompt).toContain(
+      'Completed or reconciled sessions are reusable by alias',
+    );
+    expect(prompt).toContain(
+      'Timed-out running sessions are recoverable by alias for safe resume after a live busy signal.',
     );
+    expect(prompt).toContain('Cancelled or errored sessions are not reusable.');
+    expect(prompt).toContain('exp-1 / ses_2 / explorer / running, timed out');
   });
 
   test('does not reconcile running jobs', () => {
@@ -570,6 +584,62 @@ describe('BackgroundJobBoard', () => {
     expect(updated?.completedAt).toBeDefined();
   });
 
+  test('live busy recovery clears timeout state on running jobs', () => {
+    const board = new BackgroundJobBoard();
+    board.registerLaunch({
+      taskID: 'ses_1',
+      parentSessionID: 'parent-1',
+      agent: 'explorer',
+      now: 100,
+    });
+    board.updateStatus({
+      taskID: 'ses_1',
+      state: 'running',
+      timedOut: true,
+      statusUncertain: true,
+      now: 150,
+    });
+
+    const updated = board.markRunningFromLiveSession('ses_1', 200);
+
+    expect(updated).toMatchObject({
+      state: 'running',
+      timedOut: false,
+      statusUncertain: false,
+      lastLiveBusyAt: 200,
+      updatedAt: 200,
+      alias: 'exp-1',
+    });
+  });
+
+  test('resolves timed-out running jobs for safe recovery only', () => {
+    const board = new BackgroundJobBoard();
+    board.registerLaunch({
+      taskID: 'ses_1',
+      parentSessionID: 'parent-1',
+      agent: 'explorer',
+    });
+    board.updateStatus({
+      taskID: 'ses_1',
+      state: 'running',
+      timedOut: true,
+    });
+
+    expect(
+      board.resolveReusable('parent-1', 'exp-1', 'explorer'),
+    ).toBeUndefined();
+    expect(
+      board.resolveRecoverable('parent-1', 'exp-1', 'explorer'),
+    ).toMatchObject({
+      taskID: 'ses_1',
+      state: 'running',
+      timedOut: true,
+    });
+    expect(
+      board.resolveRecoverable('parent-1', 'exp-1', 'oracle'),
+    ).toBeUndefined();
+  });
+
   test('stale status updates cannot reopen already reconciled jobs', () => {
     const board = new BackgroundJobBoard();
     board.registerLaunch({

+ 18 - 1
src/utils/background-job-board.ts

@@ -220,6 +220,8 @@ export class BackgroundJobBoard {
       ...existing,
       updatedAt: now,
       lastLiveBusyAt: now,
+      timedOut: false,
+      statusUncertain: false,
     };
 
     this.jobs.set(taskID, updated);
@@ -314,6 +316,18 @@ export class BackgroundJobBoard {
     return job;
   }
 
+  resolveRecoverable(
+    parentSessionID: string,
+    taskIDOrAlias: string,
+    agent?: string,
+  ): BackgroundJobRecord | undefined {
+    const job = this.resolve(parentSessionID, taskIDOrAlias);
+    if (!job) return undefined;
+    if (agent && job.agent !== agent) return undefined;
+    if (job.state !== 'running' || !job.timedOut) return undefined;
+    return job;
+  }
+
   markUsed(parentSessionID: string, key: string, now = Date.now()): void {
     const job = this.resolve(parentSessionID, key);
     if (!job) return;
@@ -381,7 +395,10 @@ export class BackgroundJobBoard {
     return [
       '### Background Job Board',
       'SENTINEL: background-job-board-v2',
-      'Do not poll running jobs. Wait for hook-driven completion, or use cancel_task only for explicit cancellation. Reconcile terminal jobs before final response. Reuse only completed sessions for the same specialist/context; never reuse cancelled or errored sessions.',
+      'Do not poll running jobs. Wait for hook-driven completion, or use cancel_task only for explicit cancellation. Reconcile terminal jobs before final response.',
+      'Completed or reconciled sessions are reusable by alias for the same specialist/context.',
+      'Timed-out running sessions are recoverable by alias for safe resume after a live busy signal.',
+      'Cancelled or errored sessions are not reusable.',
       '',
       '#### Active / Unreconciled',
       ...(active.length > 0