diff --git a/packages/core/src/sessions/cloudTaskUpdateNotifications.test.ts b/packages/core/src/sessions/cloudTaskUpdateNotifications.test.ts index d80e6f20dc..75a61e4f83 100644 --- a/packages/core/src/sessions/cloudTaskUpdateNotifications.test.ts +++ b/packages/core/src/sessions/cloudTaskUpdateNotifications.test.ts @@ -6,9 +6,10 @@ import { SessionService, type SessionServiceDeps } from "./sessionService"; const TASK_ID = "task-1"; const RUN_ID = "run-1"; -function turnComplete(): StoredLogEntry { +function turnComplete(timestamp?: string): StoredLogEntry { return { type: "notification", + timestamp, notification: { method: "_posthog/turn_complete", params: { sessionId: RUN_ID, stopReason: "end_turn" }, @@ -19,9 +20,10 @@ function turnComplete(): StoredLogEntry { // 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`. -function sessionPrompt(id: number): StoredLogEntry { +function sessionPrompt(id: number, timestamp?: string): StoredLogEntry { return { type: "notification", + timestamp, notification: { id, method: "session/prompt", @@ -233,15 +235,23 @@ describe("cloud task update notifications", () => { }, ); - it("notifies with the task title and stop reason, and marks activity", () => { + it("notifies with the task title, stop reason and turn duration, and marks activity", () => { const harness = createHarness(); - harness.sendUpdate(logsUpdate([sessionPrompt(1), turnComplete()], 2)); + harness.sendUpdate( + logsUpdate( + [ + sessionPrompt(1, "2026-01-01T00:00:00Z"), + turnComplete("2026-01-01T00:00:45Z"), + ], + 2, + ), + ); expect(harness.notifyPromptComplete).toHaveBeenCalledWith( "Cloud Task", "end_turn", TASK_ID, - undefined, + 45_000, ); expect(harness.markActivity).toHaveBeenCalledTimes(1); }); diff --git a/packages/core/src/sessions/sessionEventBatching.test.ts b/packages/core/src/sessions/sessionEventBatching.test.ts index 93c436edfd..e9af1af617 100644 --- a/packages/core/src/sessions/sessionEventBatching.test.ts +++ b/packages/core/src/sessions/sessionEventBatching.test.ts @@ -37,6 +37,7 @@ function createHarness() { [RUN_ID]: { taskRunId: RUN_ID, taskId: TASK_ID, + taskTitle: "Local Task", events: [], messageQueue: [], pendingPermissions: new Map(), @@ -62,7 +63,7 @@ function createHarness() { }, updateSession: (taskRunId: string, updates: Partial) => { const session = sessions[taskRunId]; - if (session) Object.assign(session, updates); + if (session) sessions[taskRunId] = { ...session, ...updates }; }, appendEvents, replaceOptimisticWithEvent: vi.fn(), @@ -80,10 +81,11 @@ function createHarness() { debug: vi.fn(), }; + const notifyPromptComplete = vi.fn(); const deps = { store, log: noopLog, - notifyPromptComplete: vi.fn(), + notifyPromptComplete, notifyPermissionRequest: vi.fn(), taskViewedApi: { markActivity: vi.fn() }, getPersistedConfigOptions: () => undefined, @@ -120,11 +122,36 @@ function createHarness() { return { service, appendEvents, + notifyPromptComplete, + updateSession: store.updateSession, emit: (event: AcpMessage) => onEvent?.(event), events: () => sessions[RUN_ID].events, }; } +function promptEcho(id: number, ts: number): AcpMessage { + return { + ts, + message: { + jsonrpc: "2.0", + id, + method: "session/prompt", + params: { sessionId: RUN_ID, prompt: [] }, + }, + } as unknown as AcpMessage; +} + +function promptResponse(id: number, ts: number): AcpMessage { + return { + ts, + message: { + jsonrpc: "2.0", + id, + result: { stopReason: "end_turn" }, + }, + } as unknown as AcpMessage; +} + describe("streamed event batching", () => { beforeEach(() => { vi.useFakeTimers(); @@ -164,4 +191,22 @@ describe("streamed event batching", () => { vi.advanceTimersByTime(FLUSH_MS); expect(h.events()).toHaveLength(2); }); + + it("keeps the turn duration when the prompt mutation clears state before the response flushes", () => { + const h = createHarness(); + + h.emit(promptEcho(1, 1_000)); + vi.advanceTimersByTime(FLUSH_MS); + + h.emit(promptResponse(1, 6_000)); + h.updateSession(RUN_ID, { isPromptPending: false, promptStartedAt: null }); + vi.advanceTimersByTime(FLUSH_MS); + + expect(h.notifyPromptComplete).toHaveBeenCalledWith( + "Local Task", + "end_turn", + TASK_ID, + 5_000, + ); + }); }); diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 36b18a1bd6..edd07c7734 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -1781,6 +1781,9 @@ export class SessionService { // above. Cloud sessions never see that response. const session = this.getSessionByRunId(taskRunId); if (session?.isCloud) { + const turnStartedAtTs = + this.liveTurnContent.get(taskRunId)?.startedAtTs ?? + session.promptStartedAt; this.d.store.updateSession(taskRunId, { isPromptPending: false, promptStartedAt: null, @@ -1793,9 +1796,7 @@ export class SessionService { session.taskTitle, "end_turn", session.taskId, - session.promptStartedAt - ? acpMsg.ts - session.promptStartedAt - : undefined, + turnStartedAtTs ? acpMsg.ts - turnStartedAtTs : undefined, ); } this.d.taskViewedApi.markActivity(session.taskId); @@ -1886,6 +1887,9 @@ export class SessionService { } else { this.d.store.appendEvents(taskRunId, [acpMsg]); } + const turnStartedAtTs = + this.liveTurnContent.get(taskRunId)?.startedAtTs ?? + session.promptStartedAt; this.updatePromptStateFromEvents(taskRunId, [acpMsg], { isLive: true }); const msg = acpMsg.message; @@ -1917,9 +1921,7 @@ export class SessionService { session.taskTitle, stopReason, session.taskId, - session.promptStartedAt - ? acpMsg.ts - session.promptStartedAt - : undefined, + turnStartedAtTs ? acpMsg.ts - turnStartedAtTs : undefined, ); }