From ee902c45738fe7f08ad4a02164df6ac6d6a17326 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Mon, 6 Jul 2026 14:05:11 +0300 Subject: [PATCH 1/2] fix(sessions): recover and resend prompts after "session not found" Local agent sessions are idle-killed after 15 minutes. When the renderer misses the idle-kill event (sleep, dropped subscription), replying to the task hit the dead session, surfaced "Session not found", and the composer had already cleared the typed prompt. Now a fatal prompt failure awaits the existing in-place recovery and resends the prompt once on the recovered session, so stale replies land instead of erroring. If a send still fails for any reason, the typed prompt is restored into the composer instead of being lost. Generated-By: PostHog Code Task-Id: 03ebf8b9-1bd2-49a9-ae91-d8f5e86ed683 --- packages/core/src/sessions/sessionService.ts | 91 ++++++++++-- .../sessionServicePromptRecovery.test.ts | 137 ++++++++++++++++++ .../sessions/hooks/useSessionCallbacks.ts | 13 ++ 3 files changed, 232 insertions(+), 9 deletions(-) create mode 100644 packages/core/src/sessions/sessionServicePromptRecovery.test.ts diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 1ea7c02b8a..6c791ce64e 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -2402,7 +2402,7 @@ export class SessionService { session: AgentSession, blocks: ContentBlock[], promptText: string, - options: { optimisticApplied?: boolean } = {}, + options: { optimisticApplied?: boolean; isRecoveryResend?: boolean } = {}, ): Promise<{ stopReason: string }> { if (!options.optimisticApplied) { this.applyOptimisticPrompt(session.taskRunId, blocks, promptText); @@ -2444,14 +2444,33 @@ export class SessionService { errorMessage, errorDetails, }); - this.startAutoRecoverLocalSession( - session.taskId, - session.taskRunId, - session.taskTitle, - errorDetails || errorMessage, - errorDetails || - "Session connection lost. Please retry or start a new session.", - ); + if (!options.isRecoveryResend) { + const resent = await this.recoverAndResendPrompt( + session, + blocks, + promptText, + errorDetails || errorMessage, + ); + if (resent) { + return resent; + } + } + // Recovery failed (or this already was the post-recovery resend): + // surface the error state so the user can retry manually. + const latest = this.d.store.getSessionByTaskId(session.taskId); + if ( + latest?.taskRunId === session.taskRunId && + latest.status !== "error" + ) { + this.setErrorSession( + session.taskId, + session.taskRunId, + session.taskTitle, + errorDetails || + "Session connection lost. Please retry or start a new session.", + "Connection lost", + ); + } } else { this.d.store.updateSession(session.taskRunId, { isPromptPending: false, @@ -2464,6 +2483,60 @@ export class SessionService { } } + /** + * A fatal prompt failure (e.g. "Session not found") usually means the + * backend agent was idle-killed or the host process restarted while the + * renderer still shows the session as connected. Recover the session in + * place and resend the prompt once, so a reply to a stale session lands + * instead of erroring and losing the user's message. + * + * Returns the resend result, or null when recovery (or the refreshed + * session lookup) failed and the caller should surface the original error. + */ + private async recoverAndResendPrompt( + session: AgentSession, + blocks: ContentBlock[], + promptText: string, + reason: string, + ): Promise<{ stopReason: string } | null> { + let recovered = false; + try { + recovered = await this.tryAutoRecoverLocalSession( + session.taskId, + session.taskRunId, + reason, + ); + } catch (recoveryError) { + this.d.log.warn("Session recovery threw while resending prompt", { + taskId: session.taskId, + taskRunId: session.taskRunId, + error: + recoveryError instanceof Error + ? recoveryError.message + : String(recoveryError), + }); + return null; + } + if (!recovered) return null; + + const refreshed = this.d.store.getSessionByTaskId(session.taskId); + if ( + !refreshed || + refreshed.taskRunId !== session.taskRunId || + refreshed.status !== "connected" + ) { + return null; + } + + this.d.log.info("Resending prompt after session recovery", { + taskId: session.taskId, + taskRunId: session.taskRunId, + }); + return this.sendLocalPrompt(refreshed, blocks, promptText, { + isRecoveryResend: true, + }); + } + /** * Steer a single queued message into the running turn now: drop it from the * queue and resend it as a steer. Native (Claude, local) injects at the next diff --git a/packages/core/src/sessions/sessionServicePromptRecovery.test.ts b/packages/core/src/sessions/sessionServicePromptRecovery.test.ts new file mode 100644 index 0000000000..dd6853b18b --- /dev/null +++ b/packages/core/src/sessions/sessionServicePromptRecovery.test.ts @@ -0,0 +1,137 @@ +import type { AgentSession } from "@posthog/shared"; +import { describe, expect, it, vi } from "vitest"; +import { SessionService, type SessionServiceDeps } from "./sessionService"; + +const TASK_ID = "task-1"; +const TASK_RUN_ID = `run-${TASK_ID}`; + +function makeSession(): AgentSession { + return { + taskRunId: TASK_RUN_ID, + taskId: TASK_ID, + taskTitle: "Test task", + channel: "", + events: [], + startedAt: 1, + status: "connected", + isPromptPending: false, + isCompacting: false, + promptStartedAt: null, + pendingPermissions: new Map(), + pausedDurationMs: 0, + messageQueue: [], + optimisticItems: [], + } as AgentSession; +} + +function createHarness() { + const sessions: Record = { + [TASK_RUN_ID]: makeSession(), + }; + const store = { + getSessions: () => sessions, + getSessionByTaskId: (taskId: string) => + Object.values(sessions).find((s) => s.taskId === taskId), + setSession: vi.fn((session: AgentSession) => { + sessions[session.taskRunId] = session; + }), + updateSession: vi.fn((taskRunId: string, patch: Partial) => { + const existing = sessions[taskRunId]; + if (existing) sessions[taskRunId] = { ...existing, ...patch }; + }), + appendOptimisticItem: vi.fn(), + clearOptimisticItems: vi.fn(), + }; + const promptMutate = vi.fn(); + const deps = { + store, + h: { extractSkillButtonId: () => undefined }, + log: { info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug: vi.fn() }, + toast: { error: vi.fn(), info: vi.fn() }, + track: vi.fn(), + getIsOnline: () => true, + addDirectoryDialog: { open: false }, + usageLimit: { show: vi.fn() }, + trpc: { + agent: { + prompt: { mutate: promptMutate }, + cancel: { mutate: vi.fn().mockResolvedValue(undefined) }, + onSessionIdleKilled: { + subscribe: () => ({ unsubscribe: vi.fn() }), + }, + }, + }, + } as unknown as SessionServiceDeps; + + const service = new SessionService(deps); + const recoverSpy = vi.spyOn( + service as unknown as { + tryAutoRecoverLocalSession: ( + taskId: string, + taskRunId: string, + reason: string, + ) => Promise; + }, + "tryAutoRecoverLocalSession", + ); + return { service, sessions, store, promptMutate, recoverSpy }; +} + +describe("SessionService prompt recovery on fatal session errors", () => { + it("recovers the session and resends the prompt once", async () => { + const { service, promptMutate, recoverSpy } = createHarness(); + promptMutate + .mockRejectedValueOnce(new Error(`Session not found: ${TASK_RUN_ID}`)) + .mockResolvedValueOnce({ stopReason: "end_turn" }); + recoverSpy.mockResolvedValue(true); + + const result = await service.sendPrompt(TASK_ID, "hello again"); + + expect(result).toEqual({ stopReason: "end_turn" }); + expect(recoverSpy).toHaveBeenCalledTimes(1); + expect(promptMutate).toHaveBeenCalledTimes(2); + }); + + it("surfaces the error state when recovery fails", async () => { + const { service, store, promptMutate, recoverSpy } = createHarness(); + promptMutate.mockRejectedValue( + new Error(`Session not found: ${TASK_RUN_ID}`), + ); + recoverSpy.mockResolvedValue(false); + + await expect(service.sendPrompt(TASK_ID, "hello again")).rejects.toThrow( + "Session not found", + ); + expect(promptMutate).toHaveBeenCalledTimes(1); + expect(store.setSession).toHaveBeenCalledWith( + expect.objectContaining({ taskRunId: TASK_RUN_ID, status: "error" }), + ); + }); + + it("does not loop when the resend hits another fatal error", async () => { + const { service, promptMutate, recoverSpy } = createHarness(); + promptMutate.mockRejectedValue( + new Error(`Session not found: ${TASK_RUN_ID}`), + ); + recoverSpy.mockResolvedValue(true); + + await expect(service.sendPrompt(TASK_ID, "hello again")).rejects.toThrow( + "Session not found", + ); + expect(recoverSpy).toHaveBeenCalledTimes(1); + expect(promptMutate).toHaveBeenCalledTimes(2); + }); + + it("swallows recovery exceptions and rethrows the original error", async () => { + const { service, promptMutate, recoverSpy } = createHarness(); + promptMutate.mockRejectedValue( + new Error(`Session not found: ${TASK_RUN_ID}`), + ); + recoverSpy.mockRejectedValue(new Error("auth restoring")); + + await expect(service.sendPrompt(TASK_ID, "hello again")).rejects.toThrow( + "Session not found", + ); + expect(promptMutate).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts index 53b0c06f93..e0ce20a63d 100644 --- a/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts +++ b/packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts @@ -1,3 +1,7 @@ +import { + isContentEmpty, + xmlToContent, +} from "@posthog/core/message-editor/content"; import { combineQueuedCloudPrompts, promptToQueuedEditorContent, @@ -106,6 +110,13 @@ export function useSessionCallbacks({ } catch (error) { const message = error instanceof Error ? error.message : "Failed to send message"; + // The composer clears optimistically on submit, so a failed send + // would otherwise lose the typed prompt. Restore it — unless the + // user has already started typing a new message. + if (isContentEmpty(useDraftStore.getState().drafts[taskId] ?? null)) { + setPendingContent(taskId, xmlToContent(text)); + requestFocus(taskId); + } toast.error(message); log.error("Failed to send prompt", error); } @@ -119,6 +130,8 @@ export function useSessionCallbacks({ sessionService, hostClient, messagingMode, + setPendingContent, + requestFocus, ], ); From 7edba26c94c41fe401599db5a7daa0fcc608d990 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Mon, 6 Jul 2026 14:36:26 +0300 Subject: [PATCH 2/2] test(sessions): parameterise prompt recovery failure cases The three failure cases exercise the same fatal-error path with different recovery behaviour, so express them as one it.each table per review feedback. All failure cases now also assert the session is surfaced in the error state. Generated-By: PostHog Code Task-Id: 03ebf8b9-1bd2-49a9-ae91-d8f5e86ed683 --- .../sessionServicePromptRecovery.test.ts | 78 +++++++++---------- 1 file changed, 37 insertions(+), 41 deletions(-) diff --git a/packages/core/src/sessions/sessionServicePromptRecovery.test.ts b/packages/core/src/sessions/sessionServicePromptRecovery.test.ts index dd6853b18b..1549401663 100644 --- a/packages/core/src/sessions/sessionServicePromptRecovery.test.ts +++ b/packages/core/src/sessions/sessionServicePromptRecovery.test.ts @@ -77,6 +77,8 @@ function createHarness() { return { service, sessions, store, promptMutate, recoverSpy }; } +type RecoverSpy = ReturnType["recoverSpy"]; + describe("SessionService prompt recovery on fatal session errors", () => { it("recovers the session and resends the prompt once", async () => { const { service, promptMutate, recoverSpy } = createHarness(); @@ -92,46 +94,40 @@ describe("SessionService prompt recovery on fatal session errors", () => { expect(promptMutate).toHaveBeenCalledTimes(2); }); - it("surfaces the error state when recovery fails", async () => { - const { service, store, promptMutate, recoverSpy } = createHarness(); - promptMutate.mockRejectedValue( - new Error(`Session not found: ${TASK_RUN_ID}`), - ); - recoverSpy.mockResolvedValue(false); - - await expect(service.sendPrompt(TASK_ID, "hello again")).rejects.toThrow( - "Session not found", - ); - expect(promptMutate).toHaveBeenCalledTimes(1); - expect(store.setSession).toHaveBeenCalledWith( - expect.objectContaining({ taskRunId: TASK_RUN_ID, status: "error" }), - ); - }); - - it("does not loop when the resend hits another fatal error", async () => { - const { service, promptMutate, recoverSpy } = createHarness(); - promptMutate.mockRejectedValue( - new Error(`Session not found: ${TASK_RUN_ID}`), - ); - recoverSpy.mockResolvedValue(true); - - await expect(service.sendPrompt(TASK_ID, "hello again")).rejects.toThrow( - "Session not found", - ); - expect(recoverSpy).toHaveBeenCalledTimes(1); - expect(promptMutate).toHaveBeenCalledTimes(2); - }); - - it("swallows recovery exceptions and rethrows the original error", async () => { - const { service, promptMutate, recoverSpy } = createHarness(); - promptMutate.mockRejectedValue( - new Error(`Session not found: ${TASK_RUN_ID}`), - ); - recoverSpy.mockRejectedValue(new Error("auth restoring")); + it.each([ + { + case: "recovery fails", + setupRecovery: (spy: RecoverSpy) => spy.mockResolvedValue(false), + expectedPromptCalls: 1, + }, + { + case: "the resend hits another fatal error", + setupRecovery: (spy: RecoverSpy) => spy.mockResolvedValue(true), + expectedPromptCalls: 2, + }, + { + case: "recovery throws", + setupRecovery: (spy: RecoverSpy) => + spy.mockRejectedValue(new Error("auth restoring")), + expectedPromptCalls: 1, + }, + ])( + "surfaces the error state and rethrows when $case", + async ({ setupRecovery, expectedPromptCalls }) => { + const { service, store, promptMutate, recoverSpy } = createHarness(); + promptMutate.mockRejectedValue( + new Error(`Session not found: ${TASK_RUN_ID}`), + ); + setupRecovery(recoverSpy); - await expect(service.sendPrompt(TASK_ID, "hello again")).rejects.toThrow( - "Session not found", - ); - expect(promptMutate).toHaveBeenCalledTimes(1); - }); + await expect(service.sendPrompt(TASK_ID, "hello again")).rejects.toThrow( + "Session not found", + ); + expect(recoverSpy).toHaveBeenCalledTimes(1); + expect(promptMutate).toHaveBeenCalledTimes(expectedPromptCalls); + expect(store.setSession).toHaveBeenCalledWith( + expect.objectContaining({ taskRunId: TASK_RUN_ID, status: "error" }), + ); + }, + ); });