fix(chat): deliver background job notifications at step boundaries

packages/cli/src/lib/__tests__/background-job-notification-delivery.test.ts

@@ -0,0 +1,98 @@
+import type { BackgroundJobTerminalEvent } from "@getpochi/common";
+import type { Message } from "@getpochi/livekit";
+import { describe, expect, it } from "vitest";
+import {
+  deliverBackgroundJobNotifications,
+  takeBackgroundJobNotificationMessage,
+} from "../background-job-notification-delivery";
+
+describe("takeBackgroundJobNotificationMessage", () => {
+  it("drains every pending event into one message", () => {
+    const pending = [event("bgjob-cmd-1"), event("bgjob-cmd-2")];
+
+    const message = takeBackgroundJobNotificationMessage(pending);
+
+    expect(pending).toHaveLength(0);
+    expect(message?.parts).toHaveLength(2);
+  });
+
+  it("returns nothing when no job finished", () => {
+    expect(takeBackgroundJobNotificationMessage([])).toBeUndefined();
+  });
+});
+
+describe("deliverBackgroundJobNotifications", () => {
+  it("delivers at a step boundary that continues the loop", () => {
+    const pending = [event("bgjob-cmd-1")];
+    const chat = fakeChat();
+
+    const delivered = deliverBackgroundJobNotifications("next", pending, chat);
+
+    expect(delivered).toBe(true);
+    expect(chat.messages).toHaveLength(1);
+    expect(chat.messages[0].role).toBe("user");
+    expect(chat.messages[0].parts).toEqual([
+      expect.objectContaining({
+        type: "data-background-job-notification",
+        data: expect.objectContaining({ backgroundJobId: "bgjob-cmd-1" }),
+      }),
+    ]);
+    // Drained, so the drain at the end of the task does not repeat it.
+    expect(pending).toHaveLength(0);
+  });
+
+  it("delivers nothing when no job finished during the step", () => {
+    const chat = fakeChat();
+
+    expect(deliverBackgroundJobNotifications("next", [], chat)).toBe(false);
+    expect(chat.messages).toHaveLength(0);
+  });
+
+  it("keeps notifications pending while the step is retried", () => {
+    const pending = [event("bgjob-cmd-1")];
+    const chat = fakeChat();
+
+    const delivered = deliverBackgroundJobNotifications("retry", pending, chat);
+
+    expect(delivered).toBe(false);
+    expect(chat.messages).toHaveLength(0);
+    expect(pending).toHaveLength(1);
+  });
+
+  it("leaves the finished step to drain notifications itself", () => {
+    const pending = [event("bgjob-cmd-1")];
+    const chat = fakeChat();
+
+    const delivered = deliverBackgroundJobNotifications(
+      "finished",
+      pending,
+      chat,
+    );
+
+    expect(delivered).toBe(false);
+    expect(chat.messages).toHaveLength(0);
+    expect(pending).toHaveLength(1);
+  });
+});
+
+function fakeChat() {
+  const messages: Message[] = [];
+  return {
+    messages,
+    appendOrReplaceMessage(message: Message) {
+      messages.push(message);
+    },
+  };
+}
+
+function event(backgroundJobId: string): BackgroundJobTerminalEvent {
+  return {
+    taskId: "task-1",
+    backgroundJobId,
+    outputFile: `/tmp/${backgroundJobId}.log`,
+    status: "completed",
+    command: `run ${backgroundJobId}`,
+    exitCode: 0,
+    finishedAt: 1,
+  };
+}

packages/cli/src/lib/background-job-notification-delivery.ts

@@ -0,0 +1,47 @@
+import type { BackgroundJobTerminalEvent } from "@getpochi/common";
+import {
+  type Message,
+  createBackgroundJobNotificationMessage,
+} from "@getpochi/livekit";
+
+export type StepResult = "finished" | "next" | "retry";
+
+export interface BackgroundJobNotificationTarget {
+  appendOrReplaceMessage(message: Message): void;
+}
+
+/**
+ * Drains every pending event into one user message, so a notification is
+ * never delivered twice.
+ */
+export function takeBackgroundJobNotificationMessage(
+  pending: BackgroundJobTerminalEvent[],
+): Message | undefined {
+  const events = pending.splice(0);
+  return events.length > 0
+    ? createBackgroundJobNotificationMessage(events)
+    : undefined;
+}
+
+/**
+ * Delivers jobs that finished mid loop as part of the continuation request
+ * the loop is about to send, instead of only when the task ends.
+ */
+export function deliverBackgroundJobNotifications(
+  stepResult: StepResult,
+  pending: BackgroundJobTerminalEvent[],
+  chat: BackgroundJobNotificationTarget,
+): boolean {
+  // "retry" resends the last message as is, and "finished" has its own drain.
+  if (stepResult !== "next") {
+    return false;
+  }
+
+  const message = takeBackgroundJobNotificationMessage(pending);
+  if (!message) {
+    return false;
+  }
+
+  chat.appendOrReplaceMessage(message);
+  return true;
+}

