fix(sessions): recover and resend prompts after "session not found"

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

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<string, AgentSession> = {
+    [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<AgentSession>) => {
+      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<boolean>;
+    },
+    "tryAutoRecoverLocalSession",
+  );
+  return { service, sessions, store, promptMutate, recoverSpy };
+}
+
+type RecoverSpy = ReturnType<typeof createHarness>["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" }),
+      );
+    },
+  );
+});

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,
     ],
   );