From e72ba82a0c7c30e6b026cc93a3268b6a5ec47361 Mon Sep 17 00:00:00 2001 From: Charles Vien Date: Wed, 22 Jul 2026 01:19:35 -0700 Subject: [PATCH] keep cloud turn armed until turn_complete --- .../cloudTaskUpdateNotifications.test.ts | 87 ++++++++++++++++++- packages/core/src/sessions/sessionService.ts | 24 ++++- .../sessions/sessionServiceHost.test.ts | 11 ++- 3 files changed, 114 insertions(+), 8 deletions(-) diff --git a/packages/core/src/sessions/cloudTaskUpdateNotifications.test.ts b/packages/core/src/sessions/cloudTaskUpdateNotifications.test.ts index 5ad145a0c0..cb4335841a 100644 --- a/packages/core/src/sessions/cloudTaskUpdateNotifications.test.ts +++ b/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 = () => diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 4bda00fe47..65830d943c 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/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, ); diff --git a/packages/ui/src/features/sessions/sessionServiceHost.test.ts b/packages/ui/src/features/sessions/sessionServiceHost.test.ts index cf9ebc06dc..88deaf7e04 100644 --- a/packages/ui/src/features/sessions/sessionServiceHost.test.ts +++ b/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 | undefined)?.currentPromptId === + null, + ); + expect(disarmed).toBe(false); }); it("flushes queued cloud messages on _posthog/turn_complete", async () => {