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..1549401663 --- /dev/null +++ b/packages/core/src/sessions/sessionServicePromptRecovery.test.ts @@ -0,0 +1,133 @@ +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 }; +} + +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(); + 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.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(recoverSpy).toHaveBeenCalledTimes(1); + expect(promptMutate).toHaveBeenCalledTimes(expectedPromptCalls); + expect(store.setSession).toHaveBeenCalledWith( + expect.objectContaining({ taskRunId: TASK_RUN_ID, status: "error" }), + ); + }, + ); +}); 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, ], );