Merge pull request #1049 from HeZ2z/fix/task-result-running-status

Author: alvinunreal

src/tools/task-result.test.ts

@@ -2,6 +2,7 @@ import { describe, expect, mock, test } from 'bun:test';
 import { BackgroundJobBoard } from '../utils/background-job-board';
 import { buildPluginInput } from '../v2/client-shim';
 import { createTaskResultTool } from './task-result';
+import { createTaskStatusTool } from './task-status';
 
 let mockClient: Record<string, any>;
 
@@ -26,11 +27,22 @@ function createTool() {
   const status = mock(async () => ({ data: {} }));
   mockClient = { session: { get, messages, status } };
 
+  const input = { directory: '/test/project' } as any;
   const tools = createTaskResultTool({
-    input: { directory: '/test/project' } as any,
+    input,
     backgroundJobBoard: board,
   });
-  return { board, get, messages, tool: tools.task_result };
+  const statusTools = createTaskStatusTool({
+    input,
+    backgroundJobBoard: board,
+  });
+  return {
+    board,
+    get,
+    messages,
+    statusTool: statusTools.task_status,
+    tool: tools.task_result,
+  };
 }
 
 describe('task_result', () => {
@@ -85,21 +97,56 @@ describe('task_result', () => {
     ).resolves.toBe('final findings');
   });
 
-  test('rejects a still-running tracked task', async () => {
-    const { board, tool, messages } = createTool();
+  test('returns a non-error status for a still-running tracked task', async () => {
+    const { board, tool, statusTool, messages } = createTool();
     board.registerLaunch({
       taskID: 'ses_child1',
       parentSessionID: 'parent-1',
       agent: 'explorer',
       description: 'trace bug',
     });
 
+    const output = await tool.execute({ task_id: 'exp-1' }, {
+      sessionID: 'parent-1',
+      agent: 'orchestrator',
+    } as any);
+
+    expect(output).toBe(
+      [
+        'task_id: ses_child1',
+        'state: running',
+        'message: Task is still running. Wait for its terminal result.',
+        'next: use task_status to inspect the task',
+      ].join('\n'),
+    );
     await expect(
-      tool.execute({ task_id: 'exp-1' }, {
+      statusTool.execute({ task_id: 'exp-1' }, {
         sessionID: 'parent-1',
         agent: 'orchestrator',
       } as any),
-    ).rejects.toThrow('still running');
+    ).resolves.toContain('state: running');
+    expect(messages).not.toHaveBeenCalled();
+  });
+
+  test('preserves a live retry state for a tracked running task', async () => {
+    const { board, tool, messages } = createTool();
+    board.registerLaunch({
+      taskID: 'ses_child1',
+      parentSessionID: 'parent-1',
+      agent: 'explorer',
+      description: 'trace bug',
+    });
+    mockClient.session.status.mockResolvedValue({
+      data: { ses_child1: { type: 'retry' } },
+    });
+
+    const output = await tool.execute({ task_id: 'exp-1' }, {
+      sessionID: 'parent-1',
+      agent: 'orchestrator',
+    } as any);
+
+    expect(output).toContain('state: retry');
+    expect(output).toContain('next: use task_status to inspect the task');
     expect(messages).not.toHaveBeenCalled();
   });
 
@@ -122,13 +169,13 @@ describe('task_result', () => {
       data: { ses_child1: { type: 'busy' } },
     });
 
-    await expect(
-      tool.execute({ task_id: 'exp-1' }, {
-        sessionID: 'parent-1',
-        agent: 'orchestrator',
-      } as any),
-    ).rejects.toThrow('still running');
+    const output = await tool.execute({ task_id: 'exp-1' }, {
+      sessionID: 'parent-1',
+      agent: 'orchestrator',
+    } as any);
 
+    expect(output).toContain('state: running');
+    expect(output).toContain('task_status');
     expect(board.get('ses_child1')).toMatchObject({
       state: 'running',
       statusUncertain: false,
@@ -235,8 +282,7 @@ describe('task_result', () => {
         sessionID: 'parent-1',
         agent: 'orchestrator',
       } as any),
-    ).rejects.toThrow('still running');
-
+    ).rejects.toThrow('changed generation');
     expect(board.get('ses_child1')).toMatchObject({
       generation: first.generation + 1,
       state: 'running',
@@ -280,8 +326,7 @@ describe('task_result', () => {
         sessionID: 'parent-1',
         agent: 'orchestrator',
       } as any),
-    ).rejects.toThrow('still running');
-
+    ).rejects.toThrow('changed generation');
     expect(board.get('ses_child1')?.state).toBe('running');
   });
 
@@ -309,33 +354,58 @@ describe('task_result', () => {
     expect(messages).not.toHaveBeenCalled();
   });
 
-  test('rejects an untracked child whose live session is busy', async () => {
+  test('returns a status for an untracked child whose live session is busy', async () => {
     const { tool, messages } = createTool();
     mockClient.session.status.mockImplementation(async () => ({
       data: { ses_child1: { type: 'busy' } },
     }));
 
+    const output = await tool.execute({ task_id: 'ses_child1' }, {
+      sessionID: 'parent-1',
+      agent: 'orchestrator',
+    } as any);
+
+    expect(output).toBe(
+      [
+        'task_id: ses_child1',
+        'state: running',
+        'message: Task is still running. Wait for its terminal result.',
+        'next: retry task_result after the task finishes',
+      ].join('\n'),
+    );
+    expect(messages).not.toHaveBeenCalled();
+
+    mockClient.session.status.mockResolvedValue({ data: {} });
+    mockClient.session.messages.mockResolvedValue({
+      data: [
+        {
+          info: { role: 'assistant', time: { completed: 100 } },
+          parts: [{ type: 'text', text: 'final findings' }],
+        },
+      ],
+    });
     await expect(
       tool.execute({ task_id: 'ses_child1' }, {
         sessionID: 'parent-1',
         agent: 'orchestrator',
       } as any),
-    ).rejects.toThrow('still running');
-    expect(messages).not.toHaveBeenCalled();
+    ).resolves.toBe('final findings');
   });
 
-  test('rejects an untracked child whose live session is retrying', async () => {
+  test('returns a status for an untracked child whose live session is retrying', async () => {
     const { tool, messages } = createTool();
     mockClient.session.status.mockImplementation(async () => ({
       data: { ses_child1: { type: 'retry' } },
     }));
 
-    await expect(
-      tool.execute({ task_id: 'ses_child1' }, {
-        sessionID: 'parent-1',
-        agent: 'orchestrator',
-      } as any),
-    ).rejects.toThrow('still running');
+    const output = await tool.execute({ task_id: 'ses_child1' }, {
+      sessionID: 'parent-1',
+      agent: 'orchestrator',
+    } as any);
+
+    expect(output).toContain('state: retry');
+    expect(output).toContain('next: retry task_result after the task finishes');
+    expect(output).not.toContain('task_status');
     expect(messages).not.toHaveBeenCalled();
   });
 

src/tools/task-result.ts

@@ -26,22 +26,17 @@ interface TaskResultToolOptions {
 /**
  * Gate tracked task retrieval on the tracked terminal outcome. Only a
  * `completed` state (or a reconciled job whose terminal outcome was
- * `completed`) may yield a successful result; running, errored and cancelled
- * jobs are rejected explicitly and accurately instead of leaking partial
- * output as a final result.
+ * `completed`) may yield a successful result. Errored, cancelled and otherwise
+ * unconfirmed jobs are rejected instead of leaking partial output as a final
+ * result.
  */
 function assertRetrievableState(
   requested: string,
   job: BackgroundJobRecord,
 ): void {
-  const terminalState =
-    job.state === 'reconciled' ? job.terminalState : job.state;
+  const terminalState = getTrackedTerminalState(job);
 
-  if (terminalState === 'running') {
-    throw new Error(
-      `Task ${requested} is still running. Wait for its terminal result instead of retrieving or duplicating it.`,
-    );
-  }
+  if (terminalState === 'running') return;
   if (terminalState === 'error') {
     throw new Error(
       `Task ${requested} ended in error: ${job.lastStatusError ?? job.resultSummary ?? 'no error details available'}`,
@@ -58,6 +53,25 @@ function assertRetrievableState(
   }
 }
 
+function getTrackedTerminalState(
+  job: BackgroundJobRecord,
+): BackgroundJobRecord['state'] | BackgroundJobRecord['terminalState'] {
+  return job.state === 'reconciled' ? job.terminalState : job.state;
+}
+
+function formatRunningTaskStatus(
+  taskID: string,
+  state: 'running' | 'retry',
+  tracked: boolean,
+): string {
+  return [
+    `task_id: ${taskID}`,
+    `state: ${state}`,
+    'message: Task is still running. Wait for its terminal result.',
+    `next: ${tracked ? 'use task_status to inspect the task' : 'retry task_result after the task finishes'}`,
+  ].join('\n');
+}
+
 function assertStableTrackedGeneration(
   requested: string,
   parentSessionID: string,
@@ -85,13 +99,11 @@ export function createTaskResultTool(
   options: TaskResultToolOptions,
 ): Record<string, ToolDefinition> {
   const task_result = tool({
-    description: `Retrieve the final text already produced by a specialist task without resuming or re-running it.
+    description: `Retrieve the final text already produced by a specialist task, or inspect its active state without resuming or re-running it.
 
-Use this when the user asks to see a prior task's full result, or before retrying work whose completed output may already answer the request. Accepts either the native task_id/session ID or the parent-scoped alias shown in the Background Job Board. This tool is read-only and never sends a new prompt to the specialist.`,
+Use this when the user asks to see a prior task's full result, or before retrying work whose completed output may already answer the request. If the task is still running, this returns a status message; only a completed task returns its final text. Accepts either the native task_id/session ID or the parent-scoped alias shown in the Background Job Board. This tool is read-only and never sends a new prompt to the specialist.`,
     args: {
-      task_id: z
-        .string()
-        .describe('Completed task ID or Background Job Board alias'),
+      task_id: z.string().describe('Task ID or Background Job Board alias'),
     },
     async execute(args, toolContext) {
       const parentSessionID = toolContext?.sessionID;
@@ -147,6 +159,19 @@ Use this when the user asks to see a prior task's full result, or before retryin
         throw new Error(`Unknown task ID or alias: ${requested}`);
       }
 
+      if (tracked && getTrackedTerminalState(tracked) === 'running') {
+        liveSnapshot ??= await getRuntimeSessionStatusSnapshot(options.input);
+        revalidateTracked();
+        if (tracked && getTrackedTerminalState(tracked) === 'running') {
+          const status = runtimeSessionStatus(liveSnapshot, taskID);
+          return formatRunningTaskStatus(
+            taskID,
+            status === 'retry' ? 'retry' : 'running',
+            true,
+          );
+        }
+      }
+
       const client = getClient(options.input);
       const sessionClient = client.session as typeof client.session & {
         get?: typeof client.session.get;
@@ -171,16 +196,19 @@ Use this when the user asks to see a prior task's full result, or before retryin
       revalidateTracked();
       const status = runtimeSessionStatus(liveSnapshot, taskID);
       const trackedTerminalState = tracked
-        ? tracked.state === 'reconciled'
-          ? tracked.terminalState
-          : tracked.state
+        ? getTrackedTerminalState(tracked)
         : undefined;
-      if (
-        (status === 'busy' || status === 'retry') &&
-        trackedTerminalState !== 'completed'
-      ) {
-        throw new Error(
-          `Task ${requested} is still running. Wait for its terminal result instead of retrieving or duplicating it.`,
+      const activeState =
+        status === 'retry'
+          ? 'retry'
+          : status === 'busy' || trackedTerminalState === 'running'
+            ? 'running'
+            : undefined;
+      if (activeState !== undefined && trackedTerminalState !== 'completed') {
+        return formatRunningTaskStatus(
+          taskID,
+          activeState,
+          tracked !== undefined,
         );
       }