Răsfoiți Sursa

Merge pull request #1008 from alvinunreal/fix/background-job-restart-recovery

fix: recover background jobs after restart
Alvin 1 lună în urmă
părinte
comite
282d5f26a4

+ 6 - 0
docs/background-orchestration.md

@@ -168,6 +168,12 @@ finalizing. Separately, the default-on orchestrator wake scheduler may prompt an
 idle parent with incomplete todos after continuous idle time; it does not depend
 on the local job board.
 
+After a full OpenCode or plugin restart, persisted running background-task
+history is rehydrated into the local job board and immediately reconciled against
+live host session status. A missing or idle child is therefore surfaced as
+`stopped, unreconciled`, while a busy child remains running; status lookup
+failures remain uncertain rather than being treated as completion.
+
 Specialist outputs are inputs, not final truth. The orchestrator reconciles them
 against each other and the original user goal.
 

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

@@ -77,6 +77,26 @@ function taskLaunchOutput(taskID: string): string {
   ].join('\n');
 }
 
+function historicalRunningTaskPart(
+  taskID: string,
+  input: Record<string, unknown> = {
+    background: true,
+    subagent_type: 'explorer',
+    description: 'recover scheduler task',
+    prompt: 'inspect scheduler state',
+  },
+) {
+  return {
+    type: 'tool',
+    tool: 'task',
+    state: {
+      status: 'running',
+      input,
+      output: taskLaunchOutput(taskID),
+    },
+  };
+}
+
 type HookOptions = {
   shouldManageSession?: (sessionID: string) => boolean;
   registerSessionAsOrchestrator?: (sessionID: string) => void;
@@ -253,6 +273,236 @@ describe('task-session-manager hook', () => {
     );
   });
 
