diff --git a/ui/src/api/companies-query.test.ts b/ui/src/api/companies-query.test.ts index 2d114d39a5..865f0718f8 100644 --- a/ui/src/api/companies-query.test.ts +++ b/ui/src/api/companies-query.test.ts @@ -1,22 +1,129 @@ +import { QueryClient } from "@tanstack/react-query"; import { describe, expect, it, vi } from "vitest"; -import { companiesListQueryOptions } from "./companies-query"; +import { companyListQueryOptions, resolveAccountUserId } from "./companies-query"; import { ApiError } from "./client"; +import { queryKeys } from "../lib/queryKeys"; + +const mockAuthApi = vi.hoisted(() => ({ + getSession: vi.fn(), +})); + +vi.mock("./auth", () => ({ + authApi: mockAuthApi, +})); const mockCompaniesApi = vi.hoisted(() => ({ list: vi.fn(), + detachInflightList: vi.fn(), })); vi.mock("./companies", () => ({ companiesApi: mockCompaniesApi, })); -describe("companiesListQueryOptions", () => { +describe("companyListQueryOptions", () => { it.each([401, 403])("treats %s company-list failures as unauthorized bootstrap state", async (status) => { mockCompaniesApi.list.mockRejectedValueOnce(new ApiError("Board access required", status, { error: "Board access required" })); - await expect(companiesListQueryOptions.queryFn()).resolves.toEqual({ + await expect(companyListQueryOptions("user-1").queryFn()).resolves.toEqual({ companies: [], unauthorized: true, }); }); }); + +describe("company list cache keys", () => { + it("gives each account its own entry, and a stable one for signed-out", () => { + expect(companyListQueryOptions("user-1").queryKey).not.toEqual( + companyListQueryOptions("user-2").queryKey, + ); + expect(companyListQueryOptions(null).queryKey).toEqual(companyListQueryOptions(null).queryKey); + }); + + it("stays under the `companies` prefix so existing invalidations still reach it", () => { + expect(companyListQueryOptions("user-1").queryKey.slice(0, 1)).toEqual(queryKeys.companies.all); + }); + + // Coalescing keys on the request path alone, so a `/companies` request issued + // under the previous session would otherwise answer this one and land the + // wrong account's list in an account-keyed entry. + it("refuses to join an in-flight request issued for another account", async () => { + mockCompaniesApi.list.mockResolvedValueOnce([]); + await companyListQueryOptions("user-2").queryFn(); + expect(mockCompaniesApi.detachInflightList).toHaveBeenCalled(); + }); +}); + +// 19 call sites invalidate `queryKeys.companies.all` after mutating a company. +// Keying the list deeper only stays safe while that prefix still matches it, so +// this asserts the matching itself rather than the shape of the key. +describe("existing invalidations still reach the account-keyed list", () => { + it("marks the list stale when the `companies` prefix is invalidated", async () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + const key = companyListQueryOptions("user-1").queryKey; + queryClient.setQueryData(key, { companies: [], unauthorized: false }); + expect(queryClient.getQueryState(key)?.isInvalidated).toBe(false); + + await queryClient.invalidateQueries({ queryKey: queryKeys.companies.all }); + + expect(queryClient.getQueryState(key)?.isInvalidated).toBe(true); + }); +}); + +// A session request that fails answers nothing. `authApi.getSession` returns +// null only for a 401 — a real "signed out" — and throws otherwise, so treating +// a failure as null would key a credentialed response to `anonymous` and hand +// it to every signed-out reader. The original bug, rebuilt for anonymous. +describe("an unavailable session is not an anonymous session", () => { + it("does not resolve an identity from an errored session entry", async () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + mockAuthApi.getSession.mockRejectedValueOnce(new Error("network down")); + await queryClient + .fetchQuery({ queryKey: queryKeys.auth.session, queryFn: () => mockAuthApi.getSession(), retry: false }) + .catch(() => undefined); + expect(queryClient.getQueryState(queryKeys.auth.session)?.status).toBe("error"); + + // Asked again rather than assumed anonymous. + mockAuthApi.getSession.mockResolvedValueOnce({ user: { id: "user-9" } }); + await expect(resolveAccountUserId(queryClient)).resolves.toBe("user-9"); + }); + + it("fetches the session when the cache has no answer, instead of defaulting to anonymous", async () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + mockAuthApi.getSession.mockResolvedValueOnce({ user: { id: "user-7" } }); + + await expect(resolveAccountUserId(queryClient)).resolves.toBe("user-7"); + expect(mockAuthApi.getSession).toHaveBeenCalled(); + }); + + it("still treats a confirmed 401 as the anonymous identity", async () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + // `getSession` resolves null for a 401 — signed out is an answer. + mockAuthApi.getSession.mockResolvedValueOnce(null); + + await expect(resolveAccountUserId(queryClient)).resolves.toBeNull(); + }); +}); + + +// Invalidation is the app saying "that answer belongs to the previous account". +// The entry stays `success` with the old user until a refetch lands, so a status +// check alone would key the new account's list to the old account's id. +describe("an invalidated session is not a current identity", () => { + it("re-reads the session rather than trusting an invalidated success", async () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + mockAuthApi.getSession.mockResolvedValueOnce({ user: { id: "user-1" } }); + await queryClient.fetchQuery({ + queryKey: queryKeys.auth.session, + queryFn: () => mockAuthApi.getSession(), + retry: false, + }); + + // Signed in as somebody else; the invite page invalidates and asks at once. + await queryClient.invalidateQueries({ queryKey: queryKeys.auth.session }); + expect(queryClient.getQueryState(queryKeys.auth.session)?.status).toBe("success"); + mockAuthApi.getSession.mockResolvedValueOnce({ user: { id: "user-2" } }); + + await expect(resolveAccountUserId(queryClient)).resolves.toBe("user-2"); + }); +}); diff --git a/ui/src/api/companies-query.ts b/ui/src/api/companies-query.ts index 4250193d04..057cdb89f5 100644 --- a/ui/src/api/companies-query.ts +++ b/ui/src/api/companies-query.ts @@ -1,25 +1,139 @@ +import { useQuery, type QueryClient } from "@tanstack/react-query"; import type { Company } from "@paperclipai/shared"; +import { authApi } from "./auth"; import { companiesApi } from "./companies"; import { ApiError } from "./client"; import { queryKeys } from "../lib/queryKeys"; export type CompanyListResult = { companies: Company[]; unauthorized: boolean }; -// Single source of truth for the `["companies"]` query. Both CompanyProvider and -// the invite landing page read this cache entry, so they must agree on the shape — -// returning a bare `Company[]` from one and this wrapped object from the other -// silently corrupts the shared cache and crashes whichever reads the other's shape. -export const companiesListQueryOptions = { - queryKey: queryKeys.companies.all, - queryFn: async (): Promise => { - try { - return { companies: await companiesApi.list(), unauthorized: false }; - } catch (err) { - if (err instanceof ApiError && (err.status === 401 || err.status === 403)) { - return { companies: [], unauthorized: true }; +// Single source of truth for the company-list query. Every consumer reads the +// same cache entry, so they must agree on the shape — returning a bare +// `Company[]` from one and this wrapped object from another silently corrupts +// the entry and crashes whichever reads the other's shape. +// +// The entry is keyed by account. Callers therefore cannot ask for "the company +// list" without saying whose, which is the point: the previous design let a +// list fetched for one account answer a question about another, and every +// consumer had to defend against that individually. +export function companyListQueryOptions(userId: string | null) { + return { + queryKey: queryKeys.companies.list(userId), + queryFn: async (): Promise => { + // Request coalescing keys on the path alone, so it would happily answer a + // fetch for this account with a `/companies` request issued under the + // previous one — putting the wrong account's list in an account-keyed + // entry. Coalescing has nothing to offer here anyway: React Query already + // dedupes concurrent fetches within a key, and *across* keys the accounts + // differ by construction, which is exactly when sharing is wrong. + companiesApi.detachInflightList(); + try { + return { companies: await companiesApi.list(), unauthorized: false }; + } catch (err) { + if (err instanceof ApiError && (err.status === 401 || err.status === 403)) { + return { companies: [], unauthorized: true }; + } + throw err; } - throw err; - } - }, + }, + retry: false, + } as const; +} + +const sessionQueryOptions = { + queryKey: queryKeys.auth.session, + queryFn: () => authApi.getSession(), retry: false, } as const; + +/** + * The signed-in account, as the cache currently understands it. + * + * `settled` means the session query *answered*, not merely that it stopped + * being pending. The distinction is the whole point of this helper. `null` is a + * real answer — `authApi.getSession` returns it for a 401, meaning "signed out" + * — but a session request that *failed* answers nothing, and it throws rather + * than returning null. Treating that failure as `null` would key an account's + * list to `anonymous` while the request still carried their cookie, writing a + * credentialed response into the entry every signed-out reader trusts. Which is + * the original bug, rebuilt for anonymous. + * + * So: settled on success only. A failed session lookup leaves the list unfetched + * until the query recovers, which it does on the next refetch. + */ +export function useAccountIdentity(): { userId: string | null; settled: boolean } { + const { data: session, isSuccess } = useQuery(sessionQueryOptions); + return { userId: session?.user.id ?? null, settled: isSuccess }; +} + +/** + * Observe the company list for the account signed in now. + * + * Pass `enabled: false` to hold for a caller's own reasons; the session gate is + * applied on top of it either way. + */ +export function useCompanyListQuery( + options: { enabled?: boolean; staleTime?: number; retry?: number | boolean } = {}, +) { + const { userId, settled } = useAccountIdentity(); + const query = useQuery({ + ...companyListQueryOptions(userId), + ...options, + enabled: settled && (options.enabled ?? true), + }); + + // While the account is unknown this query is disabled, and a disabled query + // reports `isLoading: false` with `data: undefined` — which consumers default + // to an empty list and read as "asked, and owns nothing". That is the + // destructive reading: `shouldClearStoredCompanySelection` would throw away + // the customer's stored company on every cold boot, before the session had + // even landed. + // + // Waiting for the account is part of getting the list, so it is reported as + // part of getting the list. Consumers gate on these two, and both must mean + // "no answer yet" for the whole window in which there is no answer. + const waitingForAccount = !settled && (options.enabled ?? true); + return { + ...query, + isLoading: query.isLoading || waitingForAccount, + isFetching: query.isFetching || waitingForAccount, + }; +} + +/** + * Resolve the account identity for an imperative path, fetching the session if + * the cache has no answer yet. + * + * Not `getQueryData` with a `?? null` fallback: an empty or errored session + * entry would read as "signed out" and key an authenticated response to + * `anonymous`. There is no safe default here — the identity is either known or + * must be obtained, and if it cannot be obtained the caller must fail rather + * than guess. + */ +export async function resolveAccountUserId(queryClient: QueryClient): Promise { + const state = queryClient.getQueryState>>( + queryKeys.auth.session, + ); + // `isInvalidated` matters as much as `status` here. Invalidation is how the + // app says "this is the previous account's answer" — InviteLanding invalidates + // the session immediately after a sign-in and resolves an identity in the next + // breath. The entry is still `success`, still holding the *old* user, until a + // refetch lands; trusting it would key the new account's list to the old + // account's id, which is the misattribution this whole file exists to prevent. + if (state?.status === "success" && !state.isInvalidated) { + return state.data?.user.id ?? null; + } + const session = await queryClient.fetchQuery({ ...sessionQueryOptions, staleTime: 0 }); + return session?.user.id ?? null; +} + +/** Fetch the company list for the account signed in now, ignoring what is cached. */ +export async function fetchCompanyListForCurrentAccount( + queryClient: QueryClient, +): Promise { + const userId = await resolveAccountUserId(queryClient); + return queryClient.fetchQuery({ + ...companyListQueryOptions(userId), + staleTime: 0, + }); +} diff --git a/ui/src/components/OnboardingWizard.adapters.test.tsx b/ui/src/components/OnboardingWizard.adapters.test.tsx index 942cb67341..77e25a9734 100644 --- a/ui/src/components/OnboardingWizard.adapters.test.tsx +++ b/ui/src/components/OnboardingWizard.adapters.test.tsx @@ -48,8 +48,22 @@ vi.mock("../context/CompanyContext", () => ({ // here. An empty list is fine: the drafts in this file carry no // `createdCompanyId`, so there is no ownership question — but the fetch has to // succeed for the gate to treat the draft as decidable at all. +// The company list is keyed by account, so it holds until the session query +// *succeeds*. A seeded entry is stale under the test client and refetches, so +// the refetch has to answer too — otherwise the identity errors and the list +// never runs. +const mockAuthApi = vi.hoisted(() => ({ getSession: vi.fn() })); +vi.mock("../api/auth", async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, authApi: { ...actual.authApi, getSession: mockAuthApi.getSession } }; +}); + vi.mock("../api/companies", () => ({ - companiesApi: { create: vi.fn(), list: vi.fn().mockResolvedValue([]) }, + companiesApi: { + create: vi.fn(), + list: vi.fn().mockResolvedValue([]), + detachInflightList: vi.fn(), + }, })); vi.mock("../adapters", () => ({ listUIAdapters: () => mockAdapterRegistry.list, @@ -83,6 +97,7 @@ vi.mock("./AsciiArtAnimation", () => ({ AsciiArtAnimation: () => null })); vi.mock("./FrontDoor", () => ({ FrontDoor: () => null })); vi.mock("./AgentCapsule", () => ({ AgentCapsule: () => null })); +import { queryKeys } from "../lib/queryKeys"; import { OnboardingWizard } from "./OnboardingWizard"; // eslint-disable-next-line @typescript-eslint/no-explicit-any @@ -102,6 +117,12 @@ async function mount() { const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } }, }); + // The company list is keyed by account, so it holds until the session is + // known. Seeding it is how this test says "signed in". + queryClient.setQueryData(queryKeys.auth.session, { + session: { id: "session-1", userId: "user-1" }, + user: { id: "user-1", name: "Example", email: "user-1@example.com", image: null }, + }); await act(async () => { root.render( @@ -115,6 +136,10 @@ async function mount() { describe("OnboardingWizard adapter selection", () => { beforeEach(() => { + mockAuthApi.getSession.mockResolvedValue({ + session: { id: "session-1", userId: "user-1" }, + user: { id: "user-1", name: "Example", email: "user-1@example.com", image: null }, + }); window.localStorage.clear(); mockDialog.onboardingOpen = true; mockDialog.onboardingOptions = {}; diff --git a/ui/src/components/OnboardingWizard.test.tsx b/ui/src/components/OnboardingWizard.test.tsx index 47e8afd88c..8fe95e65bc 100644 --- a/ui/src/components/OnboardingWizard.test.tsx +++ b/ui/src/components/OnboardingWizard.test.tsx @@ -7,6 +7,16 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; // --- Mocks (hoisted so vi.mock factories can close over them) ---------------- +// The company list is keyed by account, so it holds until the session query +// *succeeds*. A seeded entry is stale under the test client and refetches, so +// the refetch has to answer too — otherwise the identity errors and the list +// never runs. +const mockAuthApi = vi.hoisted(() => ({ getSession: vi.fn() })); +vi.mock("../api/auth", async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, authApi: { ...actual.authApi, getSession: mockAuthApi.getSession } }; +}); + const mockDialog = vi.hoisted(() => ({ onboardingOpen: true, onboardingOptions: {} as { initialStep?: number; companyId?: string }, @@ -23,6 +33,7 @@ const mockCompany = vi.hoisted(() => ({ })); const mockCompaniesApi = vi.hoisted(() => ({ + detachInflightList: vi.fn(), create: vi.fn(), update: vi.fn(), // The gate fetches the list itself now, rather than reading the shared @@ -131,6 +142,8 @@ async function flushReact() { }); } +const SESSION_USER_ID = "user-b"; + function render() { const container = document.createElement("div"); document.body.appendChild(container); @@ -138,11 +151,21 @@ function render() { const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } }, }); + // The company list is keyed by account, so it stays disabled until the + // session is known. Seeding it is how these tests say "signed in as B". + queryClient.setQueryData(queryKeys.auth.session, { + session: { id: "session-b", userId: SESSION_USER_ID }, + user: { id: SESSION_USER_ID, name: "B", email: "b@example.com", image: null }, + }); return { container, root, queryClient }; } describe("OnboardingWizard restore-gate (stale localStorage across accounts)", () => { beforeEach(() => { + mockAuthApi.getSession.mockResolvedValue({ + session: { id: "session-b", userId: SESSION_USER_ID }, + user: { id: SESSION_USER_ID, name: "B", email: "b@example.com", image: null }, + }); window.localStorage.clear(); mockDialog.onboardingOpen = true; mockDialog.onboardingOptions = {}; @@ -644,7 +667,7 @@ describe("OnboardingWizard restore-gate (stale localStorage across accounts)", ( ); const { root, queryClient } = render(); // A's list, already in the cache from their session. - queryClient.setQueryData(queryKeys.companies.all, { + queryClient.setQueryData(queryKeys.companies.list(SESSION_USER_ID), { companies: [{ id: "company-a", name: "Account A Co", issuePrefix: "AAC" }], unauthorized: false, }); @@ -687,7 +710,7 @@ describe("OnboardingWizard restore-gate (stale localStorage across accounts)", ( }); window.localStorage.setItem(ONBOARDING_STORAGE_KEY, draft); const { root, queryClient } = render(); - queryClient.setQueryData(queryKeys.companies.all, { + queryClient.setQueryData(queryKeys.companies.list(SESSION_USER_ID), { companies: [{ id: "c1", name: "Saved Co", issuePrefix: "SC" }], unauthorized: false, }); diff --git a/ui/src/components/OnboardingWizard.tsx b/ui/src/components/OnboardingWizard.tsx index 98b00435f4..b5b8d6c525 100644 --- a/ui/src/components/OnboardingWizard.tsx +++ b/ui/src/components/OnboardingWizard.tsx @@ -5,7 +5,7 @@ import { useLocation, useNavigate, useParams } from "@/lib/router"; import { useDialog } from "../context/DialogContext"; import { useCompany } from "../context/CompanyContext"; import { companiesApi } from "../api/companies"; -import { companiesListQueryOptions } from "../api/companies-query"; +import { useCompanyListQuery } from "../api/companies-query"; import { goalsApi } from "../api/goals"; import { agentsApi } from "../api/agents"; import { approvalsApi } from "../api/approvals"; @@ -220,8 +220,7 @@ export function OnboardingWizard() { // is what makes that reachable while the shared entry is still fresh. It // shares the query key, so the result populates the same cache entry the // rest of the app reads. - const companiesQuery = useQuery({ - ...companiesListQueryOptions, + const companiesQuery = useCompanyListQuery({ staleTime: 0, // Only a *parseable* saved draft poses the question. Without one there is // nothing to authorize, and this must not add a request to every wizard diff --git a/ui/src/context/CompanyContext.test.tsx b/ui/src/context/CompanyContext.test.tsx index b84be34bc7..26e16aa515 100644 --- a/ui/src/context/CompanyContext.test.tsx +++ b/ui/src/context/CompanyContext.test.tsx @@ -264,7 +264,11 @@ describe("CompanyProvider", () => { it("replaces a stale stored company id with the first loaded company", async () => { localStorage.setItem("paperclip.selectedCompanyId", "stale-company"); - queryClient.setQueryData(queryKeys.companies.all, { + // Signed out, and *known* to be: the list is keyed by account, so it cannot + // resolve before the session does. Seeding the answer is how this test says + // the account lookup has already happened. + queryClient.setQueryData(queryKeys.auth.session, null); + queryClient.setQueryData(queryKeys.companies.list(null), { companies: [makeCompany("company-1")], unauthorized: false, }); @@ -292,7 +296,7 @@ describe("CompanyProvider", () => { describe("when the account changes in this tab", () => { async function bootWithFirstAccount(seen: Array) { mockAuthApi.getSession.mockResolvedValue(sessionFor("user-1")); - queryClient.setQueryData(queryKeys.companies.all, { + queryClient.setQueryData(queryKeys.companies.list("user-1"), { companies: [makeCompany("company-1")], unauthorized: false, }); @@ -332,7 +336,12 @@ describe("CompanyProvider", () => { // Nothing from the previous account is exposed while the new list loads. expect(seen).toEqual([null, "company-1", null]); - expect(queryClient.getQueryData(queryKeys.companies.all)).toBeUndefined(); + expect(queryClient.getQueryData(queryKeys.companies.list("user-2"))).toBeUndefined(); + // The previous account's entry may still sit in the cache; keying by + // account is what makes that harmless, since nothing reads it now. + expect(queryClient.getQueryData(queryKeys.companies.list("user-1"))).toMatchObject({ + companies: [{ id: "company-1" }], + }); expect(resolveSecondList).not.toBeNull(); // The replacement fetch must not be coalesced into a `/companies` request // issued under the previous session. @@ -436,7 +445,7 @@ describe("CompanyProvider", () => { // This account genuinely owns no companies, and the request says so. mockCompaniesApi.list.mockResolvedValue([]); await act(async () => { - await queryClient.refetchQueries({ queryKey: queryKeys.companies.all }); + await queryClient.refetchQueries({ queryKey: queryKeys.companies.list("user-2") }); }); await flushReact(); @@ -444,6 +453,65 @@ describe("CompanyProvider", () => { expect(captured?.companyListUnavailable).toBe(false); }); + // The guarantee the account-keyed entry buys, stated directly: a list sitting + // in the cache for another account is not merely distrusted, it is not + // reachable. No gate, no freshness check, no refetch stands between the two — + // they are different entries. + it("cannot read another account's cached list at all", async () => { + const seen: Array = []; + await bootWithFirstAccount(seen); + + // user-2's list is slow. If the entries were shared, the cached user-1 + // list would answer in the meantime — which is the whole bug. + mockCompaniesApi.list.mockImplementation(() => new Promise(() => {})); + + await act(async () => { + queryClient.setQueryData(queryKeys.auth.session, sessionFor("user-2")); + }); + await flushReact(); + + expect(captured?.companies).toEqual([]); + expect(captured?.selectedCompanyId).toBeNull(); + expect(captured?.loading).toBe(true); + // Still cached, still keyed to whom it belongs, still unread. + expect(queryClient.getQueryData(queryKeys.companies.list("user-1"))).toMatchObject({ + companies: [{ id: "company-1" }], + }); + }); + + // A session request that failed answers nothing. Keying the list to + // `anonymous` on that basis would write an authenticated response — the + // request still carries the cookie — into the entry every signed-out reader + // trusts, which is the original bug rebuilt for anonymous. + // + // Cold cache on purpose: after a prior success React Query retains the old + // session `data` through a failed refetch, so the identity survives anyway + // and the bug does not show. It needs a session that never answered. + it("does not treat a never-answered session as a signed-out session", async () => { + mockAuthApi.getSession.mockRejectedValue(new Error("network down")); + localStorage.setItem("paperclip.selectedCompanyId", "company-1"); + mockCompaniesApi.list.mockResolvedValue([makeCompany("company-1")]); + + await act(async () => { + root.render( + + + undefined} /> + + , + ); + }); + await flushReact(); + await flushReact(); + + // Nothing asked for a list it could not attribute, and nothing landed in + // the anonymous entry. + expect(mockCompaniesApi.list).not.toHaveBeenCalled(); + expect(queryClient.getQueryData(queryKeys.companies.list(null))).toBeUndefined(); + // And the stored company survives: unavailable is not "owns nothing". + expect(localStorage.getItem("paperclip.selectedCompanyId")).toBe("company-1"); + }); + it("leaves the selection alone when the same account is observed again", async () => { const seen: Array = []; await bootWithFirstAccount(seen); diff --git a/ui/src/context/CompanyContext.tsx b/ui/src/context/CompanyContext.tsx index 8cc5534b65..64e41ae37c 100644 --- a/ui/src/context/CompanyContext.tsx +++ b/ui/src/context/CompanyContext.tsx @@ -8,11 +8,14 @@ import { useState, type ReactNode, } from "react"; -import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query"; +import { useMutation, useQueryClient } from "@tanstack/react-query"; import type { Company } from "@paperclipai/shared"; -import { authApi } from "../api/auth"; import { companiesApi } from "../api/companies"; -import { companiesListQueryOptions, type CompanyListResult } from "../api/companies-query"; +import { + fetchCompanyListForCurrentAccount, + useAccountIdentity, + useCompanyListQuery, +} from "../api/companies-query"; import { queryKeys } from "../lib/queryKeys"; import type { CompanySelectionSource } from "../lib/company-selection"; type CompanySelectionOptions = { source?: CompanySelectionSource }; @@ -103,8 +106,20 @@ export function CompanyProvider({ children }: { children: ReactNode }) { const [selectionSource, setSelectionSource] = useState("bootstrap"); const [selectedCompanyId, setSelectedCompanyIdState] = useState(null); + // Keyed by account, so there is no such thing here as "the list, but whose?". + // A change of account changes the key, which leaves this observer pending + // rather than holding the previous account's answer. const { data: companiesResult = { companies: [], unauthorized: false }, isLoading, error } = - useQuery(companiesListQueryOptions); + useCompanyListQuery({ + // `retry: 1` against the shared `retry: false`, and it is load-bearing + // here even though an earlier measurement found retries pointless. That + // measurement was taken against an implementation that removed the cache + // entry on an account change, which made the observer rebind and issue a + // second request for free. Keying by account replaced that mechanism, and + // the free attempt went with it: without this, one blip during a switch + // leaves the customer with no companies until they find Try again. + retry: 1, + }); const companies = companiesResult.companies; const companyListUnauthorized = companiesResult.unauthorized; const sidebarCompanies = useMemo( @@ -112,86 +127,31 @@ export function CompanyProvider({ children }: { children: ReactNode }) { [companies], ); - // The `["companies"]` entry is shared app-wide and carries no account identity, - // so it outlives a change of account in this tab: the previous account's list - // keeps being served, and the effect below would auto-select from it and write - // that company id to localStorage. Signing in through the app drops the entry, - // but nothing does when the session lapses server-side or a second account - // signs in on another tab — so watch the account itself. - const { data: session, isPending: isSessionPending } = useQuery({ - queryKey: queryKeys.auth.session, - queryFn: () => authApi.getSession(), - retry: false, - }); - const sessionUserId = session?.user.id ?? null; + // The list is account-keyed, but the *selection* is component state and does + // not change key with it. It belongs to the account that just went away, so it + // is dropped here. The stored id deliberately survives: + // resolveBootstrapCompanySelection re-validates it against the incoming list, + // so an account signing back in keeps its company while an unrelated account + // cannot inherit it. + const { userId: sessionUserId, settled: isSessionSettled } = useAccountIdentity(); const observedUserIdRef = useRef(undefined); - const [awaitingAccountScopedList, setAwaitingAccountScopedList] = useState(false); useEffect(() => { // Until the session settles the account is unknown, not changed. - if (isSessionPending) return; + if (!isSessionSettled) return; const previousUserId = observedUserIdRef.current; observedUserIdRef.current = sessionUserId; // First settled observation is this tab's boot: no previous account to leave. if (previousUserId === undefined || previousUserId === sessionUserId) return; - // The live selection belongs to the account that just went away. The stored - // one deliberately survives: resolveBootstrapCompanySelection re-validates it - // against the incoming list, so an account signing back in keeps its company - // while an unrelated account cannot inherit it. setSelectedCompanyIdState(null); setSelectionSource("bootstrap"); - setAwaitingAccountScopedList(true); - // `removeQueries`, not `resetQueries`: reset rewinds the update counters - // `isFetchedAfterMount` is derived from while mounted observers keep their - // pre-reset baseline, which strands consumers gating on that flag (see - // AppsConnect.tsx). Removal is what fits *here*, for a reason worth saying - // out loud: removal notifies nobody, so on its own it would leave mounted - // observers serving the previous account. What rebinds them is the render - // the state updates above just scheduled — every observer re-binds to a - // fresh query on the next render. Do not lift this call anywhere that lacks - // that guarantee; the sign-out sweep is exactly such a place. - queryClient.removeQueries({ queryKey: queryKeys.companies.all, exact: true }); - - // Coalescing keys on the request path alone, so without this the fetch below - // can join a `/companies` request issued under the previous session and be - // answered with that account's companies. - companiesApi.detachInflightList(); - - // Drive the replacement fetch here rather than leaning on observer refetch - // semantics, so the gate below lifts exactly when a list for this account - // has landed. - // - // No `retry` override, though `companiesListQueryOptions` sets - // `retry: false`. A transient blip already gets a second attempt: the - // observer above rebinds to a fresh query on the render these state updates - // schedule and issues its own request, so a single failure self-heals - // (measured: two attempts either way). Adding retries here only buys extra - // failed round trips before a real outage is reported, and the outage is - // what needs a way out — see `companyListUnavailable` below. - // The rejection is caught only to keep it from going unhandled — it is not - // lost. `fetchQuery` records it on the query itself, so it arrives as - // `error` below, which is where `companyListUnavailable` reads it from. - // Carrying a second copy in component state is what let the two fall out of - // step: the copy outlived the failure and kept reporting "couldn't load" - // over a later, honest empty list. - let cancelled = false; - void queryClient - .fetchQuery({ ...companiesListQueryOptions, staleTime: 0 }) - .catch(() => undefined) - .finally(() => { - if (!cancelled) setAwaitingAccountScopedList(false); - }); - return () => { - cancelled = true; - }; - }, [isSessionPending, queryClient, sessionUserId]); + }, [isSessionSettled, sessionUserId]); // Auto-select first company when list loads useEffect(() => { - // Nothing may be derived from the list until one has been fetched for the - // account signed in now. - if (awaitingAccountScopedList) return; + // `isLoading` covers the account change too: the key moved, so this observer + // is pending on a list for the new account rather than holding the old one. if (isLoading) return; // An errored list says nothing about which companies this account has, and // `retry: false` makes a single network blip stick. Treat it as undecided @@ -223,7 +183,6 @@ export function CompanyProvider({ children }: { children: ReactNode }) { setSelectionSource("bootstrap"); localStorage.setItem(STORAGE_KEY, next); }, [ - awaitingAccountScopedList, companies, companyListUnauthorized, error, @@ -247,13 +206,12 @@ export function CompanyProvider({ children }: { children: ReactNode }) { // its old error, so the recovery affordance would keep telling the customer // the list is unavailable after a retry had already succeeded. const retryCompanies = useCallback(async () => { - setAwaitingAccountScopedList(true); try { - await queryClient.fetchQuery({ ...companiesListQueryOptions, staleTime: 0 }); + await fetchCompanyListForCurrentAccount(queryClient); } catch { - // Recorded on the query, same as the replacement fetch above. - } finally { - setAwaitingAccountScopedList(false); + // Not swallowed — `fetchQuery` records it on the query, which is where + // `companyListUnavailable` below reads it from. Caught only so the + // rejection does not go unhandled. } }, [queryClient]); diff --git a/ui/src/lib/queryKeys.ts b/ui/src/lib/queryKeys.ts index 9bfaee76a7..b4680ed189 100644 --- a/ui/src/lib/queryKeys.ts +++ b/ui/src/lib/queryKeys.ts @@ -1,6 +1,20 @@ export const queryKeys = { companies: { + /** + * Prefix for everything company-shaped. Matches the list, details and stats + * below, so an `invalidateQueries` against it still reaches all of them. + * Not a cache entry of its own — the list lives under `list()`. + */ all: ["companies"] as const, + /** + * The company list, scoped to the account it was fetched for. Which + * companies you belong to is an answer about *you*, and a single shared + * entry made it look like an answer about the app: for 30 seconds after an + * account change the previous account's list was served to anyone who + * asked. A different account is a different entry, so there is nothing to + * mistake. `null` is the signed-out/local_trusted key. + */ + list: (userId: string | null) => ["companies", "list", userId ?? "anonymous"] as const, detail: (id: string) => ["companies", id] as const, stats: ["companies", "stats"] as const, exportFidelity: (companyId: string) => ["companies", companyId, "export-fidelity"] as const, diff --git a/ui/src/pages/InviteLanding.test.tsx b/ui/src/pages/InviteLanding.test.tsx index 8dd88e1723..891dbfecc0 100644 --- a/ui/src/pages/InviteLanding.test.tsx +++ b/ui/src/pages/InviteLanding.test.tsx @@ -41,6 +41,7 @@ vi.mock("../api/health", () => ({ vi.mock("../api/companies", () => ({ companiesApi: { list: () => listCompaniesMock(), + detachInflightList: () => undefined, }, })); @@ -495,7 +496,7 @@ describe("InviteLandingPage", () => { expect(acceptInviteMock).toHaveBeenCalledWith("pcp_invite_test", { requestType: "human" }); expect(setSelectedCompanyIdMock).toHaveBeenCalledWith("company-1", { source: "manual" }); expect(queryClient.getQueryState(queryKeys.access.currentBoardAccess)?.isInvalidated).toBe(true); - expect(queryClient.getQueryData(queryKeys.companies.all)).toMatchObject({ + expect(queryClient.getQueryData(queryKeys.companies.list("user-1"))).toMatchObject({ companies: [], unauthorized: false, }); @@ -755,7 +756,7 @@ describe("InviteLandingPage", () => { }); expect(acceptInviteMock).not.toHaveBeenCalled(); expect(setSelectedCompanyIdMock).toHaveBeenCalledWith("company-1", { source: "manual" }); - expect(queryClient.getQueryData(queryKeys.companies.all)).toMatchObject({ + expect(queryClient.getQueryData(queryKeys.companies.list("user-1"))).toMatchObject({ companies: [{ id: "company-1", name: "Acme Robotics" }], unauthorized: false, }); @@ -876,7 +877,7 @@ describe("InviteLandingPage", () => { const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } }, }); - queryClient.setQueryData(queryKeys.companies.all, { + queryClient.setQueryData(queryKeys.companies.list("user-1"), { companies: [], unauthorized: false, }); diff --git a/ui/src/pages/InviteLanding.tsx b/ui/src/pages/InviteLanding.tsx index f732085361..abf46f297f 100644 --- a/ui/src/pages/InviteLanding.tsx +++ b/ui/src/pages/InviteLanding.tsx @@ -8,7 +8,7 @@ import { useCompany } from "@/context/CompanyContext"; import { Link, useNavigate, useParams } from "@/lib/router"; import { accessApi } from "../api/access"; import { authApi } from "../api/auth"; -import { companiesListQueryOptions } from "../api/companies-query"; +import { fetchCompanyListForCurrentAccount, useCompanyListQuery } from "../api/companies-query"; import { healthApi } from "../api/health"; import { getAdapterLabel } from "../adapters/adapter-display-registry"; import { clearPendingInviteToken, rememberPendingInviteToken } from "../lib/invite-memory"; @@ -242,8 +242,7 @@ export function InviteLandingPage() { retry: false, }); - const companiesQuery = useQuery({ - ...companiesListQueryOptions, + const companiesQuery = useCompanyListQuery({ enabled: !!sessionQuery.data && !!inviteQuery.data?.companyId, }); const companyList = companiesQuery.data?.companies ?? []; @@ -370,7 +369,10 @@ export function InviteLandingPage() { await queryClient.invalidateQueries({ queryKey: queryKeys.auth.session }); await queryClient.invalidateQueries({ queryKey: queryKeys.health }); await queryClient.invalidateQueries({ queryKey: queryKeys.access.currentBoardAccess }); - const { companies: freshCompanies } = await queryClient.fetchQuery(companiesListQueryOptions); + // Keyed to the account that just signed in, and forced past whatever is + // cached for it, so this cannot be answered with the previous account's + // membership. + const { companies: freshCompanies } = await fetchCompanyListForCurrentAccount(queryClient); if (invite?.companyId && freshCompanies.some((company) => company.id === invite.companyId)) { clearPendingInviteToken(token);