fix(sessions): notify cloud turn completions in either log order (#3682)

Author: charlesvien

packages/core/src/sessions/cloudTaskUpdateNotifications.test.ts

@@ -6,17 +6,35 @@ import { SessionService, type SessionServiceDeps } from "./sessionService";
 const TASK_ID = "task-1";
 const RUN_ID = "run-1";
 
-function turnComplete(timestamp?: string): StoredLogEntry {
+function turnComplete(
+  timestamp?: string,
+  stopReason = "end_turn",
+): StoredLogEntry {
   return {
     type: "notification",
     timestamp,
     notification: {
       method: "_posthog/turn_complete",
-      params: { sessionId: RUN_ID, stopReason: "end_turn" },
+      params: { sessionId: RUN_ID, stopReason },
     },
   };
 }
 
+// The agent's JSON-RPC response to `session/prompt`. Real logs carry it next
+// to `_posthog/turn_complete` in either order (the two writes race in the
+// agent's log stream), so completion cases must ring for both orderings.
+function promptResponse(
+  id: number,
+  stopReason = "end_turn",
+  timestamp?: string,
+): StoredLogEntry {
+  return {
+    type: "notification",
+    timestamp,
+    notification: { id, result: { stopReason } },
+  };
+}
+
 // The `session/prompt` request that opens a turn. Its arrival is what arms the
 // turn's single completion notification, so realistic sequences pair it with a
 // later `turn_complete`.
@@ -235,6 +253,50 @@ describe("cloud task update notifications", () => {
       ],
       expected: 1,
     },
+    {
+      label: "a turn whose response precedes turn_complete",
+      updates: [
+        logsUpdate([sessionPrompt(1), promptResponse(1), turnComplete()], 3),
+      ],
+      expected: 1,
+    },
+    {
+      label: "a turn whose response follows turn_complete",
+      updates: [
+        logsUpdate([sessionPrompt(1), turnComplete(), promptResponse(1)], 3),
+      ],
+      expected: 1,
+    },
+    {
+      label: "a duplicate turn_complete after a response-first turn",
+      updates: [
+        logsUpdate([sessionPrompt(1), promptResponse(1), turnComplete()], 3),
+        logsUpdate([turnComplete()], 4),
+      ],
+      expected: 1,
+    },
+    {
+      label: "several turns with responses on both sides of turn_complete",
+      updates: [
+        logsUpdate([sessionPrompt(1), promptResponse(1), turnComplete()], 3),
+        logsUpdate([sessionPrompt(2), turnComplete(), promptResponse(2)], 6),
+      ],
+      expected: 2,
+    },
+    {
+      label: "a cancelled turn",
+      updates: [
+        logsUpdate(
+          [
+            sessionPrompt(1),
+            promptResponse(1, "cancelled"),
+            turnComplete(undefined, "cancelled"),
+          ],
+          3,
+        ),
+      ],
+      expected: 0,
+    },
     {
       // Opening a task mid-turn: its session/prompt is already in history and
       // only the turn_complete arrives live. The completion must still ring.
@@ -278,6 +340,27 @@ describe("cloud task update notifications", () => {
     expect(harness.markActivity).toHaveBeenCalledTimes(1);
   });
 
+  it("keeps the turn duration when the response precedes turn_complete", () => {
+    const harness = createHarness();
+    harness.sendUpdate(
+      logsUpdate(
+        [
+          sessionPrompt(1, "2026-01-01T00:00:00Z"),
+          promptResponse(1, "end_turn", "2026-01-01T00:00:44Z"),
+          turnComplete("2026-01-01T00:00:45Z"),
+        ],
+        3,
+      ),
+    );
+
+    expect(harness.notifyPromptComplete).toHaveBeenCalledWith(
+      "Cloud Task",
+      "end_turn",
+      TASK_ID,
+      45_000,
+    );
+  });
+
   it("notifies a pending permission once across repeated snapshots", () => {
     const harness = createHarness();
     const snapshot = () =>

packages/core/src/sessions/sessionService.ts

@@ -2863,6 +2863,16 @@ export class SessionService {
         if (session && session.currentPromptId !== msg.id) {
           continue;
         }
+        if (session?.isCloud) {
+          // Cloud logs carry both this response and `_posthog/turn_complete`,
+          // in either order (they race in the agent's log stream). Only
+          // turn_complete may disarm the turn; disarming here would make the
+          // completion notification fire or vanish based on log line order.
+          this.d.store.updateSession(taskRunId, {
+            isPromptPending: false,
+          });
+          continue;
+        }
         this.d.store.updateSession(taskRunId, {
           isPromptPending: false,
           promptStartedAt: null,
@@ -2874,13 +2884,15 @@ export class SessionService {
       }
       if (isTurnCompleteEvent(acpMsg)) {
         // Local sessions use the JSON-RPC response as the canonical turn-done
-        // signal; clearing currentPromptId here would race the id-match guard
-        // above. Cloud sessions never see that response.
+        // signal; turn_complete is the cloud one, so only cloud disarms here.
         const session = this.getSessionByRunId(taskRunId);
         if (session?.isCloud) {
           const completedActiveTurn =
             session.currentPromptId !== null &&
             session.currentPromptId !== undefined;
+          const stopReason =
+            (msg as { params?: { stopReason?: string } }).params?.stopReason ??
+            "end_turn";
           const turnStartedAtTs =
             this.liveTurnContent.get(taskRunId)?.startedAtTs ??
             session.promptStartedAt;
@@ -2891,10 +2903,14 @@ export class SessionService {
           });
           if (isLive) {
             // Queued messages will start a new turn — suppress the "done" notification in that case.
-            if (completedActiveTurn && session.messageQueue.length === 0) {
+            if (
+              completedActiveTurn &&
+              stopReason === "end_turn" &&
+              session.messageQueue.length === 0
+            ) {
               this.d.notifyPromptComplete(
                 session.taskTitle,
-                "end_turn",
+                stopReason,
                 session.taskId,
                 turnStartedAtTs ? acpMsg.ts - turnStartedAtTs : undefined,
               );

packages/ui/src/features/sessions/sessionServiceHost.test.ts

@@ -1644,11 +1644,18 @@ describe("SessionService", () => {
           "run-123",
           expect.objectContaining({
             isPromptPending: false,
-            promptStartedAt: null,
-            currentPromptId: null,
           }),
         );
       });
+      // The response must not disarm the turn: `_posthog/turn_complete` is the
+      // cloud turn-done signal and still needs currentPromptId to notify.
+      const disarmed = mockSessionStoreSetters.updateSession.mock.calls.some(
+        (call) =>
+          call[0] === "run-123" &&
+          (call[1] as Record<string, unknown> | undefined)?.currentPromptId ===
+            null,
+      );
+      expect(disarmed).toBe(false);
     });
 
     it("flushes queued cloud messages on _posthog/turn_complete", async () => {