+  test('rehydrates historical background tasks and reconciles stopped children immediately', async () => {
+    const board = new BackgroundJobBoard();
+    const status = mock(async () => ({ data: {} }));
+    const { hook } = createHook({
+      backgroundJobBoard: board,
+      sessionClient: { status },
+      runtimeStatusReconcileDelayMs: 60_000,
+    });
+    const messages = {
+      messages: [
+        {
+          info: {
+            role: 'assistant',
+            sessionID: 'parent-1',
+          },
+          parts: [historicalRunningTaskPart('historical-child')],
+        },
+        ...createMessages('parent-1', 'continue').messages,
+      ],
+    };
+
+    await transformMessages(hook, messages as never);
+
+    expect(board.get('historical-child')).toMatchObject({
+      state: 'stopped',
+      terminalUnreconciled: true,
+      background: true,
+      agent: 'explorer',
+      description: 'recover scheduler task',
+      objective: 'recover scheduler task',
+    });
+    expect(boardText(messages)).toContain(
+      'historical-child / explorer / stopped, unreconciled',
+    );
+  });
+
+  test('rehydrates a completed tool call when its child output is still running', async () => {
+    const board = new BackgroundJobBoard();
+    const { hook } = createHook({
+      backgroundJobBoard: board,
+      sessionStatus: {},
+      runtimeStatusReconcileDelayMs: 60_000,
+    });
+    const taskPart = historicalRunningTaskPart('completed-call-child');
+    taskPart.state.status = 'completed';
+    const messages = {
+      messages: [
+        {
+          info: { role: 'assistant', sessionID: 'parent-1' },
+          parts: [taskPart],
+        },
+        ...createMessages('parent-1', 'continue').messages,
+      ],
+    };
+
+    await transformMessages(hook, messages as never);
+
+    expect(board.get('completed-call-child')).toMatchObject({
+      state: 'stopped',
+      terminalUnreconciled: true,
+    });
+  });
+
+  test('consumes historical terminal completion before restart reconciliation', async () => {
+    const board = new BackgroundJobBoard();
+    const status = mock(async () => ({ data: {} }));
+    const { hook } = createHook({
+      backgroundJobBoard: board,
+      sessionClient: { status },
+      runtimeStatusReconcileDelayMs: 60_000,
+    });
+    const messages = {
+      messages: [
+        {
+          info: { role: 'assistant', sessionID: 'parent-1' },
+          parts: [historicalRunningTaskPart('completed-child')],
+        },
+        {
+          info: {
+            role: 'user',
+            agent: 'orchestrator',
+            sessionID: 'parent-1',
+          },
+          parts: [
+            {
+              type: 'text',
+              synthetic: true,
+              text: [
+                '<task id="completed-child" state="completed">',
+                '<summary>Background task completed: recovered</summary>',
+                '<task_result>',
+                'historical result',
+                '</task_result>',
+                '</task>',
+              ].join('\n'),
+            },
+          ],
+        },
+        ...createMessages('parent-1', 'continue').messages,
+      ],
+    };
+
+    await transformMessages(hook, messages as never);
+
+    expect(board.get('completed-child')).toMatchObject({
+      state: 'completed',
+      terminalUnreconciled: true,
+      resultSummary: 'historical result',
+    });
+    expect(status).not.toHaveBeenCalled();
+  });
+
+  test('keeps a rehydrated child running when live status is busy', async () => {
+    const board = new BackgroundJobBoard();
+    const { hook } = createHook({
+      backgroundJobBoard: board,
+      sessionStatus: { 'historical-child': { type: 'busy' } },
+      runtimeStatusReconcileDelayMs: 60_000,
+    });
+    const messages = {
+      messages: [
+        {
+          info: {
+            role: 'assistant',
+            agent: 'orchestrator',
+            sessionID: 'parent-1',
+          },
+          parts: [historicalRunningTaskPart('historical-child')],
+        },
+        ...createMessages('parent-1', 'continue').messages,
+      ],
+    };
+
+    await transformMessages(hook, messages as never);
+    expect(board.get('historical-child')).toMatchObject({
+      state: 'running',
+      statusUncertain: false,
+    });
+    expect(boardText(messages)).toContain(
+      'historical-child / explorer / running',
+    );
+  });
+
+  test('ignores foreground, terminal, and malformed historical task parts', async () => {
+    const board = new BackgroundJobBoard();
+    const { hook } = createHook({
+      backgroundJobBoard: board,
+      runtimeStatusReconcileDelayMs: 60_000,
+    });
+    const messages = {
+      messages: [
+        {
+          info: {
+            role: 'assistant',
+            agent: 'orchestrator',
+            sessionID: 'parent-1',
+          },
+          parts: [
+            historicalRunningTaskPart('foreground-child', {
+              background: false,
+              subagent_type: 'explorer',
+            }),
+            {
+              ...historicalRunningTaskPart('terminal-child'),
+              state: {
+                ...historicalRunningTaskPart('terminal-child').state,
+                status: 'completed',
+                output: [
+                  'task_id: terminal-child',
+                  'state: completed',
+                  '',
+                  '<task_result>',
+                  'done',
+                  '</task_result>',
+                ].join('\n'),
+              },
+            },
+            historicalRunningTaskPart('missing-id', {
+              background: true,
+              subagent_type: 'explorer',
+            }),
+          ],
+        },
+        ...createMessages('parent-1', 'continue').messages,
+      ],
+    };
+    (
+      messages.messages[0].parts[2] as { state: { output: string } }
+    ).state.output = 'state: running\nmalformed output';
+
+    await transformMessages(hook, messages as never);
+
+    expect(board.list()).toHaveLength(0);
+  });
+
+  test('rehydration is idempotent across repeated transforms', async () => {
+    const board = new BackgroundJobBoard();
+    const status = mock(async () => ({ data: {} }));
+    const { hook } = createHook({
+      backgroundJobBoard: board,
+      sessionClient: { status },
+      runtimeStatusReconcileDelayMs: 60_000,
+    });
+    const messages = {
+      messages: [
+        {
+          info: {
+            role: 'assistant',
+            agent: 'orchestrator',
+            sessionID: 'parent-1',
+          },
+          parts: [historicalRunningTaskPart('historical-child')],
+        },
+        ...createMessages('parent-1', 'continue').messages,
+      ],
+    };
+
+    await transformMessages(hook, messages as never);
+    const first = board.get('historical-child');
+    await transformMessages(hook, messages as never);
+    const second = board.get('historical-child');
+
+    expect(second).toMatchObject({
+      alias: first?.alias,
+      generation: first?.generation,
+      state: 'stopped',
+    });
+    expect(status).toHaveBeenCalledTimes(1);
+  });
+
   test('stores background task launches in job board prompt context', async () => {
     const board = new BackgroundJobBoard();
     const { hook } = createHook({ backgroundJobBoard: board });

+ 103 - 1
src/hooks/task-session-manager/index.ts

@@ -4,11 +4,14 @@ import {
   type BackgroundJobExecution,
   type BackgroundJobStore,
   type BackgroundJobSupervisor,
+  deriveTaskSessionLabel,
   isInternalInitiatorPart,
+  parseTaskIdFromTaskOutput,
+  parseTaskStateFromOutput,
 } from '../../utils';
 import { isRecord as isObjectRecord } from '../../utils/guards';
 import type { SessionLifecycle } from '../session-lifecycle';
-import { isUserMessageWithParts } from '../types';
+import { isMessageWithParts, isUserMessageWithParts } from '../types';
 import {
   BACKGROUND_JOB_BOARD_METADATA_KEY,
   type InjectedTerminalJobs,
@@ -42,6 +45,94 @@ export { BACKGROUND_JOB_BOARD_METADATA_KEY } from './board-injection';
  */
 const IDLE_RECONCILE_DELAY_MS = 2_000;
 
+const RECOVERED_TASK_AGENT_FALLBACK = 'unknown';
+
+function rehydrateHistoricalRunningTasks(
+  messages: unknown[],
+  backgroundJobBoard: BackgroundJobStore,
+  shouldManageSession: (sessionID: string) => boolean,
+  registerSessionAsOrchestrator?: (sessionID: string) => void,
+): number {
+  let rehydrated = 0;
+  const managedOrchestratorSessionIDs = new Set<string>();
+
+  for (const message of messages) {
+    if (!isMessageWithParts(message)) continue;
+    if (message.info.agent !== 'orchestrator') continue;
+
+    const parentSessionID = message.info.sessionID;
+    if (!parentSessionID) continue;
+    if (!shouldManageSession(parentSessionID)) {
+      registerSessionAsOrchestrator?.(parentSessionID);
+      if (!shouldManageSession(parentSessionID)) continue;
+    }
+    managedOrchestratorSessionIDs.add(parentSessionID);
+  }
+
+  for (const message of messages) {
+    if (!isMessageWithParts(message)) continue;
+
+    const parentSessionID = message.info.sessionID;
+    if (
+      !parentSessionID ||
+      !managedOrchestratorSessionIDs.has(parentSessionID)
+    ) {
+      continue;
+    }
+
+    for (const part of message.parts) {
+      if (part.type !== 'tool' || part.tool !== 'task') continue;
+      if (!isObjectRecord(part.state)) continue;
+
+      const state = part.state;
+      if (typeof state.output !== 'string') continue;
+      if (!isObjectRecord(state.input) || state.input.background !== true) {
+        continue;
+      }
+
+      const taskID = parseTaskIdFromTaskOutput(state.output);
+      if (!taskID || parseTaskStateFromOutput(state.output) !== 'running') {
+        continue;
+      }
+      if (backgroundJobBoard.get(taskID)) continue;
+
+      const agent =
+        typeof state.input.subagent_type === 'string' &&
+        state.input.subagent_type.trim() !== ''
+          ? state.input.subagent_type.trim()
+          : RECOVERED_TASK_AGENT_FALLBACK;
+      const description =
+        typeof state.input.description === 'string'
+          ? state.input.description
+          : undefined;
+      const prompt =
+        typeof state.input.prompt === 'string' ? state.input.prompt : undefined;
+      const label = deriveTaskSessionLabel({
+        description,
+        prompt,
+        agentType: agent,
+      });
+
+      backgroundJobBoard.registerLaunch({
+        taskID,
+        parentSessionID,
+        agent,
+        description: label,
+        objective: label,
+        background: true,
+        preserveRun: true,
+        // Historical parts do not carry a trustworthy launch timestamp. Zero
+        // also prevents this registration from looking like a live observation
+        // to the first runtime-status reconciliation.
+        now: 0,
+      });
+      rehydrated += 1;
+    }
+  }
+
+  return rehydrated;
+}
+
 export function createTaskSessionManagerHook(
   _ctx: PluginInput,
   options: {
@@ -295,6 +386,13 @@ export function createTaskSessionManagerHook(
       // cache. Terminal results are left untouched (they materialize once).
       stabilizeRunningTaskParts(messages);
 
+      const rehydratedCount = rehydrateHistoricalRunningTasks(
+        messages,
+        backgroundJobBoard,
+        options.shouldManageSession,
+        options.registerSessionAsOrchestrator,
+      );
+
       for (const [messageIndex, message] of messages.entries()) {
         if (!isUserMessageWithParts(message)) continue;
         if (message.info.agent && message.info.agent !== 'orchestrator') {
@@ -322,6 +420,10 @@ export function createTaskSessionManagerHook(
           );
         }
       }
+
+      if (rehydratedCount > 0) {
+        await runtimeStatusReconciler.reconcile();
+      }
     },
 
     injectBackgroundJobBoard: (

+ 38 - 0
src/hooks/task-session-manager/runtime-status-reconciliation.test.ts

@@ -150,6 +150,44 @@ describe('runtime status reconciliation', () => {
     expect(board.get('child-1')).toMatchObject({ state: 'running' });
   });
 
+  test('serializes overlapping reconciliation and observes jobs added in-flight', async () => {
+    const firstResponse = deferred<unknown>();
+    let lookupCount = 0;
+    const status = mock(() => {
+      lookupCount += 1;
+      if (lookupCount === 1) return firstResponse.promise;
+      return Promise.resolve({
+        data: {
+          'child-1': { type: 'busy' },
+          'child-2': { type: 'idle' },
+        },
+      });
+    });
+    const { board, reconciler } = createReconciler(status);
+
+    const firstReconciliation = reconciler.reconcile();
+    await Promise.resolve();
+    board.registerLaunch({
+      taskID: 'child-2',
+      parentSessionID: 'parent-1',
+      agent: 'fixer',
+      description: 'second reconciliation job',
+      now: 0,
+    });
+    const secondReconciliation = reconciler.reconcile();
+
+    expect(status).toHaveBeenCalledTimes(1);
+    firstResponse.resolve({ data: { 'child-1': { type: 'busy' } } });
+    await Promise.all([firstReconciliation, secondReconciliation]);
+
+    expect(status).toHaveBeenCalledTimes(2);
+    expect(board.get('child-2')).toMatchObject({
+      state: 'stopped',
+      terminalUnreconciled: true,
+    });
+    reconciler.dispose();
+  });
+
   test('does not apply an old status response to a relaunched generation', async () => {
     const response = deferred<unknown>();
     const { board, reconciler } = createReconciler(() => response.promise);

+ 90 - 67
src/hooks/task-session-manager/runtime-status-reconciliation.ts

@@ -22,10 +22,16 @@ export function createRuntimeStatusReconciler(options: {
   const delayMs = options.delayMs ?? RUNTIME_STATUS_RECONCILE_DELAY_MS;
   let timer: ReturnType<typeof setTimeout> | undefined;
   let disposed = false;
-  let reconciling = false;
+  let activeReconcile: Promise<void> | undefined;
+  let rerunRequested = false;
 
   function schedule(): void {
-    if (disposed || timer || reconciling) return;
+    if (disposed) return;
+    if (activeReconcile) {
+      rerunRequested = true;
+      return;
+    }
+    if (timer) return;
     if (
       !options.backgroundJobBoard.list().some((job) => job.state === 'running')
     ) {
@@ -38,86 +44,103 @@ export function createRuntimeStatusReconciler(options: {
     timer.unref?.();
   }
 
-  async function reconcile(): Promise<void> {
-    if (disposed || reconciling) return;
+  async function reconcilePass(): Promise<void> {
+    if (disposed) return;
     const running = options.backgroundJobBoard
       .list()
       .filter((job) => job.state === 'running');
     if (running.length === 0) return;
 
-    reconciling = true;
-    try {
-      const requestStartedAt = Date.now();
-      const snapshot = await getRuntimeSessionStatusSnapshot(options.input, {
-        timeoutMs: options.statusTimeoutMs,
+    const requestStartedAt = Date.now();
+    const snapshot = await getRuntimeSessionStatusSnapshot(options.input, {
+      timeoutMs: options.statusTimeoutMs,
+    });
+    if (disposed) return;
+    const observedAt = Date.now();
+    if (snapshot.error) {
+      for (const job of running) {
+        options.backgroundJobBoard.markStatusUncertain(
+          job.taskID,
+          `Runtime status lookup failed: ${snapshot.error}`,
+          job.generation,
+        );
+      }
+      log('[task-session-manager] runtime status reconciliation uncertain', {
+        activeJobs: running.length,
+        error: snapshot.error,
       });
+      return;
+    }
+
+    for (const job of running) {
       if (disposed) return;
-      const observedAt = Date.now();
-      if (snapshot.error) {
-        for (const job of running) {
-          options.backgroundJobBoard.markStatusUncertain(
-            job.taskID,
-            `Runtime status lookup failed: ${snapshot.error}`,
-            job.generation,
-          );
-        }
-        log('[task-session-manager] runtime status reconciliation uncertain', {
-          activeJobs: running.length,
-          error: snapshot.error,
-        });
-        return;
+      const current = options.backgroundJobBoard.get(job.taskID);
+      if (
+        current?.state !== 'running' ||
+        current.generation !== job.generation
+      ) {
+        continue;
       }
-
-      for (const job of running) {
-        if (disposed) return;
-        const current = options.backgroundJobBoard.get(job.taskID);
-        if (
-          current?.state !== 'running' ||
-          current.generation !== job.generation
-        ) {
-          continue;
-        }
-        const status = runtimeSessionStatus(snapshot, job.taskID);
-        if (status === undefined) {
-          options.backgroundJobBoard.markStatusUncertain(
-            job.taskID,
-            'Runtime status response did not contain a recognized session state.',
-            job.generation,
-          );
-          continue;
-        }
-        if (status === 'busy' || status === 'retry') {
-          options.backgroundJobBoard.markRunningFromLiveSession(
-            job.taskID,
-            observedAt,
-            job.generation,
-          );
-          continue;
-        }
-
-        const stopped = options.backgroundJobBoard.markStopped(
+      const status = runtimeSessionStatus(snapshot, job.taskID);
+      if (status === undefined) {
+        options.backgroundJobBoard.markStatusUncertain(
           job.taskID,
-          'Background session stopped before a terminal task result was received.',
-          requestStartedAt,
+          'Runtime status response did not contain a recognized session state.',
           job.generation,
         );
-        if (stopped?.state !== 'stopped') continue;
-        options.taskContextTracker.pendingManagedTaskIds.delete(job.taskID);
-        options.backgroundJobBoard.addContext(
+        continue;
+      }
+      if (status === 'busy' || status === 'retry') {
+        options.backgroundJobBoard.markRunningFromLiveSession(
           job.taskID,
-          options.taskContextTracker.contextFilesForPrompt(job.taskID),
+          observedAt,
+          job.generation,
         );
-        options.taskContextTracker.prune(options.backgroundJobBoard);
-        log('[task-session-manager] reconciled runtime-stopped job', {
-          taskID: stopped.taskID,
-          alias: stopped.alias,
-          parentSessionID: stopped.parentSessionID,
-        });
+        continue;
       }
-    } finally {
-      reconciling = false;
-      schedule();
+
+      const stopped = options.backgroundJobBoard.markStopped(
+        job.taskID,
+        'Background session stopped before a terminal task result was received.',
+        requestStartedAt,
+        job.generation,
+      );
+      if (stopped?.state !== 'stopped') continue;
+      options.taskContextTracker.pendingManagedTaskIds.delete(job.taskID);
+      options.backgroundJobBoard.addContext(
+        job.taskID,
+        options.taskContextTracker.contextFilesForPrompt(job.taskID),
+      );
+      options.taskContextTracker.prune(options.backgroundJobBoard);
+      log('[task-session-manager] reconciled runtime-stopped job', {
+        taskID: stopped.taskID,
+        alias: stopped.alias,
+        parentSessionID: stopped.parentSessionID,
+      });
+    }
+  }
+
+  async function reconcile(): Promise<void> {
+    if (disposed) return;
+    if (activeReconcile) {
+      rerunRequested = true;
+      await activeReconcile;
+      return;
     }
+
+    const run = (async () => {
+      try {
+        do {
+          rerunRequested = false;
+          await reconcilePass();
+        } while (!disposed && rerunRequested);
+      } finally {
+        activeReconcile = undefined;
+        schedule();
+      }
+    })();
+    activeReconcile = run;
+    await run;
   }
 
   function dispose(): void {