From 892b0b3606f066dd7344d6397daf25c8b52f598b Mon Sep 17 00:00:00 2001 From: Devin Foley Date: Tue, 6 Oct 2026 18:20:37 -0700 Subject: [PATCH] fix: scope quota reports to authorized subscription accounts (#14994) Co-Authored-By: Paperclip --- doc/connections/AI-CONNECTIONS.md | 51 +- .../adapters/claude-local/src/server/quota.ts | 5 +- .../src/server/codex-auth-cache.test.ts | 36 +- .../src/server/codex-auth-cache.ts | 19 +- .../src/server/quota-spawn-error.test.ts | 185 ++++- .../adapters/codex-local/src/server/quota.ts | 109 ++- packages/shared/src/types/quota.ts | 4 + .../agent-hire-ai-connections.test.ts | 34 + .../agent-test-environment-routes.test.ts | 4 + server/src/__tests__/ai-connections.test.ts | 646 +++++++++++++++++- .../__tests__/company-quota-windows.test.ts | 256 +++++++ .../src/__tests__/durable-directory.test.ts | 243 +++++++ .../__tests__/quota-windows-service.test.ts | 42 +- server/src/__tests__/quota-windows.test.ts | 40 +- server/src/lib/durable-directory.ts | 71 ++ server/src/routes/agents.ts | 19 +- server/src/routes/costs.ts | 4 +- server/src/services/ai-connection-runtime.ts | 201 ++++-- server/src/services/ai-connections.ts | 300 ++++++-- .../src/services/quota-credential-recovery.ts | 54 ++ server/src/services/quota-windows.ts | 144 +++- ui/src/components/ClaudeSubscriptionPanel.tsx | 7 +- ui/src/components/CodexSubscriptionPanel.tsx | 7 +- ui/src/lib/quota-refresh.ts | 30 + 24 files changed, 2329 insertions(+), 182 deletions(-) create mode 100644 server/src/__tests__/company-quota-windows.test.ts create mode 100644 server/src/__tests__/durable-directory.test.ts create mode 100644 server/src/lib/durable-directory.ts create mode 100644 server/src/services/quota-credential-recovery.ts create mode 100644 ui/src/lib/quota-refresh.ts diff --git a/doc/connections/AI-CONNECTIONS.md b/doc/connections/AI-CONNECTIONS.md index 2e74412af1b..3c379dc9107 100644 --- a/doc/connections/AI-CONNECTIONS.md +++ b/doc/connections/AI-CONNECTIONS.md @@ -193,6 +193,17 @@ and contain neither credentials nor raw provider errors. Expired credentials report `authentication_required`; the probe does not exchange refresh tokens or change connection health. +The Costs dashboard uses the company quota endpoint. It reads each authorized +managed subscription separately, retaining the last successful observation on +transient errors. For managed Codex accounts, a 401 can trigger one persisted +OAuth refresh; it defers while potentially competing OpenAI work is active. +Refresh takes the company secret-mutation lock before database row locks and +requires those row locks immediately, before exchanging a single-use token. +Reconnect and runtime credential write-back follow the same lock order. +Credential resolution checks its version after any database wait and reloads a +rotated value before giving it to a new run. Unknown utilization has no progress +bar, but any provider-reported reset time remains visible. + Verification: ```sh @@ -498,9 +509,10 @@ existing account name. Connecting installs access for that agent and resumes the work automatically. Explicit incompatible bindings and shared-account permission denials still fail; hiring never expands a restricted shared account's audience. -Concurrent runs of one subscription do not wait for each other. No credential -lease exists to hold them, so a fresh task execution cannot enter a contention -wait. A run that already entered this wait keeps its scheduled retries. It does +Concurrent runs of one subscription do not hold a credential lease. A fresh +task can briefly wait when a credential rotation holds the company file lock or +a grant/secret database row lock. Lock timeouts become `ai_connection_busy`, +keeping the task on its automatic pre-provider scheduled retry path. It does not request new credentials, and it does not consume the provider-failure retry allowance. Each retry revalidates the account, and existing run-dispatch rules still suppress cancelled, reassigned, or otherwise ineligible work. An assignee @@ -524,6 +536,39 @@ 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. + +Managed runtime token write-back retries a company-lock timeout twice (three +30-second acquisition attempts), so a slow quota exchange does not discard a +different account's replacement tokens. Any remaining write-back error preserves +the private runtime home for retry; cleanup deletes it only after a successful +transaction or an intentional freshness/authorization discard. These retained +homes are recovery evidence, not a background replay queue; persistent database +or lock failures still require operator intervention. + +Quota OAuth replacement tokens are encrypted with the instance secrets master +key and fsynced under `/quota-credential-recovery//` +before vault, grant, or activity writes. Storage and encryption are checked before +exchanging a single-use refresh token. Failed saves retry without another OAuth +exchange; persistent failures keep the encrypted record for the next quota poll, +including after a restart. New OpenAI subscription runtimes serialize credential +reads with quota exchanges and recover a matching pending replacement before +materializing an auth home. Authentication-failure handling also recovers a +matching replacement before marking the grant invalid. A failed recovery save +defers these actions and leaves both the active grant and journal intact. +The journal is removed only after database commit (or +when an authorized newer credential makes it obsolete). Replay checks the grant, +connection, secret, and original credential fingerprint and cannot reactivate a +revoked grant or overwrite a reconnect. Back up this directory and the instance +secrets key with the instance data. Loss of durable storage while receiving a +provider token can still require reconnecting the account. ## Advanced provider routing (2026-10-02) diff --git a/packages/adapters/claude-local/src/server/quota.ts b/packages/adapters/claude-local/src/server/quota.ts index 44d4cc633f6..4774acbd710 100644 --- a/packages/adapters/claude-local/src/server/quota.ts +++ b/packages/adapters/claude-local/src/server/quota.ts @@ -265,14 +265,15 @@ export async function fetchWithTimeout(url: string, init: RequestInit, ms = 8000 const controller = new AbortController(); const timer = setTimeout(() => controller.abort(), ms); try { - return await fetch(url, { ...init, signal: controller.signal }); + return await fetch(url, { ...init, signal: init.signal ? AbortSignal.any([init.signal, controller.signal]) : controller.signal }); } finally { clearTimeout(timer); } } -export async function fetchClaudeQuota(token: string): Promise { +export async function fetchClaudeQuota(token: string, signal?: AbortSignal): Promise { const resp = await fetchWithTimeout("https://api.anthropic.com/api/oauth/usage", { + signal, headers: { Authorization: `Bearer ${token}`, "anthropic-beta": "oauth-2025-04-20", diff --git a/packages/adapters/codex-local/src/server/codex-auth-cache.test.ts b/packages/adapters/codex-local/src/server/codex-auth-cache.test.ts index 0ca1bee28ca..2cd6ba7c173 100644 --- a/packages/adapters/codex-local/src/server/codex-auth-cache.test.ts +++ b/packages/adapters/codex-local/src/server/codex-auth-cache.test.ts @@ -463,13 +463,15 @@ describe("codex auth cache store", () => { releaseFirstCaller = resolve; }); + const entered = deferredSignal(); const firstCall = withAccountHomeSecretMutationLock(env, "company-shared", async () => { events.push("first-enter"); + entered.resolve(); await firstCallerGate; events.push("first-exit"); return "first"; }); - await new Promise((resolve) => setTimeout(resolve, 20)); + await entered.promise; const secondCall = withAccountHomeSecretMutationLock(env, "company-shared", async () => { events.push("second-enter"); events.push("second-exit"); @@ -497,12 +499,14 @@ describe("codex auth cache store", () => { releaseCompanyA = resolve; }); + const entered = deferredSignal(); const companyACall = withAccountHomeSecretMutationLock(env, "company-a", async () => { events.push("a-enter"); + entered.resolve(); await companyAGate; events.push("a-exit"); }); - await new Promise((resolve) => setTimeout(resolve, 20)); + await entered.promise; const companyBCall = withAccountHomeSecretMutationLock(env, "company-b", async () => { events.push("b-enter"); events.push("b-exit"); @@ -536,6 +540,28 @@ describe("codex auth cache store", () => { expect(events).toEqual(["promotion-enter", "mutation-enter", "mutation-exit", "promotion-exit"]); }); + it("allows awaited nested writes but expires inherited ownership after release", async () => { + const env = envFor(await makeInstanceRoot()); + const detachedGate = deferredSignal(); + let detached!: Promise; + let enteredDetached = false; + await withAccountHomeSecretMutationLock(env, "nested", async () => { + await expect(withAccountHomeSecretMutationLock(env, "nested", async () => "nested write")).resolves.toBe("nested write"); + detached = detachedGate.promise.then(() => withAccountHomeSecretMutationLock(env, "nested", async () => { enteredDetached = true; })); + }); + const holding = deferredSignal(); + const release = deferredSignal(); + const other = withAccountHomeSecretMutationLock(env, "nested", async () => { holding.resolve(); await release.promise; }); + await holding.promise; + detachedGate.resolve(); + try { + await new Promise(resolve => setTimeout(resolve, 30)); + expect(enteredDetached).toBe(false); + } finally { release.resolve(); } + await Promise.all([other, detached]); + expect(enteredDetached).toBe(true); + }); + }); describe("Phase 8: account-home directory validity before a secret write commits", () => { @@ -581,3 +607,9 @@ describe("codex auth cache store", () => { }); }); }); + +function deferredSignal() { + let resolve!: () => void; + const promise = new Promise(done => { resolve = done; }); + return { promise, resolve }; +} diff --git a/packages/adapters/codex-local/src/server/codex-auth-cache.ts b/packages/adapters/codex-local/src/server/codex-auth-cache.ts index c51921fff6b..d6dd757313a 100644 --- a/packages/adapters/codex-local/src/server/codex-auth-cache.ts +++ b/packages/adapters/codex-local/src/server/codex-auth-cache.ts @@ -1,4 +1,5 @@ -import { lstat, mkdir, open, readFile, rename, rm } from "node:fs/promises"; +import { AsyncLocalStorage } from "node:async_hooks"; +import { lstat, mkdir, open, readFile, realpath, rename, rm } from "node:fs/promises"; import path from "node:path"; import { randomUUID } from "node:crypto"; import { resolvePaperclipInstanceRootForAdapter } from "@paperclipai/adapter-utils/server-utils"; @@ -334,6 +335,11 @@ export async function withCodexAccountHomePromotionLock( * create) can still call into a secrets-service write that takes this lock * without deadlocking on itself. */ +// An outer credential transaction acquires this lock before database row locks. +// Nested secret-service writes must share that ownership, not acquire it again. +// Expire ownership on return so detached work cannot inherit a released lock. +const secretMutationOwnership = new AsyncLocalStorage>(); + export async function withAccountHomeSecretMutationLock( env: NodeJS.ProcessEnv = process.env, companyId: string, @@ -341,7 +347,16 @@ export async function withAccountHomeSecretMutationLock( ): Promise { const lockDir = resolveCodexAuthCacheNamedLockDir(env, companyId, "codex-account-home-secret-mutation-lock"); await ensurePrivateDir(lockDir); - return withDirectoryMergeLock(lockDir, fn, env); + const key = await realpath(lockDir); + const inherited = secretMutationOwnership.getStore(); + if (inherited?.get(key)?.active) return fn(); + return withDirectoryMergeLock(lockDir, async () => { + const ownership = { active: true }; + const scope = new Map(inherited); + scope.set(key, ownership); + try { return await secretMutationOwnership.run(scope, fn); } + finally { ownership.active = false; } + }, env); } /** diff --git a/packages/adapters/codex-local/src/server/quota-spawn-error.test.ts b/packages/adapters/codex-local/src/server/quota-spawn-error.test.ts index a660dae0a6c..76621357f23 100644 --- a/packages/adapters/codex-local/src/server/quota-spawn-error.test.ts +++ b/packages/adapters/codex-local/src/server/quota-spawn-error.test.ts @@ -17,7 +17,7 @@ vi.mock("node:child_process", async (importOriginal) => { }; }); -import { fetchCodexQuota, getQuotaWindows } from "./quota.js"; +import { fetchCodexQuota, fetchCodexRpcQuota, getQuotaWindows } from "./quota.js"; function createChildThatErrorsOnMicrotask(err: Error): ChildProcess { const child = new EventEmitter() as ChildProcess; @@ -27,7 +27,7 @@ function createChildThatErrorsOnMicrotask(err: Error): ChildProcess { Object.assign(child, { stdout: stream, stderr: Object.assign(new EventEmitter(), { setEncoding: () => {} }), - stdin: { write: vi.fn(), end: vi.fn() }, + stdin: Object.assign(new EventEmitter(), { write: vi.fn(), end: vi.fn() }), kill: vi.fn(), }); queueMicrotask(() => { @@ -36,6 +36,31 @@ function createChildThatErrorsOnMicrotask(err: Error): ChildProcess { return child; } +type RpcRequest = { id?: number; method: string }; + +function createRpcChild(response?: (request: RpcRequest) => unknown) { + const child = new EventEmitter() as ChildProcess; + const stdout = Object.assign(new EventEmitter(), { setEncoding: () => {} }); + const stderr = Object.assign(new EventEmitter(), { setEncoding: () => {} }); + const stdin = Object.assign(new EventEmitter(), { + write: vi.fn((line: string) => { + const request = JSON.parse(line) as RpcRequest; + if (request.id == null) return; + const message = response ? response(request) : { id: request.id, result: {} }; + if (message !== undefined) { + queueMicrotask(() => stdout.emit("data", JSON.stringify(message) + "\n")); + } + }), + end: vi.fn(), + }); + const kill = vi.fn((_signal?: string | number) => { + queueMicrotask(() => child.emit("close", 0)); + return true; + }); + Object.assign(child, { pid: 12345, stdout, stderr, stdin, kill }); + return { child, stdout, stderr, stdin, kill }; +} + describe("CodexRpcClient spawn failures", () => { let previousCodexHome: string | undefined; let isolatedCodexHome: string | undefined; @@ -51,6 +76,7 @@ describe("CodexRpcClient spawn failures", () => { }); afterEach(() => { + vi.useRealTimers(); vi.unstubAllGlobals(); if (isolatedCodexHome) { try { @@ -67,6 +93,159 @@ describe("CodexRpcClient spawn failures", () => { } }); + it("reads account quota with a supported approval policy and never starts a model turn", async () => { + const { child, stdin, kill } = createRpcChild((request) => { + const result = request.method === "account/rateLimits/read" + ? { rateLimits: { primary: { usedPercent: 37, windowDurationMins: 300 } } } + : {}; + return { id: request.id, result }; + }); + mockSpawn.mockImplementation((_command, args: string[]) => { + if (args.includes("untrusted")) return createChildThatErrorsOnMicrotask(new Error("invalid approval policy")); + return child; + }); + + const result = await fetchCodexRpcQuota(); + expect(result.windows).toEqual([expect.objectContaining({ usedPercent: 37 })]); + expect(mockSpawn).toHaveBeenCalledWith("codex", ["-s", "read-only", "-a", "on-request", "app-server"], expect.any(Object)); + expect(stdin.write.mock.calls.map(([line]) => JSON.parse(line).method)).toEqual([ + "initialize", "initialized", "account/rateLimits/read", "account/read", + ]); + expect(stdin.end).toHaveBeenCalledOnce(); + expect(kill).toHaveBeenCalledWith("SIGTERM"); + expect(kill).not.toHaveBeenCalledWith("SIGKILL"); + }); + + it("rejects RPC error envelopes and preserves their authentication classification", async () => { + const { child, kill } = createRpcChild((request) => request.method === "account/rateLimits/read" + ? { id: request.id, error: { code: -32000, message: "OAuth failed: refresh token has expired" } } + : { id: request.id, result: {} }); + mockSpawn.mockReturnValue(child); + + const result = await getQuotaWindows(); + + expect(result).toMatchObject({ ok: false, source: "codex-rpc", errorFamily: "refresh_token_expired", windows: [] }); + expect(kill).toHaveBeenCalledWith("SIGTERM"); + }); + + it.each([null, [], "invalid", 17, undefined])("rejects an invalid result envelope (%s)", async (result) => { + const { child } = createRpcChild((request) => ({ id: request.id, result })); + mockSpawn.mockReturnValue(child); + await expect(fetchCodexRpcQuota()).rejects.toThrow("invalid quota response"); + }); + + it("tolerates non-object lines and parses fragmented, out-of-order replies", async () => { + const probe = createRpcChild(); + mockSpawn.mockReturnValue(probe.child); + probe.stdin.write.mockImplementation((line) => { + const request = JSON.parse(line) as RpcRequest; + if (request.method === "initialize") { + queueMicrotask(() => probe.stdout.emit("data", 'null\n[]\nnot-json\n{"id":1,"result":{}}\n')); + } else if (request.method === "account/read") { + queueMicrotask(() => { + probe.stdout.emit("data", '{"id":3,"result":{"account":{"email":"test@example.com"}}}\n{"id":2,'); + probe.stdout.emit("data", '"result":{"rateLimits":{"primary":{"usedPercent":37}}}}\n'); + }); + } + }); + const result = await fetchCodexRpcQuota(); + expect(result.email).toBe("test@example.com"); + expect(result.windows[0].usedPercent).toBe(37); + }); + + it.each([false, true])("bounds stdout buffering with and without a completed frame (newline: %s)", async (newline) => { + const probe = createRpcChild(() => undefined); + mockSpawn.mockReturnValue(probe.child); + const result = fetchCodexRpcQuota(); + probe.stdout.emit("data", "x".repeat(64 * 1024)); + probe.stdout.emit("data", "x" + (newline ? "\n" : "")); + await expect(result).rejects.toThrow("response exceeded the quota probe limit"); + expect(probe.kill).toHaveBeenCalledWith("SIGTERM"); + }); + + it("bounds stderr diagnostics when the process exits", async () => { + const probe = createRpcChild(() => undefined); + mockSpawn.mockReturnValue(probe.child); + const result = fetchCodexRpcQuota().catch((error: Error) => error); + probe.stderr.emit("data", "x".repeat(128 * 1024)); + probe.stderr.emit("data", " end of diagnostic"); + probe.child.emit("exit", 1); + const error = await result as Error; + expect(error.message.length).toBeLessThanOrEqual(4_000); + expect(error.message).toMatch(/end of diagnostic$/); + }); + + it("handles a broken stdin without an uncaught stream error", async () => { + const probe = createRpcChild(() => undefined); + mockSpawn.mockReturnValue(probe.child); + const result = fetchCodexRpcQuota(); + probe.stdin.emit("error", new Error("write EPIPE")); + await expect(result).rejects.toThrow("write EPIPE"); + expect(probe.kill).toHaveBeenCalledWith("SIGTERM"); + }); + + it("clears other pending requests when a probe times out and kills a child that ignores SIGTERM", async () => { + vi.useFakeTimers(); + const probe = createRpcChild((request) => request.method === "initialize" + ? { id: request.id, result: {} } + : undefined); + probe.kill.mockImplementation((signal) => { + if (signal === "SIGKILL") queueMicrotask(() => probe.child.emit("close", null)); + return true; + }); + mockSpawn.mockReturnValue(probe.child); + const result = fetchCodexRpcQuota(); + const rejected = expect(result).rejects.toThrow("timed out on account/rateLimits/read"); + await vi.advanceTimersByTimeAsync(6_000); + expect(probe.stdin.end).toHaveBeenCalledOnce(); + expect(probe.kill.mock.calls).toEqual([["SIGTERM"]]); + await vi.advanceTimersByTimeAsync(250); + await rejected; + expect(probe.kill.mock.calls).toEqual([["SIGTERM"], ["SIGKILL"]]); + expect(vi.getTimerCount()).toBe(0); + expect(probe.child.listenerCount("close")).toBe(1); + }); + + it("bounds shutdown even if a killed child never emits close", async () => { + vi.useFakeTimers(); + const probe = createRpcChild(); + probe.kill.mockImplementation(() => true); + mockSpawn.mockReturnValue(probe.child); + const result = fetchCodexRpcQuota(); + await vi.advanceTimersByTimeAsync(500); + await expect(result).resolves.toMatchObject({ windows: [] }); + expect(probe.kill.mock.calls).toEqual([["SIGTERM"], ["SIGKILL"]]); + expect(vi.getTimerCount()).toBe(0); + expect(probe.child.listenerCount("close")).toBe(1); + }); + + it("reaps a real subprocess that ignores SIGTERM", async () => { + const { spawn: spawnReal } = await vi.importActual("node:child_process"); + const child = spawnReal(process.execPath, ["--input-type=commonjs", "-e", ` + process.on("SIGTERM", () => {}); + setInterval(() => {}, 60_000); + require("node:readline").createInterface({ input: process.stdin }).on("line", (line) => { + const request = JSON.parse(line); + if (request.id == null) return; + const result = request.method === "account/rateLimits/read" + ? { rateLimits: { primary: { usedPercent: 37 } } } + : {}; + process.stdout.write(JSON.stringify({ id: request.id, result }) + "\\n"); + }); + `], { stdio: ["pipe", "pipe", "pipe"] }); + mockSpawn.mockReturnValue(child); + try { + const result = await fetchCodexRpcQuota(); + expect(result.windows[0].usedPercent).toBe(37); + expect(child.signalCode).toBe("SIGKILL"); + expect(child.stdin.destroyed).toBe(true); + expect(child.stdout.destroyed).toBe(true); + expect(child.stderr.destroyed).toBe(true); + } finally { + if (child.exitCode === null && child.signalCode === null) child.kill("SIGKILL"); + } + }); + it("classifies app-server refresh-token failures as quota probe auth errors", async () => { mockSpawn.mockImplementation(() => createChildThatErrorsOnMicrotask(new Error("OAuth failed: refresh token has expired"))); @@ -110,7 +289,7 @@ describe("CodexRpcClient spawn failures", () => { expect(result.windows).toEqual([ expect.objectContaining({ label: "5h limit", - usedPercent: 50, + usedPercent: 0.5, resetsAt: "2024-03-22T12:38:31.000Z", }), ]); diff --git a/packages/adapters/codex-local/src/server/quota.ts b/packages/adapters/codex-local/src/server/quota.ts index 354a5a7d9b8..cb2407057ff 100644 --- a/packages/adapters/codex-local/src/server/quota.ts +++ b/packages/adapters/codex-local/src/server/quota.ts @@ -11,6 +11,8 @@ import { const CODEX_USAGE_SOURCE_RPC = "codex-rpc"; const CODEX_USAGE_SOURCE_WHAM = "codex-wham"; const MAX_QUOTA_ERROR_BODY_BYTES = 4_000; +const MAX_RPC_FRAME_BYTES = 64 * 1024; +const RPC_SHUTDOWN_GRACE_MS = 250; export function codexHomeDir(): string { const fromEnv = process.env.CODEX_HOME; @@ -227,7 +229,7 @@ export async function fetchWithTimeout( const controller = new AbortController(); const timer = setTimeout(() => controller.abort(), ms); try { - return await fetch(url, { ...init, signal: controller.signal }); + return await fetch(url, { ...init, signal: init.signal ? AbortSignal.any([init.signal, controller.signal]) : controller.signal }); } finally { clearTimeout(timer); } @@ -272,20 +274,22 @@ async function readResponseTextPrefix( } function normalizeCodexUsedPercent(rawPct: number | null | undefined): number | null { - if (rawPct == null) return null; - return Math.min(100, Math.round(rawPct < 1 ? rawPct * 100 : rawPct)); + if (typeof rawPct !== "number" || !Number.isFinite(rawPct)) return null; + // Both RPC and WHAM report percentage points, including values below 1%. + return Math.max(0, Math.min(100, rawPct)); } export async function fetchCodexQuota( token: string, accountId: string | null, + signal?: AbortSignal, ): Promise { const headers: Record = { Authorization: `Bearer ${token}`, }; if (accountId) headers["ChatGPT-Account-Id"] = accountId; - const resp = await fetchWithTimeout("https://chatgpt.com/backend-api/wham/usage", { headers }); + const resp = await fetchWithTimeout("https://chatgpt.com/backend-api/wham/usage", { headers, signal }); if (!resp.ok) { const message = `chatgpt wham api returned ${resp.status}`; const responseText = await readResponseTextPrefix(resp); @@ -470,7 +474,9 @@ type PendingRequest = { class CodexRpcClient { private proc = spawn( "codex", - ["-s", "read-only", "-a", "untrusted", "app-server"], + // Only account-read RPCs are sent; no agent turn is started. Keep the + // read-only sandbox and use the approval policy supported by current CLIs. + ["-s", "read-only", "-a", "on-request", "app-server"], { stdio: ["pipe", "pipe", "pipe"], env: process.env }, ); @@ -478,69 +484,101 @@ class CodexRpcClient { private buffer = ""; private pending = new Map(); private stderr = ""; + private failure: Error | null = null; + private closed = false; constructor() { this.proc.stdout.setEncoding("utf8"); this.proc.stderr.setEncoding("utf8"); this.proc.stdout.on("data", (chunk: string) => this.onStdout(chunk)); this.proc.stderr.on("data", (chunk: string) => { - this.stderr += chunk; + this.stderr = (this.stderr + chunk).slice(-MAX_QUOTA_ERROR_BODY_BYTES); }); this.proc.on("exit", () => { - for (const request of this.pending.values()) { - clearTimeout(request.timer); - request.reject(new Error(this.stderr.trim() || "codex app-server closed unexpectedly")); - } - this.pending.clear(); + this.fail(new Error(this.stderr.trim() || "codex app-server closed unexpectedly")); }); - this.proc.on("error", (err: Error) => { - for (const request of this.pending.values()) { - clearTimeout(request.timer); - request.reject(err); - } - this.pending.clear(); + this.proc.on("close", () => { + this.closed = true; }); + this.proc.on("error", (err: Error) => this.fail(err)); + // A child may close its input between writing a request and exiting. + this.proc.stdin.on("error", (err: Error) => this.fail(err)); + } + + private fail(error: Error) { + this.failure ??= error; + this.buffer = ""; + for (const request of this.pending.values()) { + clearTimeout(request.timer); + request.reject(this.failure); + } + this.pending.clear(); } private onStdout(chunk: string) { + if (this.failure) return; this.buffer += chunk; while (true) { const newlineIndex = this.buffer.indexOf("\n"); + const frame = newlineIndex < 0 ? this.buffer : this.buffer.slice(0, newlineIndex); + if (Buffer.byteLength(frame, "utf8") > MAX_RPC_FRAME_BYTES) { + this.fail(new Error("codex app-server response exceeded the quota probe limit")); + return; + } if (newlineIndex < 0) break; const line = this.buffer.slice(0, newlineIndex).trim(); this.buffer = this.buffer.slice(newlineIndex + 1); if (!line) continue; - let parsed: Record; + let value: unknown; try { - parsed = JSON.parse(line) as Record; + value = JSON.parse(line); } catch { continue; } + if (!value || typeof value !== "object" || Array.isArray(value)) continue; + const parsed = value as Record; const id = typeof parsed.id === "number" ? parsed.id : null; if (id == null) continue; const pending = this.pending.get(id); if (!pending) continue; this.pending.delete(id); clearTimeout(pending.timer); - pending.resolve(parsed); + if (parsed.error != null) { + const error = parsed.error as { message?: unknown }; + pending.reject(new Error(typeof error.message === "string" + ? error.message.slice(0, MAX_QUOTA_ERROR_BODY_BYTES) + : "codex app-server rejected the quota request")); + } else if (!parsed.result || typeof parsed.result !== "object" || Array.isArray(parsed.result)) { + pending.reject(new Error("codex app-server returned an invalid quota response")); + } else { + pending.resolve(parsed); + } } } private request(method: string, params: Record = {}, timeoutMs = 6_000): Promise> { + if (this.failure) return Promise.reject(this.failure); const id = this.nextId++; const payload = JSON.stringify({ id, method, params }) + "\n"; return new Promise>((resolve, reject) => { const timer = setTimeout(() => { - this.pending.delete(id); - reject(new Error(`codex app-server timed out on ${method}`)); + this.fail(new Error(`codex app-server timed out on ${method}`)); }, timeoutMs); this.pending.set(id, { resolve, reject, timer }); - this.proc.stdin.write(payload); + this.write(payload); }); } + private write(payload: string) { + try { + this.proc.stdin.write(payload); + } catch (error) { + this.fail(error instanceof Error ? error : new Error("codex app-server input failed")); + } + } + private notify(method: string, params: Record = {}) { - this.proc.stdin.write(JSON.stringify({ method, params }) + "\n"); + this.write(JSON.stringify({ method, params }) + "\n"); } async initialize() { @@ -568,7 +606,30 @@ class CodexRpcClient { } async shutdown() { + this.fail(new Error("codex app-server quota probe closed")); + this.proc.stdin.end(); + // A failed spawn has no process to signal. Never leave a live probe behind + // when it ignores SIGTERM, including after a request or transport failure. + if (this.closed || this.proc.pid == null) return; this.proc.kill("SIGTERM"); + await this.waitForClose(); + if (!this.closed) { + this.proc.kill("SIGKILL"); + await this.waitForClose(); + } + } + + private async waitForClose(): Promise { + if (this.closed) return; + await new Promise((resolve) => { + const done = () => { + clearTimeout(timer); + this.proc.removeListener("close", done); + resolve(); + }; + const timer = setTimeout(done, RPC_SHUTDOWN_GRACE_MS); + this.proc.once("close", done); + }); } } diff --git a/packages/shared/src/types/quota.ts b/packages/shared/src/types/quota.ts index 78ff1026c9e..a0686a9afb3 100644 --- a/packages/shared/src/types/quota.ts +++ b/packages/shared/src/types/quota.ts @@ -16,6 +16,10 @@ export interface QuotaWindow { export interface ProviderQuotaResult { /** provider slug, e.g. "anthropic", "openai" */ provider: string; + /** Connection, grant and credential revision; never a credential value. */ + accountKey?: string; + accountLabel?: string; + capturedAt?: string; /** source label when the provider reports where the quota data came from */ source?: string | null; /** true when the fetch succeeded and windows is populated */ diff --git a/server/src/__tests__/agent-hire-ai-connections.test.ts b/server/src/__tests__/agent-hire-ai-connections.test.ts index fb30faa12fe..c8fcab742b5 100644 --- a/server/src/__tests__/agent-hire-ai-connections.test.ts +++ b/server/src/__tests__/agent-hire-ai-connections.test.ts @@ -1,3 +1,5 @@ +import * as codexAdapter from "@paperclipai/adapter-codex-local/server"; +import { WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE } from "@paperclipai/adapter-utils/workspace-restore-merge"; import { randomUUID } from "node:crypto"; import { mkdtemp, rm } from "node:fs/promises"; import os from "node:os"; @@ -470,6 +472,38 @@ describe("agent-created hires use managed AI connections", () => { }); describe("hired agents sharing a subscription", () => { + it("keeps a credential-lock timeout on automatic retry without blocking the task or starting a provider", async () => { + const f = await fixture("openai", "subscription"); + const agent = hired(await request(f.app).post(`/api/companies/${f.companyId}/agent-hires`).send({ name: "Waiting teammate", role: "engineer", adapterType: f.adapterType, reportsTo: f.agentId, adapterConfig: { engine: "cli" }, runtimeConfig: { heartbeat: { enabled: false } } })); + // Only an operator may set host execution paths after the agent is hired. + await db.update(agents).set({ adapterConfig: { ...agent.adapterConfig, cwd: home } }).where(eq(agents.id, agent.id)); + const [issue] = await db.insert(issues).values({ companyId: f.companyId, title: "Wait for credential rotation", status: "todo", assigneeAgentId: agent.id, responsibleUserId: f.userId, createdByUserId: f.userId }).returning(); + const execute = vi.fn(async () => ({ exitCode: 0, signal: null, timedOut: false, resultJson: {} })); + registerServerAdapter({ ...getServerAdapter(f.adapterType), execute }); + const lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock") + .mockRejectedValueOnce(Object.assign(new Error("quota refresh holds company lock"), { code: WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE })); + const heartbeat = heartbeatService(db); + try { + const run = await heartbeat.invoke(agent.id, "assignment", { issueId: issue.id, wakeReason: "issue_assigned", responsibleUserId: f.userId }, "system"); + expect(run).not.toBeNull(); + await expect.poll(async () => (await heartbeat.getRun(run!.id))?.status, { timeout: 20_000 }).toBe("cancelled"); + await heartbeat.drainActiveRunExecutions(); + expect(await heartbeat.getRun(run!.id)).toMatchObject({ errorCode: "ai_connection_busy", resultJson: { executionRecovery: { providerWorkStarted: false } } }); + const retries = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.retryOfRunId, run!.id)); + expect(retries).toHaveLength(1); + expect(retries[0]).toMatchObject({ status: "scheduled_retry", scheduledRetryReason: "ai_connection_busy" }); + const [savedIssue] = await db.select().from(issues).where(eq(issues.id, issue.id)); + expect(savedIssue.status).not.toBe("blocked"); + expect(savedIssue.executionRunId).toBe(retries[0].id); + expect(execute).not.toHaveBeenCalled(); + expect(await aiConnectionService(db).list(f.companyId, f.userId)).toEqual([expect.objectContaining({ status: "connected" })]); + } finally { + lock.mockRestore(); + await heartbeat.drainActiveRunExecutions(); + unregisterServerAdapter(f.adapterType); + } + }); + it.each(["openai", "anthropic"] as const)("runs the %s child alongside a live parent and inherits its connection", async (provider) => { const f = await fixture(provider, "subscription"); const agent = hired(await request(f.app).post(`/api/companies/${f.companyId}/agent-hires`).send({ name: "Concurrent teammate", role: "engineer", adapterType: f.adapterType, reportsTo: f.agentId, adapterConfig: { engine: "cli" }, runtimeConfig: { heartbeat: { enabled: false } } })); diff --git a/server/src/__tests__/agent-test-environment-routes.test.ts b/server/src/__tests__/agent-test-environment-routes.test.ts index 7d5e4255053..2c2499e4f34 100644 --- a/server/src/__tests__/agent-test-environment-routes.test.ts +++ b/server/src/__tests__/agent-test-environment-routes.test.ts @@ -110,6 +110,10 @@ const mockPrepareManagedAiRuntime = vi.hoisted(() => vi.fn()); vi.mock("../services/ai-connection-runtime.js", async (importOriginal) => ({ ...(await importOriginal()), prepareManagedAiRuntime: mockPrepareManagedAiRuntime, + withManagedAiProbe: async (_db: unknown, input: unknown, probe: (runtime: unknown) => Promise) => { + const runtime = await mockPrepareManagedAiRuntime(_db, input); + try { return await probe(runtime); } finally { await runtime.cleanup(); } + }, })); const mockValidateAiApiKey = vi.hoisted(() => vi.fn(async () => undefined)); vi.mock("../routes/ai-connections.js", async (importOriginal) => ({ diff --git a/server/src/__tests__/ai-connections.test.ts b/server/src/__tests__/ai-connections.test.ts index 14550483b72..718344d09c7 100644 --- a/server/src/__tests__/ai-connections.test.ts +++ b/server/src/__tests__/ai-connections.test.ts @@ -7,15 +7,21 @@ import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; import { createHash, randomUUID } from "node:crypto"; import { spawnSync } from "node:child_process"; import { localEncryptedProvider } from "../secrets/local-encrypted-provider.js"; +import { promises as fs } from "node:fs"; import { mkdtemp, realpath, rm, access, readFile, writeFile, stat } from "node:fs/promises"; import os from "node:os"; import path from "node:path"; import { and, eq, sql } from "drizzle-orm"; -import { createDb, companies, agents, agentTaskSessions, heartbeatRuns, companyMemberships, connectionGrants, toolApplications, connectionGrantDelegations, connectionGrantMembers, toolConnections, toolConnectionInstalls, aiConnectionDefaults, aiProviderDefaults, adapterAuthSessions, environments, issues, issueThreadInteractions, issueRecoveryActions, connectionIntentDeliveries, agentWakeupRequests, companySecrets, principalPermissionGrants } from "@paperclipai/db"; +import { createDb, companies, agents, agentTaskSessions, heartbeatRuns, companyMemberships, connectionGrants, toolApplications, connectionGrantDelegations, connectionGrantMembers, toolConnections, toolConnectionInstalls, aiConnectionDefaults, aiProviderDefaults, adapterAuthSessions, environments, issues, issueThreadInteractions, issueRecoveryActions, connectionIntentDeliveries, agentWakeupRequests, companySecrets, companySecretVersions, userSecretDefinitions, principalPermissionGrants } from "@paperclipai/db"; import { startEmbeddedPostgresTestDatabase } from "@paperclipai/db/test-embedded-postgres"; +import { fetchCompanyQuotaWindows } from "../services/quota-windows.js"; +import { quotaCredentialRecoveryPath, quotaCredentialRecovery, quotaCredentialHash } from "../services/quota-credential-recovery.js"; import { aiConnectionService } from "../services/ai-connections.js"; +import { syncConnectionCredentialBindings } from "../services/connection-credential-bindings.js"; +import * as codexAdapter from "@paperclipai/adapter-codex-local/server"; +import { WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE } from "@paperclipai/adapter-utils/workspace-restore-merge"; import * as executionTarget from "@paperclipai/adapter-utils/execution-target"; -import { prepareManagedAiRuntime, assertManagedAiProjectAuth } from "../services/ai-connection-runtime.js"; +import { prepareManagedAiRuntime, withManagedAiProbe, assertManagedAiProjectAuth, isAiConnectionBusy } from "../services/ai-connection-runtime.js"; import { execute as executeGemini, testEnvironment as testGeminiEnvironment } from "@paperclipai/adapter-gemini-local/server"; import { toolAccessService } from "../services/tool-access.js"; import { secretService } from "../services/secrets.js"; @@ -54,6 +60,625 @@ beforeAll(async () => { afterAll(async () => { await database?.cleanup(); vi.unstubAllEnvs(); if (home) await rm(home, { recursive: true, force: true }); }); describe("managed AI connections", () => { + it("persists first-time recovery directories before consuming a single-use token", async () => { + const owner = "quota-durable-path"; + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const original = JSON.stringify({ tokens: { account_id: owner, access_token: "expired", refresh_token: "single-use" } }); + await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, original); + const [row] = await service.quotaAccounts(companyId, owner); + const file = quotaCredentialRecoveryPath(companyId, row.grant.id); + const recoveryRoot = path.dirname(path.dirname(file)); + await expect(access(recoveryRoot)).rejects.toMatchObject({ code: "ENOENT" }); + const synced: string[] = []; + let rejectRoot = true; + const realOpen = fs.open.bind(fs); + const open = vi.spyOn(fs, "open").mockImplementation(async (...args: Parameters) => { + const handle = await realOpen(...args); + const sync = handle.sync.bind(handle); + vi.spyOn(handle, "sync").mockImplementation(async () => { + if (rejectRoot && args[0] === recoveryRoot) throw new Error("Injected recovery directory flush failure"); + await sync(); synced.push(String(args[0])); + }); + return handle; + }); + const request = vi.spyOn(globalThis, "fetch").mockImplementation(async () => { + expect(synced).toContain(recoveryRoot); + expect(synced).toContain(path.dirname(recoveryRoot)); + return Response.json({ access_token: "fresh-durable", refresh_token: "replacement-durable" }); + }); + try { + await expect(service.refreshQuotaCredential(row, original, new AbortController().signal)).rejects.toThrow("Injected recovery directory flush failure"); + expect(request).not.toHaveBeenCalled(); + expect(await service.credential(row)).toBe(original); + rejectRoot = false; + await service.refreshQuotaCredential(row, original, new AbortController().signal); + expect(request).toHaveBeenCalledTimes(1); + expect(JSON.parse(await service.credential(row))).toMatchObject({ tokens: { access_token: "fresh-durable", refresh_token: "replacement-durable" } }); + } finally { open.mockRestore(); request.mockRestore(); } + }); + + it.each(["timeout", "unexpected"] as const)("classifies a manual probe's %s lock failure before starting the provider", async (kind) => { + const owner = `manual-probe-lock-${kind}`; + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const auth = JSON.stringify({ tokens: { account_id: owner, access_token: "valid-access", refresh_token: "valid-refresh" } }); + await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth); + const [row] = await service.quotaAccounts(companyId, owner); + const probe = vi.fn().mockResolvedValue("ok"); + const input = { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local", + binding: { provider: "openai", method: "subscription", mode: "responsible_user" } as const, config: {} }; + const failure = Object.assign(new Error("injected private lock diagnostic"), { code: kind === "timeout" ? WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE : "EACCES" }); + const lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock").mockRejectedValueOnce(failure); + try { + const error = await withManagedAiProbe(db, input, probe).then(() => { throw new Error("Probe must not start"); }, error => error); + if (kind === "timeout") { + expect(isAiConnectionBusy(error)).toBe(true); + expect(error).toMatchObject({ status: 422, message: "AI credentials are being updated. Retry shortly.", details: { code: "ai_connection_busy" } }); + } else expect(error).toBe(failure); + expect(probe).not.toHaveBeenCalled(); + expect((await db.select().from(connectionGrants).where(eq(connectionGrants.id, row.grant.id)))[0].status).toBe("active"); + } finally { lock.mockRestore(); } + expect(await withManagedAiProbe(db, input, probe)).toBe("ok"); + expect(probe).toHaveBeenCalledTimes(1); + expect(await service.credential(row)).toBe(auth); + }); + + it("keeps quota refresh outside a manual probe until its rotated credential is saved", async () => { + const owner = "manual-probe-race"; + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const auth = (hour: string) => JSON.stringify({ tokens: { account_id: owner, id_token: "identity", access_token: `access-${hour}`, refresh_token: `refresh-${hour}` }, last_refresh: `2026-09-10T${hour}:00:00Z` }); + await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth("10")); + const [row] = await service.quotaAccounts(companyId, owner); + const entered = deferredSignal(), release = deferredSignal(); + const request = vi.spyOn(globalThis, "fetch").mockRejectedValue(new Error("Quota must not exchange the probe's token")); + const probing = withManagedAiProbe(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local", + binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} }, async runtime => { + entered.resolve(); await release.promise; + await writeFile(path.join(String(runtime.config.env.CODEX_HOME), "auth.json"), auth("11")); + }); + let refreshing: Promise | undefined; + let lock: ReturnType | undefined; + try { + await entered.promise; + lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock"); + // Start outside the probe's async scope; it is a separate HTTP request. + refreshing = service.refreshQuotaCredential(row, auth("10"), new AbortController().signal); + await vi.waitFor(() => expect(lock).toHaveBeenCalled()); + expect(request).not.toHaveBeenCalled(); + release.resolve(); await probing; + expect(await refreshing).toBe(auth("11")); + expect(await service.credential(row)).toBe(auth("11")); + expect(request).not.toHaveBeenCalled(); + } finally { release.resolve(); await Promise.allSettled([probing, refreshing]); lock?.mockRestore(); request.mockRestore(); } + }); + + it("rejects preparation when reconnect replaces the selected credential reference", async () => { + const owner = "replaced-ref-owner"; + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const auth = (id: string) => JSON.stringify({ tokens: { account_id: id, access_token: id, refresh_token: id } }); + await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth("old-account")); + const [row] = await service.quotaAccounts(companyId, owner); + const definition = await secretService(db).createUserSecretDefinition(companyId, { key: "replacement_credential", name: "Replacement", provider: "local_encrypted" }, { userId: owner }); + const replacement = await secretService(db).createCurrentUserSecretValue(companyId, owner, { definitionId: definition.id, value: auth("new-account") }, { userId: owner }); + const prepare = () => prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local", + binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} }); + const original = codexAdapter.withAccountHomeSecretMutationLock; + const lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock").mockImplementationOnce(async (env, company, callback) => { + await db.update(connectionGrants).set({ credentialSecretRefs: row.grant.credentialSecretRefs.map(ref => ref.configPath === "ai.credential" ? { ...ref, secretId: replacement.id } : ref) }).where(eq(connectionGrants.id, row.grant.id)); + await syncConnectionCredentialBindings(db, row.connection); + return original(env, company, callback); + }); + try { await expect(prepare()).rejects.toThrow("AI credential changed during preparation"); } + finally { lock.mockRestore(); } + const runtime = await prepare(); + try { + expect(runtime.sessionIdentity).toContain(replacement.id); + expect(await readFile(path.join(String(runtime.config.env.CODEX_HOME), "auth.json"), "utf8")).toBe(auth("new-account")); + } finally { await runtime.cleanup(); } + }); + + it("caches quota after a real secret read but invalidates credential rotation and revocation", async () => { + const owner = "quota-cache-owner"; + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const account = await service.save(companyId, owner, { provider: "anthropic", method: "subscription", ownership: "personal", name: "Quota cache account", loginSessionId: "fixture", agentIds: [], allAgents: true }, "quota-cache-original"); + const [grant] = await db.select().from(connectionGrants).where(eq(connectionGrants.id, account.grantId)); + const secretId = grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!.secretId; + const [before] = await db.select().from(companySecrets).where(eq(companySecrets.id, secretId)); + const fetchSpy = vi.spyOn(globalThis, "fetch").mockImplementation(async () => new Response(JSON.stringify({ five_hour: { utilization: 42 } }))); + try { + const first = await fetchCompanyQuotaWindows(db, companyId, owner); + const [after] = await db.select().from(companySecrets).where(eq(companySecrets.id, secretId)); + expect(after.lastResolvedAt).not.toBeNull(); + expect(after.updatedAt.getTime()).toBeGreaterThan(before.updatedAt.getTime()); + expect(first[0].ok).toBe(true); + expect(await fetchCompanyQuotaWindows(db, companyId, owner)).toEqual(first); + expect(fetchSpy).toHaveBeenCalledTimes(1); + await secretService(db).rotate(secretId, { value: "quota-cache-rotated" }); + const rotated = await fetchCompanyQuotaWindows(db, companyId, owner); + expect(rotated[0].ok).toBe(true); + expect(rotated[0].accountKey).not.toBe(first[0].accountKey); + expect(fetchSpy).toHaveBeenCalledTimes(2); + expect(fetchSpy).toHaveBeenLastCalledWith(expect.any(String), expect.objectContaining({ headers: expect.objectContaining({ Authorization: "Bearer quota-cache-rotated" }) })); + await db.update(userSecretDefinitions).set({ status: "disabled" }).where(eq(userSecretDefinitions.id, before.userSecretDefinitionId!)); + const disabled = await fetchCompanyQuotaWindows(db, companyId, owner); + expect(disabled[0]).toMatchObject({ ok: false, errorFamily: "credentials_unavailable", windows: [] }); + expect(disabled[0].accountKey).not.toBe(rotated[0].accountKey); + expect(fetchSpy).toHaveBeenCalledTimes(2); + await db.update(userSecretDefinitions).set({ status: "active" }).where(eq(userSecretDefinitions.id, before.userSecretDefinitionId!)); + expect((await fetchCompanyQuotaWindows(db, companyId, owner))[0].ok).toBe(true); + await db.update(companySecretVersions).set({ status: "disabled", revokedAt: new Date() }).where(and(eq(companySecretVersions.secretId, secretId), eq(companySecretVersions.version, 2))); + const revoked = await fetchCompanyQuotaWindows(db, companyId, owner); + expect(revoked[0].ok).toBe(false); + expect(revoked[0].windows).toEqual([]); + expect(fetchSpy).toHaveBeenCalledTimes(2); + } 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.each(["rotation", "secret_metadata", "activity", "commit"] as const)("recovers exchanged OAuth tokens after %s persistence failure without another exchange", async (failure) => { + const owner = `quota-rollback-${failure}`; + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const original = JSON.stringify({ tokens: { access_token: `expired-${failure}`, refresh_token: "single-use", account_id: owner } }); + await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, original); + const [row] = await service.quotaAccounts(companyId, owner); + const ref = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!; + const table = failure === "secret_metadata" ? "company_secrets" : failure === "activity" ? "activity_log" : "company_secret_versions"; + const condition = failure === "secret_metadata" ? `new.id = '${ref.secretId}'::uuid and new.latest_version > old.latest_version` : failure === "activity" ? `new.action = 'ai_connection.credential_refreshed'` : `new.secret_id = '${ref.secretId}'::uuid`; + await db.execute(sql.raw(`create function quota_save_failure() returns trigger language plpgsql as $$ begin + if ${condition} then raise exception 'injected quota save failure'; end if; return new; end $$; + create ${failure === "commit" ? "constraint " : ""}trigger quota_save_failure ${failure === "commit" ? "after" : "before"} ${failure === "secret_metadata" ? "update" : "insert"} on ${table} + ${failure === "commit" ? "deferrable initially deferred" : ""} for each row execute function quota_save_failure();`)); + const request = vi.spyOn(globalThis, "fetch").mockResolvedValue(Response.json({ access_token: `fresh-${failure}`, refresh_token: `replacement-${failure}` })); + const file = quotaCredentialRecoveryPath(companyId, row.grant.id); + let triggerDropped = false; + try { + await expect(service.refreshQuotaCredential(row, original, new AbortController().signal)).rejects.toThrow(); + expect(request).toHaveBeenCalledTimes(1); + expect(await service.credential(row)).toBe(original); + const encrypted = await readFile(file, "utf8"); + expect(encrypted).not.toContain(`fresh-${failure}`); + expect(encrypted).not.toContain(`replacement-${failure}`); + await db.execute(sql.raw(`drop trigger quota_save_failure on ${table}; drop function quota_save_failure();`)); + triggerDropped = true; + // A new service instance recovers from disk, without an in-memory cache. + const recovered = aiConnectionService(db); + await recovered.refreshQuotaCredential(row, original, new AbortController().signal); + expect(JSON.parse(await recovered.credential(row))).toMatchObject({ tokens: { access_token: `fresh-${failure}`, refresh_token: `replacement-${failure}` } }); + expect(request).toHaveBeenCalledTimes(1); + await expect(access(file)).rejects.toThrow(); + expect((await db.select().from(companySecretVersions).where(eq(companySecretVersions.secretId, ref.secretId))).map(version => version.version).sort()).toEqual([1, 2]); + } finally { + request.mockRestore(); + if (!triggerDropped) await db.execute(sql.raw(`drop trigger quota_save_failure on ${table}; drop function quota_save_failure();`)); + } + }); + + it.each(["company_file", "grant_row", "secret_row"] as const)("defers runtime preparation on a contended %s lock and succeeds once it clears", async (kind) => { + const owner = `runtime-lock-${kind}`; + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const auth = JSON.stringify({ tokens: { access_token: "valid-access", refresh_token: "valid-refresh", account_id: owner } }); + await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth); + const [row] = await service.quotaAccounts(companyId, owner); + const ref = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!; + const prepare = () => prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local", + binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} }); + const held = deferredSignal(), release = deferredSignal(); + const lock = kind === "company_file" ? vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock") + .mockRejectedValueOnce(Object.assign(new Error("company lock held by another account"), { code: WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE })) : undefined; + const holding = kind === "company_file" ? Promise.resolve() : db.transaction(async tx => { + if (kind === "grant_row") await tx.select().from(connectionGrants).where(eq(connectionGrants.id, row.grant.id)).for("update"); + else await tx.select().from(companySecrets).where(eq(companySecrets.id, ref.secretId)).for("update"); + held.resolve(); await release.promise; + }); + try { + if (kind !== "company_file") await held.promise; + const error = await prepare().then(() => { throw new Error("Runtime must defer"); }, error => error); + expect(isAiConnectionBusy(error)).toBe(true); + expect(error).toMatchObject({ status: 422, details: { code: "ai_connection_busy" } }); + expect((await db.select().from(connectionGrants).where(eq(connectionGrants.id, row.grant.id)))[0].status).toBe("active"); + } finally { release.resolve(); await holding; lock?.mockRestore(); } + const run = await prepare(); + try { expect(await readFile(path.join(String(run.config.env.CODEX_HOME), "auth.json"), "utf8")).toBe(auth); } + finally { await run.cleanup(); } + }); + + it.each([ + ["runtime", "rotation"], ["runtime", "commit"], ["failed_run", "rotation"], ["failed_run", "commit"], + ] as const)("recovers pending quota credentials before %s can reuse or invalidate the consumed token after %s failure", async (consumer, failure) => { + const owner = `quota-consumer-${consumer}-${failure}`; + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const auth = (marker: string) => JSON.stringify({ tokens: { access_token: `${marker}-access`, refresh_token: `${marker}-refresh`, account_id: owner } }); + await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth("old")); + const [row] = await service.quotaAccounts(companyId, owner); + const prepare = () => prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local", + binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} }); + const oldRun = await prepare(); + const ref = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!; + const journal = quotaCredentialRecovery(companyId, row.grant.id); + await journal.write({ companyId, grantId: row.grant.id, connectionId: row.connection.id, secretId: ref.secretId, + baseHash: quotaCredentialHash(auth("old")), value: auth("new") }); + const fail = () => service.markAuthenticationFailed({ companyId, agentId, runStartedAt: new Date(), + attribution: { ...oldRun.attribution, identity: oldRun.identity } }); + await db.execute(sql.raw(`create function reject_consumer_recovery() returns trigger language plpgsql as $$ begin + if new.secret_id = '${ref.secretId}'::uuid then raise exception 'recovery unavailable'; end if; return new; end $$; + create ${failure === "commit" ? "constraint " : ""}trigger reject_consumer_recovery ${failure === "commit" ? "after" : "before"} insert on company_secret_versions + ${failure === "commit" ? "deferrable initially deferred" : ""} for each row execute function reject_consumer_recovery();`)); + try { + if (consumer === "runtime") { + const error = await prepare().then(() => { throw new Error("Runtime must defer"); }, error => error); + expect(isAiConnectionBusy(error)).toBe(true); + expect(error).toMatchObject({ status: 422, details: { code: "ai_connection_busy" } }); + } else await expect(fail()).rejects.toThrow(); + expect((await db.select().from(connectionGrants).where(eq(connectionGrants.id, row.grant.id)))[0].status).toBe("active"); + expect(await service.credential(row)).toBe(auth("old")); + expect(await journal.read()).not.toBeNull(); + } finally { + await db.execute(sql`drop trigger reject_consumer_recovery on company_secret_versions; drop function reject_consumer_recovery()`); + } + const request = vi.spyOn(globalThis, "fetch").mockRejectedValue(new Error("Recovery must not exchange again")); + let nextRun: Awaited> | undefined; + try { + if (consumer === "failed_run") await fail(); + nextRun = await prepare(); + expect(await readFile(path.join(String(nextRun.config.env.CODEX_HOME), "auth.json"), "utf8")).toBe(auth("new")); + expect(nextRun.sessionIdentity).toBe(oldRun.sessionIdentity); + await fail(); // A late failure of the old run cannot invalidate the recovered identity. + expect((await db.select().from(connectionGrants).where(eq(connectionGrants.id, row.grant.id)))[0].status).toBe("active"); + expect(await service.credential(row)).toBe(auth("new")); + expect(await journal.read()).toBeNull(); + expect(request).not.toHaveBeenCalled(); + expect(await db.select().from(companySecretVersions).where(eq(companySecretVersions.secretId, ref.secretId))).toHaveLength(2); + } finally { request.mockRestore(); await oldRun.cleanup(); await nextRun?.cleanup(); await journal.clear(); } + }); + + it.each(["rollback", "lost_commit_reply"] as const)("automatically saves one exchanged token across %s", async (failure) => { + const owner = `quota-once-${failure}`; + 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: owner } }); + await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, original); + const [row] = await service.quotaAccounts(companyId, owner); + const transaction = db.transaction.bind(db); + const interrupted = vi.spyOn(db, "transaction").mockImplementationOnce(async (work, config) => { + await transaction(async (tx) => { + await work(tx); + if (failure === "rollback") throw new Error("Injected rollback after writes"); + }, config); + throw new Error("Injected lost commit response"); + }); + const request = vi.spyOn(globalThis, "fetch").mockResolvedValue(Response.json({ access_token: "fresh-once", refresh_token: "replacement-once" })); + try { + await service.refreshQuotaCredential(row, original, new AbortController().signal); + expect(JSON.parse(await service.credential(row))).toMatchObject({ tokens: { access_token: "fresh-once", refresh_token: "replacement-once" } }); + expect(request).toHaveBeenCalledTimes(1); + await expect(access(quotaCredentialRecoveryPath(companyId, row.grant.id))).rejects.toThrow(); + } finally { interrupted.mockRestore(); request.mockRestore(); } + }); + + it("recovers the current credential's pending replacement before answering a delayed older poll", async () => { + const owner = "delayed-quota-poll"; + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const original = JSON.stringify({ tokens: { access_token: "A", refresh_token: "refresh-A", account_id: owner } }); + await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, original); + const [row] = await service.quotaAccounts(companyId, owner); + const ref = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!; + const journal = quotaCredentialRecovery(companyId, row.grant.id); + const request = vi.spyOn(globalThis, "fetch") + .mockResolvedValueOnce(Response.json({ access_token: "B", refresh_token: "refresh-B" })) + .mockResolvedValueOnce(Response.json({ access_token: "C", refresh_token: "refresh-C" })); + let triggerInstalled = false; + try { + const current = await service.refreshQuotaCredential(row, original, new AbortController().signal); + await db.execute(sql.raw(`create function delayed_quota_commit_failure() returns trigger language plpgsql as $$ begin + if new.secret_id = '${ref.secretId}'::uuid then raise exception 'injected delayed quota commit failure'; end if; return new; end $$; + create constraint trigger delayed_quota_commit_failure after insert on company_secret_versions deferrable initially deferred + for each row execute function delayed_quota_commit_failure();`)); + triggerInstalled = true; + await expect(service.refreshQuotaCredential(row, current, new AbortController().signal)).rejects.toThrow(); + expect(await service.credential(row)).toBe(current); + expect(JSON.parse((await journal.read())!.value).tokens.access_token).toBe("C"); + await db.execute(sql`drop trigger delayed_quota_commit_failure on company_secret_versions; drop function delayed_quota_commit_failure()`); + triggerInstalled = false; + const recovered = await service.refreshQuotaCredential(row, original, new AbortController().signal); + expect(JSON.parse(recovered).tokens).toMatchObject({ access_token: "C", refresh_token: "refresh-C" }); + expect(await service.credential(row)).toBe(recovered); + expect(await journal.read()).toBeNull(); + expect(request).toHaveBeenCalledTimes(2); + } finally { + if (triggerInstalled) await db.execute(sql`drop trigger delayed_quota_commit_failure on company_secret_versions; drop function delayed_quota_commit_failure()`); + request.mockRestore(); await journal.clear(); + } + }); + + it.each(["reconnected", "expired_reconnect", "revoked"] as const)("does not apply a pending refresh over a %s identity", async (change) => { + const owner = `quota-replay-${change}`; + 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: owner } }); + await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, original); + const [row] = await service.quotaAccounts(companyId, owner); + const ref = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!; + const journal = quotaCredentialRecovery(companyId, row.grant.id); + await journal.write({ companyId, grantId: row.grant.id, connectionId: row.connection.id, secretId: ref.secretId, baseHash: quotaCredentialHash(original), value: "old-pending-replacement" }); + const request = vi.spyOn(globalThis, "fetch"); + try { + if (change === "reconnected") { + await secretService(db).rotate(ref.secretId, { value: "new-authorized-identity" }); + expect(await service.refreshQuotaCredential(row, original, new AbortController().signal)).toBe("new-authorized-identity"); + expect(await service.credential(row)).toBe("new-authorized-identity"); + expect(await journal.read()).toBeNull(); + } else if (change === "expired_reconnect") { + const reconnected = JSON.stringify({ tokens: { access_token: "expired-new-account", refresh_token: "new-account-refresh", account_id: owner } }); + await secretService(db).rotate(ref.secretId, { value: reconnected }); + request.mockResolvedValueOnce(Response.json({ access_token: "new-account-fresh", refresh_token: "new-account-replacement" })); + await service.refreshQuotaCredential(row, reconnected, new AbortController().signal); + expect(request).toHaveBeenCalledOnce(); + expect(JSON.parse(String(request.mock.calls[0][1]?.body))).toMatchObject({ refresh_token: "new-account-refresh" }); + expect(JSON.parse(await service.credential(row))).toMatchObject({ tokens: { access_token: "new-account-fresh", refresh_token: "new-account-replacement" } }); + expect(await journal.read()).toBeNull(); + } else { + await db.update(connectionGrants).set({ status: "revoked" }).where(eq(connectionGrants.id, row.grant.id)); + await expect(service.refreshQuotaCredential(row, original, new AbortController().signal)).rejects.toThrow("credentials_unavailable"); + expect(await journal.read()).not.toBeNull(); + } + if (change !== "expired_reconnect") expect(request).not.toHaveBeenCalled(); + } finally { request.mockRestore(); await journal.clear(); } + }); + + 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("retries another account's token write-back while a slow quota exchange holds the company lock", async () => { + const auth = (account: string, marker: string) => JSON.stringify({ tokens: { account_id: account, id_token: `id-${marker}`, access_token: `access-${marker}`, refresh_token: `refresh-${marker}` }, last_refresh: marker === "old" ? "2026-09-10T10:00:00Z" : "2026-09-10T11:00:00Z" }); + const quotaOwner = "slow-quota-owner", runOwner = "slow-quota-other-run"; + await db.insert(companyMemberships).values([quotaOwner, runOwner].map(principalId => ({ companyId, principalId, principalType: "user" as const, status: "active" as const, membershipRole: "member" as const }))); + for (const owner of [quotaOwner, runOwner]) await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth(owner, "old")); + const [quotaRow] = await service.quotaAccounts(companyId, quotaOwner); + // Runtime credential reads now serialize with quota exchanges. Prepare this + // already-running, different account before testing its contended write-back. + const preparedRun = await prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: runOwner, adapterType: "codex_local", binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} }); + const started = deferredSignal(), release = deferredSignal(); + const request = vi.spyOn(globalThis, "fetch").mockImplementation(async () => { + started.resolve(); await release.promise; + return Response.json({ access_token: "quota-fresh", refresh_token: "quota-replacement" }); + }); + const refreshing = service.refreshQuotaCredential(quotaRow, auth(quotaOwner, "old"), new AbortController().signal); + let run: Awaited> | undefined; + let cleaning: Promise | undefined; + let lock: ReturnType | undefined; + try { + await started.promise; + run = preparedRun; + const authFile = path.join(String(run.config.env.CODEX_HOME), "auth.json"); + await writeFile(authFile, auth(runOwner, "replacement")); + // Inject the two 30-second acquisition expirations, then use the real + // still-held lock and real Postgres write-back for the final attempt. + lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock") + .mockRejectedValueOnce(Object.assign(new Error("lock timeout"), { code: WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE })) + .mockRejectedValueOnce(Object.assign(new Error("lock timeout"), { code: WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE })); + cleaning = run.cleanup(); + await vi.waitFor(() => expect(lock).toHaveBeenCalledTimes(3)); + expect(await readFile(authFile, "utf8")).toBe(auth(runOwner, "replacement")); + release.resolve(); + await Promise.all([refreshing, cleaning]); + const [runRow] = await service.quotaAccounts(companyId, runOwner); + expect(await service.credential(runRow)).toBe(auth(runOwner, "replacement")); + expect(JSON.parse(await service.credential(quotaRow))).toMatchObject({ tokens: { access_token: "quota-fresh", refresh_token: "quota-replacement" } }); + await expect(access(run.home!)).rejects.toMatchObject({ code: "ENOENT" }); + } finally { release.resolve(); await Promise.allSettled([refreshing, cleaning]); lock?.mockRestore(); request.mockRestore(); } + }); + + it.each(["new", "reconnect"])("retains a completed %s sign-in across two lock waits behind a quota refresh", async kind => { + const owner = `slow-quota-login-${kind}`, loginOwner = kind === "reconnect" ? owner : "contended-login-owner"; + await db.insert(companyMemberships).values([...new Set([owner, loginOwner])].map(principalId => ({ companyId, principalId, principalType: "user", status: "active", membershipRole: "member" }))); + const original = JSON.stringify({ tokens: { account_id: owner, access_token: "old", refresh_token: "old-refresh" } }); + await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, original); + const [quotaRow] = await service.quotaAccounts(companyId, owner); + const [environment] = await db.insert(environments).values({ name: "Contended AI login", driver: "sandbox" }).returning(); + const intent = { provider: "openai", method: "subscription", ownership: "personal", name: "Contended sign-in", agentIds: [], allAgents: true, ...(kind === "reconnect" ? { connectionId: quotaRow.connection.id } : {}) } as const; + const sessionId = randomUUID(); + await db.insert(adapterAuthSessions).values({ id: sessionId, publicSessionId: randomUUID(), companyId, environmentId: environment.id, adapterType: "codex_local", startedByUserId: loginOwner, status: "promoting", aiConnection: { ...intent, agentIds: [] }, expiresAt: new Date(Date.now() + 600_000) }); + const started = deferredSignal(), release = deferredSignal(); + const request = vi.spyOn(globalThis, "fetch").mockImplementation(async () => { + started.resolve(); await release.promise; + return Response.json({ access_token: "quota-fresh", refresh_token: "quota-replacement" }); + }); + const refreshing = service.refreshQuotaCredential(quotaRow, original, new AbortController().signal); + let saving: ReturnType | undefined; + let lock: ReturnType | undefined; + const credential = JSON.stringify({ tokens: { account_id: "signed-in", access_token: "login-token", refresh_token: "login-refresh" } }); + try { + await started.promise; + // Represent two expired 30-second acquisitions, then wait on the real + // held lock. The credential must survive until that final save commits. + lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock") + .mockRejectedValueOnce(Object.assign(new Error("lock timeout"), { code: WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE })) + .mockRejectedValueOnce(Object.assign(new Error("lock timeout"), { code: WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE })); + saving = service.save(companyId, loginOwner, { ...intent, agentIds: [] }, credential, sessionId); + void saving.catch(() => {}); + await vi.waitFor(() => expect(lock).toHaveBeenCalledTimes(3)); + expect((await db.select().from(adapterAuthSessions).where(eq(adapterAuthSessions.id, sessionId)))[0].connectionId).toBeNull(); + release.resolve(); + const [, saved] = await Promise.all([refreshing, saving]); + const [row] = (await service.quotaAccounts(companyId, loginOwner)).filter(row => row.connection.id === saved.connectionId); + expect(await service.credential(row)).toBe(credential); + expect(await service.save(companyId, loginOwner, { ...intent, agentIds: [] }, credential, sessionId)).toEqual(saved); + if (kind === "new") expect(JSON.parse(await service.credential(quotaRow))).toMatchObject({ tokens: { access_token: "quota-fresh", refresh_token: "quota-replacement" } }); + } finally { + release.resolve(); await Promise.allSettled([refreshing, saving]); lock?.mockRestore(); request.mockRestore(); + await db.delete(adapterAuthSessions).where(eq(adapterAuthSessions.id, sessionId)); + await db.delete(environments).where(eq(environments.id, environment.id)); + } + }); + + it.each(["quota", "recovery", "runtime"])("allows an in-flight reconnect after automatic %s credential rotation", async kind => { + const owner = `reconnect-refresh-${kind}`; + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const auth = (label: string, hour: number) => JSON.stringify({ tokens: { account_id: owner, id_token: "identity", access_token: `access-${label}`, refresh_token: `refresh-${label}` }, last_refresh: `2026-09-10T${hour}:00:00Z` }); + const intent = { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true } as const; + const account = await service.save(companyId, owner, { ...intent, agentIds: [] }, auth("old", 10)); + const [row] = await service.quotaAccounts(companyId, owner); + const ref = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!; + const before = (await db.select().from(companySecrets).where(eq(companySecrets.id, ref.secretId)))[0]; + const started = new Date(); + if (kind === "quota") { + const request = vi.spyOn(globalThis, "fetch").mockResolvedValue(Response.json({ access_token: "access-fresh", refresh_token: "refresh-fresh" })); + try { await service.refreshQuotaCredential(row, auth("old", 10), new AbortController().signal); } + finally { request.mockRestore(); } + } else if (kind === "recovery") { + const recovery = quotaCredentialRecovery(companyId, row.grant.id); + await recovery.write({ companyId, grantId: row.grant.id, connectionId: row.connection.id, secretId: ref.secretId, baseHash: quotaCredentialHash(auth("old", 10)), value: auth("fresh", 11) }); + expect(await service.runtimeCredential(row)).toBe(auth("fresh", 11)); + } else { + const runtime = await prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local", binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} }); + await writeFile(path.join(String(runtime.config.env.CODEX_HOME), "auth.json"), auth("fresh", 11)); + await runtime.cleanup(); + } + const after = (await db.select().from(companySecrets).where(eq(companySecrets.id, ref.secretId)))[0]; + expect(after.latestVersion).toBeGreaterThan(before.latestVersion); + expect(after.aiSessionEpoch).toBe(before.aiSessionEpoch); + const saved = await service.save(companyId, owner, { ...intent, agentIds: [], connectionId: account.connectionId }, auth("reconnected", 12), undefined, started); + expect(saved).toEqual(account); + expect(await service.credential(row)).toBe(auth("reconnected", 12)); + await expect(service.save(companyId, owner, { ...intent, agentIds: [], connectionId: account.connectionId }, auth("older-login", 13), undefined, started)).rejects.toThrow("changed"); + // A real later revocation still fences this same old reconnect attempt. + await db.update(connectionGrants).set({ status: "revoked", updatedAt: new Date() }).where(eq(connectionGrants.id, account.grantId)); + await expect(service.save(companyId, owner, { ...intent, agentIds: [], connectionId: account.connectionId }, auth("stale", 13), undefined, started)).rejects.toThrow("changed"); + }); + + it.each(["lock", "database"])("preserves replacement tokens after exhausted %s write-back and permits a later retry", async failure => { + const owner = `retained-writeback-${failure}`; + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const auth = (hour: string) => JSON.stringify({ tokens: { account_id: owner, id_token: "identity", access_token: `access-${hour}`, refresh_token: `refresh-${hour}` }, last_refresh: `2026-09-10T${hour}:00:00Z` }); + await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth("10")); + const run = await prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local", binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} }); + const authFile = path.join(String(run.config.env.CODEX_HOME), "auth.json"); + await writeFile(authFile, auth("11")); + const injected = failure === "lock" + ? vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock").mockRejectedValue(Object.assign(new Error("lock timeout"), { code: WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE })) + : vi.spyOn(db, "transaction").mockRejectedValue(new Error("database unavailable")); + try { + await expect(run.cleanup()).rejects.toThrow(failure === "lock" ? "lock timeout" : "database unavailable"); + expect(injected).toHaveBeenCalledTimes(failure === "lock" ? 3 : 1); + expect(await readFile(authFile, "utf8")).toBe(auth("11")); + } finally { injected.mockRestore(); } + await run.cleanup(); + const [row] = await service.quotaAccounts(companyId, owner); + expect(await service.credential(row)).toBe(auth("11")); + await expect(access(run.home!)).rejects.toMatchObject({ code: "ENOENT" }); + }); + + 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" }); + const original = JSON.stringify({ tokens: { access_token: "old-access", refresh_token: "single-use", id_token: "identity", account_id: "race-account" } }); + const account = await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: "Racing refresh", loginSessionId: "fixture", agentIds: [], allAgents: true }, original); + const [row] = await service.quotaAccounts(companyId, owner); + const secretId = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!.secretId; + const exchangeStarted = deferredSignal(); + const releaseExchange = deferredSignal(); + const request = vi.spyOn(globalThis, "fetch").mockImplementation(async () => { + exchangeStarted.resolve(); await releaseExchange.promise; + return Response.json({ access_token: "new-access", refresh_token: "new-refresh" }); + }); + let runtime: Awaited> | undefined; + let lock: ReturnType | undefined; + const refreshing = service.refreshQuotaCredential(row, original, new AbortController().signal); + try { + await exchangeStarted.promise; + const edit = secretService(db).update(secretId, { description: "Changed during refresh" }); + lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock"); + const preparing = prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local", binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} }); + // Runtime reads now wait at the company file lock before touching the + // old credential. The returned home must contain the committed rotation. + let prepared = false; + void preparing.then(() => { prepared = true; }); + await vi.waitFor(() => expect(lock).toHaveBeenCalled()); + expect(prepared).toBe(false); + releaseExchange.resolve(); + const results = await Promise.all([refreshing, edit, preparing]); + runtime = results[2]; + expect(JSON.parse(await readFile(path.join(String(runtime.config.env.CODEX_HOME), "auth.json"), "utf8"))).toMatchObject({ tokens: { access_token: "new-access", refresh_token: "new-refresh" } }); + expect(JSON.parse(await service.credential(row))).toMatchObject({ tokens: { access_token: "new-access", refresh_token: "new-refresh" } }); + expect((await secretService(db).getById(secretId))?.description).toBe("Changed during refresh"); + } finally { releaseExchange.resolve(); await runtime?.cleanup(); request.mockRestore(); lock?.mockRestore(); } + }, 10_000); + + 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 { + // A pre-existing transaction may own a secret row before asking for the + // company file lock. Refresh must defer before exchanging its token. + await db.transaction(async tx => { + const secretId = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!.secretId; + await tx.select().from(companySecrets).where(eq(companySecrets.id, secretId)).for("update"); + await expect(service.refreshQuotaCredential(row, original, signal)).rejects.toThrow(); + expect(request).not.toHaveBeenCalled(); + }); + 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([ ["openai", "codex_local", "responses"], ["anthropic", "claude_local", "messages"], @@ -380,6 +1005,17 @@ console.log(JSON.stringify({ type: "result", subtype: "success", result: "hello" expect(await intents.requestForRunAuthFailure(runId)).toBeNull(); }); + it("lists quotas only for the current member's accessible subscription accounts", async () => { + const user = `quota-${randomUUID()}`; + await db.insert(companyMemberships).values({ companyId, principalId: user, principalType: "user", status: "active", membershipRole: "member" }); + const subscription = await service.save(companyId, user, { provider: "openai", method: "subscription", ownership: "personal", name: "Private quota", loginSessionId: "fixture", allAgents: true, agentIds: [] }, JSON.stringify({ tokens: { access_token: "fixture", account_id: "fixture", refresh_token: "fixture" } })); + expect((await service.quotaAccounts(companyId, user)).some(row => row.connection.id === subscription.connectionId)).toBe(true); + expect((await service.quotaAccounts(companyId, "bob")).some(row => row.connection.id === subscription.connectionId)).toBe(false); + expect(await service.quotaAccounts(otherCompanyId, user)).toEqual([]); + await db.update(companyMemberships).set({ status: "inactive" }).where(and(eq(companyMemberships.companyId, companyId), eq(companyMemberships.principalId, user))); + expect(await service.quotaAccounts(companyId, user)).toEqual([]); + }); + it("repairs a teammate's missing onboarding key with their own inline AI connection", async () => { const userId = `new-teammate-${randomUUID()}`; const id = randomUUID(); @@ -1273,3 +1909,9 @@ describe("AI connection recovery delivery", () => { }, 30000, ); }); + +function deferredSignal() { + let resolve!: () => void; + const promise = new Promise(done => { resolve = done; }); + return { promise, resolve }; +} diff --git a/server/src/__tests__/company-quota-windows.test.ts b/server/src/__tests__/company-quota-windows.test.ts new file mode 100644 index 00000000000..646f821657d --- /dev/null +++ b/server/src/__tests__/company-quota-windows.test.ts @@ -0,0 +1,256 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +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(), +})); +vi.mock("../services/ai-connections.js", () => ({ + aiConnectionService: () => ({ + quotaAccounts: mocks.accounts, + credential: mocks.credential, + refreshQuotaCredential: mocks.refresh, + }), +})); +vi.mock("../adapters/registry.js", () => ({ + listServerAdapters: mocks.adapters, +})); +vi.mock("@paperclipai/adapter-codex-local/server", () => ({ + fetchCodexQuota: mocks.codex, +})); +vi.mock("@paperclipai/adapter-claude-local/server", () => ({ + fetchClaudeQuota: mocks.claude, +})); +import { fetchCompanyQuotaWindows } from "../services/quota-windows.js"; + +let revision = 1; +const db = { + select: () => ({ + from: () => ({ + leftJoin: () => ({ + leftJoin: () => ({ + where: async () => [ + { + id: "secret", + updatedAt: new Date(0), + latestVersion: revision, + status: "active", + }, + ], + }), + }), + }), + }), +} as unknown as Db; +const account = (id: string, provider = "openai", status = "connected") => ({ + connection: { id, updatedAt: new Date(0) }, + grant: { + id, + updatedAt: new Date(0), + credentialSecretRefs: [{ secretId: "secret" }], + }, + summary: { provider, name: id, status }, +}); +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" }, + }), + ); + mocks.codex.mockResolvedValue([ + { label: "5h", usedPercent: 0.5, resetsAt: null, valueLabel: null }, + ]); + }); + afterEach(() => vi.useRealTimers()); + it("uses the selected account without invoking a local CLI or fallback login", async () => { + mocks.accounts.mockResolvedValue([account("remote-codex")]); + const result = await fetchCompanyQuotaWindows(db, "company-1", "user-1"); + expect(mocks.accounts).toHaveBeenCalledWith("company-1", "user-1"); + expect(mocks.codex).toHaveBeenCalledWith("private-token", "actual-account", expect.any(AbortSignal)); + expect(mocks.adapters).not.toHaveBeenCalled(); + expect(result[0]).toMatchObject({ + ok: true, + accountLabel: "remote-codex", + source: "managed-connection", + }); + expect(result[0].capturedAt).toBeDefined(); + expect(JSON.stringify(result)).not.toContain("private-token"); + }); + it("deduplicates account requests and invalidates on credential revision", async () => { + mocks.accounts.mockResolvedValue([account("cached-codex")]); + const [a, b] = await Promise.all([ + fetchCompanyQuotaWindows(db, "company-2", "user-2"), + fetchCompanyQuotaWindows(db, "company-2", "user-2"), + ]); + expect(mocks.codex).toHaveBeenCalledTimes(1); + expect(a[0].capturedAt).toBe(b[0].capturedAt); + revision++; + const next = await fetchCompanyQuotaWindows(db, "company-2", "user-2"); + expect(mocks.codex).toHaveBeenCalledTimes(2); + expect(next[0].accountKey).not.toBe(a[0].accountKey); + }); + it("keeps two accounts for the same provider separate", async () => { + mocks.accounts.mockResolvedValue([account("first"), account("second")]); + const result = await fetchCompanyQuotaWindows(db, "company-3", "user-3"); + expect(result).toHaveLength(2); + expect(result[0].accountKey).not.toBe(result[1].accountKey); + }); + it("does not expose failures and marks authentication failures for stale-window removal", async () => { + mocks.accounts.mockResolvedValue([account("expired")]); + mocks.codex.mockRejectedValue( + new Error("401 private-token raw provider diagnostics"), + ); + const [result] = await fetchCompanyQuotaWindows(db, "company-4", "user-4"); + expect(result.errorFamily).toBe("authentication_required"); + expect(result.windows).toEqual([]); + expect(JSON.stringify(result)).not.toContain("private-token"); + }); + it("does not read revoked credentials or use the server login for an unconnected remote account", async () => { + mocks.accounts.mockResolvedValue([account("revoked", "openai", "revoked")]); + const [result] = await fetchCompanyQuotaWindows(db, "company-5", "user-5"); + expect(result.errorFamily).toBe("credentials_unavailable"); + expect(mocks.credential).not.toHaveBeenCalled(); + mocks.accounts.mockResolvedValue([]); + const empty = await fetchCompanyQuotaWindows(db, "company-5", "user-5"); + expect( + empty.every((row) => row.errorFamily === "credentials_unavailable"), + ).toBe(true); + expect(mocks.adapters).not.toHaveBeenCalled(); + }); + it("uses a managed Claude OAuth credential and keeps failed reads private", async () => { + mocks.accounts.mockResolvedValue([account("claude-account", "anthropic")]); + mocks.credential.mockResolvedValue("claude-private-oauth"); + mocks.claude.mockResolvedValue([ + { label: "Weekly", usedPercent: 12, resetsAt: null, valueLabel: null }, + ]); + const [result] = await fetchCompanyQuotaWindows(db, "company-6", "user-6"); + expect(result.ok).toBe(true); + expect(mocks.claude).toHaveBeenCalledWith("claude-private-oauth", expect.any(AbortSignal)); + mocks.accounts.mockResolvedValue([account("secret-failure")]); + mocks.credential.mockRejectedValue( + new Error("private credential store failure"), + ); + const [failed] = await fetchCompanyQuotaWindows(db, "company-6", "user-6"); + expect(failed.errorFamily).toBe("provider_unavailable"); + expect(JSON.stringify(failed)).not.toContain("private credential"); + }); + 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.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.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")]); + mocks.credential.mockReturnValue(new Promise(() => {})); + const pending = fetchCompanyQuotaWindows(db, "company-8", "user-8"); + await vi.advanceTimersByTimeAsync(20_001); + const [result] = await pending; + expect(result).toMatchObject({ + ok: false, + accountLabel: "stalled", + windows: [], + }); + expect(result.accountKey).toBeTruthy(); + expect(result.error).not.toContain("timed out"); + }); + it("does not start a provider request when a credential resolves after its deadline", async () => { + vi.useFakeTimers(); + mocks.accounts.mockResolvedValue([account("late-credential")]); + let resolveCredential!: (value: string) => void; + mocks.credential.mockReturnValue(new Promise((resolve) => { resolveCredential = resolve; })); + const pending = fetchCompanyQuotaWindows(db, "late-company", "late-user"); + await vi.advanceTimersByTimeAsync(20_001); + expect((await pending)[0].ok).toBe(false); + resolveCredential(JSON.stringify({ accessToken: "private", accountId: "actual" })); + await vi.advanceTimersByTimeAsync(0); + expect(mocks.codex).not.toHaveBeenCalled(); + }); + + it("cancels an active provider read when the account deadline expires", async () => { + vi.useFakeTimers(); + mocks.accounts.mockResolvedValue([account("slow-provider")]); + let signal!: AbortSignal; + mocks.codex.mockImplementation((_token, _id, abortSignal) => { + signal = abortSignal; + return new Promise((_resolve, reject) => signal.addEventListener("abort", () => reject(signal.reason), { once: true })); + }); + const pending = fetchCompanyQuotaWindows(db, "slow-company", "slow-user"); + await vi.advanceTimersByTimeAsync(20_001); + expect((await pending)[0].ok).toBe(false); + expect(signal.aborted).toBe(true); + }); + + it("limits account probes to four concurrent requests and preserves account order", async () => { + mocks.accounts.mockResolvedValue( + Array.from({ length: 9 }, (_, i) => account(`parallel-${i}`)), + ); + mocks.credential.mockResolvedValue( + JSON.stringify({ accessToken: "private", accountId: "actual" }), + ); + let active = 0, + maximum = 0; + mocks.codex.mockImplementation(async () => { + active++; + maximum = Math.max(maximum, active); + await new Promise((resolve) => setTimeout(resolve, 5)); + active--; + return []; + }); + const results = await fetchCompanyQuotaWindows(db, "company-9", "user-9"); + expect(maximum).toBe(4); + expect(results.map((row) => row.accountLabel)).toEqual( + Array.from({ length: 9 }, (_, i) => `parallel-${i}`), + ); + }); +}); diff --git a/server/src/__tests__/durable-directory.test.ts b/server/src/__tests__/durable-directory.test.ts new file mode 100644 index 00000000000..e6da0592966 --- /dev/null +++ b/server/src/__tests__/durable-directory.test.ts @@ -0,0 +1,243 @@ +import { createRequire } from "node:module"; +import { spawnSync } from "node:child_process"; +import { promises as fs } from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { ensureDurableDirectory } from "../lib/durable-directory.js"; + +const roots: string[] = []; +afterEach(async () => { + vi.restoreAllMocks(); + await Promise.all(roots.splice(0).map(root => fs.rm(root, { recursive: true, force: true }))); +}); + +describe.skipIf(process.platform === "win32")("durable directory creation", () => { + async function fixture() { + const root = await fs.realpath(await fs.mkdtemp(path.join(os.tmpdir(), "paperclip-durable-directory-"))); + roots.push(root); + return { root, directory: path.join(root, "instance", "recovery", "company") }; + } + + function observeSyncs(failOn?: string) { + const synced: string[] = []; + const open = fs.open.bind(fs); + vi.spyOn(fs, "open").mockImplementation(async (...args: Parameters) => { + const handle = await open(...args); + const sync = handle.sync.bind(handle); + vi.spyOn(handle, "sync").mockImplementation(async () => { + if (args[0] === failOn) throw new Error("Injected directory fsync failure"); + await sync(); + synced.push(String(args[0])); + }); + return handle; + }); + return synced; + } + + it.each([false, true])("flushes symlink targets and the directories containing each link (chained: %s)", async (chained) => { + const { root } = await fixture(); + const targetParent = path.join(root, "storage"); + const target = path.join(targetParent, "new-home"); + const writtenParent = path.join(root, "configured"); + const links = path.join(root, "links"); + await fs.mkdir(target, { recursive: true }); + await fs.mkdir(writtenParent); + await fs.mkdir(links); + await fs.symlink("../storage/new-home", path.join(links, "middle")); + const alias = path.join(writtenParent, "home"); + await fs.symlink(chained ? "../links/middle" : target, alias); + const directory = path.join(alias, "recovery"); + const synced = observeSyncs(); + await ensureDurableDirectory(directory); + expect(synced).toContain(targetParent); + expect(synced).toContain(writtenParent); + if (chained) expect(synced).toContain(links); + expect(synced).toContain(path.join(target, "recovery")); + synced.length = 0; + await ensureDurableDirectory(directory); + expect(synced).toEqual([]); + }); + + it("resolves parent components after following a symlink target", async () => { + const { root } = await fixture(); + const storage = path.join(root, "storage"); + await fs.mkdir(path.join(storage, "nested"), { recursive: true }); + await fs.mkdir(path.join(storage, "home")); + await fs.symlink(path.join(storage, "nested"), path.join(root, "through")); + await fs.symlink("through/../home", path.join(root, "alias")); + const synced = observeSyncs(); + await ensureDurableDirectory(path.join(root, "alias", "recovery")); + expect(synced).toContain(storage); + expect(synced).toContain(path.join(storage, "home", "recovery")); + }); + + it("invalidates the proof when a symlink moves the same leaf under a new parent", async () => { + const { root } = await fixture(); + const original = path.join(root, "original"); + const parent = path.join(root, "new-parent"); + const alias = path.join(root, "alias"); + await fs.mkdir(original); + await fs.symlink(original, alias); + const directory = path.join(alias, "recovery"); + await ensureDurableDirectory(directory); + const inode = (await fs.stat(directory)).ino; + await fs.mkdir(parent); + await fs.rename(original, path.join(parent, "moved")); + await fs.unlink(alias); + await fs.symlink(path.join(parent, "moved"), alias); + expect((await fs.stat(directory)).ino).toBe(inode); + const synced = observeSyncs(); + await ensureDurableDirectory(directory); + expect(synced).toContain(parent); + }); + + it("retries failed real-parent flushes through a symlink", async () => { + const { root } = await fixture(); + const parent = path.join(root, "storage"); + await fs.mkdir(path.join(parent, "home"), { recursive: true }); + const alias = path.join(root, "alias"); + await fs.symlink(path.join(parent, "home"), alias); + const directory = path.join(alias, "recovery"); + observeSyncs(parent); + await expect(ensureDurableDirectory(directory)).rejects.toThrow("Injected directory fsync failure"); + await expect(ensureDurableDirectory(directory)).rejects.toThrow("Injected directory fsync failure"); + vi.restoreAllMocks(); + const synced = observeSyncs(); + await ensureDurableDirectory(directory); + expect(synced).toContain(parent); + }); + + it("flushes every created path entry and avoids ancestor fsyncs on subsequent writes", async () => { + const { root, directory } = await fixture(); + const synced = observeSyncs(); + await ensureDurableDirectory(directory); + expect(synced.slice(0, 4)).toEqual([directory, path.dirname(directory), path.join(root, "instance"), root]); + synced.length = 0; + await ensureDurableDirectory(directory); + expect(synced).toEqual([]); + await fs.rm(directory, { recursive: true }); + await ensureDurableDirectory(directory); + expect(synced.slice(0, 4)).toEqual([directory, path.dirname(directory), path.join(root, "instance"), root]); + }); + + it("fails closed for existing storage beneath an unflushable ancestor", async () => { + const { root, directory } = await fixture(); + const storage = path.join(root, "instance"); + await fs.mkdir(storage); + const open = fs.open.bind(fs); + vi.spyOn(fs, "open").mockImplementation(async (...args: Parameters) => { + if (args[0] === root) throw Object.assign(new Error("execute-only ancestor"), { code: "EACCES" }); + return open(...args); + }); + await expect(ensureDurableDirectory(directory)).rejects.toMatchObject({ code: "EACCES" }); + expect((await fs.stat(directory)).isDirectory()).toBe(true); + }); + + it("requires permission to flush a newly created entry's parent, including on retry", async () => { + const { root, directory } = await fixture(); + const open = fs.open.bind(fs); + vi.spyOn(fs, "open").mockImplementation(async (...args: Parameters) => { + if (args[0] === root) throw Object.assign(new Error("write-only creation parent"), { code: "EACCES" }); + return open(...args); + }); + await expect(ensureDurableDirectory(directory)).rejects.toMatchObject({ code: "EACCES" }); + // mkdir now succeeds without creating anything. Its first attempt still + // owes a parent flush, so a retry cannot silently acknowledge that path. + await expect(ensureDurableDirectory(directory)).rejects.toMatchObject({ code: "EACCES" }); + vi.restoreAllMocks(); + const synced = observeSyncs(); + await ensureDurableDirectory(directory); + expect(synced).toContain(root); + }); + + it.each([ + { sibling: false, symlink: false }, { sibling: true, symlink: false }, + { sibling: false, symlink: true }, { sibling: true, symlink: true }, + ])("retains a failed parent flush across fresh processes ($sibling, symlink: $symlink)", async ({ sibling, symlink }) => { + const { root, directory: original } = await fixture(); + let directory = original; + let parent = root; + if (symlink) { + parent = path.join(root, "storage"); + await fs.mkdir(path.join(parent, "home"), { recursive: true }); + await fs.symlink(path.join(parent, "home"), path.join(root, "alias")); + directory = path.join(root, "alias", "recovery", "company"); + } + const source = new URL("../lib/durable-directory.ts", import.meta.url).href; + const child = (target: string, fail: boolean) => spawnSync(process.execPath, ["--import", createRequire(import.meta.url).resolve("tsx"), "--input-type=module", "-e", ` + import { promises as fs } from "node:fs"; + import { ensureDurableDirectory } from ${JSON.stringify(source)}; + const [directory, parent, fail] = process.argv.slice(1); + const open = fs.open.bind(fs); + fs.open = async (...args) => { + if (fail === "true" && args[0] === parent) throw Object.assign(new Error("unflushed parent"), { code: "EACCES" }); + return open(...args); + }; + try { await ensureDurableDirectory(directory); } + catch (error) { console.error(error.code); process.exitCode = error.code === "EACCES" ? 23 : 1; } + `, target, parent, String(fail)], { encoding: "utf8", timeout: 10000 }); + const first = child(directory, true); + expect(first.stderr).toContain("EACCES"); + expect(first.status).toBe(23); + expect((await fs.stat(directory)).isDirectory()).toBe(true); + const retryTarget = sibling ? path.join(path.dirname(directory), "other-company") : directory; + const retry = child(retryTarget, true); + expect(retry.status, retry.stderr).toBe(23); + expect(retry.stderr).toContain("EACCES"); + const recovered = child(retryTarget, false); + expect(recovered.status, recovered.stderr).toBe(0); + }); + + it("retries all ancestors after a failed flush even though mkdir already succeeded", async () => { + const { root, directory } = await fixture(); + observeSyncs(root); + await expect(ensureDurableDirectory(directory)).rejects.toThrow("Injected directory fsync failure"); + expect((await fs.stat(directory)).isDirectory()).toBe(true); + vi.restoreAllMocks(); + const synced = observeSyncs(); + await ensureDurableDirectory(directory); + expect(synced).toContain(root); + expect(synced).toContain(path.dirname(root)); + }); + + it.each([false, true])("keeps a shared parent flush required for sibling writes (first call pending: %s)", async (concurrent) => { + const { root, directory } = await fixture(); + const sibling = path.join(path.dirname(directory), "other-company"); + const open = fs.open.bind(fs); + let releaseFirst!: () => void; + let reachedParent!: () => void; + const held = new Promise(resolve => { releaseFirst = resolve; }); + const reached = new Promise(resolve => { reachedParent = resolve; }); + let firstParent = true; + vi.spyOn(fs, "open").mockImplementation(async (...args: Parameters) => { + if (args[0] === root) { + if (firstParent) { + firstParent = false; + reachedParent(); + if (concurrent) await held; + } + throw Object.assign(new Error("unflushed shared parent"), { code: "EACCES" }); + } + return open(...args); + }); + const first = ensureDurableDirectory(directory).then(() => null, error => error); + try { + await reached; + if (!concurrent) expect(await first).toMatchObject({ code: "EACCES" }); + await expect(ensureDurableDirectory(sibling)).rejects.toMatchObject({ code: "EACCES" }); + } finally { + releaseFirst(); + await first; + } + vi.restoreAllMocks(); + const synced = observeSyncs(); + await ensureDurableDirectory(sibling); + expect(synced).toContain(root); + await ensureDurableDirectory(directory); + synced.length = 0; + await ensureDurableDirectory(sibling); + await ensureDurableDirectory(directory); + expect(synced).toEqual([]); + }); +}); diff --git a/server/src/__tests__/quota-windows-service.test.ts b/server/src/__tests__/quota-windows-service.test.ts index 94d068ad819..8e6dfcc382c 100644 --- a/server/src/__tests__/quota-windows-service.test.ts +++ b/server/src/__tests__/quota-windows-service.test.ts @@ -3,9 +3,11 @@ import { describe, expect, it, vi, beforeEach, afterEach } from "vitest"; vi.mock("../adapters/registry.js", () => ({ listServerAdapters: vi.fn(), })); +vi.mock("../middleware/logger.js", () => ({ logger: { warn: vi.fn() } })); import { listServerAdapters } from "../adapters/registry.js"; import { fetchAllQuotaWindows } from "../services/quota-windows.js"; +import { logger } from "../middleware/logger.js"; describe("fetchAllQuotaWindows", () => { beforeEach(() => { @@ -48,9 +50,47 @@ describe("fetchAllQuotaWindows", () => { { provider: "anthropic", ok: false, - error: "quota polling timed out after 20s", + error: "Subscription quota is currently unavailable. Check usage with your provider.", windows: [], }, ]); }); + + it("keeps command diagnostics out of API responses and redacts them in server logs", async () => { + const diagnostic = 'Command failed: sh -c probe --token fixture-private-token'; + vi.mocked(listServerAdapters).mockReturnValue([ + { type: "claude_local", getQuotaWindows: vi.fn().mockResolvedValue({ + provider: "anthropic", ok: false, windows: [], error: diagnostic, + }) }, + { type: "codex_local", getQuotaWindows: vi.fn().mockRejectedValue(new Error(diagnostic)) }, + ] as never); + + const results = await fetchAllQuotaWindows(); + expect(results).toHaveLength(2); + for (const result of results) { + expect(result.ok).toBe(false); + expect(result.error).toBe("Subscription quota is currently unavailable. Check usage with your provider."); + } + expect(JSON.stringify(results)).not.toContain("Command failed"); + expect(logger.warn).toHaveBeenCalledWith(expect.objectContaining({ + adapterType: "claude_local", diagnostic: expect.stringContaining("Command failed"), + }), "Provider subscription quota unavailable"); + expect(JSON.stringify(vi.mocked(logger.warn).mock.calls)).not.toContain("fixture-private-token"); + }); + + it("isolates synchronous probe failures and preserves structured auth failure information", async () => { + vi.mocked(listServerAdapters).mockReturnValue([ + { type: "claude_local", getQuotaWindows: () => { throw new Error("local command unavailable"); } }, + { type: "codex_local", getQuotaWindows: vi.fn().mockResolvedValue({ + provider: "openai", source: "codex-rpc", ok: false, + errorFamily: "refresh_token_expired", error: "private diagnostic", windows: [], + }) }, + ] as never); + const results = await fetchAllQuotaWindows(); + expect(results[0]).toMatchObject({ provider: "anthropic", ok: false }); + expect(results[1]).toMatchObject({ + provider: "openai", source: "codex-rpc", ok: false, errorFamily: "refresh_token_expired", + }); + expect(JSON.stringify(results)).not.toContain("private diagnostic"); + }); }); diff --git a/server/src/__tests__/quota-windows.test.ts b/server/src/__tests__/quota-windows.test.ts index 156926ced6e..b56a99f7321 100644 --- a/server/src/__tests__/quota-windows.test.ts +++ b/server/src/__tests__/quota-windows.test.ts @@ -137,7 +137,7 @@ describe("WHAM used_percent normalization via fetchCodexQuota", () => { expect(windows[0]!.usedPercent).toBe(50); }); - it("treats values < 1 as fraction and multiplies by 100 (0.5 → 50%)", async () => { + it("preserves fractional percentage points (0.5 → 0.5%)", async () => { mockFetch({ rate_limit: { primary_window: { @@ -148,11 +148,10 @@ describe("WHAM used_percent normalization via fetchCodexQuota", () => { }, }); const windows = await fetchCodexQuota("token", null); - expect(windows[0]!.usedPercent).toBe(50); + expect(windows[0]!.usedPercent).toBe(0.5); }); - it("treats value exactly 1.0 as 1% (not 100%) — the < 1 heuristic boundary", async () => { - // 1.0 is NOT < 1, so it is treated as already-percentage → 1% + it("treats value exactly 1.0 as 1%", async () => { mockFetch({ rate_limit: { primary_window: { @@ -166,6 +165,16 @@ describe("WHAM used_percent normalization via fetchCodexQuota", () => { expect(windows[0]!.usedPercent).toBe(1); }); + it.each([ + [0.01, 0.01], [0.99, 0.99], [1.01, 1.01], [-1, 0], + [Number.NaN, null], [Number.POSITIVE_INFINITY, null], [undefined, null], + ])("handles percentage boundary %s consistently in RPC and WHAM", async (raw, expected) => { + mockFetch({ rate_limit: { primary_window: { used_percent: raw } } }); + const windows = await fetchCodexQuota("token", null); + expect(windows[0].usedPercent).toBe(expected); + expect(mapCodexRpcQuota({ rateLimits: { primary: { usedPercent: raw } } }).windows[0].usedPercent).toBe(expected); + }); + it("treats value 0 as 0%", async () => { mockFetch({ rate_limit: { @@ -820,6 +829,29 @@ describe("mapCodexRpcQuota", () => { // fetchWithTimeout — abort on timeout // --------------------------------------------------------------------------- +describe("managed quota cancellation", () => { + afterEach(() => vi.unstubAllGlobals()); + + it.each([ + ["Claude", (signal: AbortSignal) => fetchClaudeQuota("token", signal)], + ["Codex", (signal: AbortSignal) => fetchCodexQuota("token", "account", signal)], + ] as const)("forwards cancellation to the %s provider", async (_provider, read) => { + const controller = new AbortController(); + let requestSignal!: AbortSignal; + vi.stubGlobal("fetch", vi.fn(async (_url: string, init: RequestInit) => { + requestSignal = init.signal!; + return new Promise((_resolve, reject) => { + requestSignal.addEventListener("abort", () => reject(requestSignal.reason), { once: true }); + }); + })); + const pending = read(controller.signal); + const rejected = expect(pending).rejects.toThrow("account deadline"); + controller.abort(new Error("account deadline")); + await rejected; + expect(requestSignal.aborted).toBe(true); + }); +}); + describe("fetchWithTimeout", () => { afterEach(() => { vi.unstubAllGlobals(); diff --git a/server/src/lib/durable-directory.ts b/server/src/lib/durable-directory.ts new file mode 100644 index 00000000000..614f6777a50 --- /dev/null +++ b/server/src/lib/durable-directory.ts @@ -0,0 +1,71 @@ +import { promises as fs } from "node:fs"; +import path from "node:path"; + +// Revalidate the complete path on reuse, including symlink identities/targets. +// Moving an unchanged leaf to a new parent must invalidate its durability proof. +const durableDirectories = new Map(); + +export async function syncDirectory(directory: string) { + if (process.platform === "win32") return; + const handle = await fs.open(directory, "r"); + try { await handle.sync(); } finally { await handle.close(); } +} + +async function directoryAncestry(directory: string) { + // The configured path must also be traversable by the OS for later writes, + // including its native symlink-depth limit. + await fs.stat(directory); + let current = path.parse(directory).root; + let pending = directory.slice(current.length).split(path.sep); + const directories = new Set([current]); + const identities: string[] = []; + let links = 0; + while (pending.length > 0) { + const part = pending.shift()!; + if (!part || part === ".") continue; + if (part === "..") { current = path.dirname(current); continue; } + const next = path.join(current, part); + const stat = await fs.lstat(next, { bigint: true }); + identities.push(`${next}:${stat.dev}:${stat.ino}:${stat.birthtimeNs}`); + if (stat.isSymbolicLink()) { + if (++links > 40) throw Object.assign(new Error("Too many storage symlinks"), { code: "ELOOP" }); + const target = await fs.readlink(next); + identities.push(target); + // Resolve target components in filesystem order. Normalizing a target + // like "link/../home" before following link would select the wrong home. + if (path.isAbsolute(target)) current = path.parse(target).root; + pending = target.split(path.sep).concat(pending); + } else { + if (!stat.isDirectory()) throw Object.assign(new Error("Storage path is not a directory"), { code: "ENOTDIR" }); + current = next; + directories.add(current); + } + } + return { directories: [...directories].reverse(), identity: JSON.stringify(identities) }; +} + +/** An existing directory may be left by an interrupted mkdir/flush sequence. + * A fresh process must flush its complete ancestry before trusting it. Never + * waive a permission error based on an in-memory creation boundary: a restart + * would forget that boundary and could acknowledge a non-durable token write. */ +export async function ensureDurableDirectory(directory: string) { + const resolved = path.resolve(directory); + if (process.platform === "win32") { + await fs.mkdir(resolved, { recursive: true, mode: 0o700 }); + return; + } + let ancestry = await directoryAncestry(resolved).catch(error => { + if ((error as NodeJS.ErrnoException).code !== "ENOENT") throw error; + return null; + }); + if (!ancestry) { + await fs.mkdir(resolved, { recursive: true, mode: 0o700 }); + ancestry = await directoryAncestry(resolved); + } + if (durableDirectories.get(resolved) === ancestry.identity) return; + for (const parent of ancestry.directories) await syncDirectory(parent); + // Cache only a fully flushed path. Failures, siblings and fresh processes + // each retry the entire chain rather than inferring durability from existence. + if (durableDirectories.size >= 1024) durableDirectories.delete(durableDirectories.keys().next().value!); + durableDirectories.set(resolved, ancestry.identity); +} diff --git a/server/src/routes/agents.ts b/server/src/routes/agents.ts index 347a21a5e6e..7e9df47655b 100644 --- a/server/src/routes/agents.ts +++ b/server/src/routes/agents.ts @@ -8,7 +8,7 @@ import { agentFileStore, agentFileTokenFromHash } from "../services/agent-file-s import { pipeline } from "node:stream/promises"; import { resolveAgentAppearance, agentAvatarUrl } from "@paperclipai/shared"; import { listOpenRouterModels } from "../services/openrouter-models.js"; -import { prepareManagedAiRuntime, assertManagedAiProjectAuth, stripAiAuthBindings } from "../services/ai-connection-runtime.js"; +import { prepareManagedAiRuntime, withManagedAiProbe, assertManagedAiProjectAuth, stripAiAuthBindings } from "../services/ai-connection-runtime.js"; import { ADAPTER_AUTH_MISSING_CHECK_CODE, AI_CONNECTION_CAPABILITIES, aiConnectionBindingSchema, type AiConnectionBinding, aiRuntimeConnectionBindingSchema, type AiRuntimeConnectionBinding } from "@paperclipai/shared"; import { toolConnections } from "@paperclipai/db"; import { aiConnectionService } from "../services/ai-connections.js"; @@ -3506,17 +3506,17 @@ export function agentRoutes( const testEnvironmentId = await resolveAdapterTestEnvironmentId(companyId, environmentId); if (testEnvironmentId) await assertAdapterTestEnvironmentForCompany(companyId, testEnvironmentId); const target = await resolveAdapterTestExecutionContext({ companyId, adapterType, environmentId: testEnvironmentId }); - let managed: Awaited> | undefined; try { if (!target.executionTarget && target.fallbackChecks.length > 0) throw unprocessable("The agent environment is not available for adoption"); - managed = await prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: userId, adapterType, binding, config, allowUninstalledPersonal: newAgent, allowUninstalledShared, allowLegacyValidation: true }); + await withManagedAiProbe(db, { companyId, agentId, responsibleUserId: userId, adapterType, binding, config, allowUninstalledPersonal: newAgent, allowUninstalledShared, allowLegacyValidation: true }, async managed => { const result = await testManagedEnvironment(adapterType, { companyId, adapterType, config: managed.config, executionTarget: target.executionTarget, environmentName: target.environmentName }, binding, managed, agentId); if (result.status === "fail" || result.checks.some(check => check.code === ADAPTER_AUTH_MISSING_CHECK_CODE)) throw unprocessable("The selected AI connection failed validation in this agent’s environment. Run the agent test to see the failing checks.", { code: "ai_connection_validation_failed", checks: result.checks.filter(check => check.level === "error" || check.code === ADAPTER_AUTH_MISSING_CHECK_CODE).map(check => ({ code: check.code, level: check.level })), }); if (selection.connection.config.aiLegacyAdoption === true) await db.update(toolConnections).set({ healthStatus: "ok", config: { ...selection.connection.config, aiLegacyAdoption: false }, updatedAt: new Date() }).where(eq(toolConnections.id, selection.connection.id)); - } finally { try { await managed?.cleanup(); } finally { await target.release("released"); } } + }); + } finally { await target.release("released"); } } return { connectionId: selection.connection.id }; } @@ -3709,12 +3709,11 @@ export function agentRoutes( effectiveAdapterConfig.apiKey = req.body.testCredentials.API_SERVER_KEY; } } - const managed = aiBinding ? await prepareManagedAiRuntime(db, { companyId, agentId: req.body.agentId ?? "", responsibleUserId: responsibleUserForAiRequest(req), adapterType: type, binding: aiBinding, config: effectiveAdapterConfig, allowUninstalledPersonal: !req.body.agentId, allowUninstalledShared: !req.body.agentId && await canInstallSharedAiConnectionForNewAgent(db, req, companyId, aiBinding) }) : null; - let result; - try { - result = managed && aiBinding ? await testManagedEnvironment(type, { companyId, adapterType: type, config: managed.config, executionTarget, environmentName }, aiBinding, managed, savedAgentId ?? undefined) : await adapter.testEnvironment({ companyId, adapterType: type, config: effectiveAdapterConfig, executionTarget, environmentName }); - if (managed) result.checks.unshift({ code: "ai_connection_tested", level: "info", message: `Tested ${managed.accountName} — ${managed.accountOwnerUserId ? managed.accountOwnerUserId === responsibleUserForAiRequest(req) ? "your personal account" : "the owner’s account authorized for this agent" : "company-shared account"}. Responsible user: ${req.actor.type === "agent" ? responsibleUserForAiRequest(req) ?? "unavailable" : "the signed-in user"}.` }); - } finally { await managed?.cleanup(); } + const result = aiBinding ? await withManagedAiProbe(db, { companyId, agentId: req.body.agentId ?? "", responsibleUserId: responsibleUserForAiRequest(req), adapterType: type, binding: aiBinding, config: effectiveAdapterConfig, allowUninstalledPersonal: !req.body.agentId, allowUninstalledShared: !req.body.agentId && await canInstallSharedAiConnectionForNewAgent(db, req, companyId, aiBinding) }, async managed => { + const tested = await testManagedEnvironment(type, { companyId, adapterType: type, config: managed.config, executionTarget, environmentName }, aiBinding, managed, savedAgentId ?? undefined); + tested.checks.unshift({ code: "ai_connection_tested", level: "info", message: `Tested ${managed.accountName} — ${managed.accountOwnerUserId ? managed.accountOwnerUserId === responsibleUserForAiRequest(req) ? "your personal account" : "the owner’s account authorized for this agent" : "company-shared account"}. Responsible user: ${req.actor.type === "agent" ? responsibleUserForAiRequest(req) ?? "unavailable" : "the signed-in user"}.` }); + return tested; + }) : await adapter.testEnvironment({ companyId, adapterType: type, config: effectiveAdapterConfig, executionTarget, environmentName }); const prefixChecks = [ ...(sandboxIdentityCheck ? [sandboxIdentityCheck] : []), diff --git a/server/src/routes/costs.ts b/server/src/routes/costs.ts index 373adf06141..b66abc0fc0d 100644 --- a/server/src/routes/costs.ts +++ b/server/src/routes/costs.ts @@ -21,7 +21,7 @@ import { logActivity, } from "../services/index.js"; import { assertBoard, assertCompanyAccess, getAccessibleResource, getActorInfo } from "./authz.js"; -import { fetchAllQuotaWindows } from "../services/quota-windows.js"; +import { fetchCompanyQuotaWindows } from "../services/quota-windows.js"; import { badRequest } from "../errors.js"; import type { PluginWorkerManager } from "../services/plugin-worker-manager.js"; @@ -281,7 +281,7 @@ export function costRoutes( res.status(404).json({ error: "Company not found" }); return; } - const results = await fetchAllQuotaWindows(); + const results = await fetchCompanyQuotaWindows(db, companyId, getActorInfo(req).actorId); res.json(results); }); diff --git a/server/src/services/ai-connection-runtime.ts b/server/src/services/ai-connection-runtime.ts index bce54ad3d22..a43e29024d0 100644 --- a/server/src/services/ai-connection-runtime.ts +++ b/server/src/services/ai-connection-runtime.ts @@ -13,7 +13,8 @@ import { import { managedProviderRouting } from "./ai-provider-routing.js"; import { aiConnectionService } from "./ai-connections.js"; import { secretService } from "./secrets.js"; -import { decideCodexAuthMerge } from "@paperclipai/adapter-codex-local/server"; +import { decideCodexAuthMerge, withAccountHomeSecretMutationLock } from "@paperclipai/adapter-codex-local/server"; +import { WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE } from "@paperclipai/adapter-utils/workspace-restore-merge"; import type { AdapterExecutionTarget } from "@paperclipai/adapter-utils/execution-target"; import { runAdapterExecutionTargetProcess } from "@paperclipai/adapter-utils/execution-target"; import { decideGrokAuthMerge } from "@paperclipai/adapter-grok-local/server"; @@ -271,12 +272,31 @@ export async function prepareManagedAiRuntime( const routing = aiConnectionMetadataSchema.parse(selection.connection.config.ai).routing; const noAuth = routing?.auth === "none"; if (!noAuth && !credentialRef) throw unprocessable("The selected AI credential is unavailable"); - const readFreshness = async () => credentialRef ? (await db.select({ epoch: companySecrets.aiSessionEpoch, version: companySecrets.latestVersion }) - .from(companySecrets).where(and(eq(companySecrets.companyId, input.companyId), eq(companySecrets.id, credentialRef.secretId))).limit(1))[0] : undefined; - const freshness = noAuth ? undefined : await readFreshness(); - const value = noAuth ? "" : await service.credential(selection); - const afterRead = noAuth ? undefined : await readFreshness(); - if (!noAuth && (!freshness || freshness.version !== afterRead?.version || freshness.epoch !== afterRead.epoch)) throw unprocessable("The AI credential changed during preparation; retry this execution"); + const readFreshness = async () => { + if (!credentialRef) return undefined; + const [currentGrant] = await db.select({ refs: connectionGrants.credentialSecretRefs, status: connectionGrants.status }) + .from(connectionGrants).where(and(eq(connectionGrants.companyId, input.companyId), eq(connectionGrants.id, selection.grant.id))); + const currentRef = currentGrant?.refs.find(ref => ref.configPath === "ai.credential"); + if (currentGrant?.status !== "active" || currentRef?.secretId !== credentialRef.secretId || currentRef.versionSelector !== credentialRef.versionSelector) return undefined; + return (await db.select({ epoch: companySecrets.aiSessionEpoch, version: companySecrets.latestVersion }) + .from(companySecrets).where(and(eq(companySecrets.companyId, input.companyId), eq(companySecrets.id, credentialRef.secretId))).limit(1))[0]; + }; + const { value, freshness } = await (async () => { + if (noAuth) return { value: "", freshness: undefined }; + // Recovering a rotated quota token can advance the secret version during + // the first read. Re-read the saved credential, retaining the epoch guard + // against a concurrent reconnect or explicit rotation. + for (let attempt = 0; attempt < 2; attempt += 1) { + const before = await readFreshness(); + const value = await service.runtimeCredential(selection); + const after = await readFreshness(); + if (before && after && before.version === after.version && before.epoch === after.epoch) { + return { value, freshness: after }; + } + if (!before || before.epoch !== after?.epoch) break; + } + throw unprocessable("The AI credential changed during preparation; retry this execution"); + })(); home = await mkdtemp( path.join( os.tmpdir(), @@ -355,75 +375,85 @@ export async function prepareManagedAiRuntime( identity, home, cleanup: async () => { - try { - if (subscriptionFile) { - const refreshed = await readFile(authFile, "utf8"); - if (refreshed !== value) - await db.transaction(async (tx) => { - const [grant] = await tx - .select() - .from(connectionGrants) - .where( - and( - eq(connectionGrants.id, selection.grant.id), - eq(connectionGrants.companyId, input.companyId), - ), - ) - .for("update"); - // A missing or revoked grant blocks the write-back. Among - // active copies, the merge decision below keeps the - // credential with the newest provider freshness field. - if (!grant || grant.status !== "active") return; - const ref = grant.credentialSecretRefs.find( - (r) => r.configPath === "ai.credential", - ); - if (!ref) return; - // Lock the referenced secret row for the rest of this - // transaction. The grant-row lock above does not cover it, - // so an authorized rotation of this secret could otherwise - // land between the read and the write below and be - // overwritten by this stale write-back. - await tx - .select({ id: companySecrets.id }) - .from(companySecrets) - .where( - and( - eq(companySecrets.id, ref.secretId), - eq(companySecrets.companyId, input.companyId), - ), - ) - .for("update"); - const current = await aiConnectionService( - tx as unknown as Db, - ).credential({ ...selection, grant }); - const destination = path.join( - providerHome, - "current-auth.json", - ); - await writeFile(destination, current, { mode: 0o600 }); - const decision = - input.binding.provider === "openai" - ? await decideCodexAuthMerge(authFile, destination, { - errorLabel: "AI account refresh", - }) - : await decideGrokAuthMerge(authFile, destination, { - errorLabel: "AI account refresh", - }); - if (decision !== 10) return; - await secretService(tx).rotate( - ref.secretId, - { value: refreshed, preserveAiSessionEpoch: true }, - { userId: grant.subjectUserId }, - ); - await tx - .update(connectionGrants) - .set({ updatedAt: new Date() }) - .where(eq(connectionGrants.id, grant.id)); - }); + if (subscriptionFile) { + const refreshed = await readFile(authFile, "utf8"); + if (refreshed !== value) { + // A quota exchange can hold this company lock for 60 seconds. + // Retry its 30-second acquisition timeout without discarding the + // provider's only copy of a rotated, single-use refresh token. + const writeBack = () => withAccountHomeSecretMutationLock(undefined, input.companyId, () => db.transaction(async (tx) => { + const [grant] = await tx + .select() + .from(connectionGrants) + .where( + and( + eq(connectionGrants.id, selection.grant.id), + eq(connectionGrants.companyId, input.companyId), + ), + ) + .for("update"); + // A missing or revoked grant blocks the write-back. Among + // active copies, the merge decision below keeps the + // credential with the newest provider freshness field. + if (!grant || grant.status !== "active") return; + const ref = grant.credentialSecretRefs.find( + (r) => r.configPath === "ai.credential", + ); + if (!ref) return; + // Lock the referenced secret row for the rest of this + // transaction. The grant-row lock above does not cover it, + // so an authorized rotation of this secret could otherwise + // land between the read and the write below and be + // overwritten by this stale write-back. + await tx + .select({ id: companySecrets.id }) + .from(companySecrets) + .where( + and( + eq(companySecrets.id, ref.secretId), + eq(companySecrets.companyId, input.companyId), + ), + ) + .for("update"); + const current = await aiConnectionService( + tx as unknown as Db, + ).credential({ ...selection, grant }); + const destination = path.join( + providerHome, + "current-auth.json", + ); + await writeFile(destination, current, { mode: 0o600 }); + const decision = + input.binding.provider === "openai" + ? await decideCodexAuthMerge(authFile, destination, { + errorLabel: "AI account refresh", + }) + : await decideGrokAuthMerge(authFile, destination, { + errorLabel: "AI account refresh", + }); + if (decision !== 10) return; + await secretService(tx).rotate( + ref.secretId, + { value: refreshed, preserveAiSessionEpoch: true }, + { userId: grant.subjectUserId }, + ); + // Keep grant.updatedAt for explicit account/access changes. + // The rotated secret revision invalidates quota caches without + // rejecting an otherwise valid in-flight reconnect. + })); + for (let attempt = 0; ; attempt++) { + try { await writeBack(); break; } + catch (error) { + if (attempt >= 2 || (error as NodeJS.ErrnoException)?.code !== WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE) throw error; + } + } } - } finally { - if (home) await rm(home, { recursive: true, force: true }); } + // Preserve the private home on any failed write-back, including an + // exhausted lock retry or database failure. The caller can retry the + // same cleanup; only a committed save or an intentional discard + // (revoked grant, older/different identity) permits deletion. + if (home) await rm(home, { recursive: true, force: true }); }, }; } catch (error) { @@ -431,3 +461,26 @@ export async function prepareManagedAiRuntime( throw error; } } + +/** Manual tests/adoption have no heartbeat row for quota's active-run guard. + * Hold the shared credential lock from resolution through the provider probe + * and write-back. Nested credential writes reuse this lock's async ownership. */ +export async function withManagedAiProbe( + db: Db, + input: Parameters[1], + probe: (runtime: Awaited>) => Promise, +): Promise { + const run = async () => { + const runtime = await prepareManagedAiRuntime(db, input); + try { return await probe(runtime); } + finally { await runtime.cleanup(); } + }; + return input.binding.provider === "openai" + ? withAccountHomeSecretMutationLock(undefined, input.companyId, run).catch(error => { + if ((error as NodeJS.ErrnoException)?.code === WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE) { + throw unprocessable("AI credentials are being updated. Retry shortly.", { code: "ai_connection_busy" }); + } + throw error; + }) + : run(); +} diff --git a/server/src/services/ai-connections.ts b/server/src/services/ai-connections.ts index b7f0f53174c..1fb561a9206 100644 --- a/server/src/services/ai-connections.ts +++ b/server/src/services/ai-connections.ts @@ -1,7 +1,9 @@ +import { WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE } from "@paperclipai/adapter-utils/workspace-restore-merge"; +import { withAccountHomeSecretMutationLock } from "@paperclipai/adapter-codex-local/server"; import { syncConnectionCredentialBindings } from "./connection-credential-bindings.js"; import { createHash, randomUUID } from "node:crypto"; import { isDeepStrictEqual } from "node:util"; -import { and, eq, inArray, or } from "drizzle-orm"; +import { and, eq, inArray, or, sql } from "drizzle-orm"; import { type Db, authUsers, @@ -9,6 +11,7 @@ import { aiConnectionDefaults, aiProviderDefaults, agents, + heartbeatRuns, companyMemberships, companySecrets, userSecretDefinitions, @@ -35,6 +38,8 @@ import { type AiConnectionLoginIntent, } from "@paperclipai/shared"; import { forbidden, notFound, unprocessable } from "../errors.js"; +import { quotaCredentialRecovery, quotaCredentialHash } from "./quota-credential-recovery.js"; +import { logger } from "../middleware/logger.js"; import { logActivity } from "./activity-log.js"; import { secretService } from "./secrets.js"; import { probeAiConnectionUsage } from "./ai-connection-usage.js"; @@ -402,7 +407,7 @@ export function aiConnectionService(db: Db) { } satisfies AiConnectionAttribution, }; } - async function credential(row: Pick>, "connection" | "grant">) { + async function credential(row: Pick>, "connection" | "grant">, retry = 0): Promise { const metadata = aiConnectionMetadataSchema.parse(row.connection.config.ai); if (metadata.routing?.auth === "none") return ""; const ref = row.grant.credentialSecretRefs.find( @@ -434,6 +439,7 @@ export function aiConnectionService(db: Db) { responsibleUserId: row.grant.subjectUserId, actorType: "system" as const, }; + let value: string; if (secret.scope === "user") { if ( secret.ownerUserId !== row.grant.subjectUserId || @@ -451,14 +457,21 @@ export function aiConnectionService(db: Db) { context, ); if (!result) throw unprocessable("Reconnect this AI account"); - return result.value; + value = result.value; + } else { + value = await secrets.resolveSecretValue(row.connection.companyId, secret.id, "latest", context); } - return secrets.resolveSecretValue( - row.connection.companyId, - secret.id, - "latest", - context, - ); + // Resolution records lastResolvedAt and can wait behind a rotating writer. + // Re-read after that wait instead of handing a run the pre-rotation token. + const [latest] = await db.select().from(companySecrets).where(and( + eq(companySecrets.id, secret.id), eq(companySecrets.companyId, row.connection.companyId), + )); + if (!latest || latest.status !== "active") throw unprocessable("Reconnect this AI account"); + if (latest.latestVersion !== secret.latestVersion) { + if (retry >= 2) throw unprocessable("AI credentials are changing. Retry this execution.", { code: "ai_connection_busy" }); + return credential(row, retry + 1); + } + return value; } async function save( companyId: string, @@ -504,7 +517,7 @@ export function aiConnectionService(db: Db) { ); const id = reconnect?.connection.id ?? randomUUID(); const grantId = reconnect?.grant.id ?? randomUUID(); - return db.transaction(async (tx) => { + const persist = () => withAccountHomeSecretMutationLock(undefined, companyId, () => db.transaction(async (tx) => { const secrets = secretService(tx); if (sessionId) { const [session] = await tx @@ -812,8 +825,89 @@ export function aiConnectionService(db: Db) { details: { provider: input.provider, method: input.method, grantId }, }); return { connectionId: id, grantId }; + })); + // Quota refresh can hold this lock through a 60-second token exchange. + // Keep the verified sign-in in memory across two 30-second acquisition + // timeouts instead of making its caller discard a successful login. + for (let attempt = 0; ; attempt++) { + try { return await persist(); } + catch (error) { + if (attempt >= 2 || (error as NodeJS.ErrnoException)?.code !== WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE) throw error; + } + } + } + async function recoverQuotaCredentialInTransaction( + tx: Db, row: Parameters[0], current: string, onRecoveryStarted?: () => void, + ) { + const metadata = aiConnectionMetadataSchema.safeParse(row.connection.config.ai); + if (!metadata.success || metadata.data.provider !== "openai" || metadata.data.method !== "subscription") return null; + const recovery = quotaCredentialRecovery(row.connection.companyId, row.grant.id); + const pending = await recovery.read(); + if (!pending) return null; + const ref = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential"); + if (!ref || pending.secretId !== ref.secretId || pending.connectionId !== row.connection.id || pending.baseHash !== quotaCredentialHash(current)) { + await recovery.clear(); + return null; + } + onRecoveryStarted?.(); + await secretService(tx).rotate(ref.secretId, { value: pending.value, preserveAiSessionEpoch: true }, { userId: row.grant.subjectUserId }); + // Automatic rotation changes secret revision, not the grant mutation marker + // used to fence reconnect attempts. Ownership and access are unchanged. + await logActivity(tx, { companyId: row.connection.companyId, actorType: "system", actorId: "quota_recovery", + action: "ai_connection.credential_refreshed", entityType: "tool_connection", entityId: row.connection.id, + details: { provider: "openai", grantId: row.grant.id } }); + return pending.value; + } + + /** Never hand a consumed token to a new provider while its replacement awaits commit. */ + async function runtimeCredential(row: Parameters[0]) { + const metadata = aiConnectionMetadataSchema.safeParse(row.connection.config.ai); + if (!metadata.success || metadata.data.provider !== "openai" || metadata.data.method !== "subscription") return credential(row); + let savingReplacement = false; + return withAccountHomeSecretMutationLock(undefined, row.connection.companyId, async () => { + let recovered = false; + const value = await db.transaction(async tx => { + const [grant] = await tx.select().from(connectionGrants).where(and( + eq(connectionGrants.id, row.grant.id), eq(connectionGrants.companyId, row.connection.companyId), + eq(connectionGrants.connectionId, row.connection.id), + )).for("update", { noWait: true }); + 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, row.connection.companyId), + )).for("update", { noWait: true }); + const selected = { connection: row.connection, grant }; + const current = await aiConnectionService(tx as unknown as Db).credential(selected); + const pending = await recoverQuotaCredentialInTransaction(tx as unknown as Db, selected, current, () => { + savingReplacement = true; + }); + recovered = pending !== null; + return pending ?? current; + }); + if (recovered) await quotaCredentialRecovery(row.connection.companyId, row.grant.id).clear() + .catch(error => logger.warn({ err: error, grantId: row.grant.id }, "Saved runtime credential recovery cleanup pending")); + return value; + }).catch(error => { + // Once a matching replacement is found, a failed write or commit must + // defer execution. Keep its journal so the next attempt can save it. + if (savingReplacement) { + throw unprocessable("AI credentials are being updated. Retry shortly.", { code: "ai_connection_busy" }); + } + // Drizzle wraps Postgres errors in `cause`. Contention is a temporary + // pre-provider wait, not evidence that this account needs reconnecting. + let current: unknown = error; + for (let depth = 0; depth < 4 && current && typeof current === "object"; depth += 1) { + const value = current as { code?: unknown; cause?: unknown }; + if (value.code === WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE || value.code === "55P03") { + throw unprocessable("AI credentials are being updated. Retry shortly.", { code: "ai_connection_busy" }); + } + current = value.cause; + } + throw error; }); } + /** A late failure must never invalidate credentials that were refreshed or reconnected meanwhile. */ async function markAuthenticationFailed(input: { companyId: string; @@ -822,41 +916,159 @@ export function aiConnectionService(db: Db) { runStartedAt: Date; attribution: AiConnectionAttribution & { identity: string }; }) { - return db.transaction(async (tx) => { - const { attribution } = input; - const [grant] = await tx.select().from(connectionGrants).where(and( - eq(connectionGrants.companyId, input.companyId), - eq(connectionGrants.id, attribution.grantId), - eq(connectionGrants.connectionId, attribution.connectionId), - )).for("update"); - if (!grant || grant.status !== "active" || grant.updatedAt > input.runStartedAt) return; - const [connection] = await tx.select().from(toolConnections).where(and( - eq(toolConnections.companyId, input.companyId), - eq(toolConnections.id, attribution.connectionId), - eq(toolConnections.connectionPurpose, "ai"), - )); - if (!connection) return; - const metadata = aiConnectionMetadataSchema.safeParse(connection.config.ai); - if (!metadata.success || metadata.data.provider !== attribution.provider || metadata.data.method !== attribution.method) return; - const ref = grant.credentialSecretRefs.find((candidate) => candidate.configPath === "ai.credential"); - if (!ref) return; - await tx.select({ id: companySecrets.id }).from(companySecrets).where(and( - eq(companySecrets.companyId, input.companyId), eq(companySecrets.id, ref.secretId), - )).for("update"); - const value = await aiConnectionService(tx as unknown as Db).credential({ connection, grant }); - const generation = createHash("sha256").update(value).digest("hex").slice(0, 16); - if (attribution.identity !== `${grant.id}:${attribution.responsibleUserId ?? "shared"}:${generation}`) return; - await tx.update(connectionGrants).set({ status: "needs_reauthorization", updatedAt: new Date() }) - .where(eq(connectionGrants.id, grant.id)); - await tx.update(toolConnections).set({ healthStatus: "error", healthMessage: "Sign in again to restore this AI connection.", updatedAt: new Date() }) - .where(eq(toolConnections.id, connection.id)); - await logActivity(tx as unknown as Db, { - companyId: input.companyId, actorType: "system", actorId: input.runId ? "heartbeat" : "adapter_test", - agentId: input.agentId, runId: input.runId, action: "ai_connection.authentication_failed", - entityType: "tool_connection", entityId: connection.id, - details: { provider: attribution.provider, grantId: grant.id }, + return withAccountHomeSecretMutationLock(undefined, input.companyId, async () => { + let recovered = false; + await db.transaction(async (tx) => { + const { attribution } = input; + const [grant] = await tx.select().from(connectionGrants).where(and( + eq(connectionGrants.companyId, input.companyId), + eq(connectionGrants.id, attribution.grantId), + eq(connectionGrants.connectionId, attribution.connectionId), + )).for("update"); + if (!grant || grant.status !== "active" || grant.updatedAt > input.runStartedAt) return; + const [connection] = await tx.select().from(toolConnections).where(and( + eq(toolConnections.companyId, input.companyId), + eq(toolConnections.id, attribution.connectionId), + eq(toolConnections.connectionPurpose, "ai"), + )); + if (!connection) return; + const metadata = aiConnectionMetadataSchema.safeParse(connection.config.ai); + if (!metadata.success || metadata.data.provider !== attribution.provider || metadata.data.method !== attribution.method) return; + const ref = grant.credentialSecretRefs.find((candidate) => candidate.configPath === "ai.credential"); + if (!ref) return; + await tx.select({ id: companySecrets.id }).from(companySecrets).where(and( + eq(companySecrets.companyId, input.companyId), eq(companySecrets.id, ref.secretId), + )).for("update"); + const value = await aiConnectionService(tx as unknown as Db).credential({ connection, grant }); + const generation = createHash("sha256").update(value).digest("hex").slice(0, 16); + if (attribution.identity !== `${grant.id}:${attribution.responsibleUserId ?? "shared"}:${generation}`) return; + // A quota exchange may have consumed this run's token before a failed + // database commit. Save its working replacement instead of revoking it. + recovered = await recoverQuotaCredentialInTransaction(tx as unknown as Db, { connection, grant }, value) !== null; + if (recovered) return; + await tx.update(connectionGrants).set({ status: "needs_reauthorization", updatedAt: new Date() }) + .where(eq(connectionGrants.id, grant.id)); + await tx.update(toolConnections).set({ healthStatus: "error", healthMessage: "Sign in again to restore this AI connection.", updatedAt: new Date() }) + .where(eq(toolConnections.id, connection.id)); + await logActivity(tx as unknown as Db, { + companyId: input.companyId, actorType: "system", actorId: input.runId ? "heartbeat" : "adapter_test", + agentId: input.agentId, runId: input.runId, action: "ai_connection.authentication_failed", + entityType: "tool_connection", entityId: connection.id, + details: { provider: attribution.provider, grantId: grant.id }, + }); }); + if (recovered) await quotaCredentialRecovery(input.companyId, input.attribution.grantId).clear() + .catch(error => logger.warn({ err: error, grantId: input.attribution.grantId }, "Saved failed-run credential recovery cleanup pending")); }); } - return { list, select, credential, probeUsage, save, setDefault, membership, markAuthenticationFailed }; + /** Refresh only the selected vaulted identity, serializing rotation and reconnect. */ + async function refreshQuotaCredential(row: Parameters[0], failedValue: string, signal: AbortSignal) { + return withAccountHomeSecretMutationLock(undefined, row.connection.companyId, async () => { + const recovery = quotaCredentialRecovery(row.connection.companyId, row.grant.id); + let savingReplacement = false; + const save = () => db.transaction(async (tx) => { + // Do not wait on a transaction that might itself be waiting for the file + // lock. All database locks must be available before consuming OAuth tokens. + 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", { noWait: true }); + 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", { noWait: true }); + const current = await aiConnectionService(tx as unknown as Db).credential({ connection: row.connection, grant }); + signal.throwIfAborted(); + let pending = await recovery.read(); + if (pending && (pending.secretId !== ref.secretId || pending.connectionId !== row.connection.id || pending.baseHash !== quotaCredentialHash(current))) { + // Reconnect may reuse this grant. Once the current authorized token + // itself expires, an older journal must not prevent refreshing it. + await recovery.clear(); + pending = null; + } + // A delayed poll may have read A before B was saved and B's refresh + // produced an uncommitted C. Recover C before returning or clearing its + // journal; a newer database credential alone does not prove it is usable. + if (current !== failedValue && !pending) 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 && !pending) throw new Error("provider_unavailable"); + let refreshed = pending?.value; + if (!refreshed) { + await recovery.prepare(); + 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: refreshSignal, + 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"); + 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() }); + // This fsynced encrypted journal survives transaction rollback and host + // restart. Never exchange the consumed token again after a failed save. + await recovery.write({ companyId, grantId: grant.id, connectionId: row.connection.id, + secretId: ref.secretId, baseHash: quotaCredentialHash(current), value: refreshed }); + } + savingReplacement = true; + // 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, preserveAiSessionEpoch: true }, { userId: grant.subjectUserId }); + 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; + }); + for (let attempt = 0; ; attempt++) { + try { + const value = await save(); + // Clear only after commit. A committed-but-lost response is detected + // by the current credential comparison on the next attempt. + await recovery.clear().catch((error) => logger.warn({ err: error, grantId: row.grant.id }, "Saved quota credential recovery cleanup pending")); + return value; + } catch (error) { + if (!savingReplacement || attempt >= 2) throw error; + // Saving retries have no provider side effect. Persistent failures + // retain the journal for the next quota poll, including after restart. + } + } + }); + } + // 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) { + if (!(await membership(companyId, userId))) return []; + const visible = await list(companyId, userId); + const accounts = await rows(companyId); + return accounts.flatMap(row => { + const summary = visible.find(item => item.id === row.connection.id && item.grantId === row.grant.id); + if (!summary || summary.method !== "subscription" || !["openai", "anthropic"].includes(summary.provider)) return []; + return [{ ...row, summary }]; + }); + } + return { list, quotaAccounts, refreshQuotaCredential, select, credential, runtimeCredential, probeUsage, save, setDefault, membership, markAuthenticationFailed }; } diff --git a/server/src/services/quota-credential-recovery.ts b/server/src/services/quota-credential-recovery.ts new file mode 100644 index 00000000000..d666641c059 --- /dev/null +++ b/server/src/services/quota-credential-recovery.ts @@ -0,0 +1,54 @@ +import { createHash, randomUUID } from "node:crypto"; +import { promises as fs } from "node:fs"; +import path from "node:path"; +import { z } from "zod"; +import { resolvePaperclipInstanceRoot } from "../home-paths.js"; +import { localEncryptedProvider } from "../secrets/local-encrypted-provider.js"; +import { ensureDurableDirectory, syncDirectory } from "../lib/durable-directory.js"; + +const recoverySchema = z.object({ + companyId: z.string().uuid(), grantId: z.string().uuid(), connectionId: z.string().uuid(), secretId: z.string().uuid(), + baseHash: z.string().length(64), value: z.string().min(1), +}); +type Recovery = z.infer; +export const quotaCredentialHash = (value: string) => createHash("sha256").update(value).digest("hex"); +export const quotaCredentialRecoveryPath = (companyId: string, grantId: string) => + path.join(resolvePaperclipInstanceRoot(), "quota-credential-recovery", companyId, `${grantId}.json`); + +/** Called only while holding the company credential mutation lock. The local + * encrypted provider uses the instance secrets master key, including when the + * destination vault is remote. No plaintext token enters this journal. */ +export function quotaCredentialRecovery(companyId: string, grantId: string) { + const file = quotaCredentialRecoveryPath(companyId, grantId); + const directory = path.dirname(file); + async function writeEncrypted(target: string, value: string) { + const encrypted = await localEncryptedProvider.createVersion({ value }); + await ensureDurableDirectory(directory); + const temporary = `${target}.${randomUUID()}.tmp`; + const handle = await fs.open(temporary, "wx", 0o600); + try { + try { await handle.writeFile(JSON.stringify(encrypted.material)); await handle.sync(); } + finally { await handle.close(); } + await fs.rename(temporary, target); + await syncDirectory(directory); + } catch (error) { await fs.rm(temporary, { force: true }); throw error; } + } + return { + // Verify encryption, storage, and fsync before exchanging a single-use token. + prepare: async () => { + const probe = `${file}.${randomUUID()}.probe`; + try { await writeEncrypted(probe, "credential-recovery-probe"); } + finally { await fs.rm(probe, { force: true }); } + }, + read: async (): Promise => { + let material: string; + try { material = await fs.readFile(file, "utf8"); } + catch (error) { if ((error as NodeJS.ErrnoException).code === "ENOENT") return null; throw error; } + const pending = recoverySchema.parse(JSON.parse(await localEncryptedProvider.resolveVersion({ material: JSON.parse(material), externalRef: null }))); + if (pending.companyId !== companyId || pending.grantId !== grantId) throw new Error("credentials_unavailable"); + return pending; + }, + write: (input: Recovery) => writeEncrypted(file, JSON.stringify(recoverySchema.parse(input))), + clear: async () => { await fs.rm(file, { force: true }); }, + }; +} diff --git a/server/src/services/quota-windows.ts b/server/src/services/quota-windows.ts index 664406abdec..2d363969b26 100644 --- a/server/src/services/quota-windows.ts +++ b/server/src/services/quota-windows.ts @@ -1,5 +1,135 @@ import type { ProviderQuotaResult } from "@paperclipai/shared"; +import { redactDiagnosticText } from "@paperclipai/adapter-utils"; import { listServerAdapters } from "../adapters/registry.js"; +import { logger } from "../middleware/logger.js"; +import { createHash } from "node:crypto"; +import { eq, inArray, and } from "drizzle-orm"; +import { companySecrets, companySecretVersions, userSecretDefinitions, type Db } from "@paperclipai/db"; +import { fetchCodexQuota } from "@paperclipai/adapter-codex-local/server"; +import { fetchClaudeQuota } from "@paperclipai/adapter-claude-local/server"; +import { aiConnectionService } from "./ai-connections.js"; + +const accountRequests = new Map }>(); +const publicError = "Subscription quota is currently unavailable. Check usage with your provider."; + +/** Managed connections use their own credentials regardless of where agents run. + * Never silently substitute the control-plane host's login for a remote account. */ +export async function fetchCompanyQuotaWindows(db: Db, companyId: string, userId: string): Promise { + const service = aiConnectionService(db); + const accounts = await service.quotaAccounts(companyId, userId); + if (!accounts.length) { + return ["anthropic", "openai"].map(provider => ({ provider, ok: false, + accountKey: `unconnected:${companyId}:${provider}`, source: "managed-connection", + errorFamily: "credentials_unavailable", error: "Connect a subscription account in AI connections to view its quota.", windows: [] })); + } + // 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 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, 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 + || (revision.definitionId && (revision.definitionStatus !== "active" || revision.definitionDeletedAt !== null))); + }); + if (row.summary.status !== "connected" || definitionUnavailable) { + return { ...base, ok: false, errorFamily: "credentials_unavailable", error: publicError, windows: [] }; + } + const key = `${userId}:${accountKey}`; + let pending = accountRequests.get(key); + if (!pending || pending.expires <= Date.now()) { + const controller = new AbortController(); + const result = (async (): Promise => { + let credentialResolved = false; + try { + const value = await service.credential(row); + controller.signal.throwIfAborted(); + credentialResolved = true; + let windows; + if (base.provider === "openai") { + let auth; + try { auth = JSON.parse(value); } catch { throw new Error("credentials_unavailable"); } + const token = auth.tokens?.access_token ?? auth.accessToken; + if (typeof token !== "string" || !token) { + return { ...base, ok: false, errorFamily: "credentials_unavailable", error: publicError, windows: [] }; + } + 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)); + // 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); + } + } 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 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, + errorFamily }; + } + })(); + 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); + } + return pending.result; + } + const results: ProviderQuotaResult[] = new Array(accounts.length); + let next = 0; + await Promise.all(Array.from({ length: Math.min(4, accounts.length) }, async () => { + while (next < accounts.length) { + const index = next++; + results[index] = await readAccount(accounts[index]); + } + })); + return results; +} const QUOTA_PROVIDER_TIMEOUT_MS = 20_000; @@ -24,24 +154,31 @@ export async function fetchAllQuotaWindows(): Promise { const adapters = listServerAdapters().filter((a) => a.getQuotaWindows != null); const settled = await Promise.allSettled( - adapters.map((adapter) => withQuotaTimeout(adapter.type, adapter.getQuotaWindows!())), + adapters.map((adapter) => withQuotaTimeout(adapter.type, Promise.resolve().then(() => adapter.getQuotaWindows!()))), ); return settled.map((result, i) => { - if (result.status === "fulfilled") return result.value; const adapterType = adapters[i]!.type; - return { + const quota: ProviderQuotaResult = result.status === "fulfilled" ? result.value : { provider: providerSlugForAdapterType(adapterType), ok: false, error: String(result.reason), windows: [], }; + if (quota.ok) return { ...quota, error: undefined }; + logger.warn({ + adapterType, + errorFamily: quota.errorFamily, + diagnostic: redactDiagnosticText(quota.error ?? "Quota probe failed").slice(0, 2_000), + }, "Provider subscription quota unavailable"); + return { ...quota, error: "Subscription quota is currently unavailable. Check usage with your provider." }; }); } async function withQuotaTimeout( adapterType: string, task: Promise, + onTimeout?: () => void, ): Promise { let timeoutId: NodeJS.Timeout | null = null; try { @@ -49,6 +186,7 @@ async function withQuotaTimeout( task, new Promise((resolve) => { timeoutId = setTimeout(() => { + onTimeout?.(); resolve({ provider: providerSlugForAdapterType(adapterType), ok: false, diff --git a/ui/src/components/ClaudeSubscriptionPanel.tsx b/ui/src/components/ClaudeSubscriptionPanel.tsx index 3a950a457a0..bd814aec624 100644 --- a/ui/src/components/ClaudeSubscriptionPanel.tsx +++ b/ui/src/components/ClaudeSubscriptionPanel.tsx @@ -1,3 +1,4 @@ +import { quotaUnavailableMessage } from "@/lib/quota-refresh"; import type { QuotaWindow } from "@paperclipai/shared"; import { cn, quotaSourceDisplayName } from "@/lib/utils"; @@ -66,7 +67,7 @@ export function ClaudeSubscriptionPanel({ Anthropic subscription
- Live Claude quota windows. + Quota from the Claude subscription connected to this server.
{source ? ( @@ -77,8 +78,8 @@ export function ClaudeSubscriptionPanel({ {error ? ( -
- {error} +
+ {quotaUnavailableMessage(windows.length > 0)}
) : null} diff --git a/ui/src/components/CodexSubscriptionPanel.tsx b/ui/src/components/CodexSubscriptionPanel.tsx index 5f26314024a..d18685563e1 100644 --- a/ui/src/components/CodexSubscriptionPanel.tsx +++ b/ui/src/components/CodexSubscriptionPanel.tsx @@ -1,3 +1,4 @@ +import { quotaUnavailableMessage } from "@/lib/quota-refresh"; import type { QuotaWindow } from "@paperclipai/shared"; import { cn, quotaSourceDisplayName } from "@/lib/utils"; @@ -69,7 +70,7 @@ export function CodexSubscriptionPanel({ Codex subscription
- Live Codex quota windows. + Quota from the Codex subscription connected to this server.
{source ? ( @@ -80,8 +81,8 @@ export function CodexSubscriptionPanel({ {error ? ( -
- {error} +
+ {quotaUnavailableMessage(windows.length > 0)}
) : null} 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."; +}