feat: persist pending task prompts and recover it on failure

packages/ui/src/features/task-detail/components/PendingPromptRecovery.tsx

@@ -0,0 +1,34 @@
+import { openTaskInput } from "@posthog/ui/router/useOpenTask";
+import { logger } from "@posthog/ui/shell/logger";
+import { pendingTaskPromptStoreApi } from "@posthog/ui/shell/pendingTaskPromptStore";
+import { useEffect } from "react";
+
+const log = logger.scope("pending-prompt-recovery");
+
+let recoveryStarted = false;
+
+export function PendingPromptRecovery(): null {
+  useEffect(() => {
+    if (recoveryStarted) {
+      return;
+    }
+    recoveryStarted = true;
+    void recoverNewestPendingPrompt();
+  }, []);
+  return null;
+}
+
+async function recoverNewestPendingPrompt(): Promise<void> {
+  await pendingTaskPromptStoreApi.whenHydrated();
+  const orphans = pendingTaskPromptStoreApi.getAllNewestFirst();
+  const [newest] = orphans;
+  if (!newest) {
+    return;
+  }
+
+  log.info("Recovering an unsent prompt whose task never finished creating", {
+    remaining: orphans.length - 1,
+  });
+  pendingTaskPromptStoreApi.clear(newest.key);
+  openTaskInput({ initialPrompt: newest.prompt.promptText });
+}

packages/ui/src/features/task-detail/hooks/useTaskCreation.ts

@@ -279,6 +279,8 @@ export function useTaskCreation({
         }
       }
 
