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
;
+}
+
+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 = [];
+
+ await act(async () => {
+ root.render(
+
+
+ seen.push(companyId)} />
+
+ ,
+ );
+ });
+
+ 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 = [];
+
+ await act(async () => {
+ root.render(
+
+
+ seen.push(companyId)} />
+
+ ,
+ );
+ });
+
+ 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(null);
export function CompanyProvider({ children }: { children: ReactNode }) {
const queryClient = useQueryClient();
const [selectionSource, setSelectionSource] = useState("bootstrap");
- const [selectedCompanyId, setSelectedCompanyIdState] = useState(() => localStorage.getItem(STORAGE_KEY));
+ const [selectedCompanyId, setSelectedCompanyIdState] = useState(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);