feat(web): add single-active session gate recovery

Author: pallyoung

packages/web/src/app/providers.lifecycle.test.tsx

@@ -2,8 +2,14 @@ import type { Workspace } from "@coder-studio/core";
 import { act, render } from "@testing-library/react";
 import { createStore, Provider } from "jotai";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import {
+  activationGenerationAtom,
+  activationReasonAtom,
+  activationStatusAtom,
+} from "../atoms/activation";
 import { authenticatedAtom } from "../atoms/app-ui";
 import { authEnabledAtom, connectionStatusAtom } from "../atoms/connection";
+import { sessionsAtom } from "../atoms/sessions";
 import {
   activeWorkspaceIdAtom,
   workspaceOrderAtom,
@@ -12,9 +18,11 @@ import {
 } from "../atoms/workspaces";
 import { terminalPreferencesAtom } from "../features/terminal-panel/preferences";
 import {
+  fileTreeAtomFamily,
   fileTreeStaleAtomFamily,
   gitBranchListAtomFamily,
   gitStateAtomFamily,
+  loadedDirsAtomFamily,
   worktreeListAtomFamily,
 } from "../features/workspace/atoms";
 import { AppProviders, resetAppProvidersSingletonsForTests } from "./providers";
@@ -56,6 +64,30 @@ function renderProviders(store = createStore()) {
   return { store, ...rendered };
 }
 
+function createWsSendCommandMock(
+  handler?: (op: string, args: unknown) => Promise<unknown> | unknown
+) {
+  return vi.fn().mockImplementation(async (op: string, args: unknown) => {
+    if (op === "activation.claim") {
+      return {
+        active: true,
+        generation: 1,
+        recoveryMode: "fresh",
+      };
+    }
+
+    if (op === "activation.heartbeat" || op === "activation.release") {
+      return { ok: true };
+    }
+
+    if (handler) {
+      return await handler(op, args);
+    }
+
+    return undefined;
+  });
+}
+
 function setVisibilityState(value: "visible" | "hidden") {
   Object.defineProperty(document, "visibilityState", {
     configurable: true,
@@ -128,7 +160,7 @@ describe("AppProviders lifecycle recovery", () => {
       }),
       getStatus: vi.fn(() => "disconnected"),
       recoverConnection: vi.fn(),
-      sendCommand: vi.fn().mockResolvedValue(undefined),
+      sendCommand: createWsSendCommandMock(),
     };
   });
 
@@ -452,6 +484,151 @@ describe("AppProviders lifecycle recovery", () => {
     });
   });
 