+      let createdTaskId: string | undefined;
+
       try {
         if (!contentOverride) {
           const plainText = editor.getText()?.trim() ?? plainPromptText;
@@ -341,6 +343,7 @@ export function useTaskCreation({
             if (signalReportId) {
               clearTaskInputReportAssociation();
             }
+            createdTaskId = output.task.id;
             if (pendingTaskKey) {
               pendingTaskPromptStoreApi.move(pendingTaskKey, output.task.id);
             }
@@ -376,6 +379,14 @@ export function useTaskCreation({
         }
 
         if (result.success) {
+          if (!result.data.provisioningError) {
+            if (pendingTaskKey) {
+              pendingTaskPromptStoreApi.clear(pendingTaskKey);
+            }
+            if (createdTaskId) {
+              pendingTaskPromptStoreApi.clear(createdTaskId);
+            }
+          }
           setAdditionalDirectoriesOverride(null);
           // Guarantee the editor draft is wiped on success. editor.clear()
           // above only runs inside the onTaskReady callback (and after it
@@ -412,6 +423,9 @@ export function useTaskCreation({
           }
           if (pendingTaskKey) {
             pendingTaskPromptStoreApi.clear(pendingTaskKey);
+            if (createdTaskId) {
+              pendingTaskPromptStoreApi.clear(createdTaskId);
+            }
             openTaskInput({ initialPrompt: plainPromptText });
           }
         }
@@ -423,6 +437,9 @@ export function useTaskCreation({
         log.error("Unexpected error during task creation", { error });
         if (pendingTaskKey) {
           pendingTaskPromptStoreApi.clear(pendingTaskKey);
+          if (createdTaskId) {
+            pendingTaskPromptStoreApi.clear(createdTaskId);
+          }
           openTaskInput({ initialPrompt: plainPromptText });
         }
         return false;

packages/ui/src/shell/App.tsx

@@ -18,6 +18,7 @@ import { OnboardingFlow } from "@posthog/ui/features/onboarding/components/Onboa
 import { useOnboardingStore } from "@posthog/ui/features/onboarding/onboardingStore";
 import { SettingsDialog } from "@posthog/ui/features/settings/SettingsDialog";
 import { UpdateBanner } from "@posthog/ui/features/sidebar/components/UpdateBanner";
+import { PendingPromptRecovery } from "@posthog/ui/features/task-detail/components/PendingPromptRecovery";
 import { LoginTransition } from "@posthog/ui/primitives/LoginTransition";
 import { router } from "@posthog/ui/router/router";
 import { track } from "@posthog/ui/shell/analytics";
@@ -165,6 +166,7 @@ function App() {
             from anywhere in the app. Sibling of the router so it stays mounted
             across every route (not just the canvas space). Renders null. */}
         <CanvasGenerationToaster />
+        <PendingPromptRecovery />
       </motion.div>
     );
   };

packages/ui/src/shell/pendingTaskPromptStore.test.ts

@@ -0,0 +1,105 @@
+import {
+  pendingTaskPromptStoreApi,
+  usePendingTaskPromptStore,
+} from "@posthog/ui/shell/pendingTaskPromptStore";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+function resetStore(): void {
+  usePendingTaskPromptStore.setState({ byKey: {}, _hasHydrated: false });
+}
+
+describe("pendingTaskPromptStore", () => {
+  beforeEach(() => {
+    resetStore();
+    let now = 1_000;
+    vi.spyOn(Date, "now").mockImplementation(() => {
+      now += 1_000;
+      return now;
+    });
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+    resetStore();
+  });
+
+  it("stamps createdAt when a prompt is set", () => {
+    pendingTaskPromptStoreApi.set("k1", { promptText: "hi", attachments: [] });
+    const stored = pendingTaskPromptStoreApi.get("k1");
+    expect(stored?.promptText).toBe("hi");
+    expect(typeof stored?.createdAt).toBe("number");
+  });
+
+  it("moves an entry to a new key, preserving its contents", () => {
+    pendingTaskPromptStoreApi.set("pending-1", {
+      promptText: "draft",
+      attachments: [],
+    });
+    const before = pendingTaskPromptStoreApi.get("pending-1");
+
+    pendingTaskPromptStoreApi.move("pending-1", "task-42");
+
+    expect(pendingTaskPromptStoreApi.get("pending-1")).toBeUndefined();
+    expect(pendingTaskPromptStoreApi.get("task-42")).toEqual(before);
+  });
+
+  it("clears an entry by key", () => {
+    pendingTaskPromptStoreApi.set("k1", { promptText: "x", attachments: [] });
+    pendingTaskPromptStoreApi.clear("k1");
+    expect(pendingTaskPromptStoreApi.get("k1")).toBeUndefined();
+  });
+
+  it("returns every surviving entry newest-first for recovery", () => {
+    pendingTaskPromptStoreApi.set("old", {
+      promptText: "old",
+      attachments: [],
+    });
+    pendingTaskPromptStoreApi.set("mid", {
+      promptText: "mid",
+      attachments: [],
+    });
+    pendingTaskPromptStoreApi.set("new", {
+      promptText: "new",
+      attachments: [],
+    });
+
+    const recovered = pendingTaskPromptStoreApi.getAllNewestFirst();
+
+    expect(recovered.map((r) => r.key)).toEqual(["new", "mid", "old"]);
+  });
+
+  it("caps stored prompts to the newest, dropping the oldest", () => {
+    for (let i = 0; i < 22; i++) {
+      pendingTaskPromptStoreApi.set(`k${i}`, {
+        promptText: `p${i}`,
+        attachments: [],
+      });
+    }
+
+    const recovered = pendingTaskPromptStoreApi.getAllNewestFirst();
+
+    expect(recovered).toHaveLength(20);
+    expect(pendingTaskPromptStoreApi.get("k0")).toBeUndefined();
+    expect(pendingTaskPromptStoreApi.get("k1")).toBeUndefined();
+    expect(pendingTaskPromptStoreApi.get("k21")).toBeDefined();
+  });
+
+  it("whenHydrated resolves immediately once the store has hydrated", async () => {
+    usePendingTaskPromptStore.getState().setHasHydrated(true);
+    await expect(
+      pendingTaskPromptStoreApi.whenHydrated(),
+    ).resolves.toBeUndefined();
+  });
+
+  it("whenHydrated waits for hydration to complete", async () => {
+    let resolved = false;
+    const waiter = pendingTaskPromptStoreApi.whenHydrated().then(() => {
+      resolved = true;
+    });
+
+    expect(resolved).toBe(false);
+    usePendingTaskPromptStore.getState().setHasHydrated(true);
+    await waiter;
+    expect(resolved).toBe(true);
+  });
+});

packages/ui/src/shell/pendingTaskPromptStore.ts

@@ -1,50 +1,131 @@
 import type { UserMessageAttachment } from "@posthog/ui/features/sessions/userMessageTypes";
+import { logger } from "@posthog/ui/shell/logger";
+import { electronStorage } from "@posthog/ui/shell/rendererStorage";
 import { create } from "zustand";
+import { persist } from "zustand/middleware";
+
+const log = logger.scope("pending-task-prompts");
+
+const MAX_PENDING_PROMPTS = 20;
 
 export interface PendingTaskPrompt {
   promptText: string;
   attachments: UserMessageAttachment[];
+  createdAt: number;
+}
+
+export type PendingTaskPromptInput = Omit<PendingTaskPrompt, "createdAt">;
+
+function capToNewest(
+  byKey: Record<string, PendingTaskPrompt>,
+): Record<string, PendingTaskPrompt> {
+  const keys = Object.keys(byKey);
+  if (keys.length <= MAX_PENDING_PROMPTS) {
+    return byKey;
+  }
+  const keptKeys = keys
+    .sort((a, b) => byKey[b].createdAt - byKey[a].createdAt)
+    .slice(0, MAX_PENDING_PROMPTS);
+  log.warn("Dropping oldest unrecovered prompts beyond cap", {
+    dropped: keys.length - keptKeys.length,
+  });
+  const kept: Record<string, PendingTaskPrompt> = {};
+  for (const key of keptKeys) {
+    kept[key] = byKey[key];
+  }
+  return kept;
 }
 
 interface PendingTaskPromptStore {
   byKey: Record<string, PendingTaskPrompt>;
-  set: (key: string, prompt: PendingTaskPrompt) => void;
+  _hasHydrated: boolean;
+  setHasHydrated: (hydrated: boolean) => void;
+  set: (key: string, prompt: PendingTaskPromptInput) => void;
   get: (key: string) => PendingTaskPrompt | undefined;
   move: (fromKey: string, toKey: string) => void;
   clear: (key: string) => void;
 }
 
-export const usePendingTaskPromptStore = create<PendingTaskPromptStore>(
-  (set, get) => ({
-    byKey: {},
-    set: (key, prompt) =>
-      set((state) => ({ byKey: { ...state.byKey, [key]: prompt } })),
-    get: (key) => get().byKey[key],
-    move: (fromKey, toKey) => {
-      if (fromKey === toKey) return;
-      set((state) => {
-        const entry = state.byKey[fromKey];
-        if (!entry) return state;
-        const { [fromKey]: _removed, ...rest } = state.byKey;
-        return { byKey: { ...rest, [toKey]: entry } };
-      });
+export const usePendingTaskPromptStore = create<PendingTaskPromptStore>()(
+  persist(
+    (set, get) => ({
+      byKey: {},
+      _hasHydrated: false,
+      setHasHydrated: (hydrated) => set({ _hasHydrated: hydrated }),
+      set: (key, prompt) =>
+        set((state) => ({
+          byKey: capToNewest({
+            ...state.byKey,
+            [key]: { ...prompt, createdAt: Date.now() },
+          }),
+        })),
+      get: (key) => get().byKey[key],
+      move: (fromKey, toKey) => {
+        if (fromKey === toKey) {
+          return;
+        }
+        set((state) => {
+          const entry = state.byKey[fromKey];
+          if (!entry) {
+            return state;
+          }
+          const { [fromKey]: _removed, ...rest } = state.byKey;
+          return { byKey: { ...rest, [toKey]: entry } };
+        });
+      },
+      clear: (key) =>
+        set((state) => {
+          if (!(key in state.byKey)) {
+            return state;
+          }
+          const { [key]: _removed, ...rest } = state.byKey;
+          return { byKey: rest };
+        }),
+    }),
+    {
+      name: "pending-task-prompts",
+      storage: electronStorage,
+      partialize: (state) => ({ byKey: state.byKey }),
+      onRehydrateStorage: () => (state, error) => {
+        if (error) {
+          usePendingTaskPromptStore.getState().setHasHydrated(true);
+        } else {
+          state?.setHasHydrated(true);
+        }
+      },
     },
-    clear: (key) =>
-      set((state) => {
-        if (!(key in state.byKey)) return state;
-        const { [key]: _removed, ...rest } = state.byKey;
-        return { byKey: rest };
-      }),
-  }),
+  ),
 );
 
+export interface RecoverablePendingPrompt {
+  key: string;
+  prompt: PendingTaskPrompt;
+}
+
 export const pendingTaskPromptStoreApi = {
-  set: (key: string, prompt: PendingTaskPrompt) =>
+  set: (key: string, prompt: PendingTaskPromptInput) =>
     usePendingTaskPromptStore.getState().set(key, prompt),
   get: (key: string) => usePendingTaskPromptStore.getState().get(key),
   move: (fromKey: string, toKey: string) =>
     usePendingTaskPromptStore.getState().move(fromKey, toKey),
   clear: (key: string) => usePendingTaskPromptStore.getState().clear(key),
+  getAllNewestFirst: (): RecoverablePendingPrompt[] =>
+    Object.entries(usePendingTaskPromptStore.getState().byKey)
+      .map(([key, prompt]) => ({ key, prompt }))
+      .sort((a, b) => b.prompt.createdAt - a.prompt.createdAt),
+  whenHydrated: (): Promise<void> => {
+    if (usePendingTaskPromptStore.getState()._hasHydrated) {
+      return Promise.resolve();
+    }
+    return new Promise((resolve) => {
+      const unsubscribe = usePendingTaskPromptStore.subscribe((state) => {
+        if (state._hasHydrated) {
+          unsubscribe();
+          resolve();
+        }
+      });
+    });
+  },
 };
 
 export function usePendingTaskPrompt(