mirror of
https://github.com/paperclipai/paperclip.git
synced 2026-10-24 10:13:11 +02:00
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 <noreply@paperclip.ing>
This commit is contained in:
Nicky LeachandPaperclip authored and GitHub committed
1 parent
da77a0c28c
commit
9ed55f6931
14 files changed
+554
-514
No files matched your search
@@ -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();
|
||||
|
||||
@@ -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<void>(resolve => { holdAcquired = resolve; });
|
||||
let releaseHold!: () => void;
|
||||
const holdReleased = new Promise<void>(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);
|
||||
|
||||
@@ -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<typeof import("../adapters/index.js")>("../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<typeof import("../services/issues.js")>("../services/issues.js");
|
||||
return { ...actual, issueService: (...args: Parameters<typeof actual.issueService>) => {
|
||||
const service = actual.issueService(...args);
|
||||
return { ...service, checkout: async (...checkoutArgs: Parameters<typeof service.checkout>) => {
|
||||
const result = await service.checkout(...checkoutArgs);
|
||||
await afterCheckout();
|
||||
return result;
|
||||
} };
|
||||
} };
|
||||
});
|
||||
|
||||
describe("heartbeat AI subscription contention", () => {
|
||||
let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>>;
|
||||
let db: ReturnType<typeof createDb>;
|
||||
let heartbeat: ReturnType<typeof heartbeatService>;
|
||||
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<ReturnType<typeof fixture>>, context: Record<string, unknown> = {}) {
|
||||
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();
|
||||
}
|
||||
});
|
||||
});
|
||||
@@ -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<number>();
|
||||
let releaseHolder!: () => void;
|
||||
const holderReleased = new Promise<void>((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<unknown>;
|
||||
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
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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) {
|
||||
|
||||
Reference in new issue
Block a user