fix(ui): retain startup autosave through session recovery

Co-Authored-By: Paperclip <noreply@paperclip.ing>

tests/e2e/disconnect-write-recovery.spec.ts

@@ -178,10 +178,19 @@ test("replaying an accepted comment does not reopen completed work", async ({ pa
   }
 });
 
-test("startup work retries a failed session lookup while other APIs stay healthy", async ({ page, request }) => {
+for (const delayWorkspacePreference of [false, true]) test(delayWorkspacePreference
+  ? "startup title survives session recovery while workspace preferences remount the editor"
+  : "startup work retries a failed session lookup while other APIs stay healthy", async ({ page, request }) => {
   const company = await json(await request.post("/api/companies", { data: { name: `Session recovery ${Date.now()}` } }));
   const session = await json(await request.get("/api/auth/get-session"));
   let healthy = false, saving = false, failedAfterSave = 0;
+  let releasePreference!: () => void;
+  const preferenceGate = new Promise<void>(resolve => { releasePreference = resolve; });
+  let sawPreference!: () => void;
+  const preferenceRequested = new Promise<void>(resolve => { sawPreference = resolve; });
+  if (delayWorkspacePreference) await page.route("**/api/workspace-preferences/me", async route => {
+    sawPreference(); await preferenceGate; await route.continue();
+  });
   try {
     const issue = await json(await request.post(`/api/companies/${company.id}/issues`, {
       data: { title: "Before session recovery", status: "backlog" },
@@ -199,12 +208,24 @@ test("startup work retries a failed session lookup while other APIs stay healthy
     saving = true;
     await header.locator("textarea").press("Enter");
     await expect.poll(() => failedAfterSave).toBeGreaterThan(0);
+    if (delayWorkspacePreference) {
+      await expect(page.getByText("Couldn't update the task", { exact: true })).toBeVisible();
+      await page.evaluate(() => new Promise(resolve => requestAnimationFrame(() => resolve(null))));
+    }
     healthy = true;
+    if (delayWorkspacePreference) {
+      // Returning to the tab retries session lookup while the saved edit waits.
+      await page.evaluate(() => window.dispatchEvent(new Event("visibilitychange")));
+      await preferenceRequested;
+      await expect(page.getByTestId("task-chat-history-content")).toBeHidden();
+      releasePreference();
+    }
     await expect.poll(async () => (await json(await request.get(`/api/issues/${issue.id}`))).title,
       { timeout: 15_000 }).toBe("Saved after session recovery");
     await page.reload();
     await expect(header.getByRole("heading", { name: "Saved after session recovery", exact: true })).toBeVisible();
   } finally {
+    releasePreference();
     await request.patch(`/api/companies/${company.id}`, { data: { status: "archived" } });
   }
 });

ui/src/hooks/useAccountBoundMutation.ts

@@ -5,68 +5,12 @@ import {
   QueryClientContext,
   type DefaultError,
   type MutateOptions,
-  type QueryClient,
   type UseMutationOptions,
   type UseMutationResult,
   type MutationFunctionContext,
 } from "@tanstack/react-query";
-import { queryKeys } from "@/lib/queryKeys";
-import { authApi } from "@/api/auth";
-import { classifySendFailure, PendingSendAccountChangedError } from "@/lib/pending-send";
-
-type AccountGeneration = { owner: string | undefined; generation: number };
-const accounts = new WeakMap<QueryClient, AccountGeneration>();
-
-function currentOwner(client: QueryClient): string | undefined {
-  const session = client.getQueryData<{ user?: { id?: string }; session?: { userId?: string } } | null>(queryKeys.auth.session);
-  return session === undefined ? undefined : session?.user?.id ?? session?.session?.userId ?? "signed-out";
-}
-
-function bindAccount(client: QueryClient) {
-  let account = accounts.get(client);
-  if (!account) {
-    account = { owner: currentOwner(client), generation: 0 };
-    accounts.set(client, account);
-    const tracked = account;
-    client.getQueryCache().subscribe(event => {
-      const owner = currentOwner(client);
-      const removed = event.type === "removed" && JSON.stringify(event.query.queryKey) === JSON.stringify(queryKeys.auth.session);
-      if (owner !== tracked.owner || removed) {
-        if (removed || tracked.owner !== undefined || tracked.generation !== 0) tracked.generation += 1;
-        tracked.owner = owner;
-      }
-    });
-  }
-  const { owner, generation } = account;
-  const isCurrentScope = () => account.generation === generation && (owner === undefined ? generation === 0 : account.owner === owner);
-  const isCurrentAccount = () => account.owner !== undefined && isCurrentScope();
-  const whenAccountReady = async () => {
-    if (account.owner !== undefined || !isCurrentScope()) return;
-    let stopped = false;
-    let unsubscribe = () => {};
-    const ready = new Promise<void>(resolve => {
-      unsubscribe = client.getQueryCache().subscribe(() => {
-        if (account.owner !== undefined || !isCurrentScope()) resolve();
-      });
-    });
-    const lookup = async () => {
-      for (let attempt = 0; attempt < 3; attempt += 1) {
-        if (stopped || account.owner !== undefined || !isCurrentScope()) return;
-        try {
-          await client.fetchQuery({ queryKey: queryKeys.auth.session, queryFn: () => authApi.getSession(), retry: false });
-          return;
-        } catch (error) {
-          if (stopped || account.owner !== undefined || !isCurrentScope()) return;
-          if (attempt === 2 || classifySendFailure(error) !== "pending") throw error;
-          await new Promise(resolve => setTimeout(resolve, 250 * 2 ** attempt));
-        }
-      }
-    };
-    try { await Promise.race([ready, lookup()]); }
-    finally { stopped = true; unsubscribe(); }
-  };
-  return { isCurrentAccount, isCurrentScope, whenAccountReady };
-}
+import { PendingSendAccountChangedError } from "@/lib/pending-send";
+import { bindAccount } from "@/lib/account-scope";
 
 export function useCurrentAccountGuard(): () => boolean {
   const client = useContext(QueryClientContext);

ui/src/hooks/useDurableAutosave.test.tsx

@@ -173,6 +173,51 @@ describe("useDurableAutosave", () => {
     client.clear();
   });
 
+  it.each([false, true])("keeps an unowned startup edit when identity and unmount race (identityFirst=%s)", async identityFirst => {
+    store.setBrowserOnline(false);
+    const client = new QueryClient({ defaultOptions: { queries: { staleTime: Infinity, retry: false, enabled: false } } });
+    const save = vi.fn().mockResolvedValue(undefined);
+    act(() => root.render(<QueryClientProvider client={client}><ConnectivityProvider store={store}>
+      <Harness sourceId="startup-remount" targetKey="task-a:title" save={save} />
+    </ConnectivityProvider></QueryClientProvider>));
+    await act(async () => { await current.run("Retained startup title"); });
+    expect(current.status).toBe("waiting");
+    act(() => {
+      if (identityFirst) client.setQueryData(queryKeys.auth.session, { user: { id: "user-1" } });
+      root.unmount(); root = createRoot(container);
+      if (!identityFirst) client.setQueryData(queryKeys.auth.session, { user: { id: "user-1" } });
+    });
+    await act(async () => { store.setBrowserOnline(true); store.reportSuccess(); });
+    await advance(2_000);
+    expect(save).toHaveBeenCalledExactlyOnceWith("Retained startup title");
+    expect(store.getSnapshot().pendingWrites).toBe(0);
+    client.clear();
+  });
+
+  it.each(["different-account", "round-trip", "removed-session", "signed-out"])(
+    "does not adopt a later identity for a detached startup edit (%s)", async transition => {
+      store.setBrowserOnline(false);
+      const client = new QueryClient({ defaultOptions: { queries: { staleTime: Infinity, retry: false, enabled: false } } });
+      const save = vi.fn().mockResolvedValue(undefined);
+      act(() => root.render(<QueryClientProvider client={client}><ConnectivityProvider store={store}>
+        <Harness sourceId="private-startup" targetKey="private-title" save={save} />
+      </ConnectivityProvider></QueryClientProvider>));
+      await act(async () => { await current.run("Private startup title"); });
+      act(() => { root.unmount(); root = createRoot(container); });
+      act(() => {
+        if (transition === "removed-session") client.removeQueries({ queryKey: queryKeys.auth.session });
+        else client.setQueryData(queryKeys.auth.session, transition === "signed-out" ? null : { user: { id: "user-1" } });
+        if (transition === "round-trip") client.setQueryData(queryKeys.auth.session, { user: { id: "user-2" } });
+        client.setQueryData(queryKeys.auth.session, { user: { id: transition === "different-account" ? "user-2" : "user-1" } });
+      });
+      await act(async () => { store.setBrowserOnline(true); store.reportSuccess(); });
+      await advance(2_000);
+      expect(save).not.toHaveBeenCalled();
+      expect(store.getSnapshot().pendingWrites).toBe(0);
+      client.clear();
+    },
+  );
+
   it("bounds retries after an editor unmounts", async () => {
     const save = vi.fn().mockRejectedValue(outage());
     render({ save });

ui/src/hooks/useDurableAutosave.ts

@@ -1,4 +1,6 @@
-import { useCallback, useEffect, useId, useRef, useState } from "react";
+import { useCallback, useContext, useEffect, useId, useRef, useState } from "react";
+import { QueryClientContext } from "@tanstack/react-query";
+import { bindAccount } from "@/lib/account-scope";
 import { classifyError, describeError, type ErrorDescription } from "@/api/errors";
 import { loadStructuredDraft, saveStructuredDraft } from "@/lib/composer-draft";
 import { useConnectivityStore } from "@/lib/connectivity";
@@ -108,6 +110,7 @@ const SAVED_LINGER_MS = 1600;
  * - Errors become readable copy for "Couldn't save (reason)".
  */
 export function useDurableAutosave<T>(options: DurableAutosaveOptions<T>): DurableAutosave<T> {
+  const client = useContext(QueryClientContext);
   const owner = useDraftOwner();
   const readOwner = useDraftOwnerReader();
   const instanceId = useId();
@@ -120,7 +123,7 @@ export function useDurableAutosave<T>(options: DurableAutosaveOptions<T>): Durab
   const targetKey = options.targetKey ?? options.sourceId;
   const targetRef = useRef(targetKey);
   targetRef.current = targetKey;
-  type Job = { token: symbol; registryKey: string; resends: number; saving?: boolean; id: number; target: string; owner: string | undefined; draft: T;
+  type Job = { token: symbol; registryKey: string; resends: number; saving?: boolean; id: number; target: string; owner: string | undefined; account: ReturnType<typeof bindAccount> | null; draft: T;
     save: (draft: T) => Promise<unknown>; matchesServer?: DurableAutosaveOptions<T>["matchesServer"];
     onRetryFailed?: DurableAutosaveJob<T>["onRetryFailed"]; storageKey: string | null; action: string };
   const jobs = useRef(new Map<string, Job>());
@@ -142,6 +145,22 @@ export function useDurableAutosave<T>(options: DurableAutosaveOptions<T>): Durab
     entry?.active.delete(job.token);
     if (entry && !entry.active.size) registry.delete(job.registryKey);
   };
+  // Resolve first identity from the same generation guard used by mutations.
+  // This also runs after unmount, when React cannot migrate the job for us.
+  function owned(job: Job) {
+    if (job.account && !job.account.isCurrentScope()) return false;
+    const current = readOwner();
+    if (job.owner === undefined && current !== undefined) {
+      const nextKey = JSON.stringify([current, job.target]);
+      if (registry.get(job.registryKey)?.latest !== job.token || registry.has(nextKey)) return false;
+      releaseJob(job);
+      job.owner = current;
+      job.registryKey = nextKey;
+      registry.set(nextKey, { latest: job.token, active: new Set([job.token]) });
+    }
+    return current === job.owner && generations.current.get(job.target) === job.id
+      && registry.get(job.registryKey)?.latest === job.token;
+  }
   const mounted = useRef(true);
   const savingTimer = useRef<ReturnType<typeof setTimeout> | null>(null);
   const savedTimer = useRef<ReturnType<typeof setTimeout> | null>(null);
@@ -160,17 +179,11 @@ export function useDurableAutosave<T>(options: DurableAutosaveOptions<T>): Durab
     if (previous === undefined && owner !== undefined) {
       for (const job of jobs.current.values()) {
         if (job.owner !== undefined) continue;
-        const nextKey = JSON.stringify([owner, job.target]);
-        const superseded = registry.get(job.registryKey)?.latest !== job.token || registry.has(nextKey);
-        releaseJob(job);
-        if (superseded) {
+        if (!owned(job)) {
+          releaseJob(job);
           jobs.current.delete(job.target);
           if (generations.current.get(job.target) === job.id) generations.current.delete(job.target);
-          continue;
         }
-        job.owner = owner;
-        job.registryKey = nextKey;
-        registry.set(nextKey, { latest: job.token, active: new Set([job.token]) });
       }
       refresh(value => value + 1);
       return;
@@ -194,12 +207,12 @@ export function useDurableAutosave<T>(options: DurableAutosaveOptions<T>): Durab
     return stored?.version === 1 && stored.draft !== undefined ? stored.draft : null;
   }, [storageKey]);
   const execute = useCallback(async (job: Job, retry = false): Promise<DurableAutosaveResult> => {
-    const owned = () => readOwner() === job.owner && generations.current.get(job.target) === job.id && registry.get(job.registryKey)?.latest === job.token;
-    const visible = () => mounted.current && owned() && targetRef.current === job.target;
+    const stillOwned = () => owned(job);
+    const visible = () => mounted.current && stillOwned() && targetRef.current === job.target;
     const publish = (status: DurableAutosaveStatus, error: ErrorDescription | null = null) => {
       if (visible()) setView({ target: job.target, status, error });
     };
-    if (!owned()) {
+    if (!stillOwned()) {
       if (jobs.current.get(job.target) === job) jobs.current.delete(job.target);
       releaseJob(job);
       if (mounted.current) refresh(value => value + 1);
@@ -234,7 +247,12 @@ export function useDurableAutosave<T>(options: DurableAutosaveOptions<T>): Durab
       remove();
       return { kind: "saved" };
     };
-    try { await job.save(job.draft); return done(); }
+    try {
+      if (!mounted.current && job.account && !job.account.isCurrentAccount()) await job.account.whenAccountReady();
+      if (!stillOwned()) { remove(); return { kind: "waiting" }; }
+      await job.save(job.draft);
+      return done();
+    }
     catch (initialCause) {
       let cause = initialCause;
       if (classifyError(cause) === "conflict" && job.matchesServer) {
@@ -255,7 +273,7 @@ export function useDurableAutosave<T>(options: DurableAutosaveOptions<T>): Durab
     const current = optionsRef.current;
     const target = current.targetKey ?? current.sourceId;
     const admittedOwner = ownerRef.current ?? readOwner();
-    const job: Job = { token: Symbol(), registryKey: JSON.stringify([admittedOwner, target]), resends: 0, id: ++sequence.current, target, owner: admittedOwner, draft,
+    const job: Job = { token: Symbol(), registryKey: JSON.stringify([admittedOwner, target]), resends: 0, id: ++sequence.current, target, owner: admittedOwner, account: client ? bindAccount(client) : null, draft,
       save: overrides.save ?? current.save, matchesServer: overrides.matchesServer ?? current.matchesServer,
       onRetryFailed: overrides.onRetryFailed, storageKey: current.storageKey ?? null, action: current.action ?? "save" };
     const previous = jobs.current.get(target);
@@ -266,7 +284,7 @@ export function useDurableAutosave<T>(options: DurableAutosaveOptions<T>): Durab
     generations.current.set(target, job.id); jobs.current.set(target, job);
     refresh(value => value + 1);
     return execute(job);
-  }, [execute, readOwner]);
+  }, [client, execute, readOwner]);
   const waitingJobs = [...jobs.current.values()].filter(job => !job.saving).length;
   const auto = useAutoResend({ sourceId: options.sourceId, active: waitingJobs > 0, pendingCount: waitingJobs,
     resend: async (): Promise<SendOutcome> => {
@@ -281,7 +299,6 @@ export function useDurableAutosave<T>(options: DurableAutosaveOptions<T>): Durab
   detachedWork.current = () => {
     if (!jobs.current.size) return;
     if (auto.stalled) { for (const job of jobs.current.values()) releaseJob(job); jobs.current.clear(); return; }
-    const admittedOwner = ownerRef.current;
     const source = `${optionsRef.current.sourceId}:detached:${instanceId}`;
     let timer: ReturnType<typeof setTimeout> | null = null;
     let disposed = false, running = false;
@@ -298,7 +315,11 @@ export function useDurableAutosave<T>(options: DurableAutosaveOptions<T>): Durab
     };
     const wake = async () => {
       if (disposed || running) return;
-      if (mounted.current || readOwner() !== admittedOwner || admittedOwner === undefined) { release(); return; }
+      if (mounted.current) { release(); return; }
+      for (const job of jobs.current.values()) if (!owned(job)) {
+        releaseJob(job);
+        jobs.current.delete(job.target);
+      }
       if (!jobs.current.size) { release(); return; }
       store.setPendingWrites(source, jobs.current.size);
       if (store.getSnapshot().status !== "online") return;

ui/src/lib/account-scope.ts

@@ -0,0 +1,59 @@
+import type { QueryClient } from "@tanstack/react-query";
+import { queryKeys } from "./queryKeys";
+import { authApi } from "@/api/auth";
+import { classifySendFailure } from "./pending-send";
+
+type AccountGeneration = { owner: string | undefined; generation: number };
+const accounts = new WeakMap<QueryClient, AccountGeneration>();
+
+function currentOwner(client: QueryClient): string | undefined {
+  const session = client.getQueryData<{ user?: { id?: string }; session?: { userId?: string } } | null>(queryKeys.auth.session);
+  return session === undefined ? undefined : session?.user?.id ?? session?.session?.userId ?? "signed-out";
+}
+
+export function bindAccount(client: QueryClient) {
+  let account = accounts.get(client);
+  if (!account) {
+    account = { owner: currentOwner(client), generation: 0 };
+    accounts.set(client, account);
+    const tracked = account;
+    client.getQueryCache().subscribe(event => {
+      const owner = currentOwner(client);
+      const removed = event.type === "removed" && JSON.stringify(event.query.queryKey) === JSON.stringify(queryKeys.auth.session);
+      if (owner !== tracked.owner || removed) {
+        if (removed || tracked.owner !== undefined || tracked.generation !== 0) tracked.generation += 1;
+        tracked.owner = owner;
+      }
+    });
+  }
+  const { owner, generation } = account;
+  const isCurrentScope = () => account.generation === generation && (owner === undefined ? generation === 0 : account.owner === owner);
+  const isCurrentAccount = () => account.owner !== undefined && isCurrentScope();
+  const whenAccountReady = async () => {
+    if (account.owner !== undefined || !isCurrentScope()) return;
+    let stopped = false;
+    let unsubscribe = () => {};
+    const ready = new Promise<void>(resolve => {
+      unsubscribe = client.getQueryCache().subscribe(() => {
+        if (account.owner !== undefined || !isCurrentScope()) resolve();
+      });
+    });
+    const lookup = async () => {
+      for (let attempt = 0; attempt < 3; attempt += 1) {
+        if (stopped || account.owner !== undefined || !isCurrentScope()) return;
+        try {
+          await client.fetchQuery({ queryKey: queryKeys.auth.session, queryFn: () => authApi.getSession(), retry: false });
+          return;
+        } catch (error) {
+          if (stopped || account.owner !== undefined || !isCurrentScope()) return;
+          if (attempt === 2 || classifySendFailure(error) !== "pending") throw error;
+          await new Promise(resolve => setTimeout(resolve, 250 * 2 ** attempt));
+        }
+      }
+    };
+    try { await Promise.race([ready, lookup()]); }
+    finally { stopped = true; unsubscribe(); }
+  };
+  return { isCurrentAccount, isCurrentScope, whenAccountReady };
+}
+