From bd377f0906d234bc022e682e461a96a8d15cc96a Mon Sep 17 00:00:00 2001 From: Devin Foley Date: Sat, 3 Oct 2026 21:06:14 -0700 Subject: [PATCH] fix: recover managed quota credentials and preserve account displays Co-Authored-By: Paperclip --- server/src/__tests__/ai-connections.test.ts | 27 ++++++++ .../__tests__/company-quota-windows.test.ts | 28 +++++++-- server/src/services/ai-connections.ts | 62 ++++++++++++++++++- server/src/services/quota-windows.ts | 26 +++++--- ui/src/components/AccountQuotaPanels.test.tsx | 24 +++++++ ui/src/components/AccountQuotaPanels.tsx | 61 ++++++++++++++++++ ui/src/components/ProviderQuotaCard.tsx | 9 ++- ui/src/lib/quota-refresh.ts | 30 +++++++++ ui/src/pages/Costs.production.tsx | 4 +- ui/src/pages/Costs.tsx | 4 +- 10 files changed, 258 insertions(+), 17 deletions(-) create mode 100644 ui/src/components/AccountQuotaPanels.test.tsx create mode 100644 ui/src/components/AccountQuotaPanels.tsx create mode 100644 ui/src/lib/quota-refresh.ts diff --git a/server/src/__tests__/ai-connections.test.ts b/server/src/__tests__/ai-connections.test.ts index 3cb74ed7039..3d8632114a4 100644 --- a/server/src/__tests__/ai-connections.test.ts +++ b/server/src/__tests__/ai-connections.test.ts @@ -84,6 +84,33 @@ describe("managed AI connections", () => { } finally { fetchSpy.mockRestore(); } }); + it("serializes quota credential refresh, persists rotation, and defers while the provider is active", async () => { + const owner = "quota-refresh-owner"; + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const original = JSON.stringify({ tokens: { access_token: "expired-access", refresh_token: "single-use", id_token: "identity", account_id: null } }); + const account = await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: "Refresh account", loginSessionId: "fixture", agentIds: [], allAgents: true }, original); + const [row] = await service.quotaAccounts(companyId, owner); + const request = vi.spyOn(globalThis, "fetch").mockImplementation(async () => new Response(JSON.stringify({ access_token: "fresh-access", refresh_token: "rotated-once" }))); + const signal = new AbortController().signal; + try { + const activeId = randomUUID(); + await db.insert(heartbeatRuns).values({ id: activeId, companyId, agentId, invocationSource: "on_demand", status: "running", contextSnapshot: { aiConnection: { grantId: account.grantId } } }); + await expect(service.refreshQuotaCredential(row, original, signal)).rejects.toThrow("provider_unavailable"); + expect(request).not.toHaveBeenCalled(); + await db.update(heartbeatRuns).set({ status: "succeeded" }).where(eq(heartbeatRuns.id, activeId)); + const [first, second] = await Promise.all([ + service.refreshQuotaCredential(row, original, signal), + service.refreshQuotaCredential(row, original, signal), + ]); + expect(first).toBe(second); + expect(request).toHaveBeenCalledTimes(1); + expect(JSON.parse(await service.credential(row))).toMatchObject({ tokens: { access_token: "fresh-access", refresh_token: "rotated-once", account_id: null } }); + await db.update(connectionGrants).set({ status: "revoked" }).where(eq(connectionGrants.id, account.grantId)); + await expect(service.refreshQuotaCredential(row, first, signal)).rejects.toThrow("credentials_unavailable"); + expect(request).toHaveBeenCalledTimes(1); + } finally { request.mockRestore(); } + }); + it.each([false, true])("reports the authoritative connection-manager capability for custom grants (manager: %s)", async (manager) => { const userId = `custom-manager-${manager}`; await db.insert(companyMemberships).values({ companyId, principalType: "user", principalId: userId, 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 d848986790a..3af1d3f04dc 100644 --- a/server/src/__tests__/company-quota-windows.test.ts +++ b/server/src/__tests__/company-quota-windows.test.ts @@ -3,6 +3,7 @@ import type { Db } from "@paperclipai/db"; const mocks = vi.hoisted(() => ({ accounts: vi.fn(), credential: vi.fn(), + refresh: vi.fn(), codex: vi.fn(), claude: vi.fn(), adapters: vi.fn(), @@ -11,6 +12,7 @@ vi.mock("../services/ai-connections.js", () => ({ aiConnectionService: () => ({ quotaAccounts: mocks.accounts, credential: mocks.credential, + refreshQuotaCredential: mocks.refresh, }), })); vi.mock("../adapters/registry.js", () => ({ @@ -56,6 +58,7 @@ describe("connected account quotas", () => { beforeEach(() => { vi.clearAllMocks(); revision = 1; + mocks.refresh.mockRejectedValue(new Error("authentication_required")); mocks.credential.mockResolvedValue( JSON.stringify({ tokens: { access_token: "private-token", account_id: "actual-account" }, @@ -135,17 +138,34 @@ describe("connected account quotas", () => { new Error("private credential store failure"), ); const [failed] = await fetchCompanyQuotaWindows(db, "company-6", "user-6"); - expect(failed.errorFamily).toBe("authentication_required"); + expect(failed.errorFamily).toBe("provider_unavailable"); expect(JSON.stringify(failed)).not.toContain("private credential"); }); - it("rejects an incomplete account identity without querying a different account", async () => { + it("reads quota for a verified login without an optional account ID", async () => { mocks.accounts.mockResolvedValue([account("missing-identity")]); mocks.credential.mockResolvedValue( JSON.stringify({ accessToken: "private" }), ); const [result] = await fetchCompanyQuotaWindows(db, "company-7", "user-7"); - expect(result.errorFamily).toBe("credentials_unavailable"); - expect(mocks.codex).not.toHaveBeenCalled(); + expect(result.ok).toBe(true); + expect(mocks.codex).toHaveBeenCalledWith("private", null, expect.any(AbortSignal)); + }); + it("distinguishes permission errors from revoked credentials", async () => { + mocks.accounts.mockResolvedValue([account("permission-error")]); + mocks.codex.mockRejectedValueOnce(new Error("chatgpt wham api returned 403")); + const [result] = await fetchCompanyQuotaWindows(db, "permission-company", "user"); + expect(result.errorFamily).toBe("permission_denied"); + expect(mocks.refresh).not.toHaveBeenCalled(); + }); + 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" } })); + 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"); }); it("bounds a stalled credential lookup and preserves its account identity", async () => { vi.useFakeTimers(); diff --git a/server/src/services/ai-connections.ts b/server/src/services/ai-connections.ts index d27d19a7577..767ae7ecb0c 100644 --- a/server/src/services/ai-connections.ts +++ b/server/src/services/ai-connections.ts @@ -1,6 +1,6 @@ import { syncConnectionCredentialBindings } from "./connection-credential-bindings.js"; import { createHash, randomUUID } from "node:crypto"; -import { and, eq, inArray, or } from "drizzle-orm"; +import { and, eq, inArray, or, sql } from "drizzle-orm"; import { type Db, authUsers, @@ -8,6 +8,7 @@ import { aiConnectionDefaults, aiProviderDefaults, agents, + heartbeatRuns, companyMemberships, companySecrets, userSecretDefinitions, @@ -848,6 +849,63 @@ export function aiConnectionService(db: Db) { }); }); } + /** Refresh only the selected vaulted identity, serializing rotation and reconnect. */ + async function refreshQuotaCredential(row: Parameters[0], failedValue: string, signal: AbortSignal) { + return db.transaction(async (tx) => { + const companyId = row.connection.companyId; + const [grant] = await tx.select().from(connectionGrants).where(and( + eq(connectionGrants.id, row.grant.id), eq(connectionGrants.companyId, companyId), + eq(connectionGrants.connectionId, row.connection.id), + )).for("update"); + if (!grant || grant.status !== "active") throw new Error("credentials_unavailable"); + const ref = grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential"); + if (!ref) throw new Error("credentials_unavailable"); + await tx.select({ id: companySecrets.id }).from(companySecrets).where(and( + eq(companySecrets.id, ref.secretId), eq(companySecrets.companyId, companyId), + )).for("update"); + const current = await aiConnectionService(tx as unknown as Db).credential({ connection: row.connection, grant }); + signal.throwIfAborted(); + // Another poll, an agent, or reconnect already supplied a newer credential. + if (current !== failedValue) return current; + // An active provider may rotate this single-use token itself. Defer polling + // until its credential write-back instead of invalidating its live copy. + const [active] = await tx.select({ id: heartbeatRuns.id }).from(heartbeatRuns) + .innerJoin(agents, and(eq(agents.id, heartbeatRuns.agentId), eq(agents.companyId, companyId))).where(and( + eq(heartbeatRuns.companyId, companyId), inArray(heartbeatRuns.status, ["queued", "running"]), + or(sql`${heartbeatRuns.contextSnapshot}->'aiConnection'->>'grantId' = ${grant.id}`, + sql`${agents.runtimeConfig}->'aiConnection'->>'provider' = 'openai'`, eq(agents.adapterType, "codex_local")), + )).limit(1); + if (active) throw new Error("provider_unavailable"); + const auth = JSON.parse(current); + const refreshToken = auth.tokens?.refresh_token; + if (typeof refreshToken !== "string" || !refreshToken) throw new Error("authentication_required"); + // 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, + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ client_id: "app_EMoamEEZ73f0CkXaXp7hrann", grant_type: "refresh_token", refresh_token: refreshToken }), + }); + if (!response.ok) { + await response.body?.cancel(); + throw new Error(response.status === 400 || response.status === 401 ? "authentication_required" : "provider_unavailable"); + } + const tokens = await response.json() as Record; + if (typeof tokens.access_token !== "string" || !tokens.access_token) throw new Error("provider_unavailable"); + const refreshed = JSON.stringify({ ...auth, tokens: { ...auth.tokens, + access_token: tokens.access_token, + ...(typeof tokens.refresh_token === "string" && tokens.refresh_token ? { refresh_token: tokens.refresh_token } : {}), + ...(typeof tokens.id_token === "string" && tokens.id_token ? { id_token: tokens.id_token } : {}), + }, last_refresh: new Date().toISOString() }); + // Persist the rotated token even if the quota deadline expires after the + // successful exchange; losing it would strand all subsequent executions. + await secretService(tx).rotate(ref.secretId, { value: refreshed }, { userId: grant.subjectUserId }); + await tx.update(connectionGrants).set({ updatedAt: new Date() }).where(eq(connectionGrants.id, grant.id)); + await logActivity(tx as unknown as Db, { companyId, actorType: "system", actorId: "quota", + action: "ai_connection.credential_refreshed", entityType: "tool_connection", entityId: row.connection.id, + details: { provider: "openai", grantId: grant.id } }); + return refreshed; + }); + } // A quota read uses the same credential audience as execution. Operator status // alone never grants access to another member's personal subscription. async function quotaAccounts(companyId: string, userId: string) { @@ -860,5 +918,5 @@ export function aiConnectionService(db: Db) { return [{ ...row, summary }]; }); } - return { list, quotaAccounts, select, credential, probeUsage, save, setDefault, membership, markAuthenticationFailed }; + return { list, quotaAccounts, refreshQuotaCredential, select, credential, probeUsage, save, setDefault, membership, markAuthenticationFailed }; } diff --git a/server/src/services/quota-windows.ts b/server/src/services/quota-windows.ts index c290969c96f..f2631d87617 100644 --- a/server/src/services/quota-windows.ts +++ b/server/src/services/quota-windows.ts @@ -35,7 +35,8 @@ export async function fetchCompanyQuotaWindows(db: Db, companyId: string, userId const base = { provider: row.summary.provider, accountKey, accountLabel: row.summary.name, source: "managed-connection" }; const definitionUnavailable = row.grant.credentialSecretRefs.some(ref => { const revision = revisions.find(value => value.id === ref.secretId); - return revision?.definitionId && (revision.definitionStatus !== "active" || revision.definitionDeletedAt !== null); + return revision && (revision.status !== "active" || revision.versionStatus === "disabled" || revision.revokedAt != null + || (revision.definitionId && (revision.definitionStatus !== "active" || revision.definitionDeletedAt !== null))); }); if (row.summary.status !== "connected" || definitionUnavailable) { return { ...base, ok: false, errorFamily: "credentials_unavailable", error: publicError, windows: [] }; @@ -52,23 +53,34 @@ export async function fetchCompanyQuotaWindows(db: Db, companyId: string, userId credentialResolved = true; let windows; if (base.provider === "openai") { - const auth = JSON.parse(value); + let auth; + try { auth = JSON.parse(value); } catch { throw new Error("credentials_unavailable"); } const token = auth.tokens?.access_token ?? auth.accessToken; - const accountId = auth.tokens?.account_id ?? auth.accountId; - if (typeof token !== "string" || !token || typeof accountId !== "string" || !accountId) { + if (typeof token !== "string" || !token) { return { ...base, ok: false, errorFamily: "credentials_unavailable", error: publicError, windows: [] }; } - windows = await fetchCodexQuota(token, accountId, controller.signal); + const read = (auth: { tokens?: { access_token?: string; account_id?: string | null }; accessToken?: string; accountId?: string | null }) => + fetchCodexQuota(auth.tokens?.access_token ?? auth.accessToken!, auth.tokens?.account_id ?? auth.accountId ?? null, controller.signal); + try { windows = await read(auth); } + catch (error) { + if (!(error instanceof Error) || !/\b401\b/.test(error.message)) throw error; + auth = JSON.parse(await service.refreshQuotaCredential(row, value, controller.signal)); + controller.signal.throwIfAborted(); + windows = await read(auth); + } } else { windows = await fetchClaudeQuota(value, controller.signal); } return { ...base, ok: true, windows, capturedAt: new Date().toISOString() }; } catch (error) { const message = error instanceof Error ? error.message : "Quota unavailable"; - const invalid = !credentialResolved || /401|403|refresh_token_(?:reused|expired|invalidated)/i.test(message); + const missing = message === "credentials_unavailable"; + const invalid = message === "authentication_required" || (credentialResolved && /\b401\b|refresh_token_(?:reused|expired|invalidated)/i.test(message)); + const errorFamily = missing ? "credentials_unavailable" : invalid ? "authentication_required" + : credentialResolved && /\b403\b/.test(message) ? "permission_denied" : "provider_unavailable"; logger.warn({ companyId, accountKey, provider: base.provider, authenticationFailed: invalid }, "Connected account quota unavailable"); return { ...base, ok: false, windows: [], error: publicError, - ...(invalid ? { errorFamily: "authentication_required" } : {}) }; + errorFamily }; } })(); const bounded = withQuotaTimeout(base.provider, result, () => controller.abort()).then(quota => ({ ...base, ...quota, ...(quota.ok ? {} : { error: publicError }) })); diff --git a/ui/src/components/AccountQuotaPanels.test.tsx b/ui/src/components/AccountQuotaPanels.test.tsx new file mode 100644 index 00000000000..c87a755b924 --- /dev/null +++ b/ui/src/components/AccountQuotaPanels.test.tsx @@ -0,0 +1,24 @@ +import { renderToStaticMarkup } from "react-dom/server"; +import { describe, expect, it } from "vitest"; +import { AccountQuotaPanels } from "./AccountQuotaPanels"; + +describe("account quota panels", () => { + it("renders both accounts for one provider without merging their windows", () => { + const html = renderToStaticMarkup(); + expect(html).toContain("Personal limit"); + expect(html).toContain("Shared limit"); + expect(html).toContain("25% used"); + expect(html).toContain("75% used"); + }); + it("displays unknown utilization without inventing zero percent", () => { + const html = renderToStaticMarkup(); + expect(html).toContain("$5 remaining"); + expect(html).not.toContain("0%"); + expect(html).not.toContain('role="progressbar"'); + }); +}); diff --git a/ui/src/components/AccountQuotaPanels.tsx b/ui/src/components/AccountQuotaPanels.tsx new file mode 100644 index 00000000000..31cd43d05d7 --- /dev/null +++ b/ui/src/components/AccountQuotaPanels.tsx @@ -0,0 +1,61 @@ +import type { ProviderQuotaResult } from "@paperclipai/shared"; +import { QuotaBar } from "./QuotaBar"; +import { quotaUnavailableMessage } from "../lib/quota-refresh"; + +export function AccountQuotaPanels({ + accounts, + failed, +}: { + accounts: ProviderQuotaResult[]; + failed?: boolean; +}) { + return ( +
+ {accounts.map((account) => ( +
+

+ {account.accountLabel ?? "Subscription account"} +

+ {(failed || !account.ok) && ( +

+ {account.errorFamily === "credentials_unavailable" + ? "Connect or reconnect a subscription in AI connections to view its quota." + : quotaUnavailableMessage(account.windows.length > 0)} +

+ )} + {account.windows.map((window, index) => window.usedPercent == null ? ( +
+ {window.label} + {window.valueLabel ?? "Usage not reported"} +
+ ) : ( + + ))} + {account.capturedAt && ( +

+ Last checked {new Date(account.capturedAt).toLocaleString()} +

+ )} +
+ ))} +
+ ); +} diff --git a/ui/src/components/ProviderQuotaCard.tsx b/ui/src/components/ProviderQuotaCard.tsx index bd20aedbe5a..c3273133f82 100644 --- a/ui/src/components/ProviderQuotaCard.tsx +++ b/ui/src/components/ProviderQuotaCard.tsx @@ -1,5 +1,6 @@ +import { AccountQuotaPanels } from "./AccountQuotaPanels"; import { useMemo } from "react"; -import type { CostByProviderModel, CostWindowSpendRow, QuotaWindow } from "@paperclipai/shared"; +import type { CostByProviderModel, CostWindowSpendRow, QuotaWindow, ProviderQuotaResult } from "@paperclipai/shared"; import { Card, CardContent, CardHeader, CardTitle, CardDescription } from "@/components/ui/card"; import { Skeleton } from "@/components/ui/skeleton"; import { QuotaBar } from "./QuotaBar"; @@ -33,6 +34,7 @@ interface ProviderQuotaCardProps { quotaError?: string | null; quotaSource?: string | null; quotaLoading?: boolean; + quotaAccounts?: ProviderQuotaResult[]; } export function ProviderQuotaCard({ @@ -47,6 +49,7 @@ export function ProviderQuotaCard({ quotaError = null, quotaSource = null, quotaLoading = false, + quotaAccounts, }: ProviderQuotaCardProps) { // single-pass aggregation over rows — memoized so the 8 derived values are not // recomputed on every parent render tick (providers tab polls every 30s, and each @@ -127,7 +130,7 @@ export function ProviderQuotaCard({ const isCodexQuotaPanel = provider === "openai" && quotaSource?.startsWith("codex-"); const supportsSubscriptionQuota = provider === "anthropic" || provider === "openai"; const showSubscriptionQuotaSection = - supportsSubscriptionQuota && (quotaLoading || quotaWindows.length > 0 || quotaError != null); + supportsSubscriptionQuota && (quotaLoading || Boolean(quotaAccounts?.length) || quotaWindows.length > 0 || quotaError != null); return ( @@ -321,6 +324,8 @@ export function ProviderQuotaCard({ {quotaLoading ? ( + ) : quotaAccounts ? ( + ) : isClaudeQuotaPanel ? ( ) : isCodexQuotaPanel ? ( diff --git a/ui/src/lib/quota-refresh.ts b/ui/src/lib/quota-refresh.ts new file mode 100644 index 00000000000..75df4027e1e --- /dev/null +++ b/ui/src/lib/quota-refresh.ts @@ -0,0 +1,30 @@ +import type { ProviderQuotaResult } from "@paperclipai/shared"; + +const invalidCredentials = new Set([ + "credentials_unavailable", + "authentication_required", + "refresh_token_reused", + "refresh_token_expired", + "refresh_token_invalidated", +]); + +/** Keep last successful windows through transient failures, never revoked auth. */ +export function retainQuotaWindows( + previous: ProviderQuotaResult[] | undefined, + incoming: ProviderQuotaResult[], +): ProviderQuotaResult[] { + return incoming.map((result) => { + if (result.ok || (result.errorFamily && invalidCredentials.has(result.errorFamily))) return result; + const last = previous?.find((entry) => entry.provider === result.provider && entry.accountKey === result.accountKey); + return last?.windows.length + ? { ...result, windows: last.windows, source: last.source, capturedAt: last.capturedAt } + : result; + }); +} + +// Never render adapter diagnostics, even when talking to an older server. +export function quotaUnavailableMessage(hasPreviousWindows: boolean): string { + return hasPreviousWindows + ? "Showing the last available quota. Updates will resume automatically." + : "Subscription quota is currently unavailable. Check usage with your provider."; +} diff --git a/ui/src/pages/Costs.production.tsx b/ui/src/pages/Costs.production.tsx index bbb8b9a8043..96b8b9442b7 100644 --- a/ui/src/pages/Costs.production.tsx +++ b/ui/src/pages/Costs.production.tsx @@ -973,6 +973,7 @@ export function Costs() { weekSpendCents={weekSpendByProvider.get(provider) ?? 0} windowRows={windowSpendByProvider.get(provider) ?? []} showDeficitNotch={deficitNotchByProvider.get(provider) ?? false} + quotaAccounts={quotaData?.filter(account => account.provider === provider)} quotaWindows={quotaWindowsByProvider.get(provider) ?? []} quotaError={quotaErrorsByProvider.get(provider) ?? null} quotaSource={quotaSourcesByProvider.get(provider) ?? null} @@ -993,7 +994,8 @@ export function Costs() { weekSpendCents={weekSpendByProvider.get(provider) ?? 0} windowRows={windowSpendByProvider.get(provider) ?? []} showDeficitNotch={deficitNotchByProvider.get(provider) ?? false} - quotaWindows={quotaWindowsByProvider.get(provider) ?? []} + quotaAccounts={quotaData?.filter(account => account.provider === provider)} + quotaWindows={quotaWindowsByProvider.get(provider) ?? []} quotaError={quotaErrorsByProvider.get(provider) ?? null} quotaSource={quotaSourcesByProvider.get(provider) ?? null} quotaLoading={quotaLoading} diff --git a/ui/src/pages/Costs.tsx b/ui/src/pages/Costs.tsx index e28cf08dd9a..b2abbc7cef8 100644 --- a/ui/src/pages/Costs.tsx +++ b/ui/src/pages/Costs.tsx @@ -1001,6 +1001,7 @@ export function Costs({ weekSpendCents={weekSpendByProvider.get(provider) ?? 0} windowRows={windowSpendByProvider.get(provider) ?? []} showDeficitNotch={deficitNotchByProvider.get(provider) ?? false} + quotaAccounts={quotaData?.filter(account => account.provider === provider)} quotaWindows={quotaWindowsByProvider.get(provider) ?? []} quotaError={quotaErrorsByProvider.get(provider) ?? null} quotaSource={quotaSourcesByProvider.get(provider) ?? null} @@ -1021,7 +1022,8 @@ export function Costs({ weekSpendCents={weekSpendByProvider.get(provider) ?? 0} windowRows={windowSpendByProvider.get(provider) ?? []} showDeficitNotch={deficitNotchByProvider.get(provider) ?? false} - quotaWindows={quotaWindowsByProvider.get(provider) ?? []} + quotaAccounts={quotaData?.filter(account => account.provider === provider)} + quotaWindows={quotaWindowsByProvider.get(provider) ?? []} quotaError={quotaErrorsByProvider.get(provider) ?? null} quotaSource={quotaSourcesByProvider.get(provider) ?? null} quotaLoading={quotaLoading}