+  it("claims activation when the websocket becomes connected", async () => {
+    const store = createStore();
+    setVisibilityState("visible");
+
+    renderProviders(store);
+
+    await vi.waitFor(() => {
+      expect(wsState.client?.connect).toHaveBeenCalled();
+    });
+
+    act(() => {
+      wsState.client?.statusHandler?.("connected");
+    });
+
+    await vi.waitFor(() => {
+      const claimCalls =
+        wsState.client?.sendCommand?.mock.calls.filter(([op]) => op === "activation.claim") ?? [];
+      expect(claimCalls.length).toBeGreaterThan(0);
+      expect(claimCalls[0]?.[1]).toEqual(
+        expect.objectContaining({
+          clientInstanceId: expect.any(String),
+        })
+      );
+      expect(store.get(activationStatusAtom)).toBe("active");
+      expect(store.get(activationGenerationAtom)).toBe(1);
+      expect(store.get(activationReasonAtom)).toBeNull();
+    });
+  });
+
+  it("disconnects and gates when activation.revoked is received", async () => {
+    const store = createStore();
+    seedWorkspaces(store, ["ws-1"], "ws-1");
+    act(() => {
+      store.set(activationStatusAtom, "active");
+      store.set(activationGenerationAtom, 1);
+      store.set(activationReasonAtom, null);
+      store.set(gitStateAtomFamily("ws-1"), {
+        branch: "feature/test",
+        ahead: 1,
+        behind: 0,
+        modified: [],
+        staged: [],
+        untracked: [],
+        deleted: [],
+      });
+      store.set(gitBranchListAtomFamily("ws-1"), {
+        current: "feature/test",
+        branches: [],
+        loading: false,
+      });
+      store.set(fileTreeAtomFamily("ws-1"), new Map([[".", []]]));
+      store.set(loadedDirsAtomFamily("ws-1"), new Set(["src"]));
+      store.set(worktreeListAtomFamily("ws-1"), {
+        items: [],
+        loading: false,
+        lastLoadedAt: Date.now(),
+      });
+      store.set(fileTreeStaleAtomFamily("ws-1"), true);
+      store.set(sessionsAtom, {
+        "session-1": {
+          id: "session-1",
+          workspaceId: "ws-1",
+          terminalId: "terminal-1",
+          providerId: "codex",
+          state: "running",
+          capability: "full",
+          startedAt: Date.now(),
+          lastActiveAt: Date.now(),
+        },
+      });
+    });
+
+    renderProviders(store);
+
+    await vi.waitFor(() => {
+      expect(wsState.client?.connect).toHaveBeenCalled();
+    });
+
+    act(() => {
+      wsState.client?.eventHandler?.(
+        "activation.revoked",
+        { reason: "displaced", generation: 2 },
+        1
+      );
+    });
+
+    await vi.waitFor(() => {
+      expect(wsState.client?.disconnect).toHaveBeenCalledWith("single_active_displaced");
+      expect(store.get(activationStatusAtom)).toBe("gated");
+      expect(store.get(activationReasonAtom)).toBe("displaced");
+      expect(store.get(activationGenerationAtom)).toBe(2);
+      expect(store.get(workspacesLoadStateAtom)).toBe("idle");
+      expect(store.get(workspaceOrderAtom)).toEqual([]);
+      expect(store.get(workspacesAtom)).toEqual({});
+      expect(store.get(activeWorkspaceIdAtom)).toBeNull();
+      expect(store.get(fileTreeAtomFamily("ws-1"))).toBeNull();
+      expect(Array.from(store.get(loadedDirsAtomFamily("ws-1")))).toEqual([]);
+      expect(store.get(gitStateAtomFamily("ws-1"))).toBeNull();
+      expect(store.get(gitBranchListAtomFamily("ws-1")).current).toBe("");
+      expect(store.get(worktreeListAtomFamily("ws-1")).items).toEqual([]);
+      expect(store.get(fileTreeStaleAtomFamily("ws-1"))).toBe(false);
+      expect(store.get(sessionsAtom)).toEqual({});
+    });
+  });
+
+  it("does not auto-claim again while activation remains gated", async () => {
+    const store = createStore();
+
+    renderProviders(store);
+
+    await vi.waitFor(() => {
+      expect(wsState.client?.connect).toHaveBeenCalled();
+    });
+
+    act(() => {
+      store.set(activationStatusAtom, "gated");
+      wsState.client?.statusHandler?.("connected");
+    });
+
+    const claimCalls =
+      wsState.client?.sendCommand?.mock.calls.filter(([op]) => op === "activation.claim") ?? [];
+
+    expect(claimCalls).toHaveLength(0);
+  });
+
+  it("does not auto-recover the websocket from foreground signals while gated", async () => {
+    const store = createStore();
+    setVisibilityState("visible");
+
+    renderProviders(store);
+
+    await vi.waitFor(() => {
+      expect(wsState.client?.connect).toHaveBeenCalled();
+    });
+
+    act(() => {
+      store.set(activationStatusAtom, "gated");
+      window.dispatchEvent(new Event("focus"));
+      window.dispatchEvent(new Event("online"));
+      window.dispatchEvent(new Event("pageshow"));
+    });
+
+    expect(wsState.client?.recoverConnection).not.toHaveBeenCalled();
+  });
+
   it("hydrates terminal copy-on-select preferences from settings.get once connected", async () => {
     const store = createStore();
     setVisibilityState("visible");

packages/web/src/app/providers.tsx

@@ -33,6 +33,11 @@ import {
   workspacesLoadStateAtom,
   wsClientAtom,
 } from "../atoms";
+import {
+  activationGenerationAtom,
+  activationReasonAtom,
+  activationStatusAtom,
+} from "../atoms/activation";
 import { authenticatedAtom } from "../atoms/app-ui";
 import type { DispatchCommand } from "../atoms/connection";
 import { activeWorkspaceIdAtom } from "../atoms/workspaces";
@@ -45,11 +50,14 @@ import {
 } from "../features/terminal-panel/preferences";
 import {
   editorRefreshTokenAtomFamily,
+  fileTreeAtomFamily,
   fileTreeStaleAtomFamily,
   gitBranchListAtomFamily,
   gitStateAtomFamily,
+  loadedDirsAtomFamily,
   worktreeListAtomFamily,
 } from "../features/workspace/atoms";
+import { useActivation } from "../hooks/use-activation";
 import type { ConnectionStatus, EventListener } from "../ws";
 import { resolveWsUrl, WsClient } from "../ws";
 
@@ -107,6 +115,43 @@ function mergeRefreshHints(
   };
 }
 
+function resetServerProjectedState(store: Store): void {
+  const workspaceIds = store.get(workspaceOrderAtom);
+  const terminalIds = Object.values(store.get(sessionsAtom))
+    .map((session) => session.terminalId)
+    .filter((terminalId): terminalId is string => Boolean(terminalId));
+
+  store.set(workspacesAtom, {});
+  store.set(workspaceOrderAtom, []);
+  store.set(workspacesLoadStateAtom, "idle");
+  store.set(workspacesLoadErrorAtom, null);
+  store.set(sessionsAtom, {});
+  store.set(activeWorkspaceIdAtom, null);
+  store.set(supervisorsAtom, new Map());
+  store.set(supervisorCyclesAtom, new Map());
+
+  for (const workspaceId of workspaceIds) {
+    store.set(fileTreeAtomFamily(workspaceId), null);
+    store.set(loadedDirsAtomFamily(workspaceId), new Set());
+    store.set(gitStateAtomFamily(workspaceId), null);
+    store.set(gitBranchListAtomFamily(workspaceId), {
+      current: "",
+      branches: [],
+      loading: false,
+    });
+    store.set(worktreeListAtomFamily(workspaceId), {
+      items: [],
+      loading: false,
+    });
+    store.set(fileTreeStaleAtomFamily(workspaceId), false);
+    store.set(editorRefreshTokenAtomFamily(workspaceId), 0);
+  }
+
+  for (const terminalId of terminalIds) {
+    store.set(terminalMetaAtomFamily(terminalId), null);
+  }
+}
+
 function parseWorkspaceRefreshHint(
   topic: string,
   payload: unknown
@@ -183,6 +228,7 @@ export function AppProviders({ children }: AppProvidersProps) {
   // Get Jotai store for writing to atomFamily atoms
   const store = useStore();
   const dispatch = useAtomValue(dispatchCommandAtom);
+  const { claim } = useActivation();
 
   useSessionNotifications();
 
@@ -246,6 +292,18 @@ export function AppProviders({ children }: AppProvidersProps) {
     connectionStatusRef.current = connectionStatus;
   }, [connectionStatus]);
 
+  useEffect(() => {
+    if (connectionStatus !== "connected") {
+      return;
+    }
+
+    if (store.get(activationStatusAtom) === "gated") {
+      return;
+    }
+
+    void claim();
+  }, [claim, connectionStatus, store]);
+
   // Initialize theme from localStorage
   useEffect(() => {
     const savedTheme = localStorage.getItem("ui.theme");
@@ -391,6 +449,10 @@ export function AppProviders({ children }: AppProvidersProps) {
     };
 
     const triggerForegroundRecovery = () => {
+      if (store.get(activationStatusAtom) === "gated") {
+        return;
+      }
+
       syncWorkspaceActivity();
       if (document.visibilityState !== "visible") {
         lastForegroundRecoveryAtRef.current = null;
@@ -429,6 +491,10 @@ export function AppProviders({ children }: AppProvidersProps) {
     };
 
     const handleOnline = () => {
+      if (store.get(activationStatusAtom) === "gated") {
+        return;
+      }
+
       wsClientRef.current?.recoverConnection("network_online");
     };
 
@@ -529,6 +595,30 @@ export function AppProviders({ children }: AppProvidersProps) {
 
     // Event handler: route WS events to atoms
     const handleEvent: EventListener = (topic: string, payload: unknown, _seq: number) => {
+      if (topic === "activation.revoked") {
+        const data = (payload ?? {}) as {
+          reason?: string;
+          generation?: number;
+        };
+
+        store.set(activationStatusAtom, "gated");
+        store.set(
+          activationReasonAtom,
+          typeof data.reason === "string" && data.reason.length > 0 ? data.reason : "displaced"
+        );
+        store.set(
+          activationGenerationAtom,
+          typeof data.generation === "number" ? data.generation : null
+        );
+        resetServerProjectedState(store);
+        workspaceActivityRef.current = {
+          mode: "inactive",
+          workspaceId: null,
+        };
+        wsClientRef.current?.disconnect("single_active_displaced");
+        return;
+      }
+
       const refreshInfo = parseWorkspaceRefreshHint(topic, payload);
       if (refreshInfo) {
         queueWorkspaceRefresh(refreshInfo.workspaceId, refreshInfo.hint);
@@ -544,6 +634,7 @@ export function AppProviders({ children }: AppProvidersProps) {
     // Subscribe to all topics we care about
     const topics = [
       "connection.*", // Connection-level events
+      "activation.*",
       "workspace.*", // All workspace events (glob pattern)
     ];