diff --git a/doc/connections/AI-CONNECTIONS.md b/doc/connections/AI-CONNECTIONS.md index c0ca25fa2c6..fc160934d22 100644 --- a/doc/connections/AI-CONNECTIONS.md +++ b/doc/connections/AI-CONNECTIONS.md @@ -523,3 +523,12 @@ identity, a different grant or responsible user, or a changed credential generat requires a fresh session. The metadata is removed before passing session params to an adapter. Temporary authentication-home paths do not change the configuration fingerprint. These checks do not relax current connection authorization. + +Quota polling has a 20-second response deadline. Once a managed OAuth refresh +starts, it has its own 60-second request lifetime so the replacement token body +can still be read and committed after the dashboard stops waiting. A successful +refreshed observation uses the saved grant/secret revision as its cache identity. +A reconnect during credential resolution defers the poll instead of associating +one account's quota with another revision. Both Costs surfaces retain matching +successful observations on transient failures and clear them on authentication +failure or credential rotation. diff --git a/server/src/__tests__/ai-connections.test.ts b/server/src/__tests__/ai-connections.test.ts index dbd467740bf..ccc5df5b29b 100644 --- a/server/src/__tests__/ai-connections.test.ts +++ b/server/src/__tests__/ai-connections.test.ts @@ -84,6 +84,53 @@ describe("managed AI connections", () => { } finally { fetchSpy.mockRestore(); } }); + it("caches refreshed quota under the saved grant and secret revision", async () => { + const owner = "quota-refreshed-cache-owner"; + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const original = JSON.stringify({ tokens: { access_token: "expired-cache", refresh_token: "single-use", account_id: "cache-account" } }); + await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: "Refreshed cache", loginSessionId: "fixture", agentIds: [], allAgents: true }, original); + const request = vi.spyOn(globalThis, "fetch").mockImplementation(async (url, options) => { + if (String(url).endsWith("/oauth/token")) return Response.json({ access_token: "fresh-cache", refresh_token: "replacement" }); + if (new Headers(options?.headers).get("authorization") === "Bearer expired-cache") return new Response("Expired", { status: 401 }); + return Response.json({ rate_limit: { primary_window: { used_percent: 42 } } }); + }); + try { + const first = await fetchCompanyQuotaWindows(db, companyId, owner); + expect(first[0]).toMatchObject({ ok: true, windows: [{ usedPercent: 42 }] }); + expect(await fetchCompanyQuotaWindows(db, companyId, owner)).toEqual(first); + expect(request).toHaveBeenCalledTimes(3); // rejected quota, OAuth, refreshed quota + } finally { request.mockRestore(); } + }); + + it("saves exchanged tokens when the quota deadline expires during response-body reading", async () => { + const owner = "quota-body-deadline-owner"; + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const original = JSON.stringify({ tokens: { access_token: "expired", refresh_token: "single-use", account_id: "deadline-account" } }); + await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: "Slow token body", loginSessionId: "fixture", agentIds: [], allAgents: true }, original); + const [row] = await service.quotaAccounts(companyId, owner); + const reading = deferredSignal(), releaseBody = deferredSignal(), quotaDeadline = new AbortController(); + let refreshSignal: AbortSignal | null | undefined; + const request = vi.spyOn(globalThis, "fetch").mockImplementation(async (_url, options) => { + refreshSignal = options?.signal; + return { ok: true, json: async () => { + reading.resolve(); await releaseBody.promise; + refreshSignal?.throwIfAborted(); + return { access_token: "fresh-after-deadline", refresh_token: "replacement" }; + } } as Response; + }); + const saving = service.refreshQuotaCredential(row, original, quotaDeadline.signal); + try { + await reading.promise; + quotaDeadline.abort(); + expect(refreshSignal?.aborted).toBe(false); + releaseBody.resolve(); + await saving; + expect(JSON.parse(await service.credential(row))).toMatchObject({ tokens: { access_token: "fresh-after-deadline", refresh_token: "replacement" } }); + await expect(service.refreshQuotaCredential(row, original, quotaDeadline.signal)).rejects.toThrow(); + expect(request).toHaveBeenCalledTimes(1); + } finally { releaseBody.resolve(); await saving.catch(() => {}); request.mockRestore(); } + }); + it("serializes quota exchange with a concurrent secret edit and runtime credential read", async () => { const owner = "quota-race-owner"; await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); diff --git a/server/src/__tests__/company-quota-windows.test.ts b/server/src/__tests__/company-quota-windows.test.ts index 3af1d3f04dc..646f821657d 100644 --- a/server/src/__tests__/company-quota-windows.test.ts +++ b/server/src/__tests__/company-quota-windows.test.ts @@ -160,13 +160,35 @@ describe("connected account quotas", () => { it("refreshes an expired selected credential once before retrying quota", async () => { mocks.accounts.mockResolvedValue([account("refresh-account")]); mocks.codex.mockRejectedValueOnce(new Error("chatgpt wham api returned 401")); - mocks.refresh.mockResolvedValueOnce(JSON.stringify({ tokens: { access_token: "new-token", account_id: "same-account" } })); + mocks.refresh.mockImplementationOnce(async () => { + const value = JSON.stringify({ tokens: { access_token: "new-token", account_id: "same-account" } }); + revision++; + mocks.credential.mockResolvedValue(value); + return value; + }); const [result] = await fetchCompanyQuotaWindows(db, "refresh-company", "user"); expect(result.ok).toBe(true); expect(mocks.refresh).toHaveBeenCalledTimes(1); expect(mocks.codex).toHaveBeenLastCalledWith("new-token", "same-account", expect.any(AbortSignal)); expect(JSON.stringify(result)).not.toContain("new-token"); + const again = await fetchCompanyQuotaWindows(db, "refresh-company", "user"); + expect(again[0]).toEqual(result); + expect(mocks.codex).toHaveBeenCalledTimes(2); // initial 401 and one refreshed read + }); + it("does not cache a quota read under a credential revision that changed during resolution", async () => { + mocks.accounts.mockResolvedValue([account("rotating-during-refresh")]); + mocks.codex.mockRejectedValueOnce(new Error("401")); + mocks.refresh.mockResolvedValue(JSON.stringify({ accessToken: "refreshed" })); + mocks.credential.mockResolvedValueOnce(JSON.stringify({ accessToken: "expired" })).mockImplementationOnce(async () => { + revision++; + return JSON.stringify({ accessToken: "changed-again" }); + }); + const [result] = await fetchCompanyQuotaWindows(db, "changing-revision", "user"); + expect(result).toMatchObject({ ok: false, errorFamily: "provider_unavailable" }); + expect(mocks.codex).toHaveBeenCalledTimes(1); + }); + it("bounds a stalled credential lookup and preserves its account identity", async () => { vi.useFakeTimers(); mocks.accounts.mockResolvedValue([account("stalled")]); diff --git a/server/src/services/ai-connections.ts b/server/src/services/ai-connections.ts index 3b452933784..89ee73554bf 100644 --- a/server/src/services/ai-connections.ts +++ b/server/src/services/ai-connections.ts @@ -890,9 +890,13 @@ export function aiConnectionService(db: Db) { const auth = JSON.parse(current); const refreshToken = auth.tokens?.refresh_token; if (typeof refreshToken !== "string" || !refreshToken) throw new Error("authentication_required"); + // Once exchange starts, its lifetime is independent of the dashboard's + // shorter deadline. A consumed single-use token must still be read/saved. + signal.throwIfAborted(); + const refreshSignal = AbortSignal.timeout(60_000); // Matches the Codex OAuth client (openai/codex, login/src/auth/manager.rs). const response = await fetch("https://auth.openai.com/oauth/token", { - method: "POST", redirect: "error", signal, + method: "POST", redirect: "error", signal: refreshSignal, headers: { "Content-Type": "application/json" }, body: JSON.stringify({ client_id: "app_EMoamEEZ73f0CkXaXp7hrann", grant_type: "refresh_token", refresh_token: refreshToken }), }); diff --git a/server/src/services/quota-windows.ts b/server/src/services/quota-windows.ts index f2631d87617..2d363969b26 100644 --- a/server/src/services/quota-windows.ts +++ b/server/src/services/quota-windows.ts @@ -25,14 +25,20 @@ export async function fetchCompanyQuotaWindows(db: Db, companyId: string, userId // Resolve revision metadata before consulting the cache, so rotation and // revocation cannot serve windows associated with an old credential. const secretIds = accounts.flatMap(row => row.grant.credentialSecretRefs.map(ref => ref.secretId)); - const revisions = secretIds.length ? await db.select({ id: companySecrets.id, latestVersion: companySecrets.latestVersion, status: companySecrets.status, versionStatus: companySecretVersions.status, revokedAt: companySecretVersions.revokedAt, definitionId: companySecrets.userSecretDefinitionId, definitionStatus: userSecretDefinitions.status, definitionDeletedAt: userSecretDefinitions.deletedAt }) - .from(companySecrets).leftJoin(companySecretVersions, and(eq(companySecretVersions.secretId, companySecrets.id), eq(companySecretVersions.version, companySecrets.latestVersion))).leftJoin(userSecretDefinitions, and(eq(userSecretDefinitions.id, companySecrets.userSecretDefinitionId), eq(userSecretDefinitions.companyId, companyId))).where(and(eq(companySecrets.companyId, companyId), inArray(companySecrets.id, secretIds))) : []; - async function readAccount(row: (typeof accounts)[number]): Promise { - const accountKey = createHash("sha256").update(JSON.stringify([companyId, row.connection.id, row.grant.id, + const readRevisions = async (ids: string[]) => ids.length ? await db.select({ id: companySecrets.id, latestVersion: companySecrets.latestVersion, status: companySecrets.status, versionStatus: companySecretVersions.status, revokedAt: companySecretVersions.revokedAt, definitionId: companySecrets.userSecretDefinitionId, definitionStatus: userSecretDefinitions.status, definitionDeletedAt: userSecretDefinitions.deletedAt }) + .from(companySecrets).leftJoin(companySecretVersions, and(eq(companySecretVersions.secretId, companySecrets.id), eq(companySecretVersions.version, companySecrets.latestVersion))).leftJoin(userSecretDefinitions, and(eq(userSecretDefinitions.id, companySecrets.userSecretDefinitionId), eq(userSecretDefinitions.companyId, companyId))).where(and(eq(companySecrets.companyId, companyId), inArray(companySecrets.id, ids))) : []; + const revisions = await readRevisions(secretIds); + const identity = (row: (typeof accounts)[number], metadata: typeof revisions) => ({ + provider: row.summary.provider, + accountKey: createHash("sha256").update(JSON.stringify([companyId, row.connection.id, row.grant.id, row.connection.updatedAt, row.grant.updatedAt, - row.grant.credentialSecretRefs.map(ref => [ref.secretId, revisions.find(r => r.id === ref.secretId)]), - ])).digest("hex"); - const base = { provider: row.summary.provider, accountKey, accountLabel: row.summary.name, source: "managed-connection" }; + row.grant.credentialSecretRefs.map(ref => [ref.secretId, metadata.find(r => r.id === ref.secretId)]), + ])).digest("hex"), + accountLabel: row.summary.name, source: "managed-connection", + }); + async function readAccount(row: (typeof accounts)[number]): Promise { + let base = identity(row, revisions); + const { accountKey } = base; const definitionUnavailable = row.grant.credentialSecretRefs.some(ref => { const revision = revisions.find(value => value.id === ref.secretId); return revision && (revision.status !== "active" || revision.versionStatus === "disabled" || revision.revokedAt != null @@ -65,6 +71,22 @@ export async function fetchCompanyQuotaWindows(db: Db, companyId: string, userId catch (error) { if (!(error instanceof Error) || !/\b401\b/.test(error.message)) throw error; auth = JSON.parse(await service.refreshQuotaCredential(row, value, controller.signal)); + // Rotation changes both the grant and secret revision. Return and + // cache the observation under the identity the next poll will see. + const refreshed = (await service.quotaAccounts(companyId, userId)).find(candidate => + candidate.connection.id === row.connection.id && candidate.grant.id === row.grant.id); + if (!refreshed || refreshed.summary.status !== "connected") throw new Error("credentials_unavailable"); + base = identity(refreshed, await readRevisions(refreshed.grant.credentialSecretRefs.map(ref => ref.secretId))); + // A reconnect can race the completed refresh. Bind the value and + // cache identity to one revision, rather than label old quota with + // a newer account's key. A concurrent change defers this poll. + const refreshedValue = await service.credential(refreshed); + const verified = (await service.quotaAccounts(companyId, userId)).find(candidate => + candidate.connection.id === row.connection.id && candidate.grant.id === row.grant.id); + if (!verified || verified.summary.status !== "connected") throw new Error("credentials_unavailable"); + const verifiedIdentity = identity(verified, await readRevisions(verified.grant.credentialSecretRefs.map(ref => ref.secretId))); + if (verifiedIdentity.accountKey !== base.accountKey) throw new Error("provider_unavailable"); + auth = JSON.parse(refreshedValue); controller.signal.throwIfAborted(); windows = await read(auth); } @@ -83,7 +105,15 @@ export async function fetchCompanyQuotaWindows(db: Db, companyId: string, userId errorFamily }; } })(); - const bounded = withQuotaTimeout(base.provider, result, () => controller.abort()).then(quota => ({ ...base, ...quota, ...(quota.ok ? {} : { error: publicError }) })); + const bounded: Promise = withQuotaTimeout(base.provider, result, () => controller.abort()).then(quota => { + const observation = { ...base, ...quota, ...(quota.ok ? {} : { error: publicError }) }; + if (observation.ok && observation.accountKey !== accountKey) { + if (accountRequests.get(key)?.result === bounded) accountRequests.delete(key); + if (accountRequests.size >= 500) accountRequests.delete(accountRequests.keys().next().value!); + accountRequests.set(`${userId}:${observation.accountKey}`, { result: Promise.resolve(observation), expires: Date.now() + 30_000 }); + } + return observation; + }); pending = { result: bounded, expires: Date.now() + 30_000 }; if (accountRequests.size >= 500) accountRequests.delete(accountRequests.keys().next().value!); accountRequests.set(key, pending); diff --git a/ui/src/pages/Costs.production.tsx b/ui/src/pages/Costs.production.tsx index 96b8b9442b7..0dba6b087e5 100644 --- a/ui/src/pages/Costs.production.tsx +++ b/ui/src/pages/Costs.production.tsx @@ -9,6 +9,7 @@ import type { CostWindowSpendRow, FinanceEvent, QuotaWindow, + ProviderQuotaResult, } from "@paperclipai/shared"; import { ArrowDownLeft, ArrowUpRight, ChevronDown, ChevronRight, Coins, DollarSign, ReceiptText } from "lucide-react"; import { budgetsApi } from "../api/budgets"; @@ -28,6 +29,7 @@ import { StatusBadge } from "../components/StatusBadge"; import { useBreadcrumbs } from "../context/BreadcrumbContext"; import { useCompany } from "../context/CompanyContext"; import { useDateRange, PRESET_KEYS, PRESET_LABELS } from "../hooks/useDateRange"; +import { retainQuotaWindows } from "../lib/quota-refresh"; import { queryKeys } from "../lib/queryKeys"; import { billingTypeDisplayName, cn, formatCents, formatTokens, providerDisplayName } from "../lib/utils"; import { Button } from "@/components/ui/button"; @@ -333,6 +335,10 @@ export function Costs() { const { data: quotaData, isLoading: quotaLoading } = useQuery({ queryKey: queryKeys.usageQuotaWindows(companyId), queryFn: () => costsApi.quotaWindows(companyId), + structuralSharing: (previous, incoming) => retainQuotaWindows( + previous as ProviderQuotaResult[] | undefined, + incoming as ProviderQuotaResult[], + ), enabled: !!selectedCompanyId && mainTab === "providers", refetchInterval: 300_000, staleTime: 60_000, diff --git a/ui/src/pages/Costs.quota.test.tsx b/ui/src/pages/Costs.quota.test.tsx new file mode 100644 index 00000000000..5cd3020250e --- /dev/null +++ b/ui/src/pages/Costs.quota.test.tsx @@ -0,0 +1,55 @@ +// @vitest-environment jsdom +import { act } from "react"; +import { createRoot } from "react-dom/client"; +import { MemoryRouter } from "react-router-dom"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { Costs } from "./Costs"; +import { SidebarProvider } from "../context/SidebarContext"; + +const api = vi.hoisted(() => Object.fromEntries([ + "summary", "byAgent", "byProject", "byAgentModel", "financeSummary", "financeByBiller", + "financeByKind", "financeEvents", "byProvider", "byBiller", "windowSpend", "quotaWindows", +].map(key => [key, vi.fn()]))); +vi.mock("../api/costs", () => ({ costsApi: api })); +vi.mock("../api/budgets", () => ({ budgetsApi: { overview: async () => ({ policies: [], activeIncidents: [] }) } })); +vi.mock("../context/CompanyContext", () => ({ useCompany: () => ({ selectedCompanyId: "quota-company" }) })); +vi.mock("../context/BreadcrumbContext", () => ({ useBreadcrumbs: () => ({ setBreadcrumbs: () => {} }) })); +Object.assign(globalThis, { IS_REACT_ACT_ENVIRONMENT: true }); + +describe("Costs account quota polling", () => { + let container: HTMLDivElement, root: ReturnType, client: QueryClient; + const queryKey = ["usage-quota-windows", "quota-company"]; + const observed = { provider: "openai", accountKey: "credential-one", accountLabel: "Codie account", ok: true, + capturedAt: "2026-10-03T12:00:00Z", windows: [{ label: "5h limit", usedPercent: 37, resetsAt: null, valueLabel: null }] }; + beforeEach(() => { + vi.stubGlobal("matchMedia", vi.fn(() => ({ matches: false, addEventListener: vi.fn(), removeEventListener: vi.fn() }))); + for (const mock of Object.values(api)) mock.mockReset().mockResolvedValue([]); + api.summary.mockResolvedValue({ spendCents: 0, budgetCents: 0 }); + api.financeSummary.mockResolvedValue({ debitCents: 0, creditCents: 0, netCents: 0, eventCount: 0 }); + api.byProvider.mockResolvedValue([{ provider: "openai", biller: "openai", model: "test", billingType: "metered_api", costCents: 0, inputTokens: 1, outputTokens: 1, cachedInputTokens: 0 }]); + api.quotaWindows.mockResolvedValue([observed]); + container = document.createElement("div"); document.body.append(container); root = createRoot(container); + client = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + }); + afterEach(() => { act(() => root.unmount()); client.clear(); container.remove(); vi.unstubAllGlobals(); }); + const flush = async () => { await act(async () => { await new Promise(resolve => setTimeout(resolve, 10)); }); }; + it.each(["provider_unavailable", "permission_denied"])("retains the observation on %s but clears rejected or rotated credentials", async (errorFamily) => { + await act(async () => root.render()); + await flush(); + expect(container.textContent).toContain("37% used"); + api.quotaWindows.mockResolvedValue([{ ...observed, ok: false, windows: [], capturedAt: undefined, errorFamily }]); + await act(async () => { await client.invalidateQueries({ queryKey }); }); await flush(); + expect(container.textContent).toContain("37% used"); + expect(container.textContent).toContain("Showing the last available quota"); + expect(client.getQueryData(queryKey)).toMatchObject([{ capturedAt: observed.capturedAt }]); + api.quotaWindows.mockResolvedValue([{ ...observed, ok: false, windows: [], errorFamily: "authentication_required" }]); + await act(async () => { await client.invalidateQueries({ queryKey }); }); await flush(); + expect(container.textContent).not.toContain("37% used"); + api.quotaWindows.mockResolvedValue([observed]); + await act(async () => { await client.invalidateQueries({ queryKey }); }); await flush(); + api.quotaWindows.mockResolvedValue([{ ...observed, accountKey: "credential-two", ok: false, windows: [], errorFamily }]); + await act(async () => { await client.invalidateQueries({ queryKey }); }); await flush(); + expect(container.textContent).not.toContain("37% used"); + }); +}); diff --git a/ui/src/pages/Costs.tsx b/ui/src/pages/Costs.tsx index b2abbc7cef8..1ff4d2d4a7f 100644 --- a/ui/src/pages/Costs.tsx +++ b/ui/src/pages/Costs.tsx @@ -9,6 +9,7 @@ import type { CostWindowSpendRow, FinanceEvent, QuotaWindow, + ProviderQuotaResult, } from "@paperclipai/shared"; import { ArrowDownLeft, ArrowUpRight, ChevronDown, ChevronRight, Coins, DollarSign, ReceiptText } from "lucide-react"; import { budgetsApi } from "../api/budgets"; @@ -28,6 +29,7 @@ import { StatusBadge } from "../components/StatusBadge"; import { useBreadcrumbs } from "../context/BreadcrumbContext"; import { useCompany } from "../context/CompanyContext"; import { useDateRange, PRESET_KEYS, PRESET_LABELS } from "../hooks/useDateRange"; +import { retainQuotaWindows } from "../lib/quota-refresh"; import { queryKeys } from "../lib/queryKeys"; import { billingTypeDisplayName, cn, formatCents, formatTokens, providerDisplayName } from "../lib/utils"; import { Button } from "@/components/ui/button"; @@ -354,6 +356,10 @@ export function Costs({ const { data: quotaData, isLoading: quotaLoading } = useQuery({ queryKey: queryKeys.usageQuotaWindows(companyId), queryFn: () => costsApi.quotaWindows(companyId), + structuralSharing: (previous, incoming) => retainQuotaWindows( + previous as ProviderQuotaResult[] | undefined, + incoming as ProviderQuotaResult[], + ), enabled: !!selectedCompanyId && mainTab === "providers", refetchInterval: 300_000, staleTime: 60_000,