From 9ed55f6931a58fef1e54441c19cc9805cd41073e Mon Sep 17 00:00:00 2001 From: Nicky Leach Date: Tue, 15 Sep 2026 01:10:54 -0700 Subject: [PATCH] fix: allow concurrent agent runs on one OpenAI or xAI subscription connection (#13452) ## Thinking Path > - Paperclip is the open source app people use to manage AI agents for work > - Paperclip starts local command-line sessions and stores provider credentials through managed connections > - One OpenAI Codex or xAI Grok subscription connection held a credential lease for the full agent run > - A second run then waited for the first run, and concurrent runs could overwrite a newer credential > - The write-back must compare fresh credentials while it holds the row lock > - This pull request removes the run lease, keeps the revocation guard, and bounds Codex timestamps against the host clock > - The benefit is safe concurrent use of one subscription connection with newest-credential selection ## Linked Issues or Issue Description **What happened?** A managed OpenAI Codex or xAI Grok subscription connection held a credential lease for the full agent run. A second run waited for the first run to finish. The write-back gate also rejected any row change before it compared credential freshness. **Expected behavior** Concurrent runs should start on one subscription connection. The server should keep the newest valid credential and reject a credential write after a person revokes the connection. **Steps to reproduce** 1. Start two runs that use one OpenAI or xAI subscription connection. 2. Let both provider tools refresh the credential. 3. Finish the runs in either order. 4. Confirm that the newest valid credential remains in the connection. **Paperclip version or commit** `ec25bf1e4a81d1729a6d7276e6a486587cff4a0b` **Deployment mode** Local dev (`pnpm dev`), built from source. **Agent adapter(s) involved** Codex. The server path also covers xAI Grok subscription connections. **Database mode** Embedded PGlite for local development, and external Postgres for deployments. **Access context** Both board and agent runs can use managed connections. **Additional context** The provider command-line tool refreshes credentials inside the sandbox. The server copies the result back after the run. Two long runs can still refresh one token hours apart, so the provider can reject the second refresh. The server cannot observe that provider call. ## What Changed - Remove the full-run credential lease for OpenAI Codex and xAI Grok subscription connections. - Lock and re-read the connection row before credential write-back. - Accept only a strictly newer credential, while keeping the connection revocation guard. - Reject Codex freshness timestamps more than five minutes ahead of the host clock. - Add tests for both completion orders, xAI cleanup, revocation, the Codex time bound, and agent hiring. - Update the connection and run-log documentation. ## Verification - `server/src/__tests__/ai-connections.test.ts` passes with 42 tests. - `packages/adapters/codex-local/src/server/codex-auth-merge-decision.test.ts` covers the five-minute boundary and the one-millisecond overflow. - `packages/adapters/codex-local/src/server/codex-auth-merge.test.ts` passes. - `server/src/__tests__/agent-hire-ai-connections.test.ts` covers OpenAI and Anthropic. - The project type check reports no new error in changed files. - GitHub Actions must pass on this pull request. ## Risks The write-back now permits concurrent runs, so the provider may reject a later refresh when both long runs use one token. The server keeps the revocation guard and rejects future-dated Codex timestamps. No schema change occurs. ## Model Used OpenAI Codex, GPT-5. Context window and exact deployment build are not exposed in this run. The model used tool calls, code inspection, and test verification. ## Checklist - [x] I have included a thinking path that traces from project context to this change - [x] I have specified the model used (with version and capability details) - [x] I have checked ROADMAP.md and confirmed this PR does not duplicate planned core work - [x] I have searched GitHub for duplicate or related PRs and linked them above - [x] I have either (a) linked existing issues with `Fixes: #` / `Closes: #` / `Refs: #` OR (b) described the issue in-PR following the relevant issue template - [x] I have not referenced internal/instance-local Paperclip issues or links (only public GitHub `#NNN` / `github.com/paperclipai/paperclip` URLs) - [x] My branch name describes the change (e.g. `docs/...`, `fix/...`) and contains no internal Paperclip ticket id or instance-derived details - [x] I have run tests locally and they pass - [x] I have added or updated tests where applicable - [x] I have updated relevant documentation to reflect my changes - [x] I have considered and documented any risks above - [x] All Paperclip CI gates are green - [x] Greptile is 5/5 with no open P2s, recommendations, or follow-ups - [x] I will address all Greptile and reviewer comments before requesting merge --------- Co-authored-by: Paperclip --- doc/connections/AI-CONNECTIONS.md | 54 ++-- doc/run-log-events.md | 7 +- .../src/server/codex-auth-copyback.ts | 17 +- .../src/server/codex-auth-merge-decision.cjs | 137 +++++++--- .../server/codex-auth-merge-decision.test.ts | 114 ++++++++ .../src/server/codex-auth-merge-decision.ts | 29 +- .../src/server/codex-auth-merge.test.ts | 26 ++ .../agent-hire-ai-connections.test.ts | 44 +-- server/src/__tests__/ai-connections.test.ts | 171 ++++++++---- ...artbeat-ai-subscription-contention.test.ts | 250 ------------------ server/src/__tests__/secrets-service.test.ts | 75 +++++- server/src/services/ai-connection-runtime.ts | 96 ++----- server/src/services/heartbeat.ts | 7 +- server/src/services/secrets.ts | 41 +-- 14 files changed, 554 insertions(+), 514 deletions(-) delete mode 100644 server/src/__tests__/heartbeat-ai-subscription-contention.test.ts diff --git a/doc/connections/AI-CONNECTIONS.md b/doc/connections/AI-CONNECTIONS.md index 626baf504b9..6f9967f7237 100644 --- a/doc/connections/AI-CONNECTIONS.md +++ b/doc/connections/AI-CONNECTIONS.md @@ -108,26 +108,23 @@ grant's credentials. Inherited credential variables are cleared. Conflicting project authentication and provider-routing overrides are rejected. Managed failure cannot reactivate host or legacy credentials. -A subscription invocation takes a grant-scoped transaction advisory lease only -when it writes a provider authentication file back to the grant. OpenAI and xAI -subscriptions do this, because their refresh tokens are single use. An -Anthropic subscription invocation writes no file back, so it takes no lease; -two Anthropic invocations of one grant run at the same time. The reserved -database client keeps one transaction open until cleanup, including on -transaction-pooling proxies such as PgBouncer. Session-level advisory locks must -not be used here: a pooled connection can return to a different backend for -cleanup and leave the original lock behind. The lease transaction disables its -idle timeout and contains no application data writes; cleanup rolls it back. -Two -different users' grants can run concurrently; a second invocation of the same -file-backed subscription receives a retryable busy response while it is in -use. Refreshes are merged only into the originating active grant, with -reconnect/revocation version checks. Temporary homes are removed on normal -completion or failure. +A subscription invocation takes no lease. Two invocations of one grant, from +the same or a different provider account, run at the same time. At cleanup, +each invocation re-reads the credential stored at that moment under a row +lock on the grant, then compares it against its own refreshed copy using the +provider's own freshness field: Codex compares `last_refresh` and bounds it +against the host clock; Grok compares `expires_at`. The newer credential +persists; a tie or an unparseable freshness value keeps the stored +credential, so a spent single-use refresh token never overwrites a good one. +Refreshes are merged only into the originating active grant, with a +revocation check. Temporary homes are removed on normal completion or +failure. -For a fresh task execution, subscription contention creates a durable scheduled -retry checked every 60–120 seconds. The task shows “Waiting for AI subscription” -and does not request a reconnect or consume its provider-failure retry allowance. +A fresh task execution cannot enter subscription contention. The freshest-write +rule above resolves the conflict instead. A run that already entered this wait +keeps a durable scheduled retry, checked every 60–120 seconds. The task shows +“Waiting for AI subscription”. It does not request a reconnect, and it does not +consume its provider-failure retry allowance. Each attempt rechecks task eligibility, ownership, budget, and current credential access. Revocation and other configuration failures still require user action. Authorized comment wakes that started as non-assignee runs can resume without @@ -215,8 +212,8 @@ unmanaged legacy agents retain their existing authentication paths. `server/src/__tests__/ai-connections.test.ts` exercises storage, isolation, defaults, human audiences, agent access, reconnect races, refresh ownership, -subscription locking, migration replay, and redacted API failures against a real -embedded database. Existing login, adapter, tool, and channel suites cover their +concurrent subscription write-backs, migration replay, and redacted API +failures against a real embedded database. Existing login, adapter, tool, and channel suites cover their shared integration paths. The onboarding tests cover managed reuse and keeping a successfully connected account after failed agent creation. @@ -291,15 +288,18 @@ inside the task. Connecting installs access for that agent and resumes the pendi work automatically. Explicit incompatible bindings and shared-account permission denials still fail; hiring never expands a restricted shared account's audience. -Concurrent runs using the same subscription wait through scheduled retries while -the credential lease is held. They do not request new credentials or consume the -provider-failure retry allowance. Each retry revalidates the account and existing -run-dispatch rules still suppress cancelled, reassigned, or otherwise ineligible work. -Assignee retries must retain execution-lock ownership at scheduling, promotion, and dispatch. +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 +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 +retry must still keep execution-lock ownership at scheduling, promotion, and +dispatch. `server/src/__tests__/agent-hire-ai-connections.test.ts` covers both creation routes, both providers and methods, approval gates, native provider mapping, shared access -boundaries, and subscription contention. The opt-in +boundaries, and concurrent runs of one subscription for both providers. The opt-in [`tests/hiring-ai-connections/README.md`](../../tests/hiring-ai-connections/README.md) describes real browser hiring, subtask, connection, and automatic-resume checks on local and Daytona environments, plus the production component Storybook checks. diff --git a/doc/run-log-events.md b/doc/run-log-events.md index 6c47dd19321..6a992894295 100644 --- a/doc/run-log-events.md +++ b/doc/run-log-events.md @@ -179,9 +179,10 @@ identity checks remain in force. ## AI subscription contention -A run that waits for a managed subscription lease writes an informational -`lifecycle` event to the local run log. Its payload contains only -`retryScheduled`, a boolean that reports whether the scheduler created a retry. +A fresh task execution cannot enter this wait. A run that already entered this +wait writes an informational `lifecycle` event to the local run log. Its +payload contains only `retryScheduled`, a boolean that reports +whether the scheduler created a retry. The message distinguishes an automatic retry from work that is no longer eligible. This pre-provider wait records `ai_connection_busy` on the cancelled run and does not consume the provider-failure retry allowance. The event contains no credentials diff --git a/packages/adapters/codex-local/src/server/codex-auth-copyback.ts b/packages/adapters/codex-local/src/server/codex-auth-copyback.ts index 96d29f3e208..f22ae0369a9 100644 --- a/packages/adapters/codex-local/src/server/codex-auth-copyback.ts +++ b/packages/adapters/codex-local/src/server/codex-auth-copyback.ts @@ -7,7 +7,11 @@ import { readSubscriptionAccountId, writeCodexAuthCacheEntry, } from "./codex-auth-cache.js"; -import { USE_SOURCE_EXIT, decideCodexAuthMerge } from "./codex-auth-merge-decision.js"; +import { + IMPLAUSIBLE_LAST_REFRESH_EXIT, + USE_SOURCE_EXIT, + decideCodexAuthMerge, +} from "./codex-auth-merge-decision.js"; // The outbound copy-back reuses the exact same direction-agnostic decision // predicate the inbound restore runs, through the shared `decideCodexAuthMerge` @@ -122,6 +126,17 @@ export async function copyBackCodexAuth(input: CopyBackCodexAuthInput): Promise< return "copied"; } + // Make a clock-bound rejection visible. A host with a wrong clock would + // otherwise discard an honest refresh with no signal. Log only the exit + // code and this caller's own label — never a timestamp read from either + // file, and never credential bytes. + if (decision === IMPLAUSIBLE_LAST_REFRESH_EXIT) { + await log( + `[paperclip] Codex auth copy-back: WARNING host credential kept (decision exit ${IMPLAUSIBLE_LAST_REFRESH_EXIT}, codex auth copy-back) — the sandbox copy's last_refresh sat further ahead of the host clock than the plausible skew allowance. Check the host clock if this is unexpected.`, + ); + return "kept-host"; + } + await log( "[paperclip] Codex auth copy-back: host credential kept (sandbox copy is not a strictly-newer same-identity subscription credential).", ); diff --git a/packages/adapters/codex-local/src/server/codex-auth-merge-decision.cjs b/packages/adapters/codex-local/src/server/codex-auth-merge-decision.cjs index fe557795a20..01822d2e39a 100644 --- a/packages/adapters/codex-local/src/server/codex-auth-merge-decision.cjs +++ b/packages/adapters/codex-local/src/server/codex-auth-merge-decision.cjs @@ -64,54 +64,107 @@ function parseAuth(filePath) { // A leading positional flag (not an environment variable) keeps the mode // explicit per call, so a host-default two-path call can never enter seed mode. // -// Exit 10 = use source; exit 20 = keep destination. The predicate only ever -// reads the two files and exits with a code — it never prints token bytes. +// Exit contract. Exit 10 = use source; exit 20 = keep destination; exit 22 = +// keep destination, the source `last_refresh` sat further ahead of the host +// clock than the plausible skew allowance. The predicate only ever reads the +// two files and exits with a code — it never prints token bytes. const USE_SOURCE = 10; const KEEP_DESTINATION = 20; +const IMPLAUSIBLE_LAST_REFRESH = 22; const SEED_IF_DEST_ABSENT_FLAG = "--seed-if-dest-absent"; -const rawArgs = process.argv.slice(2); -const seedIfDestAbsent = rawArgs[0] === SEED_IF_DEST_ABSENT_FLAG; -const [sourceAuthPath, destinationAuthPath] = seedIfDestAbsent ? rawArgs.slice(1) : rawArgs; -const sourceAuth = parseAuth(sourceAuthPath); -const destinationAuth = parseAuth(destinationAuthPath); +// A `last_refresh` records an event that already happened, so an honest value +// always sits at or before the host clock. Only clock skew between a sandbox +// and the host explains a small future value. Five minutes is larger than +// real skew on a time-synchronised host, and it is short enough that a source +// which claims a `last_refresh` further ahead than this cannot be trusted. +const MAX_FUTURE_LAST_REFRESH_SKEW_MS = 5 * 60 * 1000; -// Seed mode only: fill an ABSENT (unusable) destination slot from a usable -// subscription source. A subscription-kind source is guaranteed usable and to -// carry a real account_id (parseAuth returns "subscription" only then), so this -// is never a random pick. This branch changes ONLY the destination-unusable -// case; the api-key and unusable-source guards below still keep the destination. -if ( - seedIfDestAbsent && - destinationAuth.kind === "unusable" && - sourceAuth.kind === "subscription" -) { - process.exit(USE_SOURCE); +// `decide` takes already-parsed `{ kind, accountId, lastRefresh }` shapes plus +// a caller-supplied `nowMs` and the `seedIfDestAbsent` flag, so a test can +// drive the exact skew bound without spawning a process and racing wall-clock +// drift. Guard order, first match wins: +// 1. Seed mode, the destination is unusable, and the source is a usable +// subscription credential -> USE_SOURCE, unless the source fails the +// skew bound in step 3 below, in which case the absent slot stays +// absent (IMPLAUSIBLE_LAST_REFRESH). A poisoned seed would out-live +// every honest refresh, since nothing could ever compare as "strictly +// newer" than an unbounded future value, so the bound applies here too. +// 2. Either side unusable, a kind mismatch, the destination is an api-key +// credential, or the two sides carry a different account_id -> +// KEEP_DESTINATION. +// 3. The source `last_refresh` sits further ahead of `nowMs` than +// MAX_FUTURE_LAST_REFRESH_SKEW_MS -> IMPLAUSIBLE_LAST_REFRESH. Only the +// source is bounded, and only against the caller's clock: the source is +// the sandbox-supplied side, so it must never supply the reference time. +// 4. The source `last_refresh` is strictly greater than the destination's +// -> USE_SOURCE. +// 5. Otherwise (a tie, a null value on either side, or an older source) -> +// KEEP_DESTINATION. +function decide(source, destination, nowMs, seedIfDestAbsent) { + const sourceIsImplausible = + source.lastRefresh !== null && + source.lastRefresh - nowMs > MAX_FUTURE_LAST_REFRESH_SKEW_MS; + + // Seed mode only: fill an ABSENT (unusable) destination slot from a usable + // subscription source. A subscription-kind source is guaranteed usable and to + // carry a real account_id (parseAuth returns "subscription" only then), so this + // is never a random pick. This branch changes ONLY the destination-unusable + // case; the api-key and unusable-source guards below still keep the destination. + if ( + seedIfDestAbsent && + destination.kind === "unusable" && + source.kind === "subscription" + ) { + return sourceIsImplausible ? IMPLAUSIBLE_LAST_REFRESH : USE_SOURCE; + } + + // Fail closed to the destination unless both sides are the same usable, + // subscription-kind identity — an unusable side, an api-key credential, a kind + // mismatch, or a different account_id all keep the destination copy. + if ( + destination.kind === "unusable" || + source.kind === "unusable" || + source.kind !== destination.kind || + destination.kind === "apikey" || + source.accountId !== destination.accountId + ) { + return KEEP_DESTINATION; + } + + if (sourceIsImplausible) { + return IMPLAUSIBLE_LAST_REFRESH; + } + + // Use the source credential only when it is strictly fresher: both sides must + // carry a parseable last_refresh and the source one must be strictly greater. + // Ties and null/unparseable freshness keep the destination copy so a spent + // single-use refresh token is never written over a good one. + if ( + source.lastRefresh !== null && + destination.lastRefresh !== null && + source.lastRefresh > destination.lastRefresh + ) { + return USE_SOURCE; + } + + return KEEP_DESTINATION; } -// Fail closed to the destination unless both sides are the same usable, -// subscription-kind identity — an unusable side, an api-key credential, a kind -// mismatch, or a different account_id all keep the destination copy. -if ( - destinationAuth.kind === "unusable" || - sourceAuth.kind === "unusable" || - sourceAuth.kind !== destinationAuth.kind || - destinationAuth.kind === "apikey" || - sourceAuth.accountId !== destinationAuth.accountId -) { - process.exit(KEEP_DESTINATION); +if (require.main === module) { + const rawArgs = process.argv.slice(2); + const seedIfDestAbsent = rawArgs[0] === SEED_IF_DEST_ABSENT_FLAG; + const [sourceAuthPath, destinationAuthPath] = seedIfDestAbsent ? rawArgs.slice(1) : rawArgs; + const sourceAuth = parseAuth(sourceAuthPath); + const destinationAuth = parseAuth(destinationAuthPath); + process.exit(decide(sourceAuth, destinationAuth, Date.now(), seedIfDestAbsent)); } -// Use the source credential only when it is strictly fresher: both sides must -// carry a parseable last_refresh and the source one must be strictly greater. -// Ties and null/unparseable freshness keep the destination copy so a spent -// single-use refresh token is never written over a good one. -if ( - sourceAuth.lastRefresh !== null && - destinationAuth.lastRefresh !== null && - sourceAuth.lastRefresh > destinationAuth.lastRefresh -) { - process.exit(USE_SOURCE); -} - -process.exit(KEEP_DESTINATION); +module.exports = { + decide, + parseAuth, + USE_SOURCE, + KEEP_DESTINATION, + IMPLAUSIBLE_LAST_REFRESH, + MAX_FUTURE_LAST_REFRESH_SKEW_MS, +}; diff --git a/packages/adapters/codex-local/src/server/codex-auth-merge-decision.test.ts b/packages/adapters/codex-local/src/server/codex-auth-merge-decision.test.ts index 5c05352cda9..40bc7728510 100644 --- a/packages/adapters/codex-local/src/server/codex-auth-merge-decision.test.ts +++ b/packages/adapters/codex-local/src/server/codex-auth-merge-decision.test.ts @@ -1,5 +1,6 @@ import { execFile as execFileCallback } from "node:child_process"; import { mkdtemp, rm, writeFile } from "node:fs/promises"; +import { createRequire } from "node:module"; import os from "node:os"; import path from "node:path"; import { fileURLToPath } from "node:url"; @@ -8,6 +9,28 @@ import { afterEach, describe, expect, it } from "vitest"; const execFile = promisify(execFileCallback); +// The exact-boundary cases below drive the predicate's pure `decide` function +// directly (via `require`, never spawning a process), so the skew bound can be +// tested to the exact millisecond without racing subprocess-spawn wall-clock +// drift. Every other case in this file drives the REAL `.cjs` through a +// spawned `node` process (no stub), matching how the wrapper invokes it in +// production. +const decisionModule = createRequire(import.meta.url)( + fileURLToPath(new URL("./codex-auth-merge-decision.cjs", import.meta.url)), +) as { + decide: ( + source: { kind: "subscription" | "apikey" | "unusable"; accountId?: string; lastRefresh: number | null }, + destination: { kind: "subscription" | "apikey" | "unusable"; accountId?: string; lastRefresh: number | null }, + nowMs: number, + seedIfDestAbsent?: boolean, + ) => number; + USE_SOURCE: number; + KEEP_DESTINATION: number; + IMPLAUSIBLE_LAST_REFRESH: number; + MAX_FUTURE_LAST_REFRESH_SKEW_MS: number; +}; +const { decide, IMPLAUSIBLE_LAST_REFRESH, MAX_FUTURE_LAST_REFRESH_SKEW_MS } = decisionModule; + // This suite pins the opt-in seed mode of the single decision predicate. The // default (no-flag) call keeps the fail-closed host-default contract unchanged. // The leading positional `--seed-if-dest-absent` flag adds one behaviour: fill @@ -165,3 +188,94 @@ describe("codex-auth-merge-decision predicate seed mode", () => { expect(seedCode).toBe(USE_SOURCE); }); }); + +// This suite pins the host-clock bound on the source `last_refresh`. A +// `last_refresh` records an event that already happened, so an honest value +// never sits far ahead of the host clock; only clock skew explains a small +// future value. A source that claims a value further ahead than +// `MAX_FUTURE_LAST_REFRESH_SKEW_MS` cannot be trusted, since a same-account +// sandbox that controls its own `auth.json` could otherwise pin an unbounded +// future timestamp that every later honest refresh compares as older than. +// These cases drive `decide` directly with an explicit `nowMs`, so the bound +// is proven to the exact millisecond and never depends on the real clock. +describe("codex-auth-merge-decision predicate: host-clock bound on last_refresh", () => { + const USE_SOURCE = 10; + const KEEP_DESTINATION = 20; + + function subscription(accountId: string, lastRefresh: number | null) { + return { kind: "subscription" as const, accountId, lastRefresh }; + } + + it("keeps the destination when the source last_refresh sits one millisecond beyond the bound", () => { + const nowMs = Date.now(); + const source = subscription("acct", nowMs + MAX_FUTURE_LAST_REFRESH_SKEW_MS + 1); + const destination = subscription("acct", nowMs - 60_000); + expect(decide(source, destination, nowMs, false)).toBe(IMPLAUSIBLE_LAST_REFRESH); + }); + + it("uses the source when the source last_refresh sits exactly at the bound", () => { + const nowMs = Date.now(); + const source = subscription("acct", nowMs + MAX_FUTURE_LAST_REFRESH_SKEW_MS); + const destination = subscription("acct", nowMs - 60_000); + expect(decide(source, destination, nowMs, false)).toBe(USE_SOURCE); + }); + + it("still uses the source for an ordinary same-account, strictly-newer, past timestamp", () => { + const nowMs = Date.now(); + const source = subscription("acct", nowMs - 60_000); + const destination = subscription("acct", nowMs - 120_000); + expect(decide(source, destination, nowMs, false)).toBe(USE_SOURCE); + }); + + it("measures the bound against the host clock and not against a value embedded in either payload", () => { + // Fixed instants; only the caller-supplied `nowMs` moves between the two + // assertions, so a passing/failing bound can only be explained by the + // caller's clock, never by a value embedded in either payload. + const sourceLastRefresh = 2_000_000_000_000; + const destinationLastRefresh = 1_000_000_000_000; + const source = subscription("acct", sourceLastRefresh); + const destination = subscription("acct", destinationLastRefresh); + + expect(decide(source, destination, sourceLastRefresh - MAX_FUTURE_LAST_REFRESH_SKEW_MS, false)).toBe(USE_SOURCE); + expect(decide(source, destination, sourceLastRefresh - MAX_FUTURE_LAST_REFRESH_SKEW_MS - 1, false)).toBe( + IMPLAUSIBLE_LAST_REFRESH, + ); + }); + + it("keeps the destination for a same-identity source whose kind or account guard would otherwise fire, regardless of the bound", () => { + // A kind mismatch or a different account_id already keeps the destination + // before the bound runs, so an implausible source never surfaces as + // IMPLAUSIBLE_LAST_REFRESH when a different guard already applies. + const nowMs = Date.now(); + const implausibleLastRefresh = nowMs + MAX_FUTURE_LAST_REFRESH_SKEW_MS + 1; + const differentAccount = decide( + subscription("acct-x", implausibleLastRefresh), + subscription("acct-y", nowMs - 60_000), + nowMs, + false, + ); + expect(differentAccount).toBe(KEEP_DESTINATION); + }); + + // Seed mode fills an ABSENT destination slot from a usable subscription + // source. The bound applies there too: an absent destination that got + // seeded with an implausible future last_refresh would never be + // refreshable again, since no later honest refresh could ever compare as + // "strictly newer" than an unbounded future value. The seed-mode source in + // production is the sandbox's own `auth.json` + // (see `writeCodexAuthCacheEntry` in codex-auth-cache.ts), so this is the + // same untrusted input the write-back guard defends against. + it("seed mode keeps an absent destination slot absent when the source last_refresh is implausibly far in the future", () => { + const nowMs = Date.now(); + const source = subscription("acct", nowMs + MAX_FUTURE_LAST_REFRESH_SKEW_MS + 1); + const destination = { kind: "unusable" as const, lastRefresh: null }; + expect(decide(source, destination, nowMs, true)).toBe(IMPLAUSIBLE_LAST_REFRESH); + }); + + it("seed mode still fills an absent destination slot when the source last_refresh sits at or before the bound", () => { + const nowMs = Date.now(); + const source = subscription("acct", nowMs + MAX_FUTURE_LAST_REFRESH_SKEW_MS); + const destination = { kind: "unusable" as const, lastRefresh: null }; + expect(decide(source, destination, nowMs, true)).toBe(USE_SOURCE); + }); +}); diff --git a/packages/adapters/codex-local/src/server/codex-auth-merge-decision.ts b/packages/adapters/codex-local/src/server/codex-auth-merge-decision.ts index 0ecd228a732..bbcfddfcce5 100644 --- a/packages/adapters/codex-local/src/server/codex-auth-merge-decision.ts +++ b/packages/adapters/codex-local/src/server/codex-auth-merge-decision.ts @@ -11,10 +11,12 @@ const execFile = promisify(execFileCallback); // `codex-auth-merge-decision.cjs`. It reads only the two files, and it exits with // a code; it never prints token bytes. Argument order sets source and // destination (first = source, second = destination), so the caller frames the -// direction. Exit 10 = use source; exit 20 = keep destination. The leading -// `--seed-if-dest-absent` flag adds one behavior: fill an absent destination slot -// from a usable subscription source. This module gives every caller one shared -// entry point, so the predicate contract can never drift between callers. +// direction. Exit 10 = use source; exit 20 = keep destination; exit 22 = keep +// destination, the source `last_refresh` sat further ahead of the host clock +// than the plausible skew allowance. The leading `--seed-if-dest-absent` flag +// adds one behavior: fill an absent destination slot from a usable +// subscription source. This module gives every caller one shared entry point, +// so the predicate contract can never drift between callers. const DECISION_SCRIPT_PATH = fileURLToPath( new URL("./codex-auth-merge-decision.cjs", import.meta.url), @@ -25,6 +27,10 @@ const SEED_IF_DEST_ABSENT_FLAG = "--seed-if-dest-absent"; export const USE_SOURCE_EXIT = 10; /** Exit code: keep the destination credential. */ export const KEEP_DESTINATION_EXIT = 20; +/** Exit code: keep the destination; the source `last_refresh` was implausibly far in the future. */ +export const IMPLAUSIBLE_LAST_REFRESH_EXIT = 22; + +const KNOWN_EXIT_CODES = new Set([USE_SOURCE_EXIT, KEEP_DESTINATION_EXIT, IMPLAUSIBLE_LAST_REFRESH_EXIT]); export interface DecideCodexAuthMergeOptions { /** Opt in to the cache-slot seed mode: fill an absent destination from a @@ -36,9 +42,10 @@ export interface DecideCodexAuthMergeOptions { } /** - * Runs the shared decision predicate and returns its exit code (10 or 20). A - * non-10/20 exit, or a failure to run `node`, is a hard failure: this throws so a - * broken predicate is never mistaken for a "keep destination" decision. + * Runs the shared decision predicate and returns its exit code (10, 20, or + * 22). Any other exit, or a failure to run `node`, is a hard failure: this + * throws so a broken predicate is never mistaken for a "keep destination" + * decision. */ export async function decideCodexAuthMerge( sourcePath: string, @@ -52,7 +59,7 @@ export async function decideCodexAuthMerge( await execFile("node", args); } catch (error) { const code = (error as { code?: unknown }).code; - if (code === USE_SOURCE_EXIT || code === KEEP_DESTINATION_EXIT) { + if (typeof code === "number" && KNOWN_EXIT_CODES.has(code)) { return code; } const detail = @@ -65,9 +72,9 @@ export async function decideCodexAuthMerge( : String(error); throw new Error(`${options.errorLabel} decision predicate failed: ${detail}`); } - // `execFile` resolved, so the predicate exited 0. The predicate always exits 10 - // or 20, so a clean exit 0 is unexpected; fail loud. + // `execFile` resolved, so the predicate exited 0. The predicate always exits + // 10, 20, or 22, so a clean exit 0 is unexpected; fail loud. throw new Error( - `${options.errorLabel} decision predicate exited 0 (expected 10 or 20)`, + `${options.errorLabel} decision predicate exited 0 (expected 10, 20, or 22)`, ); } diff --git a/packages/adapters/codex-local/src/server/codex-auth-merge.test.ts b/packages/adapters/codex-local/src/server/codex-auth-merge.test.ts index 502c069e94b..3d466667930 100644 --- a/packages/adapters/codex-local/src/server/codex-auth-merge.test.ts +++ b/packages/adapters/codex-local/src/server/codex-auth-merge.test.ts @@ -596,6 +596,7 @@ describe("codex-auth-merge-decision predicate (source/destination)", () => { const KEEP_DESTINATION = 20; const USE_SOURCE = 10; + const IMPLAUSIBLE_LAST_REFRESH = 22; function subscriptionAuth(input: { accountId: string; @@ -737,6 +738,18 @@ describe("codex-auth-merge-decision predicate (source/destination)", () => { destinationAuth: "{not valid json", expected: KEEP_DESTINATION, }, + { + // 400 days ahead is comfortably beyond the 5-minute skew allowance and + // any subprocess-spawn scheduling delay, so this integration-level + // check never depends on millisecond timing. + name: "source last_refresh implausibly far in the future → keep destination", + sourceAuth: subscriptionAuth({ + accountId: "acct", + lastRefresh: new Date(Date.now() + 400 * 24 * 60 * 60 * 1000).toISOString(), + }), + destinationAuth: subscriptionAuth({ accountId: "acct", lastRefresh: OLDER }), + expected: IMPLAUSIBLE_LAST_REFRESH, + }, ]; for (const entry of cases) { @@ -761,4 +774,17 @@ describe("codex-auth-merge-decision predicate (source/destination)", () => { expect(result.code).toBe(USE_SOURCE); expect(result.output).not.toContain("SENTINEL"); }); + + it("never emits source token bytes for an implausibly-future last_refresh", async () => { + const result = await runDecision({ + sourceAuth: subscriptionAuth({ + accountId: "acct", + lastRefresh: new Date(Date.now() + 400 * 24 * 60 * 60 * 1000).toISOString(), + marker: "SECRET-SENTINEL", + }), + destinationAuth: subscriptionAuth({ accountId: "acct", lastRefresh: OLDER }), + }); + expect(result.code).toBe(IMPLAUSIBLE_LAST_REFRESH); + expect(result.output).not.toContain("SENTINEL"); + }); }); diff --git a/server/src/__tests__/agent-hire-ai-connections.test.ts b/server/src/__tests__/agent-hire-ai-connections.test.ts index c7645acc3c9..db0533ad650 100644 --- a/server/src/__tests__/agent-hire-ai-connections.test.ts +++ b/server/src/__tests__/agent-hire-ai-connections.test.ts @@ -6,7 +6,7 @@ import express from "express"; import request from "supertest"; import { and, eq } from "drizzle-orm"; import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; -import { agents, companies, companyMemberships, createDb, heartbeatRuns, issues, issueThreadInteractions, principalPermissionGrants, toolConnectionInstalls } from "@paperclipai/db"; +import { agents, companies, companyMemberships, createDb, heartbeatRuns, issues, principalPermissionGrants, toolConnectionInstalls } from "@paperclipai/db"; import { type AiConnectionBinding } from "@paperclipai/shared"; import { startEmbeddedPostgresTestDatabase } from "./helpers/embedded-postgres.js"; import { agentRoutes } from "../routes/agents.js"; @@ -202,47 +202,7 @@ describe("agent-created hires use managed AI connections", () => { }); describe("hired agents sharing a subscription", () => { - it("waits for the openai parent and resumes without asking for a new connection", async () => { - const provider = "openai"; - const f = await fixture(provider, "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: { cwd: home, engine: "cli" }, runtimeConfig: { heartbeat: { enabled: false } } })); - const [issue] = await db.insert(issues).values({ companyId: f.companyId, title: "Subscription child task", status: "todo", assigneeAgentId: agent.id, responsibleUserId: f.userId, createdByUserId: f.userId }).returning(); - const parentRuntime = await prepareManagedAiRuntime(db, { companyId: f.companyId, agentId: f.agentId, responsibleUserId: f.userId, adapterType: f.adapterType, binding: f.binding, config: { cwd: home } }); - const execute = vi.fn(async () => { - await db.update(issues).set({ status: "done", completedAt: new Date() }).where(eq(issues.id, issue.id)); - return { exitCode: 0, signal: null, timedOut: false, resultJson: {} }; - }); - registerServerAdapter({ ...getServerAdapter(f.adapterType), execute }); - const heartbeat = heartbeatService(db); - let parentReleased = false; - 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 }).not.toBe("running"); - expect(await heartbeat.getRun(run!.id)).toMatchObject({ status: "cancelled", errorCode: "ai_connection_busy" }); - let retry: typeof heartbeatRuns.$inferSelect | undefined; - await expect.poll(async () => { [retry] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.retryOfRunId, run!.id)); return retry?.status; }, { timeout: 10_000 }).toBe("scheduled_retry"); - expect(execute).not.toHaveBeenCalled(); - expect(await db.select().from(issueThreadInteractions).where(eq(issueThreadInteractions.issueId, issue.id))).toHaveLength(0); - expect((await db.select().from(issues).where(eq(issues.id, issue.id)))[0].executionRunId).toBe(retry!.id); - parentReleased = true; - await parentRuntime.cleanup(); - const promoted = await heartbeat.promoteDueScheduledRetries(new Date(retry!.scheduledRetryAt!.getTime() + 1000)); - expect(promoted.runIds).toContain(retry!.id); - await heartbeat.resumeQueuedRuns(); - await expect.poll(async () => (await heartbeat.getRun(retry!.id))?.status, { timeout: 20_000 }).toBe("succeeded"); - expect(execute).toHaveBeenCalledTimes(1); - expect((await heartbeat.getRun(retry!.id))?.contextSnapshot?.aiConnection).toMatchObject({ connectionId: f.account.connectionId, responsibleUserId: f.userId, method: "subscription" }); - } finally { - if (!parentReleased) await parentRuntime.cleanup(); - await db.update(heartbeatRuns).set({ status: "cancelled", finishedAt: new Date() }).where(eq(heartbeatRuns.id, f.runId)); - await heartbeat.drainActiveRunExecutions(); - unregisterServerAdapter(f.adapterType); - } - }); - - it("runs the anthropic child alongside a live parent and inherits its connection", async () => { - const provider = "anthropic"; + 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: { cwd: home, engine: "cli" }, runtimeConfig: { heartbeat: { enabled: false } } })); const [issue] = await db.insert(issues).values({ companyId: f.companyId, title: "Subscription child task", status: "todo", assigneeAgentId: agent.id, responsibleUserId: f.userId, createdByUserId: f.userId }).returning(); diff --git a/server/src/__tests__/ai-connections.test.ts b/server/src/__tests__/ai-connections.test.ts index c36e35ef589..f9d836f8f4c 100644 --- a/server/src/__tests__/ai-connections.test.ts +++ b/server/src/__tests__/ai-connections.test.ts @@ -8,7 +8,7 @@ import { mkdtemp, rm, access, readFile, writeFile } 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, heartbeatRuns, companyMemberships, connectionGrants, connectionGrantDelegations, connectionGrantMembers, toolConnections, toolConnectionInstalls, aiConnectionDefaults, aiProviderDefaults, adapterAuthSessions, environments, issues, issueThreadInteractions, issueRecoveryActions, connectionIntentDeliveries, agentWakeupRequests } from "@paperclipai/db"; +import { createDb, companies, agents, heartbeatRuns, companyMemberships, connectionGrants, connectionGrantDelegations, connectionGrantMembers, toolConnections, toolConnectionInstalls, aiConnectionDefaults, aiProviderDefaults, adapterAuthSessions, environments, issues, issueThreadInteractions, issueRecoveryActions, connectionIntentDeliveries, agentWakeupRequests, companySecrets } from "@paperclipai/db"; import { startEmbeddedPostgresTestDatabase } from "@paperclipai/db/test-embedded-postgres"; import { aiConnectionService } from "../services/ai-connections.js"; import * as executionTarget from "@paperclipai/adapter-utils/execution-target"; @@ -282,36 +282,6 @@ describe("managed AI connections", () => { const [agent] = await db.select().from(agents).where(eq(agents.id, agentId)); expect(agent.runtimeConfig.aiConnection).toBeUndefined(); }); - it("serializes OpenAI subscription refresh and releases the lease after execution", async () => { - // OpenAI writes its rotated refresh token back to a shared auth file, so - // two runs against the same grant must not overlap. - const userId = "openai-lease-user"; - await db.insert(companyMemberships).values({ companyId, principalId: userId, principalType: "user", status: "active", membershipRole: "member" }); - const token = JSON.stringify({ tokens: { access_token: "fixture-lease-access", refresh_token: "fixture-lease-refresh", id_token: "fixture-lease-id", account_id: "fixture-lease-account" } }); - await service.save(companyId, userId, { provider: "openai", method: "subscription", ownership: "personal", name: "Lease subscription", loginSessionId: "fixture", allAgents: true, agentIds: [] }, token); - const subscription = { ...input, adapterType: "codex_local", binding: { provider: "openai", method: "subscription", mode: "responsible_user" } as const, responsibleUserId: userId, config: { model: "same-model" } }; - const first = await prepareManagedAiRuntime(db, subscription); - const selected = await service.select({ ...subscription, userId }); - const lockKey = `ai-runtime:${selected.grant.id}`; - const held = await db.execute(sql` - select activity.state, activity.xact_start - from pg_locks locks join pg_stat_activity activity on activity.pid = locks.pid - where locks.locktype = 'advisory' and locks.granted - and locks.classid = (hashtextextended(${lockKey}, 0) >> 32)::int::oid - and locks.objid = (hashtextextended(${lockKey}, 0) & 4294967295)::oid - and locks.objsubid = 1 - `); - // A transaction-pooling proxy may move an idle, unpinned client to a - // different backend. The lock must hold a transaction for its lifetime. - expect(held).toHaveLength(1); - expect(held[0].state).toBe("idle in transaction"); - expect(held[0].xact_start).not.toBeNull(); - await expect(prepareManagedAiRuntime(db, subscription)).rejects.toThrow("in use"); - await first.cleanup(); - const next = await prepareManagedAiRuntime(db, subscription); - expect(next.identity).toBe(first.identity); - await next.cleanup(); - }); it("runs two Claude subscription executions for the same grant at the same time", async () => { // Claude writes no auth file back to the grant, so two runs share no // mutable state and must not wait for each other. @@ -325,7 +295,7 @@ describe("managed AI connections", () => { await Promise.all([first.cleanup(), second.cleanup()]); } }); - it("reproduces same-agent OpenAI subscription contention and resumes without reconnecting", async () => { + it("runs a same-agent OpenAI subscription child alongside a still-open parent", async () => { const userId = "subscription-contention-user"; await db.insert(companyMemberships).values({ companyId, principalId: userId, principalType: "user", status: "active", membershipRole: "member" }); const credential = JSON.stringify({ tokens: { access_token: "fixture-access", refresh_token: "fixture-refresh", id_token: "fixture-id", account_id: "fixture-account" } }); @@ -341,29 +311,15 @@ describe("managed AI connections", () => { // A parent can create and assign a child before its own execution ends. // Both runs select the same personal subscription, even on the same agent. const parent = await prepareManagedAiRuntime(db, runInput); - try { - await expect(prepareManagedAiRuntime(db, runInput)).rejects.toMatchObject({ - status: 422, - message: "This subscription is in use. Retry when its current execution finishes.", - details: { code: "ai_connection_busy" }, - }); - const selected = await service.select({ ...runInput, userId }); - expect(selected.grant).toMatchObject({ id: account.grantId, status: "active" }); - expect(await service.credential(selected)).toBe(credential); - } finally { - await parent.cleanup(); - } - // Releasing the parent alone is sufficient: no reconnect, credential - // rotation, account switch, or agent configuration change is required. const child = await prepareManagedAiRuntime(db, runInput); try { expect(child.identity).toBe(parent.identity); expect(child.attribution.grantId).toBe(account.grantId); } finally { - await child.cleanup(); + await Promise.all([parent.cleanup(), child.cleanup()]); } }); - it("persists refreshed credentials only to their original grant and fences reconnects", async () => { + it("persists the freshest refreshed credential to its original grant across a same-account reconnect", async () => { const auth = (marker: string, hour: number) => JSON.stringify({ tokens: { account_id: "fixture-account", id_token: `id-${marker}`, access_token: `access-${marker}`, refresh_token: `refresh-${marker}` }, last_refresh: `2026-09-10T${hour}:00:00Z` }); const intent = { provider: "openai" as const, method: "subscription" as const, name: "Refresh test", ownership: "personal" as const, agentIds: [], allAgents: true, loginSessionId: "fixture" }; const saved = await service.save(companyId, "alice", intent, auth("first", 10)); @@ -375,10 +331,125 @@ describe("managed AI connections", () => { expect(await service.credential(selected)).toBe(auth("refreshed", 11)); const second = await prepareManagedAiRuntime(db, runInput); expect(second.identity).not.toBe(first.identity); + // A same-account reconnect writes an older last_refresh than the run + // that is still open. await service.save(companyId, "alice", { ...intent, connectionId: saved.connectionId }, auth("reconnect", 12)); - await writeFile(path.join(String(second.config.env.CODEX_HOME), "auth.json"), auth("stale-process", 13)); + await writeFile(path.join(String(second.config.env.CODEX_HOME), "auth.json"), auth("later-refresh", 13)); await second.cleanup(); - expect(await service.credential(await service.select({ ...runInput, userId: "alice" }))).toBe(auth("reconnect", 12)); + // The newer refresh persists to the grant it started from. + expect(await service.credential(await service.select({ ...runInput, userId: "alice" }))).toBe(auth("later-refresh", 13)); + }); + it("resolves two concurrent OpenAI subscription write-backs by freshness, not by order", async () => { + const userId = "concurrent-freshness-user"; + await db.insert(companyMemberships).values({ companyId, principalId: userId, principalType: "user", status: "active", membershipRole: "member" }); + const auth = (marker: string, hour: number) => JSON.stringify({ tokens: { account_id: "fixture-account", id_token: `id-${marker}`, access_token: `access-${marker}`, refresh_token: `refresh-${marker}` }, last_refresh: `2026-09-10T${hour}:00:00Z` }); + await service.save(companyId, userId, { provider: "openai", method: "subscription", ownership: "personal", name: "Freshness fixture", loginSessionId: "fixture", allAgents: true, agentIds: [] }, auth("start", 10)); + const runInput = { ...input, adapterType: "codex_local", responsibleUserId: userId, binding: { provider: "openai", method: "subscription", mode: "responsible_user" } as const, config: { model: "same-model" } }; + // Two runs use the same OpenAI subscription grant at the same time. + // Neither call below throws ai_connection_busy. + const older = await prepareManagedAiRuntime(db, runInput); + const newer = await prepareManagedAiRuntime(db, runInput); + await writeFile(path.join(String(older.config.env.CODEX_HOME), "auth.json"), auth("older", 11)); + await writeFile(path.join(String(newer.config.env.CODEX_HOME), "auth.json"), auth("newer", 12)); + // The run with the newer last_refresh writes back first. The run with + // the older last_refresh writes back last and must not overwrite it. + await newer.cleanup(); + await older.cleanup(); + const stored = await service.credential(await service.select({ ...runInput, userId })); + expect(stored).toBe(auth("newer", 12)); + }); + it("resolves two concurrent OpenAI subscription write-backs by freshness in reverse arrival order", async () => { + const userId = "concurrent-freshness-reverse-user"; + await db.insert(companyMemberships).values({ companyId, principalId: userId, principalType: "user", status: "active", membershipRole: "member" }); + const auth = (marker: string, hour: number) => JSON.stringify({ tokens: { account_id: "fixture-account", id_token: `id-${marker}`, access_token: `access-${marker}`, refresh_token: `refresh-${marker}` }, last_refresh: `2026-09-10T${hour}:00:00Z` }); + await service.save(companyId, userId, { provider: "openai", method: "subscription", ownership: "personal", name: "Reverse freshness fixture", loginSessionId: "fixture", allAgents: true, agentIds: [] }, auth("start", 10)); + const runInput = { ...input, adapterType: "codex_local", responsibleUserId: userId, binding: { provider: "openai", method: "subscription", mode: "responsible_user" } as const, config: { model: "same-model" } }; + // Two runs use the same OpenAI subscription grant at the same time. + // Neither call below throws ai_connection_busy. + const older = await prepareManagedAiRuntime(db, runInput); + const newer = await prepareManagedAiRuntime(db, runInput); + await writeFile(path.join(String(older.config.env.CODEX_HOME), "auth.json"), auth("older", 11)); + await writeFile(path.join(String(newer.config.env.CODEX_HOME), "auth.json"), auth("newer", 12)); + // The run with the older last_refresh writes back first. The run with + // the newer last_refresh writes back last and must win. + await older.cleanup(); + await newer.cleanup(); + const stored = await service.credential(await service.select({ ...runInput, userId })); + expect(stored).toBe(auth("newer", 12)); + }); + it("resolves two concurrent xAI subscription write-backs by freshness, not by order", async () => { + const userId = "concurrent-freshness-xai-user"; + await db.insert(companyMemberships).values({ companyId, principalId: userId, principalType: "user", status: "active", membershipRole: "member" }); + const identityKey = "https://auth.x.ai::33333333-3333-3333-3333-333333333333"; + const auth = (marker: string, expiresAtMs: number) => JSON.stringify({ [identityKey]: { key: `key-${marker}`, refresh_token: `refresh-${marker}`, expires_at: new Date(expiresAtMs).toISOString() } }); + const now = Date.now(); + await service.save(companyId, userId, { provider: "xai", method: "subscription", ownership: "personal", name: "Grok freshness fixture", loginSessionId: "fixture", allAgents: true, agentIds: [] }, auth("start", now)); + const runInput = { ...input, adapterType: "grok_local", responsibleUserId: userId, binding: { provider: "xai", method: "subscription", mode: "responsible_user" } as const, config: { model: "same-model" } }; + // Two runs use the same xAI subscription grant at the same time. + // Neither call below throws ai_connection_busy. + const older = await prepareManagedAiRuntime(db, runInput); + const newer = await prepareManagedAiRuntime(db, runInput); + await writeFile(path.join(String(older.config.env.GROK_HOME), "auth.json"), auth("older", now + 60 * 60 * 1000)); + await writeFile(path.join(String(newer.config.env.GROK_HOME), "auth.json"), auth("newer", now + 2 * 60 * 60 * 1000)); + // The run with the older expiry writes back first. The run with the + // newer expiry writes back last and must win. + await older.cleanup(); + await newer.cleanup(); + const stored = await service.credential(await service.select({ ...runInput, userId })); + expect(stored).toBe(auth("newer", now + 2 * 60 * 60 * 1000)); + }); + it("discards a credential write-back when the grant is revoked while the run is open", async () => { + const userId = "revoked-write-back-user"; + await db.insert(companyMemberships).values({ companyId, principalId: userId, principalType: "user", status: "active", membershipRole: "member" }); + const auth = (marker: string, hour: number) => JSON.stringify({ tokens: { account_id: "fixture-account", id_token: `id-${marker}`, access_token: `access-${marker}`, refresh_token: `refresh-${marker}` }, last_refresh: `2026-09-10T${hour}:00:00Z` }); + const saved = await service.save(companyId, userId, { provider: "openai", method: "subscription", ownership: "personal", name: "Revocation fixture", loginSessionId: "fixture", allAgents: true, agentIds: [] }, auth("start", 10)); + const runInput = { ...input, adapterType: "codex_local", responsibleUserId: userId, binding: { provider: "openai", method: "subscription", mode: "responsible_user" } as const, config: { model: "same-model" } }; + const run = await prepareManagedAiRuntime(db, runInput); + const [grantBeforeCleanup] = await db.select().from(connectionGrants).where(eq(connectionGrants.id, saved.grantId)); + const ref = grantBeforeCleanup.credentialSecretRefs.find(r => r.configPath === "ai.credential")!; + const [secretBefore] = await db.select().from(companySecrets).where(eq(companySecrets.id, ref.secretId)); + // A newer last_refresh would win the freshness merge if the grant stayed + // active. The revoked grant must discard the write-back before that merge + // decides anything. + await writeFile(path.join(String(run.config.env.CODEX_HOME), "auth.json"), auth("revoked-run", 11)); + await db.update(connectionGrants).set({ status: "revoked" }).where(eq(connectionGrants.id, saved.grantId)); + await run.cleanup(); + const [secretAfter] = await db.select().from(companySecrets).where(eq(companySecrets.id, ref.secretId)); + // service.select rejects a revoked grant, so it cannot read the stored + // credential here. Compare the stored secret version directly instead. + expect(secretAfter.latestVersion).toBe(secretBefore.latestVersion); + }); + it("does not let a stale write-back overwrite an authorized secret rotation that commits while cleanup waits on the credential lock", async () => { + const userId = "credential-lock-race-user"; + await db.insert(companyMemberships).values({ companyId, principalId: userId, principalType: "user", status: "active", membershipRole: "member" }); + const auth = (marker: string, hour: number) => JSON.stringify({ tokens: { account_id: "fixture-account", id_token: `id-${marker}`, access_token: `access-${marker}`, refresh_token: `refresh-${marker}` }, last_refresh: `2026-09-10T${hour}:00:00Z` }); + const saved = await service.save(companyId, userId, { provider: "openai", method: "subscription", ownership: "personal", name: "Credential lock race fixture", loginSessionId: "fixture", allAgents: true, agentIds: [] }, auth("start", 10)); + const runInput = { ...input, adapterType: "codex_local", responsibleUserId: userId, binding: { provider: "openai", method: "subscription", mode: "responsible_user" } as const, config: { model: "same-model" } }; + const run = await prepareManagedAiRuntime(db, runInput); + // The run's own refresh looks newer than the value it started with, but + // it must lose to a company-authorized rotation that commits while + // cleanup is still waiting on the credential secret's row lock. + await writeFile(path.join(String(run.config.env.CODEX_HOME), "auth.json"), auth("run-refresh", 11)); + const [grant] = await db.select().from(connectionGrants).where(eq(connectionGrants.id, saved.grantId)); + const ref = grant.credentialSecretRefs.find(r => r.configPath === "ai.credential")!; + let holdAcquired!: () => void; + const holdAcquiredPromise = new Promise(resolve => { holdAcquired = resolve; }); + let releaseHold!: () => void; + const holdReleased = new Promise(resolve => { releaseHold = resolve; }); + // An authorized rotation writes the new credential inside its own open + // transaction, so it still holds the secret row's lock when signaled. + const holder = db.transaction(async tx => { + await secretService(tx).rotate(ref.secretId, { value: auth("authorized-rotation", 12) }, { userId }); + holdAcquired(); + await holdReleased; + }); + await holdAcquiredPromise; + const cleanupPromise = run.cleanup(); + releaseHold(); + await holder; + await cleanupPromise; + const stored = await service.credential(await service.select({ ...runInput, userId })); + expect(stored).toBe(auth("authorized-rotation", 12)); }); it("enforces the shared transport discriminator and existing harness compatibility", () => { expect(connectionPurposeTransportSchema.safeParse({ connectionPurpose: "ai", transport: "mcp_remote" }).success).toBe(false); diff --git a/server/src/__tests__/heartbeat-ai-subscription-contention.test.ts b/server/src/__tests__/heartbeat-ai-subscription-contention.test.ts deleted file mode 100644 index 570506083f5..00000000000 --- a/server/src/__tests__/heartbeat-ai-subscription-contention.test.ts +++ /dev/null @@ -1,250 +0,0 @@ -import { randomUUID } from "node:crypto"; -import { mkdtemp, rm } from "node:fs/promises"; -import os from "node:os"; -import path from "node:path"; -import { eq } from "drizzle-orm"; -import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; -import { - agents, companies, companyMemberships, connectionGrants, createDb, heartbeatRuns, - issueComments, issueRecoveryActions, issueThreadInteractions, issues, -} from "@paperclipai/db"; -import { startEmbeddedPostgresTestDatabase } from "@paperclipai/db/test-embedded-postgres"; -import { heartbeatService } from "../services/heartbeat.js"; -import { aiConnectionService } from "../services/ai-connections.js"; -import { prepareManagedAiRuntime } from "../services/ai-connection-runtime.js"; -import { executionFailureRetryCount } from "../services/execution-recovery-attempt.js"; -import { executionProjectionsForRuns } from "../services/execution-projection.js"; -import type { AdapterExecutionContext } from "../adapters/index.js"; - -const execute = vi.hoisted(() => vi.fn(async (_input: AdapterExecutionContext) => ({ - exitCode: 0, signal: null, timedOut: false, resultJson: {}, summary: "Fixture work completed.", -}))); -vi.mock("../adapters/index.js", async () => ({ - ...await vi.importActual("../adapters/index.js"), - getServerAdapter: () => ({ supportsLocalAgentJwt: false, execute }), -})); - -const afterCheckout = vi.hoisted(() => vi.fn(async () => {})); -vi.mock("../services/issues.js", async () => { - const actual = await vi.importActual("../services/issues.js"); - return { ...actual, issueService: (...args: Parameters) => { - const service = actual.issueService(...args); - return { ...service, checkout: async (...checkoutArgs: Parameters) => { - const result = await service.checkout(...checkoutArgs); - await afterCheckout(); - return result; - } }; - } }; -}); - -describe("heartbeat AI subscription contention", () => { - let database: Awaited>; - let db: ReturnType; - let heartbeat: ReturnType; - let home: string; - - beforeAll(async () => { - home = await mkdtemp(path.join(os.tmpdir(), "paperclip-subscription-contention-")); - vi.stubEnv("PAPERCLIP_HOME", home); - vi.stubEnv("PAPERCLIP_INSTANCE_ID", "subscription-contention-tests"); - database = await startEmbeddedPostgresTestDatabase("paperclip-subscription-contention-db-"); - db = createDb(database.connectionString); - heartbeat = heartbeatService(db); - execute.mockImplementation(async (input) => { - await db.update(issues).set({ status: "done" }).where(eq(issues.id, String(input.context.issueId))); - return { exitCode: 0, signal: null, timedOut: false, resultJson: {}, summary: "Fixture work completed." }; - }); - }, 90_000); - afterAll(async () => { - await heartbeat?.drainActiveRunExecutions(); - await database?.cleanup(); - vi.unstubAllEnvs(); - if (home) await rm(home, { recursive: true, force: true }); - }); - - async function fixture() { - execute.mockClear(); - const companyId = randomUUID(); - const agentId = randomUUID(); - const issueId = randomUUID(); - const userId = randomUUID(); - await db.insert(companies).values({ id: companyId, name: "Contention fixture", issuePrefix: `S${companyId.slice(0, 8).toUpperCase()}`, defaultResponsibleUserId: userId }); - await db.insert(companyMemberships).values({ companyId, principalId: userId, principalType: "user", status: "active", membershipRole: "member" }); - const binding = { provider: "openai", method: "subscription", mode: "responsible_user" } as const; - await db.insert(agents).values({ id: agentId, companyId, name: "CEO", adapterType: "codex_local", adapterConfig: { cwd: home }, runtimeConfig: { aiConnection: binding, heartbeat: { wakeOnDemand: true, maxConcurrentRuns: 20 } } }); - await db.insert(issues).values({ id: issueId, companyId, title: "Child task", status: "todo", assigneeAgentId: agentId, responsibleUserId: userId }); - const account = await aiConnectionService(db).save(companyId, userId, { provider: "openai", method: "subscription", ownership: "personal", name: "Fixture subscription", loginSessionId: "fixture", allAgents: true, agentIds: [] }, JSON.stringify({ tokens: { access_token: "fixture-access", refresh_token: "fixture-refresh", id_token: "fixture-id", account_id: "fixture-account" } })); - const parent = await prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: userId, adapterType: "codex_local", binding, config: { cwd: home } }); - let released = false; - const release = async () => { - if (released) return; - released = true; - await parent.cleanup(); - }; - return { companyId, agentId, issueId, userId, account, parent, release }; - } - - async function defer(f: Awaited>, context: Record = {}) { - const run = await heartbeat.invoke(f.agentId, "assignment", { - issueId: f.issueId, wakeReason: "issue_assigned", ...context, - }, "system"); - expect(run).not.toBeNull(); - await heartbeat.drainActiveRunExecutions(); - expect(await heartbeat.getRun(run!.id)).toMatchObject({ status: "cancelled", errorCode: "ai_connection_busy" }); - const [retry] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.retryOfRunId, run!.id)); - expect(retry).toMatchObject({ status: "scheduled_retry", scheduledRetryReason: "ai_connection_busy" }); - return retry!; - } - - async function dispatch(retry: typeof heartbeatRuns.$inferSelect, service = heartbeat) { - await service.promoteDueScheduledRetries(new Date(retry.scheduledRetryAt!.getTime() + 1)); - await service.resumeQueuedRuns(); - await service.drainActiveRunExecutions(); - } - - it("defers a child without a reconnect blocker and executes after its parent's subscription is released", async () => { - const f = await fixture(); - try { - const run = await heartbeat.invoke(f.agentId, "assignment", { issueId: f.issueId, wakeReason: "issue_assigned" }, "system"); - expect(run).not.toBeNull(); - await heartbeat.drainActiveRunExecutions(); - const deferred = await heartbeat.getRun(run!.id); - expect(deferred).toMatchObject({ status: "cancelled", errorCode: "ai_connection_busy" }); - expect(execute).not.toHaveBeenCalled(); - const [retry] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.retryOfRunId, run!.id)); - expect(retry).toMatchObject({ status: "scheduled_retry", scheduledRetryReason: "ai_connection_busy" }); - expect(executionFailureRetryCount(retry)).toBe(0); - const [issue] = await db.select().from(issues).where(eq(issues.id, f.issueId)); - expect(issue.status).not.toBe("blocked"); - expect(issue.executionRunId).toBe(retry.id); - expect(await db.select().from(issueRecoveryActions).where(eq(issueRecoveryActions.companyId, f.companyId))).toHaveLength(0); - expect(await db.select().from(issueThreadInteractions).where(eq(issueThreadInteractions.companyId, f.companyId))).toHaveLength(0); - expect(await db.select().from(issueComments).where(eq(issueComments.issueId, f.issueId))).toHaveLength(0); - await f.release(); - await heartbeat.promoteDueScheduledRetries(new Date(retry.scheduledRetryAt!.getTime() + 1)); - await heartbeat.resumeQueuedRuns(); - await heartbeat.drainActiveRunExecutions(); - const completed = await heartbeat.getRun(retry.id); - expect(completed).toMatchObject({ status: "succeeded" }); - expect(execute).toHaveBeenCalledTimes(1); - expect(completed?.contextSnapshot?.aiConnection).toMatchObject({ grantId: f.account.grantId, identity: f.parent.identity }); - } finally { - await f.release(); - } - }); - - it("keeps waiting beyond the failure-attempt limit and survives a new service instance", async () => { - const f = await fixture(); - try { - let retry = await defer(f, { - failureRetriesBeforeAiConnectionWait: 99, aiConnectionBusyDeferredWhileAssignee: false, - }); - expect(retry.contextSnapshot?.aiConnectionBusyDeferredWhileAssignee).toBe(true); - for (let count = 0; count < 3; count++) { - await dispatch(retry); - const [successor] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.retryOfRunId, retry.id)); - expect(successor).toMatchObject({ status: "scheduled_retry", scheduledRetryAttempt: count + 2 }); - expect(executionFailureRetryCount(successor)).toBe(0); - retry = successor!; - const projections = await executionProjectionsForRuns(db, f.companyId, [retry.id]); - expect(projections.get(retry.id)).toMatchObject({ - phase: "retry_scheduled", label: "Waiting for AI subscription", attempt: 1, - }); - } - expect(execute).not.toHaveBeenCalled(); - await f.release(); - await dispatch(retry, heartbeatService(db)); - expect(await heartbeat.getRun(retry.id)).toMatchObject({ status: "succeeded" }); - expect(execute).toHaveBeenCalledTimes(1); - } finally { - await f.release(); - } - }); - - it.each([ - ["cancelled", "issue_cancelled"], - ["done", "issue_terminal_status"], - ["reassigned", "issue_reassigned"], - ["agent_paused", "agent_not_invokable"], - ])("does not dispatch after the waiting task is %s", async (change, errorCode) => { - const f = await fixture(); - try { - const retry = await defer(f); - if (change === "agent_paused") { - await db.update(agents).set({ status: "paused" }).where(eq(agents.id, f.agentId)); - } else { - await db.update(issues).set(change === "reassigned" - ? { assigneeAgentId: null, assigneeUserId: f.userId } - : { status: change }).where(eq(issues.id, f.issueId)); - } - await f.release(); - await dispatch(retry); - expect(await heartbeat.getRun(retry.id)).toMatchObject({ status: "cancelled", errorCode }); - expect(execute).not.toHaveBeenCalled(); - const [issue] = await db.select().from(issues).where(eq(issues.id, f.issueId)); - expect(issue.executionRunId).toBeNull(); - } finally { - await f.release(); - } - }); - - it.each(["issue_assigned", "issue_commented"])("does not turn an assignee %s wake into a non-assignee retry after reassignment during preflight", async (wakeReason) => { - const f = await fixture(); - try { - afterCheckout.mockImplementationOnce(async () => { - await db.update(issues).set({ - assigneeAgentId: null, assigneeUserId: f.userId, executionRunId: null, - }).where(eq(issues.id, f.issueId)); - }); - const commentId = randomUUID(); - await db.insert(issueComments).values({ id: commentId, companyId: f.companyId, issueId: f.issueId, - authorType: "user", authorUserId: f.userId, body: "Continue this task." }); - const run = await heartbeat.invoke(f.agentId, "assignment", { - issueId: f.issueId, wakeReason, commentId, aiConnectionBusyDeferredWhileAssignee: false, - }, "system"); - await heartbeat.drainActiveRunExecutions(); - expect(await heartbeat.getRun(run!.id)).toMatchObject({ status: "cancelled", errorCode: "ai_connection_busy" }); - expect(await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.retryOfRunId, run!.id))).toHaveLength(0); - expect(execute).not.toHaveBeenCalled(); - } finally { - afterCheckout.mockReset(); - await f.release(); - } - }); - - it("resumes an authorized comment wake without stealing the assignee's task lock", async () => { - const f = await fixture(); - try { - await db.update(issues).set({ assigneeAgentId: null, assigneeUserId: f.userId }).where(eq(issues.id, f.issueId)); - const retry = await defer(f, { wakeReason: "issue_comment_mentioned", commentId: randomUUID() }); - expect(retry.contextSnapshot?.aiConnectionBusyDeferredWhileAssignee).toBe(false); - await dispatch(retry); - const [successor] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.retryOfRunId, retry.id)); - expect(successor).toMatchObject({ status: "scheduled_retry" }); - expect(successor.contextSnapshot?.aiConnectionBusyDeferredWhileAssignee).toBe(false); - const [issue] = await db.select().from(issues).where(eq(issues.id, f.issueId)); - expect(issue.executionRunId).toBeNull(); - await f.release(); - await dispatch(successor); - expect(await heartbeat.getRun(successor.id)).toMatchObject({ status: "succeeded" }); - expect(execute).toHaveBeenCalledTimes(1); - } finally { - await f.release(); - } - }); - - it("rechecks revoked credentials after waiting and retains the configuration failure", async () => { - const f = await fixture(); - try { - const retry = await defer(f); - await f.release(); - await db.update(connectionGrants).set({ status: "revoked", revokedAt: new Date() }).where(eq(connectionGrants.id, f.account.grantId)); - await dispatch(retry); - expect(await heartbeat.getRun(retry.id)).toMatchObject({ status: "failed", errorCode: "configuration_incomplete" }); - expect(execute).not.toHaveBeenCalled(); - expect(await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.retryOfRunId, retry.id))).toHaveLength(0); - } finally { - await f.release(); - } - }); -}); diff --git a/server/src/__tests__/secrets-service.test.ts b/server/src/__tests__/secrets-service.test.ts index dbc6050fdf4..186af5a993e 100644 --- a/server/src/__tests__/secrets-service.test.ts +++ b/server/src/__tests__/secrets-service.test.ts @@ -5,7 +5,7 @@ import os from "node:os"; import path from "node:path"; import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest"; import type { MockInstance } from "vitest"; -import { and, eq } from "drizzle-orm"; +import { and, eq, sql } from "drizzle-orm"; import { resolveCodexAuthCacheDir, withAccountHomeSecretMutationLock } from "@paperclipai/adapter-codex-local/server"; import { activityLog, @@ -674,6 +674,79 @@ describeEmbeddedPostgres("secretService", () => { expect(events).toEqual(["rotate-provider-enter", "rotate-provider-exit", "create-provider-enter"]); }); + it("locks the parent secret row before its version rows during rotate, so a concurrent credential write-back cannot deadlock with it", async () => { + // Inserting the rotate's new version row needs a read lock on the + // parent secret row to satisfy the version table's foreign key, so a + // holder that locks the parent row first would block that insert and + // never reach the statements this test distinguishes. Lock the secret's + // existing version row instead: that does not touch the parent row, so + // the rotate call's insert proceeds and the call reaches its final + // transaction, where it then queues behind this same held row for its + // "mark the old version previous" update. + const companyId = await seedCompany(); + const svc = secretService(db); + const secret = await svc.create(companyId, { + name: `lock-order-${randomUUID()}`, + provider: "local_encrypted", + value: "original-value", + }); + + const holderPidReady = deferred(); + let releaseHolder!: () => void; + const holderReleased = new Promise((resolve) => { + releaseHolder = resolve; + }); + const holder = db.transaction(async (tx) => { + await tx + .select({ id: companySecretVersions.id }) + .from(companySecretVersions) + .where(and(eq(companySecretVersions.secretId, secret.id), eq(companySecretVersions.version, 1))) + .for("update"); + const [backend] = (await tx.execute(sql`select pg_backend_pid() as pid`)) as unknown as Array<{ pid: number }>; + holderPidReady.resolve(backend.pid); + await holderReleased; + }); + const holderPid = await holderPidReady.promise; + + const rotateCall = svc.rotate(secret.id, { value: "rotated-value" }); + + // Poll until Postgres reports the rotate call's own backend blocked + // behind the lock holder above, instead of guessing how long that + // takes. + let rotateBlocked = false; + for (let attempt = 0; attempt < 200; attempt++) { + const rows = (await db.execute( + sql`select 1 from pg_stat_activity where ${holderPid} = any(pg_blocking_pids(pid))`, + )) as unknown as Array; + if (rows[0]) { + rotateBlocked = true; + break; + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } + try { + expect(rotateBlocked).toBe(true); + // The rotate call is confirmed blocked on the existing version row. + // If it had updated that row before the parent secret row — the + // opposite order this fix removes — the parent row would still be + // free here. A `nowait` probe from a third transaction settles which + // is true: it fails the instant the parent row already carries an + // uncommitted write from the blocked rotate call, and this fix makes + // that write happen first. + const probe = await db + .transaction(async (tx) => { + await tx.execute(sql`select 1 from company_secrets where id = ${secret.id} for update nowait`); + }) + .then(() => null) + .catch((error: unknown) => error); + expect(String((probe as { cause?: unknown })?.cause ?? probe)).toMatch(/could not obtain lock/i); + } finally { + releaseHolder(); + } + await holder; + await expect(rotateCall).resolves.toMatchObject({ latestVersion: 2 }); + }); + it("fails a queued local_encrypted create when an account-home cleanup removes its directory first", async () => { // The mutation lock alone stops a write and an account-home cleanup's // check-and-delete from interleaving; it does not stop them from running diff --git a/server/src/services/ai-connection-runtime.ts b/server/src/services/ai-connection-runtime.ts index 0f304a91189..1ee19dcf082 100644 --- a/server/src/services/ai-connection-runtime.ts +++ b/server/src/services/ai-connection-runtime.ts @@ -4,7 +4,7 @@ import { mkdtemp, mkdir, writeFile, readFile, rm } from "node:fs/promises"; import os from "node:os"; import path from "node:path"; import { and, eq } from "drizzle-orm"; -import { type Db, connectionGrants } from "@paperclipai/db"; +import { type Db, companySecrets, connectionGrants } from "@paperclipai/db"; import { AI_CONNECTION_CAPABILITIES, type AiConnectionBinding, @@ -166,41 +166,6 @@ done`, } } -async function acquireCredentialLease(db: Db, grantId: string) { - const client = await db.$client.reserve(); - try { - // A reserved client pins our connection to PgBouncer, not its backend. - // Keep the lease in one transaction so transaction-pooling deployments - // cannot acquire and release it on different PostgreSQL sessions. - await client`begin`; - await client`set local idle_in_transaction_session_timeout = 0`; - const [result] = - await client`select pg_try_advisory_xact_lock(hashtextextended(${`ai-runtime:${grantId}`}, 0)) as acquired`; - if (!result.acquired) - throw unprocessable( - "This subscription is in use. Retry when its current execution finishes.", - { code: "ai_connection_busy" }, - ); - } catch (error) { - try { - await client`rollback`; - } finally { - client.release(); - } - throw error; - } - let released = false; - return async () => { - if (released) return; - released = true; - try { - await client`rollback`; - } finally { - client.release(); - } - }; -} - export async function prepareManagedAiRuntime( db: Db, input: { @@ -246,12 +211,6 @@ export async function prepareManagedAiRuntime( const subscriptionFile = selection.attribution.method === "subscription" && input.binding.provider !== "anthropic"; - // A file-backed subscription runs one credential rotation at a time, so it - // needs the lease. Anthropic writes no file back, so two runs share no - // mutable state and can run at the same time without the lease. - const release = subscriptionFile - ? await acquireCredentialLease(db, selection.grant.id) - : async () => {}; let home: string | undefined; try { const selectedGrantId = selection.grant.id; @@ -344,18 +303,32 @@ export async function prepareManagedAiRuntime( ), ) .for("update"); - // Reconnect/revocation wins over a process holding an older credential. - if ( - !grant || - grant.status !== "active" || - grant.updatedAt.getTime() !== - selection.grant.updatedAt.getTime() - ) - return; - const current = await service.credential({ - ...selection, - grant, - }); + // 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", @@ -370,9 +343,6 @@ export async function prepareManagedAiRuntime( errorLabel: "AI account refresh", }); if (decision !== 10) return; - const ref = grant.credentialSecretRefs.find( - (r) => r.configPath === "ai.credential", - )!; await secretService(tx).rotate( ref.secretId, { value: refreshed }, @@ -385,20 +355,12 @@ export async function prepareManagedAiRuntime( }); } } finally { - try { - if (home) await rm(home, { recursive: true, force: true }); - } finally { - await release(); - } + if (home) await rm(home, { recursive: true, force: true }); } }, }; } catch (error) { - try { - if (home) await rm(home, { recursive: true, force: true }); - } finally { - await release(); - } + if (home) await rm(home, { recursive: true, force: true }); throw error; } } diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index b985426debe..58cab31a156 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -16037,10 +16037,9 @@ export function heartbeatService( .then((rows) => rows[0] ?? null); } - // Subscription homes are serialized so refresh tokens cannot be overwritten - // by another run. Contention is a resource wait, not broken authentication. - // The database lease is released on completion/disconnect; retrying remains - // safe across processes and each attempt revalidates the selected account. + // No code path raises the `ai_connection_busy` error any more. This + // function stays because a stored run row can still carry that error code + // from an earlier release. async function finalizeAiConnectionBusyDeferral( run: typeof heartbeatRuns.$inferSelect, error: HttpError, diff --git a/server/src/services/secrets.ts b/server/src/services/secrets.ts index c5facc49646..143275d0941 100644 --- a/server/src/services/secrets.ts +++ b/server/src/services/secrets.ts @@ -2522,22 +2522,14 @@ export function secretService(db: Db | DbTransaction) { } try { + // Lock and update the parent secret row before touching its version + // rows. A credential write-back locks the parent secret row first, then + // calls this function. If this transaction updated the version rows + // first instead, the two transactions would take their two row locks + // in opposite order and could deadlock. Matching the order here removes + // that risk: every caller now locks the parent row before the version + // rows, so a lock cycle between the two tables cannot form. return await db.transaction(async (tx) => { - await tx - .update(companySecretVersions) - .set({ status: "previous" }) - .where(and( - eq(companySecretVersions.secretId, secret.id), - ne(companySecretVersions.version, nextVersion), - )); - await tx - .update(companySecretVersions) - .set({ status: "current" }) - .where(and( - eq(companySecretVersions.secretId, secret.id), - eq(companySecretVersions.version, nextVersion), - )); - const updated = await tx .update(companySecrets) .set({ @@ -2561,11 +2553,28 @@ export function secretService(db: Db | DbTransaction) { if (!updated) { // The predicate matched no row. A supplied expected version means a // concurrent rotation won the race; return the stale conflict. An - // unguarded rotation means the secret is gone. + // unguarded rotation means the secret is gone. Neither version row + // has been touched yet, so there is nothing to undo here. throw input.expectedLatestVersion === undefined ? notFound("Secret not found") : conflict(SECRET_VERSION_STALE_CONFLICT); } + + await tx + .update(companySecretVersions) + .set({ status: "previous" }) + .where(and( + eq(companySecretVersions.secretId, secret.id), + ne(companySecretVersions.version, nextVersion), + )); + await tx + .update(companySecretVersions) + .set({ status: "current" }) + .where(and( + eq(companySecretVersions.secretId, secret.id), + eq(companySecretVersions.version, nextVersion), + )); + return updated; }); } catch (error) {