From 2a9aa7d2536b1f1abd77eef09bfed15f080cc8de Mon Sep 17 00:00:00 2001
From: Dotta <bippadotta@protonmail.com>
Date: Mon, 27 Apr 2026 11:27:25 -0500
Subject: [PATCH] Ignore stale stored company selections

---
 ui/src/context/CompanyContext.test.tsx | 114 +++++++++++++++++++++++++
 ui/src/context/CompanyContext.tsx      |  20 +++--
 2 files changed, 128 insertions(+), 6 deletions(-)
 create mode 100644 ui/src/context/CompanyContext.test.tsx

diff --git a/ui/src/context/CompanyContext.test.tsx b/ui/src/context/CompanyContext.test.tsx
new file mode 100644
index 00000000000..739fde6ed38
--- /dev/null
+++ b/ui/src/context/CompanyContext.test.tsx
@@ -0,0 +1,114 @@
+// @vitest-environment jsdom
+
+import { act, useEffect } from "react";
+import { createRoot, type Root } from "react-dom/client";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { CompanyProvider, useCompany } from "./CompanyContext";
+
+const mockCompaniesApi = vi.hoisted(() => ({
+  list: vi.fn(),
+  create: vi.fn(),
+}));
+
+vi.mock("../api/companies", () => ({
+  companiesApi: mockCompaniesApi,
+}));
+
+function Probe({ onSelectedCompanyId }: { onSelectedCompanyId: (companyId: string | null) => void }) {
+  const { selectedCompanyId } = useCompany();
+  useEffect(() => {
+    onSelectedCompanyId(selectedCompanyId);
+  }, [onSelectedCompanyId, selectedCompanyId]);
+  return <div data-selected-company-id={selectedCompanyId ?? ""} />;
+}
+
+describe("CompanyProvider", () => {
+  let container: HTMLDivElement;
+  let root: Root;
+  let queryClient: QueryClient;
+
+  beforeEach(() => {
+    (globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true;
+    localStorage.clear();
+    container = document.createElement("div");
+    document.body.appendChild(container);
+    root = createRoot(container);
+    queryClient = new QueryClient({
+      defaultOptions: {
+        queries: { retry: false },
+      },
+    });
+  });
+
+  afterEach(() => {
+    act(() => {
+      root.unmount();
+    });
+    queryClient.clear();
+    container.remove();
+    vi.clearAllMocks();
+  });
+
+  it("does not expose a stale stored company id before companies load", async () => {
+    localStorage.setItem("paperclip.selectedCompanyId", "stale-company");
+    mockCompaniesApi.list.mockImplementation(() => new Promise(() => {}));
+    const seen: Array<string | null> = [];
+
+    await act(async () => {
+      root.render(
+        <QueryClientProvider client={queryClient}>
+          <CompanyProvider>
+            <Probe onSelectedCompanyId={(companyId) => seen.push(companyId)} />
+          </CompanyProvider>
+        </QueryClientProvider>,
+      );
+    });
+
+    expect(seen).toEqual([null]);
+  });
+
+  it("replaces a stale stored company id with the first loaded company", async () => {
+    localStorage.setItem("paperclip.selectedCompanyId", "stale-company");
+    mockCompaniesApi.list.mockResolvedValue([
+      {
+        id: "company-1",
+        name: "Paperclip",
+        description: null,
+        status: "active",
+        issuePrefix: "PAP",
+        issueCounter: 1,
+        budgetMonthlyCents: 0,
+        spentMonthlyCents: 0,
+        requireBoardApprovalForNewAgents: false,
+        feedbackDataSharingEnabled: false,
+        feedbackDataSharingConsentAt: null,
+        feedbackDataSharingConsentByUserId: null,
+        feedbackDataSharingTermsVersion: null,
+        brandColor: null,
+        logoAssetId: null,
+        logoUrl: null,
+        createdAt: new Date(),
+        updatedAt: new Date(),
+      },
+    ]);
+    const seen: Array<string | null> = [];
+
+    await act(async () => {
+      root.render(
+        <QueryClientProvider client={queryClient}>
+          <CompanyProvider>
+            <Probe onSelectedCompanyId={(companyId) => seen.push(companyId)} />
+          </CompanyProvider>
+        </QueryClientProvider>,
+      );
+    });
+
+    await act(async () => {
+      await vi.waitFor(() => {
+        expect(seen).toEqual([null, "company-1"]);
+        expect(localStorage.getItem("paperclip.selectedCompanyId")).toBe("company-1");
+      });
+    });
+  });
+});
diff --git a/ui/src/context/CompanyContext.tsx b/ui/src/context/CompanyContext.tsx
index 86afa175f70..2cb95912b8e 100644
--- a/ui/src/context/CompanyContext.tsx
+++ b/ui/src/context/CompanyContext.tsx
@@ -38,7 +38,7 @@ const CompanyContext = createContext<CompanyContextValue | null>(null);
 export function CompanyProvider({ children }: { children: ReactNode }) {
   const queryClient = useQueryClient();
   const [selectionSource, setSelectionSource] = useState<CompanySelectionSource>("bootstrap");
-  const [selectedCompanyId, setSelectedCompanyIdState] = useState<string | null>(() => localStorage.getItem(STORAGE_KEY));
+  const [selectedCompanyId, setSelectedCompanyIdState] = useState<string | null>(null);
 
   const { data: companies = [], isLoading, error } = useQuery({
     queryKey: queryKeys.companies.all,
@@ -61,18 +61,26 @@ export function CompanyProvider({ children }: { children: ReactNode }) {
 
   // Auto-select first company when list loads
   useEffect(() => {
-    if (companies.length === 0) return;
+    if (isLoading) return;
+    if (companies.length === 0) {
+      if (selectedCompanyId !== null) {
+        setSelectedCompanyIdState(null);
+      }
+      localStorage.removeItem(STORAGE_KEY);
+      return;
+    }
 
     const selectableCompanies = sidebarCompanies.length > 0 ? sidebarCompanies : companies;
-    const stored = localStorage.getItem(STORAGE_KEY);
-    if (stored && selectableCompanies.some((c) => c.id === stored)) return;
     if (selectedCompanyId && selectableCompanies.some((c) => c.id === selectedCompanyId)) return;
 
-    const next = selectableCompanies[0]!.id;
+    const stored = localStorage.getItem(STORAGE_KEY);
+    const next = stored && selectableCompanies.some((c) => c.id === stored)
+      ? stored
+      : selectableCompanies[0]!.id;
     setSelectedCompanyIdState(next);
     setSelectionSource("bootstrap");
     localStorage.setItem(STORAGE_KEY, next);
-  }, [companies, selectedCompanyId, sidebarCompanies]);
+  }, [companies, isLoading, selectedCompanyId, sidebarCompanies]);
 
   const setSelectedCompanyId = useCallback((companyId: string, options?: CompanySelectionOptions) => {
     setSelectedCompanyIdState(companyId);