From 274d7c7ba3442c7965b9000fcf307145c6e348d7 Mon Sep 17 00:00:00 2001
From: Richard Solomou <richard.s@posthog.com>
Date: Tue, 31 Mar 2026 12:55:56 +0300
Subject: [PATCH] fix(code): preserve initial prompt across session error
 transitions

When a new task's initial prompt was sent but the session errored (e.g.
agent crash, subscription error, or fatal prompt failure), the prompt
was lost because it was never stored on the session object. Clicking
Retry would then reconnect without re-sending the original prompt.

Two fixes:
- Store initialPrompt on the session in createNewLocalSession so it
  survives error transitions via updateSession
- Preserve initialPrompt in setErrorSession when carrying forward
  session state to a new error session
---
 .../features/sessions/service/service.ts      | 40 +++++++++++++++++++
 .../features/sessions/stores/sessionStore.ts  |  3 ++
 2 files changed, 43 insertions(+)

diff --git a/apps/code/src/renderer/features/sessions/service/service.ts b/apps/code/src/renderer/features/sessions/service/service.ts
index c9c289eed..0c2b9351b 100644
--- a/apps/code/src/renderer/features/sessions/service/service.ts
+++ b/apps/code/src/renderer/features/sessions/service/service.ts
@@ -202,6 +202,9 @@ export class SessionService {
         session.status = "error";
         session.errorMessage =
           "Authentication required. Please sign in to continue.";
+        if (initialPrompt?.length) {
+          session.initialPrompt = initialPrompt;
+        }
         sessionStoreSetters.setSession(session);
         return;
       }
@@ -296,6 +299,9 @@ export class SessionService {
       session.status = "error";
       session.errorTitle = "Failed to connect";
       session.errorMessage = message;
+      if (initialPrompt?.length) {
+        session.initialPrompt = initialPrompt;
+      }
 
       if (latestRun?.log_url) {
         try {
@@ -504,6 +510,9 @@ export class SessionService {
     if (existing?.logUrl) {
       session.logUrl = existing.logUrl;
     }
+    if (existing?.initialPrompt?.length) {
+      session.initialPrompt = existing.initialPrompt;
+    }
     sessionStoreSetters.setSession(session);
   }
 
@@ -564,6 +573,13 @@ export class SessionService {
       useSessionAdapterStore.getState().setAdapter(taskRun.id, adapter);
     }
 
+    // Store the initial prompt on the session so retry/reset flows can
+    // re-send it if the session errors after this point (e.g. subscription
+    // error, agent crash, or prompt failure).
+    if (initialPrompt?.length) {
+      session.initialPrompt = initialPrompt;
+    }
+
     sessionStoreSetters.setSession(session);
     this.subscribeToChannel(taskRun.id);
 
@@ -1786,8 +1802,32 @@ export class SessionService {
    * Retry connecting to the existing session (resume attempt using
    * the sessionId from logs). Does NOT tear down — avoids the connect
    * effect loop.
+   *
+   * If the session failed before any conversation started (has an
+   * initialPrompt saved from the original creation attempt), creates
+   * a fresh session and re-sends the prompt instead of reconnecting
+   * to an empty session.
    */
   async clearSessionError(taskId: string, repoPath: string): Promise<void> {
+    const session = sessionStoreSetters.getSessionByTaskId(taskId);
+    if (session?.initialPrompt?.length) {
+      const { taskTitle, initialPrompt } = session;
+      await this.teardownSession(session.taskRunId);
+      const auth = this.getAuthCredentials();
+      if (!auth) {
+        throw new Error(
+          "Unable to reach server. Please check your connection.",
+        );
+      }
+      await this.createNewLocalSession(
+        taskId,
+        taskTitle,
+        repoPath,
+        auth,
+        initialPrompt,
+      );
+      return;
+    }
     await this.reconnectInPlace(taskId, repoPath);
   }
 
diff --git a/apps/code/src/renderer/features/sessions/stores/sessionStore.ts b/apps/code/src/renderer/features/sessions/stores/sessionStore.ts
index 86421dec0..ed1248d0e 100644
--- a/apps/code/src/renderer/features/sessions/stores/sessionStore.ts
+++ b/apps/code/src/renderer/features/sessions/stores/sessionStore.ts
@@ -1,4 +1,5 @@
 import type {
+  ContentBlock,
   SessionConfigOption,
   SessionConfigSelectGroup,
   SessionConfigSelectOption,
@@ -68,6 +69,8 @@ export interface AgentSession {
   cloudOutput?: Record<string, unknown> | null;
   /** Cloud task error message */
   cloudErrorMessage?: string | null;
+  /** Initial prompt to re-send on retry if the first connection attempt failed */
+  initialPrompt?: ContentBlock[];
   /** Cloud task branch */
   cloudBranch?: string | null;
   /** Number of session/prompt events to skip from polled logs (set during resume) */