fix(web): handle overlapping target persistence rollbacks

Author: pallyoung

packages/web/src/features/workspace/actions/use-persist-workspace-last-viewed-target.test.tsx

@@ -164,4 +164,111 @@ describe("usePersistWorkspaceLastViewedTarget", () => {
 
     expect(store.get(lastViewedTargetAtom)?.workspaceId).toBe("ws-2");
   });
+
+  it("rolls back to the last confirmed target when overlapping writes both fail", async () => {
+    const store = createStore();
+    const firstWrite = createDeferred<{ workspaceId: string; updatedAt: number }>();
+    const secondWrite = createDeferred<{ workspaceId: string; updatedAt: number }>();
+    const sendCommand = vi
+      .fn()
+      .mockImplementationOnce(() => firstWrite.promise)
+      .mockImplementationOnce(() => secondWrite.promise);
+
+    store.set(wsClientAtom, {
+      sendCommand,
+      subscribe: vi.fn(() => () => {}),
+    } as never);
+    store.set(lastViewedTargetAtom, {
+      workspaceId: "ws-0",
+      updatedAt: 1,
+    });
+
+    const { result } = renderHook(() => usePersistWorkspaceLastViewedTarget(), {
+      wrapper: wrapperFor(store),
+    });
+
+    let firstRequest!: Promise<unknown>;
+    let secondRequest!: Promise<unknown>;
+
+    act(() => {
+      firstRequest = result.current({ workspaceId: "ws-1" });
+      secondRequest = result.current({ workspaceId: "ws-2" });
+    });
+
+    expect(store.get(lastViewedTargetAtom)?.workspaceId).toBe("ws-2");
+
+    await act(async () => {
+      firstWrite.reject(
+        new CommandResultError({
+          code: "write_failed",
+          message: "failed",
+        })
+      );
+      await firstRequest;
+    });
+
+    expect(store.get(lastViewedTargetAtom)?.workspaceId).toBe("ws-2");
+
+    await act(async () => {
+      secondWrite.reject(
+        new CommandResultError({
+          code: "write_failed",
+          message: "failed",
+        })
+      );
+      await secondRequest;
+    });
+
+    expect(store.get(lastViewedTargetAtom)?.workspaceId).toBe("ws-0");
+  });
+
+  it("rolls back to an older confirmed target when a newer overlapping write fails", async () => {
+    const store = createStore();
+    const firstWrite = createDeferred<{ workspaceId: string; updatedAt: number }>();
+    const secondWrite = createDeferred<{ workspaceId: string; updatedAt: number }>();
+    const sendCommand = vi
+      .fn()
+      .mockImplementationOnce(() => firstWrite.promise)
+      .mockImplementationOnce(() => secondWrite.promise);
+
+    store.set(wsClientAtom, {
+      sendCommand,
+      subscribe: vi.fn(() => () => {}),
+    } as never);
+    store.set(lastViewedTargetAtom, {
+      workspaceId: "ws-0",
+      updatedAt: 1,
+    });
+
+    const { result } = renderHook(() => usePersistWorkspaceLastViewedTarget(), {
+      wrapper: wrapperFor(store),
+    });
+
+    let firstRequest!: Promise<unknown>;
+    let secondRequest!: Promise<unknown>;
+
+    act(() => {
+      firstRequest = result.current({ workspaceId: "ws-1" });
+      secondRequest = result.current({ workspaceId: "ws-2" });
+    });
+
+    await act(async () => {
+      firstWrite.resolve({ workspaceId: "ws-1", updatedAt: 11 });
+      await firstRequest;
+    });
+
+    expect(store.get(lastViewedTargetAtom)?.workspaceId).toBe("ws-2");
+
+    await act(async () => {
+      secondWrite.reject(
+        new CommandResultError({
+          code: "write_failed",
+          message: "failed",
+        })
+      );
+      await secondRequest;
+    });
+
+    expect(store.get(lastViewedTargetAtom)?.workspaceId).toBe("ws-1");
+  });
 });