packages/cli/src/task-runner.ts

@@ -35,7 +35,6 @@ import {
   type LiveKitStore,
   type Message,
   type Task,
-  createBackgroundJobNotificationMessage,
   processContentOutput,
 } from "@getpochi/livekit";
 import { LiveChatKit } from "@getpochi/livekit/node";
@@ -61,6 +60,10 @@ import {
 } from "ai";
 import type z from "zod";
 import { BackgroundJobManager } from "./lib/background-job-manager";
+import {
+  deliverBackgroundJobNotifications,
+  takeBackgroundJobNotificationMessage,
+} from "./lib/background-job-notification-delivery";
 import type { FileSystem } from "./lib/file-system";
 import { readEnvironment } from "./lib/read-environment";
 import { createSpinner } from "./lib/spinner";
@@ -460,10 +463,9 @@ export class TaskRunner {
   }
 
   private takePendingBackgroundJobNotifications(): Message | undefined {
-    const events = this.pendingBackgroundJobNotifications.splice(0);
-    return events.length > 0
-      ? createBackgroundJobNotificationMessage(events)
-      : undefined;
+    return takeBackgroundJobNotificationMessage(
+      this.pendingBackgroundJobNotifications,
+    );
   }
 
   /**
@@ -538,6 +540,13 @@ export class TaskRunner {
       this.stepCount.throwIfReachedMaxRetries();
     }
 
+    // Deliver at this step boundary instead of waiting for the task to end.
+    deliverBackgroundJobNotifications(
+      result,
+      this.pendingBackgroundJobNotifications,
+      this.chat,
+    );
+
     this.abortSignal?.throwIfAborted();
     await this.chatKit.chat.sendMessage();
     return result;

packages/vscode-webui/src/features/chat/components/chat-toolbar.test.tsx

@@ -1,18 +1,38 @@
+import type { BackgroundJobNotification } from "@getpochi/common";
 import type { Message, Task } from "@getpochi/livekit";
 import type { Todo } from "@getpochi/tools";
 // @vitest-environment jsdom
-import { render, screen } from "@testing-library/react";
+import { act, render, screen } from "@testing-library/react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { ChatToolbar } from "./chat-toolbar";
 
-const chatSubmitMocks = vi.hoisted(() => ({
-  useChatSubmit: vi.fn(() => ({
-    handleSubmit: vi.fn(),
-    handleSteerSubmit: vi.fn(),
-    handleSteerQueuedMessage: vi.fn(),
-    handleStop: vi.fn(),
-    pauseQueueRef: { current: false },
-  })),
+const chatSubmitMocks = vi.hoisted(() => {
+  const sendQueuedMessage = vi.fn(() => Promise.resolve(true));
+  const setQueuedMessages = {
+    current: undefined as
+      | React.Dispatch<React.SetStateAction<unknown[]>>
+      | undefined,
+  };
+  return {
+    sendQueuedMessage,
+    setQueuedMessages,
+    useChatSubmit: vi.fn((props: { setQueuedMessages: unknown }) => {
+      setQueuedMessages.current = props.setQueuedMessages as React.Dispatch<
+        React.SetStateAction<unknown[]>
+      >;
+      return {
+        handleSubmit: vi.fn(),
+        handleSteerSubmit: vi.fn(),
+        handleSteerQueuedMessage: vi.fn(),
+        handleStop: vi.fn(),
+        sendQueuedMessage,
+      };
+    }),
+  };
+});
+const backgroundJobMocks = vi.hoisted(() => ({
+  notifications: [] as unknown[],
+  acknowledge: vi.fn(() => Promise.resolve()),
 }));
 const userEditsMocks = vi.hoisted(() => ({
   userEdits: [] as Array<{
@@ -105,8 +125,8 @@ vi.mock("@/lib/hooks/use-add-complete-tool-calls", () => ({
 }));
 vi.mock("@/lib/hooks/use-background-job-notifications", () => ({
   useBackgroundJobNotifications: () => ({
-    notifications: [],
-    acknowledge: undefined,
+    notifications: backgroundJobMocks.notifications,
+    acknowledge: backgroundJobMocks.acknowledge,
   }),
 }));
 vi.mock("@/lib/hooks/use-custom-agents", () => ({
@@ -189,7 +209,11 @@ const auditTodo: Todo = {
   priority: "medium",
 };
 
-function renderToolbar(isSubTask: boolean, lastCheckpointHash?: string) {
+function renderToolbar(
+  isSubTask: boolean,
+  lastCheckpointHash?: string,
+  deliverBackgroundJobNotificationsRef?: React.RefObject<() => boolean>,
+) {
   render(
     <ChatToolbar
       chat={
@@ -229,13 +253,33 @@ function renderToolbar(isSubTask: boolean, lastCheckpointHash?: string) {
       todoPaused={false}
       onTodoPausedChange={vi.fn()}
       taskId="task-1"
+      deliverBackgroundJobNotificationsRef={
+        deliverBackgroundJobNotificationsRef
+      }
     />,
   );
 }
 
+function notification(backgroundJobId: string): BackgroundJobNotification {
+  return {
+    notificationId: `${backgroundJobId}:terminal`,
+    backgroundJobId,
+    outputFile: `/tmp/${backgroundJobId}.log`,
+    command: `run ${backgroundJobId}`,
+    status: "completed",
+    summary: `Background command "${backgroundJobId}" completed`,
+    exitCode: 0,
+    finishedAt: 1,
+  };
+}
+
 describe("ChatToolbar", () => {
   beforeEach(() => {
     chatSubmitMocks.useChatSubmit.mockClear();
+    chatSubmitMocks.sendQueuedMessage.mockClear();
+    chatSubmitMocks.setQueuedMessages.current = undefined;
+    backgroundJobMocks.notifications = [];
+    backgroundJobMocks.acknowledge.mockClear();
     userEditsMocks.userEdits = [];
   });
 
@@ -301,4 +345,87 @@ describe("ChatToolbar", () => {
       }),
     );
   });
+
+  describe("background job notification delivery", () => {
+    it("sends a queued notification instead of a plain continuation", async () => {
+      backgroundJobMocks.notifications = [notification("bgjob-cmd-1")];
+      const deliverRef: React.RefObject<() => boolean> = {
+        current: () => false,
+      };
+
+      renderToolbar(false, undefined, deliverRef);
+
+      let delivered: boolean | undefined;
+      await act(async () => {
+        delivered = deliverRef.current();
+      });
+
+      expect(delivered).toBe(true);
+      // The guard is kept so an intentional manual approval mode is not
+      // silently turned into auto approve by a notification.
+      expect(chatSubmitMocks.sendQueuedMessage).toHaveBeenCalledWith(0, {
+        keepAutoApproveGuard: true,
+      });
+    });
+
+    it("delivers nothing when no notification is queued", async () => {
+      const deliverRef: React.RefObject<() => boolean> = {
+        current: () => false,
+      };
+
+      renderToolbar(false, undefined, deliverRef);
+
+      let delivered: boolean | undefined;
+      await act(async () => {
+        delivered = deliverRef.current();
+      });
+
+      expect(delivered).toBe(false);
+      expect(chatSubmitMocks.sendQueuedMessage).not.toHaveBeenCalled();
+    });
+
+    it("delivers nothing while a queued user message is ahead of the notification", async () => {
+      backgroundJobMocks.notifications = [notification("bgjob-cmd-1")];
+      const deliverRef: React.RefObject<() => boolean> = {
+        current: () => false,
+      };
+
+      renderToolbar(false, undefined, deliverRef);
+
+      await act(async () => {
+        chatSubmitMocks.setQueuedMessages.current?.((current) => [
+          { parts: [{ type: "text", text: "hello" }], raw: { text: "hello" } },
+          ...current,
+        ]);
+      });
+
+      let delivered: boolean | undefined;
+      await act(async () => {
+        delivered = deliverRef.current();
+      });
+
+      expect(delivered).toBe(false);
+      expect(chatSubmitMocks.sendQueuedMessage).not.toHaveBeenCalled();
+    });
+
+    it("does not send the same notification twice when the decision is evaluated again", async () => {
+      backgroundJobMocks.notifications = [notification("bgjob-cmd-1")];
+      const deliverRef: React.RefObject<() => boolean> = {
+        current: () => false,
+      };
+
+      renderToolbar(false, undefined, deliverRef);
+
+      let second: boolean | undefined;
+      await act(async () => {
+        deliverRef.current();
+        second = deliverRef.current();
+      });
+
+      // Still true: the pending delivery starts the next request, so the caller
+      // must not start a plain continuation on top of it.
+      expect(second).toBe(true);
+      expect(chatSubmitMocks.sendQueuedMessage).toHaveBeenCalledOnce();
+    });
+  });
 });