packages/web/src/features/workspace/actions/use-persist-workspace-last-viewed-target.ts

@@ -1,4 +1,5 @@
 import type { WorkspaceLastViewedTarget } from "@coder-studio/core";
+import type { Store } from "jotai";
 import { useAtomValue, useSetAtom, useStore } from "jotai";
 import { useCallback } from "react";
 import { lastViewedTargetAtom } from "../../../atoms/app-ui";
@@ -9,6 +10,32 @@ interface PersistWorkspaceLastViewedTargetInput {
   sessionId?: string;
 }
 
+const confirmedTargetByStore = new WeakMap<Store, WorkspaceLastViewedTarget | null>();
+const optimisticTargetsByStore = new WeakMap<Store, WeakSet<WorkspaceLastViewedTarget>>();
+
+function getOptimisticTargets(store: Store): WeakSet<WorkspaceLastViewedTarget> {
+  const existing = optimisticTargetsByStore.get(store);
+  if (existing) {
+    return existing;
+  }
+
+  const next = new WeakSet<WorkspaceLastViewedTarget>();
+  optimisticTargetsByStore.set(store, next);
+  return next;
+}
+
+function getConfirmedTarget(store: Store): WorkspaceLastViewedTarget | null {
+  const currentTarget = store.get(lastViewedTargetAtom);
+  const optimisticTargets = getOptimisticTargets(store);
+
+  if (currentTarget && !optimisticTargets.has(currentTarget)) {
+    confirmedTargetByStore.set(store, currentTarget);
+    return currentTarget;
+  }
+
+  return confirmedTargetByStore.get(store) ?? null;
+}
+
 export function usePersistWorkspaceLastViewedTarget() {
   const dispatch = useAtomValue(dispatchCommandAtom);
   const setLastViewedTarget = useSetAtom(lastViewedTargetAtom);
@@ -17,6 +44,8 @@ export function usePersistWorkspaceLastViewedTarget() {
   return useCallback(
     async ({ workspaceId, sessionId }: PersistWorkspaceLastViewedTargetInput) => {
       const lastViewedTarget = store.get(lastViewedTargetAtom);
+      const optimisticTargets = getOptimisticTargets(store);
+      getConfirmedTarget(store);
 
       if (!workspaceId) {
         return null;
@@ -34,25 +63,38 @@ export function usePersistWorkspaceLastViewedTarget() {
         sessionId,
         updatedAt: Date.now(),
       };
-      const previousTarget = lastViewedTarget;
+      optimisticTargets.add(optimisticTarget);
       setLastViewedTarget(optimisticTarget);
 
       const result = await dispatch<WorkspaceLastViewedTarget>("workspace.lastViewedTarget.set", {
         workspaceId,
         sessionId,
       });
 
-      if (store.get(lastViewedTargetAtom) !== optimisticTarget) {
-        return store.get(lastViewedTargetAtom);
+      const latestTarget = store.get(lastViewedTargetAtom);
+
+      if (result.ok && result.data) {
+        confirmedTargetByStore.set(store, result.data);
+
+        if (latestTarget === optimisticTarget) {
+          optimisticTargets.delete(optimisticTarget);
+          setLastViewedTarget(result.data);
+          return result.data;
+        }
+
+        optimisticTargets.delete(optimisticTarget);
+        return latestTarget;
       }
 
-      if (!result.ok || !result.data) {
-        setLastViewedTarget(previousTarget);
-        return optimisticTarget;
+      if (latestTarget === optimisticTarget) {
+        const confirmedTarget = getConfirmedTarget(store);
+        optimisticTargets.delete(optimisticTarget);
+        setLastViewedTarget(confirmedTarget);
+        return confirmedTarget;
       }
 
-      setLastViewedTarget(result.data);
-      return result.data;
+      optimisticTargets.delete(optimisticTarget);
+      return latestTarget;
     },
     [dispatch, setLastViewedTarget, store]
   );