mirror of
https://github.com/paperclipai/paperclip.git
synced 2026-10-29 12:59:13 +01:00
fix: scope quota reports to authorized subscription accounts
Co-Authored-By: Paperclip <noreply@paperclip.ing>
This commit is contained in:
Devin FoleyandPaperclip committed
1 parent
a85fa3af31
commit
e85b02f99e
24 files changed
+1983
-182
No files matched your search
@@ -1,3 +1,5 @@
|
||||
import * as codexAdapter from "@paperclipai/adapter-codex-local/server";
|
||||
import { WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE } from "@paperclipai/adapter-utils/workspace-restore-merge";
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { mkdtemp, rm } from "node:fs/promises";
|
||||
import os from "node:os";
|
||||
@@ -468,6 +470,38 @@ describe("agent-created hires use managed AI connections", () => {
|
||||
});
|
||||
|
||||
describe("hired agents sharing a subscription", () => {
|
||||
it("keeps a credential-lock timeout on automatic retry without blocking the task or starting a provider", async () => {
|
||||
const f = await fixture("openai", "subscription");
|
||||
const agent = hired(await request(f.app).post(`/api/companies/${f.companyId}/agent-hires`).send({ name: "Waiting teammate", role: "engineer", adapterType: f.adapterType, reportsTo: f.agentId, adapterConfig: { engine: "cli" }, runtimeConfig: { heartbeat: { enabled: false } } }));
|
||||
// Only an operator may set host execution paths after the agent is hired.
|
||||
await db.update(agents).set({ adapterConfig: { ...agent.adapterConfig, cwd: home } }).where(eq(agents.id, agent.id));
|
||||
const [issue] = await db.insert(issues).values({ companyId: f.companyId, title: "Wait for credential rotation", status: "todo", assigneeAgentId: agent.id, responsibleUserId: f.userId, createdByUserId: f.userId }).returning();
|
||||
const execute = vi.fn(async () => ({ exitCode: 0, signal: null, timedOut: false, resultJson: {} }));
|
||||
registerServerAdapter({ ...getServerAdapter(f.adapterType), execute });
|
||||
const lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock")
|
||||
.mockRejectedValueOnce(Object.assign(new Error("quota refresh holds company lock"), { code: WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE }));
|
||||
const heartbeat = heartbeatService(db);
|
||||
try {
|
||||
const run = await heartbeat.invoke(agent.id, "assignment", { issueId: issue.id, wakeReason: "issue_assigned", responsibleUserId: f.userId }, "system");
|
||||
expect(run).not.toBeNull();
|
||||
await expect.poll(async () => (await heartbeat.getRun(run!.id))?.status, { timeout: 20_000 }).toBe("cancelled");
|
||||
await heartbeat.drainActiveRunExecutions();
|
||||
expect(await heartbeat.getRun(run!.id)).toMatchObject({ errorCode: "ai_connection_busy", resultJson: { executionRecovery: { providerWorkStarted: false } } });
|
||||
const retries = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.retryOfRunId, run!.id));
|
||||
expect(retries).toHaveLength(1);
|
||||
expect(retries[0]).toMatchObject({ status: "scheduled_retry", scheduledRetryReason: "ai_connection_busy" });
|
||||
const [savedIssue] = await db.select().from(issues).where(eq(issues.id, issue.id));
|
||||
expect(savedIssue.status).not.toBe("blocked");
|
||||
expect(savedIssue.executionRunId).toBe(retries[0].id);
|
||||
expect(execute).not.toHaveBeenCalled();
|
||||
expect(await aiConnectionService(db).list(f.companyId, f.userId)).toEqual([expect.objectContaining({ status: "connected" })]);
|
||||
} finally {
|
||||
lock.mockRestore();
|
||||
await heartbeat.drainActiveRunExecutions();
|
||||
unregisterServerAdapter(f.adapterType);
|
||||
}
|
||||
});
|
||||
|
||||
it.each(["openai", "anthropic"] as const)("runs the %s child alongside a live parent and inherits its connection", async (provider) => {
|
||||
const f = await fixture(provider, "subscription");
|
||||
const agent = hired(await request(f.app).post(`/api/companies/${f.companyId}/agent-hires`).send({ name: "Concurrent teammate", role: "engineer", adapterType: f.adapterType, reportsTo: f.agentId, adapterConfig: { engine: "cli" }, runtimeConfig: { heartbeat: { enabled: false } } }));
|
||||
|
||||
@@ -110,6 +110,10 @@ const mockPrepareManagedAiRuntime = vi.hoisted(() => vi.fn());
|
||||
vi.mock("../services/ai-connection-runtime.js", async (importOriginal) => ({
|
||||
...(await importOriginal<typeof import("../services/ai-connection-runtime.js")>()),
|
||||
prepareManagedAiRuntime: mockPrepareManagedAiRuntime,
|
||||
withManagedAiProbe: async (_db: unknown, input: unknown, probe: (runtime: unknown) => Promise<unknown>) => {
|
||||
const runtime = await mockPrepareManagedAiRuntime(_db, input);
|
||||
try { return await probe(runtime); } finally { await runtime.cleanup(); }
|
||||
},
|
||||
}));
|
||||
const mockValidateAiApiKey = vi.hoisted(() => vi.fn(async () => undefined));
|
||||
vi.mock("../routes/ai-connections.js", async (importOriginal) => ({
|
||||
|
||||
@@ -4,15 +4,21 @@ import { issueRecoveryActionService } from "../services/issue-recovery-actions.j
|
||||
import * as localCredentials from "../services/local-ai-credentials.js";
|
||||
import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { promises as fs } from "node:fs";
|
||||
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, companySecrets, principalPermissionGrants } from "@paperclipai/db";
|
||||
import { createDb, companies, agents, heartbeatRuns, companyMemberships, connectionGrants, connectionGrantDelegations, connectionGrantMembers, toolConnections, toolConnectionInstalls, aiConnectionDefaults, aiProviderDefaults, adapterAuthSessions, environments, issues, issueThreadInteractions, issueRecoveryActions, connectionIntentDeliveries, agentWakeupRequests, companySecrets, companySecretVersions, userSecretDefinitions, principalPermissionGrants } from "@paperclipai/db";
|
||||
import { startEmbeddedPostgresTestDatabase } from "@paperclipai/db/test-embedded-postgres";
|
||||
import { fetchCompanyQuotaWindows } from "../services/quota-windows.js";
|
||||
import { quotaCredentialRecoveryPath, quotaCredentialRecovery, quotaCredentialHash } from "../services/quota-credential-recovery.js";
|
||||
import { aiConnectionService } from "../services/ai-connections.js";
|
||||
import { syncConnectionCredentialBindings } from "../services/connection-credential-bindings.js";
|
||||
import * as codexAdapter from "@paperclipai/adapter-codex-local/server";
|
||||
import { WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE } from "@paperclipai/adapter-utils/workspace-restore-merge";
|
||||
import * as executionTarget from "@paperclipai/adapter-utils/execution-target";
|
||||
import { prepareManagedAiRuntime, assertManagedAiProjectAuth } from "../services/ai-connection-runtime.js";
|
||||
import { prepareManagedAiRuntime, withManagedAiProbe, assertManagedAiProjectAuth, isAiConnectionBusy } from "../services/ai-connection-runtime.js";
|
||||
import { toolAccessService } from "../services/tool-access.js";
|
||||
import { secretService } from "../services/secrets.js";
|
||||
import { aiConnectionBindingSchema, connectionPurposeTransportSchema, isAiConnectionCompatible } from "@paperclipai/shared";
|
||||
@@ -46,6 +52,512 @@ beforeAll(async () => {
|
||||
afterAll(async () => { await database?.cleanup(); vi.unstubAllEnvs(); if (home) await rm(home, { recursive: true, force: true }); });
|
||||
|
||||
describe("managed AI connections", () => {
|
||||
it("persists first-time recovery directories before consuming a single-use token", async () => {
|
||||
const owner = "quota-durable-path";
|
||||
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
|
||||
const original = JSON.stringify({ tokens: { account_id: owner, access_token: "expired", refresh_token: "single-use" } });
|
||||
await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, original);
|
||||
const [row] = await service.quotaAccounts(companyId, owner);
|
||||
const file = quotaCredentialRecoveryPath(companyId, row.grant.id);
|
||||
const recoveryRoot = path.dirname(path.dirname(file));
|
||||
await expect(access(recoveryRoot)).rejects.toMatchObject({ code: "ENOENT" });
|
||||
const synced: string[] = [];
|
||||
let rejectRoot = true;
|
||||
const realOpen = fs.open.bind(fs);
|
||||
const open = vi.spyOn(fs, "open").mockImplementation(async (...args: Parameters<typeof fs.open>) => {
|
||||
const handle = await realOpen(...args);
|
||||
const sync = handle.sync.bind(handle);
|
||||
vi.spyOn(handle, "sync").mockImplementation(async () => {
|
||||
if (rejectRoot && args[0] === recoveryRoot) throw new Error("Injected recovery directory flush failure");
|
||||
await sync(); synced.push(String(args[0]));
|
||||
});
|
||||
return handle;
|
||||
});
|
||||
const request = vi.spyOn(globalThis, "fetch").mockImplementation(async () => {
|
||||
expect(synced).toContain(recoveryRoot);
|
||||
expect(synced).toContain(path.dirname(recoveryRoot));
|
||||
return Response.json({ access_token: "fresh-durable", refresh_token: "replacement-durable" });
|
||||
});
|
||||
try {
|
||||
await expect(service.refreshQuotaCredential(row, original, new AbortController().signal)).rejects.toThrow("Injected recovery directory flush failure");
|
||||
expect(request).not.toHaveBeenCalled();
|
||||
expect(await service.credential(row)).toBe(original);
|
||||
rejectRoot = false;
|
||||
await service.refreshQuotaCredential(row, original, new AbortController().signal);
|
||||
expect(request).toHaveBeenCalledTimes(1);
|
||||
expect(JSON.parse(await service.credential(row))).toMatchObject({ tokens: { access_token: "fresh-durable", refresh_token: "replacement-durable" } });
|
||||
} finally { open.mockRestore(); request.mockRestore(); }
|
||||
});
|
||||
|
||||
it.each(["timeout", "unexpected"] as const)("classifies a manual probe's %s lock failure before starting the provider", async (kind) => {
|
||||
const owner = `manual-probe-lock-${kind}`;
|
||||
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
|
||||
const auth = JSON.stringify({ tokens: { account_id: owner, access_token: "valid-access", refresh_token: "valid-refresh" } });
|
||||
await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth);
|
||||
const [row] = await service.quotaAccounts(companyId, owner);
|
||||
const probe = vi.fn().mockResolvedValue("ok");
|
||||
const input = { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local",
|
||||
binding: { provider: "openai", method: "subscription", mode: "responsible_user" } as const, config: {} };
|
||||
const failure = Object.assign(new Error("injected private lock diagnostic"), { code: kind === "timeout" ? WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE : "EACCES" });
|
||||
const lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock").mockRejectedValueOnce(failure);
|
||||
try {
|
||||
const error = await withManagedAiProbe(db, input, probe).then(() => { throw new Error("Probe must not start"); }, error => error);
|
||||
if (kind === "timeout") {
|
||||
expect(isAiConnectionBusy(error)).toBe(true);
|
||||
expect(error).toMatchObject({ status: 422, message: "AI credentials are being updated. Retry shortly.", details: { code: "ai_connection_busy" } });
|
||||
} else expect(error).toBe(failure);
|
||||
expect(probe).not.toHaveBeenCalled();
|
||||
expect((await db.select().from(connectionGrants).where(eq(connectionGrants.id, row.grant.id)))[0].status).toBe("active");
|
||||
} finally { lock.mockRestore(); }
|
||||
expect(await withManagedAiProbe(db, input, probe)).toBe("ok");
|
||||
expect(probe).toHaveBeenCalledTimes(1);
|
||||
expect(await service.credential(row)).toBe(auth);
|
||||
});
|
||||
|
||||
it("keeps quota refresh outside a manual probe until its rotated credential is saved", async () => {
|
||||
const owner = "manual-probe-race";
|
||||
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
|
||||
const auth = (hour: string) => JSON.stringify({ tokens: { account_id: owner, id_token: "identity", access_token: `access-${hour}`, refresh_token: `refresh-${hour}` }, last_refresh: `2026-09-10T${hour}:00:00Z` });
|
||||
await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth("10"));
|
||||
const [row] = await service.quotaAccounts(companyId, owner);
|
||||
const entered = deferredSignal(), release = deferredSignal();
|
||||
const request = vi.spyOn(globalThis, "fetch").mockRejectedValue(new Error("Quota must not exchange the probe's token"));
|
||||
const probing = withManagedAiProbe(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local",
|
||||
binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} }, async runtime => {
|
||||
entered.resolve(); await release.promise;
|
||||
await writeFile(path.join(String(runtime.config.env.CODEX_HOME), "auth.json"), auth("11"));
|
||||
});
|
||||
let refreshing: Promise<string> | undefined;
|
||||
let lock: ReturnType<typeof vi.spyOn> | undefined;
|
||||
try {
|
||||
await entered.promise;
|
||||
lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock");
|
||||
// Start outside the probe's async scope; it is a separate HTTP request.
|
||||
refreshing = service.refreshQuotaCredential(row, auth("10"), new AbortController().signal);
|
||||
await vi.waitFor(() => expect(lock).toHaveBeenCalled());
|
||||
expect(request).not.toHaveBeenCalled();
|
||||
release.resolve(); await probing;
|
||||
expect(await refreshing).toBe(auth("11"));
|
||||
expect(await service.credential(row)).toBe(auth("11"));
|
||||
expect(request).not.toHaveBeenCalled();
|
||||
} finally { release.resolve(); await Promise.allSettled([probing, refreshing]); lock?.mockRestore(); request.mockRestore(); }
|
||||
});
|
||||
|
||||
it("rejects preparation when reconnect replaces the selected credential reference", async () => {
|
||||
const owner = "replaced-ref-owner";
|
||||
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
|
||||
const auth = (id: string) => JSON.stringify({ tokens: { account_id: id, access_token: id, refresh_token: id } });
|
||||
await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth("old-account"));
|
||||
const [row] = await service.quotaAccounts(companyId, owner);
|
||||
const definition = await secretService(db).createUserSecretDefinition(companyId, { key: "replacement_credential", name: "Replacement", provider: "local_encrypted" }, { userId: owner });
|
||||
const replacement = await secretService(db).createCurrentUserSecretValue(companyId, owner, { definitionId: definition.id, value: auth("new-account") }, { userId: owner });
|
||||
const prepare = () => prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local",
|
||||
binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} });
|
||||
const original = codexAdapter.withAccountHomeSecretMutationLock;
|
||||
const lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock").mockImplementationOnce(async (env, company, callback) => {
|
||||
await db.update(connectionGrants).set({ credentialSecretRefs: row.grant.credentialSecretRefs.map(ref => ref.configPath === "ai.credential" ? { ...ref, secretId: replacement.id } : ref) }).where(eq(connectionGrants.id, row.grant.id));
|
||||
await syncConnectionCredentialBindings(db, row.connection);
|
||||
return original(env, company, callback);
|
||||
});
|
||||
try { await expect(prepare()).rejects.toThrow("AI credential changed during preparation"); }
|
||||
finally { lock.mockRestore(); }
|
||||
const runtime = await prepare();
|
||||
try {
|
||||
expect(runtime.sessionIdentity).toContain(replacement.id);
|
||||
expect(await readFile(path.join(String(runtime.config.env.CODEX_HOME), "auth.json"), "utf8")).toBe(auth("new-account"));
|
||||
} finally { await runtime.cleanup(); }
|
||||
});
|
||||
|
||||
it("caches quota after a real secret read but invalidates credential rotation and revocation", async () => {
|
||||
const owner = "quota-cache-owner";
|
||||
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
|
||||
const account = await service.save(companyId, owner, { provider: "anthropic", method: "subscription", ownership: "personal", name: "Quota cache account", loginSessionId: "fixture", agentIds: [], allAgents: true }, "quota-cache-original");
|
||||
const [grant] = await db.select().from(connectionGrants).where(eq(connectionGrants.id, account.grantId));
|
||||
const secretId = grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!.secretId;
|
||||
const [before] = await db.select().from(companySecrets).where(eq(companySecrets.id, secretId));
|
||||
const fetchSpy = vi.spyOn(globalThis, "fetch").mockImplementation(async () => new Response(JSON.stringify({ five_hour: { utilization: 42 } })));
|
||||
try {
|
||||
const first = await fetchCompanyQuotaWindows(db, companyId, owner);
|
||||
const [after] = await db.select().from(companySecrets).where(eq(companySecrets.id, secretId));
|
||||
expect(after.lastResolvedAt).not.toBeNull();
|
||||
expect(after.updatedAt.getTime()).toBeGreaterThan(before.updatedAt.getTime());
|
||||
expect(first[0].ok).toBe(true);
|
||||
expect(await fetchCompanyQuotaWindows(db, companyId, owner)).toEqual(first);
|
||||
expect(fetchSpy).toHaveBeenCalledTimes(1);
|
||||
await secretService(db).rotate(secretId, { value: "quota-cache-rotated" });
|
||||
const rotated = await fetchCompanyQuotaWindows(db, companyId, owner);
|
||||
expect(rotated[0].ok).toBe(true);
|
||||
expect(rotated[0].accountKey).not.toBe(first[0].accountKey);
|
||||
expect(fetchSpy).toHaveBeenCalledTimes(2);
|
||||
expect(fetchSpy).toHaveBeenLastCalledWith(expect.any(String), expect.objectContaining({ headers: expect.objectContaining({ Authorization: "Bearer quota-cache-rotated" }) }));
|
||||
await db.update(userSecretDefinitions).set({ status: "disabled" }).where(eq(userSecretDefinitions.id, before.userSecretDefinitionId!));
|
||||
const disabled = await fetchCompanyQuotaWindows(db, companyId, owner);
|
||||
expect(disabled[0]).toMatchObject({ ok: false, errorFamily: "credentials_unavailable", windows: [] });
|
||||
expect(disabled[0].accountKey).not.toBe(rotated[0].accountKey);
|
||||
expect(fetchSpy).toHaveBeenCalledTimes(2);
|
||||
await db.update(userSecretDefinitions).set({ status: "active" }).where(eq(userSecretDefinitions.id, before.userSecretDefinitionId!));
|
||||
expect((await fetchCompanyQuotaWindows(db, companyId, owner))[0].ok).toBe(true);
|
||||
await db.update(companySecretVersions).set({ status: "disabled", revokedAt: new Date() }).where(and(eq(companySecretVersions.secretId, secretId), eq(companySecretVersions.version, 2)));
|
||||
const revoked = await fetchCompanyQuotaWindows(db, companyId, owner);
|
||||
expect(revoked[0].ok).toBe(false);
|
||||
expect(revoked[0].windows).toEqual([]);
|
||||
expect(fetchSpy).toHaveBeenCalledTimes(2);
|
||||
} finally { fetchSpy.mockRestore(); }
|
||||
});
|
||||
|
||||
it("caches refreshed quota under the saved grant and secret revision", async () => {
|
||||
const owner = "quota-refreshed-cache-owner";
|
||||
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
|
||||
const original = JSON.stringify({ tokens: { access_token: "expired-cache", refresh_token: "single-use", account_id: "cache-account" } });
|
||||
await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: "Refreshed cache", loginSessionId: "fixture", agentIds: [], allAgents: true }, original);
|
||||
const request = vi.spyOn(globalThis, "fetch").mockImplementation(async (url, options) => {
|
||||
if (String(url).endsWith("/oauth/token")) return Response.json({ access_token: "fresh-cache", refresh_token: "replacement" });
|
||||
if (new Headers(options?.headers).get("authorization") === "Bearer expired-cache") return new Response("Expired", { status: 401 });
|
||||
return Response.json({ rate_limit: { primary_window: { used_percent: 42 } } });
|
||||
});
|
||||
try {
|
||||
const first = await fetchCompanyQuotaWindows(db, companyId, owner);
|
||||
expect(first[0]).toMatchObject({ ok: true, windows: [{ usedPercent: 42 }] });
|
||||
expect(await fetchCompanyQuotaWindows(db, companyId, owner)).toEqual(first);
|
||||
expect(request).toHaveBeenCalledTimes(3); // rejected quota, OAuth, refreshed quota
|
||||
} finally { request.mockRestore(); }
|
||||
});
|
||||
|
||||
it.each(["rotation", "grant", "activity", "commit"] as const)("recovers exchanged OAuth tokens after %s persistence failure without another exchange", async (failure) => {
|
||||
const owner = `quota-rollback-${failure}`;
|
||||
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
|
||||
const original = JSON.stringify({ tokens: { access_token: `expired-${failure}`, refresh_token: "single-use", account_id: owner } });
|
||||
await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, original);
|
||||
const [row] = await service.quotaAccounts(companyId, owner);
|
||||
const ref = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!;
|
||||
const table = failure === "grant" ? "connection_grants" : failure === "activity" ? "activity_log" : "company_secret_versions";
|
||||
const condition = failure === "grant" ? `new.id = '${row.grant.id}'::uuid` : failure === "activity" ? `new.action = 'ai_connection.credential_refreshed'` : `new.secret_id = '${ref.secretId}'::uuid`;
|
||||
await db.execute(sql.raw(`create function quota_save_failure() returns trigger language plpgsql as $$ begin
|
||||
if ${condition} then raise exception 'injected quota save failure'; end if; return new; end $$;
|
||||
create ${failure === "commit" ? "constraint " : ""}trigger quota_save_failure ${failure === "commit" ? "after" : "before"} ${failure === "grant" ? "update" : "insert"} on ${table}
|
||||
${failure === "commit" ? "deferrable initially deferred" : ""} for each row execute function quota_save_failure();`));
|
||||
const request = vi.spyOn(globalThis, "fetch").mockResolvedValue(Response.json({ access_token: `fresh-${failure}`, refresh_token: `replacement-${failure}` }));
|
||||
const file = quotaCredentialRecoveryPath(companyId, row.grant.id);
|
||||
let triggerDropped = false;
|
||||
try {
|
||||
await expect(service.refreshQuotaCredential(row, original, new AbortController().signal)).rejects.toThrow();
|
||||
expect(request).toHaveBeenCalledTimes(1);
|
||||
expect(await service.credential(row)).toBe(original);
|
||||
const encrypted = await readFile(file, "utf8");
|
||||
expect(encrypted).not.toContain(`fresh-${failure}`);
|
||||
expect(encrypted).not.toContain(`replacement-${failure}`);
|
||||
await db.execute(sql.raw(`drop trigger quota_save_failure on ${table}; drop function quota_save_failure();`));
|
||||
triggerDropped = true;
|
||||
// A new service instance recovers from disk, without an in-memory cache.
|
||||
const recovered = aiConnectionService(db);
|
||||
await recovered.refreshQuotaCredential(row, original, new AbortController().signal);
|
||||
expect(JSON.parse(await recovered.credential(row))).toMatchObject({ tokens: { access_token: `fresh-${failure}`, refresh_token: `replacement-${failure}` } });
|
||||
expect(request).toHaveBeenCalledTimes(1);
|
||||
await expect(access(file)).rejects.toThrow();
|
||||
expect((await db.select().from(companySecretVersions).where(eq(companySecretVersions.secretId, ref.secretId))).map(version => version.version).sort()).toEqual([1, 2]);
|
||||
} finally {
|
||||
request.mockRestore();
|
||||
if (!triggerDropped) await db.execute(sql.raw(`drop trigger quota_save_failure on ${table}; drop function quota_save_failure();`));
|
||||
}
|
||||
});
|
||||
|
||||
it.each(["company_file", "grant_row", "secret_row"] as const)("defers runtime preparation on a contended %s lock and succeeds once it clears", async (kind) => {
|
||||
const owner = `runtime-lock-${kind}`;
|
||||
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
|
||||
const auth = JSON.stringify({ tokens: { access_token: "valid-access", refresh_token: "valid-refresh", account_id: owner } });
|
||||
await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth);
|
||||
const [row] = await service.quotaAccounts(companyId, owner);
|
||||
const ref = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!;
|
||||
const prepare = () => prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local",
|
||||
binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} });
|
||||
const held = deferredSignal(), release = deferredSignal();
|
||||
const lock = kind === "company_file" ? vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock")
|
||||
.mockRejectedValueOnce(Object.assign(new Error("company lock held by another account"), { code: WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE })) : undefined;
|
||||
const holding = kind === "company_file" ? Promise.resolve() : db.transaction(async tx => {
|
||||
if (kind === "grant_row") await tx.select().from(connectionGrants).where(eq(connectionGrants.id, row.grant.id)).for("update");
|
||||
else await tx.select().from(companySecrets).where(eq(companySecrets.id, ref.secretId)).for("update");
|
||||
held.resolve(); await release.promise;
|
||||
});
|
||||
try {
|
||||
if (kind !== "company_file") await held.promise;
|
||||
const error = await prepare().then(() => { throw new Error("Runtime must defer"); }, error => error);
|
||||
expect(isAiConnectionBusy(error)).toBe(true);
|
||||
expect(error).toMatchObject({ status: 422, details: { code: "ai_connection_busy" } });
|
||||
expect((await db.select().from(connectionGrants).where(eq(connectionGrants.id, row.grant.id)))[0].status).toBe("active");
|
||||
} finally { release.resolve(); await holding; lock?.mockRestore(); }
|
||||
const run = await prepare();
|
||||
try { expect(await readFile(path.join(String(run.config.env.CODEX_HOME), "auth.json"), "utf8")).toBe(auth); }
|
||||
finally { await run.cleanup(); }
|
||||
});
|
||||
|
||||
it.each([
|
||||
["runtime", "rotation"], ["runtime", "commit"], ["failed_run", "rotation"], ["failed_run", "commit"],
|
||||
] as const)("recovers pending quota credentials before %s can reuse or invalidate the consumed token after %s failure", async (consumer, failure) => {
|
||||
const owner = `quota-consumer-${consumer}-${failure}`;
|
||||
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
|
||||
const auth = (marker: string) => JSON.stringify({ tokens: { access_token: `${marker}-access`, refresh_token: `${marker}-refresh`, account_id: owner } });
|
||||
await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth("old"));
|
||||
const [row] = await service.quotaAccounts(companyId, owner);
|
||||
const prepare = () => prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local",
|
||||
binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} });
|
||||
const oldRun = await prepare();
|
||||
const ref = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!;
|
||||
const journal = quotaCredentialRecovery(companyId, row.grant.id);
|
||||
await journal.write({ companyId, grantId: row.grant.id, connectionId: row.connection.id, secretId: ref.secretId,
|
||||
baseHash: quotaCredentialHash(auth("old")), value: auth("new") });
|
||||
const fail = () => service.markAuthenticationFailed({ companyId, agentId, runStartedAt: new Date(),
|
||||
attribution: { ...oldRun.attribution, identity: oldRun.identity } });
|
||||
await db.execute(sql.raw(`create function reject_consumer_recovery() returns trigger language plpgsql as $$ begin
|
||||
if new.secret_id = '${ref.secretId}'::uuid then raise exception 'recovery unavailable'; end if; return new; end $$;
|
||||
create ${failure === "commit" ? "constraint " : ""}trigger reject_consumer_recovery ${failure === "commit" ? "after" : "before"} insert on company_secret_versions
|
||||
${failure === "commit" ? "deferrable initially deferred" : ""} for each row execute function reject_consumer_recovery();`));
|
||||
try {
|
||||
if (consumer === "runtime") {
|
||||
const error = await prepare().then(() => { throw new Error("Runtime must defer"); }, error => error);
|
||||
expect(isAiConnectionBusy(error)).toBe(true);
|
||||
expect(error).toMatchObject({ status: 422, details: { code: "ai_connection_busy" } });
|
||||
} else await expect(fail()).rejects.toThrow();
|
||||
expect((await db.select().from(connectionGrants).where(eq(connectionGrants.id, row.grant.id)))[0].status).toBe("active");
|
||||
expect(await service.credential(row)).toBe(auth("old"));
|
||||
expect(await journal.read()).not.toBeNull();
|
||||
} finally {
|
||||
await db.execute(sql`drop trigger reject_consumer_recovery on company_secret_versions; drop function reject_consumer_recovery()`);
|
||||
}
|
||||
const request = vi.spyOn(globalThis, "fetch").mockRejectedValue(new Error("Recovery must not exchange again"));
|
||||
let nextRun: Awaited<ReturnType<typeof prepare>> | undefined;
|
||||
try {
|
||||
if (consumer === "failed_run") await fail();
|
||||
nextRun = await prepare();
|
||||
expect(await readFile(path.join(String(nextRun.config.env.CODEX_HOME), "auth.json"), "utf8")).toBe(auth("new"));
|
||||
expect(nextRun.sessionIdentity).toBe(oldRun.sessionIdentity);
|
||||
await fail(); // A late failure of the old run cannot invalidate the recovered identity.
|
||||
expect((await db.select().from(connectionGrants).where(eq(connectionGrants.id, row.grant.id)))[0].status).toBe("active");
|
||||
expect(await service.credential(row)).toBe(auth("new"));
|
||||
expect(await journal.read()).toBeNull();
|
||||
expect(request).not.toHaveBeenCalled();
|
||||
expect(await db.select().from(companySecretVersions).where(eq(companySecretVersions.secretId, ref.secretId))).toHaveLength(2);
|
||||
} finally { request.mockRestore(); await oldRun.cleanup(); await nextRun?.cleanup(); await journal.clear(); }
|
||||
});
|
||||
|
||||
it.each(["rollback", "lost_commit_reply"] as const)("automatically saves one exchanged token across %s", async (failure) => {
|
||||
const owner = `quota-once-${failure}`;
|
||||
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
|
||||
const original = JSON.stringify({ tokens: { access_token: "expired", refresh_token: "single-use", account_id: owner } });
|
||||
await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, original);
|
||||
const [row] = await service.quotaAccounts(companyId, owner);
|
||||
const transaction = db.transaction.bind(db);
|
||||
const interrupted = vi.spyOn(db, "transaction").mockImplementationOnce(async (work, config) => {
|
||||
await transaction(async (tx) => {
|
||||
await work(tx);
|
||||
if (failure === "rollback") throw new Error("Injected rollback after writes");
|
||||
}, config);
|
||||
throw new Error("Injected lost commit response");
|
||||
});
|
||||
const request = vi.spyOn(globalThis, "fetch").mockResolvedValue(Response.json({ access_token: "fresh-once", refresh_token: "replacement-once" }));
|
||||
try {
|
||||
await service.refreshQuotaCredential(row, original, new AbortController().signal);
|
||||
expect(JSON.parse(await service.credential(row))).toMatchObject({ tokens: { access_token: "fresh-once", refresh_token: "replacement-once" } });
|
||||
expect(request).toHaveBeenCalledTimes(1);
|
||||
await expect(access(quotaCredentialRecoveryPath(companyId, row.grant.id))).rejects.toThrow();
|
||||
} finally { interrupted.mockRestore(); request.mockRestore(); }
|
||||
});
|
||||
|
||||
it.each(["reconnected", "expired_reconnect", "revoked"] as const)("does not apply a pending refresh over a %s identity", async (change) => {
|
||||
const owner = `quota-replay-${change}`;
|
||||
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
|
||||
const original = JSON.stringify({ tokens: { access_token: "expired", refresh_token: "single-use", account_id: owner } });
|
||||
await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, original);
|
||||
const [row] = await service.quotaAccounts(companyId, owner);
|
||||
const ref = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!;
|
||||
const journal = quotaCredentialRecovery(companyId, row.grant.id);
|
||||
await journal.write({ companyId, grantId: row.grant.id, connectionId: row.connection.id, secretId: ref.secretId, baseHash: quotaCredentialHash(original), value: "old-pending-replacement" });
|
||||
const request = vi.spyOn(globalThis, "fetch");
|
||||
try {
|
||||
if (change === "reconnected") {
|
||||
await secretService(db).rotate(ref.secretId, { value: "new-authorized-identity" });
|
||||
expect(await service.refreshQuotaCredential(row, original, new AbortController().signal)).toBe("new-authorized-identity");
|
||||
expect(await service.credential(row)).toBe("new-authorized-identity");
|
||||
expect(await journal.read()).toBeNull();
|
||||
} else if (change === "expired_reconnect") {
|
||||
const reconnected = JSON.stringify({ tokens: { access_token: "expired-new-account", refresh_token: "new-account-refresh", account_id: owner } });
|
||||
await secretService(db).rotate(ref.secretId, { value: reconnected });
|
||||
request.mockResolvedValueOnce(Response.json({ access_token: "new-account-fresh", refresh_token: "new-account-replacement" }));
|
||||
await service.refreshQuotaCredential(row, reconnected, new AbortController().signal);
|
||||
expect(request).toHaveBeenCalledOnce();
|
||||
expect(JSON.parse(String(request.mock.calls[0][1]?.body))).toMatchObject({ refresh_token: "new-account-refresh" });
|
||||
expect(JSON.parse(await service.credential(row))).toMatchObject({ tokens: { access_token: "new-account-fresh", refresh_token: "new-account-replacement" } });
|
||||
expect(await journal.read()).toBeNull();
|
||||
} else {
|
||||
await db.update(connectionGrants).set({ status: "revoked" }).where(eq(connectionGrants.id, row.grant.id));
|
||||
await expect(service.refreshQuotaCredential(row, original, new AbortController().signal)).rejects.toThrow("credentials_unavailable");
|
||||
expect(await journal.read()).not.toBeNull();
|
||||
}
|
||||
if (change !== "expired_reconnect") expect(request).not.toHaveBeenCalled();
|
||||
} finally { request.mockRestore(); await journal.clear(); }
|
||||
});
|
||||
|
||||
it("saves exchanged tokens when the quota deadline expires during response-body reading", async () => {
|
||||
const owner = "quota-body-deadline-owner";
|
||||
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
|
||||
const original = JSON.stringify({ tokens: { access_token: "expired", refresh_token: "single-use", account_id: "deadline-account" } });
|
||||
await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: "Slow token body", loginSessionId: "fixture", agentIds: [], allAgents: true }, original);
|
||||
const [row] = await service.quotaAccounts(companyId, owner);
|
||||
const reading = deferredSignal(), releaseBody = deferredSignal(), quotaDeadline = new AbortController();
|
||||
let refreshSignal: AbortSignal | null | undefined;
|
||||
const request = vi.spyOn(globalThis, "fetch").mockImplementation(async (_url, options) => {
|
||||
refreshSignal = options?.signal;
|
||||
return { ok: true, json: async () => {
|
||||
reading.resolve(); await releaseBody.promise;
|
||||
refreshSignal?.throwIfAborted();
|
||||
return { access_token: "fresh-after-deadline", refresh_token: "replacement" };
|
||||
} } as Response;
|
||||
});
|
||||
const saving = service.refreshQuotaCredential(row, original, quotaDeadline.signal);
|
||||
try {
|
||||
await reading.promise;
|
||||
quotaDeadline.abort();
|
||||
expect(refreshSignal?.aborted).toBe(false);
|
||||
releaseBody.resolve();
|
||||
await saving;
|
||||
expect(JSON.parse(await service.credential(row))).toMatchObject({ tokens: { access_token: "fresh-after-deadline", refresh_token: "replacement" } });
|
||||
await expect(service.refreshQuotaCredential(row, original, quotaDeadline.signal)).rejects.toThrow();
|
||||
expect(request).toHaveBeenCalledTimes(1);
|
||||
} finally { releaseBody.resolve(); await saving.catch(() => {}); request.mockRestore(); }
|
||||
});
|
||||
|
||||
it("retries another account's token write-back while a slow quota exchange holds the company lock", async () => {
|
||||
const auth = (account: string, marker: string) => JSON.stringify({ tokens: { account_id: account, id_token: `id-${marker}`, access_token: `access-${marker}`, refresh_token: `refresh-${marker}` }, last_refresh: marker === "old" ? "2026-09-10T10:00:00Z" : "2026-09-10T11:00:00Z" });
|
||||
const quotaOwner = "slow-quota-owner", runOwner = "slow-quota-other-run";
|
||||
await db.insert(companyMemberships).values([quotaOwner, runOwner].map(principalId => ({ companyId, principalId, principalType: "user" as const, status: "active" as const, membershipRole: "member" as const })));
|
||||
for (const owner of [quotaOwner, runOwner]) await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth(owner, "old"));
|
||||
const [quotaRow] = await service.quotaAccounts(companyId, quotaOwner);
|
||||
// Runtime credential reads now serialize with quota exchanges. Prepare this
|
||||
// already-running, different account before testing its contended write-back.
|
||||
const preparedRun = await prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: runOwner, adapterType: "codex_local", binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} });
|
||||
const started = deferredSignal(), release = deferredSignal();
|
||||
const request = vi.spyOn(globalThis, "fetch").mockImplementation(async () => {
|
||||
started.resolve(); await release.promise;
|
||||
return Response.json({ access_token: "quota-fresh", refresh_token: "quota-replacement" });
|
||||
});
|
||||
const refreshing = service.refreshQuotaCredential(quotaRow, auth(quotaOwner, "old"), new AbortController().signal);
|
||||
let run: Awaited<ReturnType<typeof prepareManagedAiRuntime>> | undefined;
|
||||
let cleaning: Promise<void> | undefined;
|
||||
let lock: ReturnType<typeof vi.spyOn> | undefined;
|
||||
try {
|
||||
await started.promise;
|
||||
run = preparedRun;
|
||||
const authFile = path.join(String(run.config.env.CODEX_HOME), "auth.json");
|
||||
await writeFile(authFile, auth(runOwner, "replacement"));
|
||||
// Inject the two 30-second acquisition expirations, then use the real
|
||||
// still-held lock and real Postgres write-back for the final attempt.
|
||||
lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock")
|
||||
.mockRejectedValueOnce(Object.assign(new Error("lock timeout"), { code: WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE }))
|
||||
.mockRejectedValueOnce(Object.assign(new Error("lock timeout"), { code: WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE }));
|
||||
cleaning = run.cleanup();
|
||||
await vi.waitFor(() => expect(lock).toHaveBeenCalledTimes(3));
|
||||
expect(await readFile(authFile, "utf8")).toBe(auth(runOwner, "replacement"));
|
||||
release.resolve();
|
||||
await Promise.all([refreshing, cleaning]);
|
||||
const [runRow] = await service.quotaAccounts(companyId, runOwner);
|
||||
expect(await service.credential(runRow)).toBe(auth(runOwner, "replacement"));
|
||||
expect(JSON.parse(await service.credential(quotaRow))).toMatchObject({ tokens: { access_token: "quota-fresh", refresh_token: "quota-replacement" } });
|
||||
await expect(access(run.home!)).rejects.toMatchObject({ code: "ENOENT" });
|
||||
} finally { release.resolve(); await Promise.allSettled([refreshing, cleaning]); lock?.mockRestore(); request.mockRestore(); }
|
||||
});
|
||||
|
||||
it.each(["lock", "database"])("preserves replacement tokens after exhausted %s write-back and permits a later retry", async failure => {
|
||||
const owner = `retained-writeback-${failure}`;
|
||||
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
|
||||
const auth = (hour: string) => JSON.stringify({ tokens: { account_id: owner, id_token: "identity", access_token: `access-${hour}`, refresh_token: `refresh-${hour}` }, last_refresh: `2026-09-10T${hour}:00:00Z` });
|
||||
await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth("10"));
|
||||
const run = await prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local", binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} });
|
||||
const authFile = path.join(String(run.config.env.CODEX_HOME), "auth.json");
|
||||
await writeFile(authFile, auth("11"));
|
||||
const injected = failure === "lock"
|
||||
? vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock").mockRejectedValue(Object.assign(new Error("lock timeout"), { code: WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE }))
|
||||
: vi.spyOn(db, "transaction").mockRejectedValue(new Error("database unavailable"));
|
||||
try {
|
||||
await expect(run.cleanup()).rejects.toThrow(failure === "lock" ? "lock timeout" : "database unavailable");
|
||||
expect(injected).toHaveBeenCalledTimes(failure === "lock" ? 3 : 1);
|
||||
expect(await readFile(authFile, "utf8")).toBe(auth("11"));
|
||||
} finally { injected.mockRestore(); }
|
||||
await run.cleanup();
|
||||
const [row] = await service.quotaAccounts(companyId, owner);
|
||||
expect(await service.credential(row)).toBe(auth("11"));
|
||||
await expect(access(run.home!)).rejects.toMatchObject({ code: "ENOENT" });
|
||||
});
|
||||
|
||||
it("serializes quota exchange with a concurrent secret edit and runtime credential read", async () => {
|
||||
const owner = "quota-race-owner";
|
||||
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
|
||||
const original = JSON.stringify({ tokens: { access_token: "old-access", refresh_token: "single-use", id_token: "identity", account_id: "race-account" } });
|
||||
const account = await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: "Racing refresh", loginSessionId: "fixture", agentIds: [], allAgents: true }, original);
|
||||
const [row] = await service.quotaAccounts(companyId, owner);
|
||||
const secretId = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!.secretId;
|
||||
const exchangeStarted = deferredSignal();
|
||||
const releaseExchange = deferredSignal();
|
||||
const request = vi.spyOn(globalThis, "fetch").mockImplementation(async () => {
|
||||
exchangeStarted.resolve(); await releaseExchange.promise;
|
||||
return Response.json({ access_token: "new-access", refresh_token: "new-refresh" });
|
||||
});
|
||||
let runtime: Awaited<ReturnType<typeof prepareManagedAiRuntime>> | undefined;
|
||||
let lock: ReturnType<typeof vi.spyOn> | undefined;
|
||||
const refreshing = service.refreshQuotaCredential(row, original, new AbortController().signal);
|
||||
try {
|
||||
await exchangeStarted.promise;
|
||||
const edit = secretService(db).update(secretId, { description: "Changed during refresh" });
|
||||
lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock");
|
||||
const preparing = prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local", binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} });
|
||||
// Runtime reads now wait at the company file lock before touching the
|
||||
// old credential. The returned home must contain the committed rotation.
|
||||
let prepared = false;
|
||||
void preparing.then(() => { prepared = true; });
|
||||
await vi.waitFor(() => expect(lock).toHaveBeenCalled());
|
||||
expect(prepared).toBe(false);
|
||||
releaseExchange.resolve();
|
||||
const results = await Promise.all([refreshing, edit, preparing]);
|
||||
runtime = results[2];
|
||||
expect(JSON.parse(await readFile(path.join(String(runtime.config.env.CODEX_HOME), "auth.json"), "utf8"))).toMatchObject({ tokens: { access_token: "new-access", refresh_token: "new-refresh" } });
|
||||
expect(JSON.parse(await service.credential(row))).toMatchObject({ tokens: { access_token: "new-access", refresh_token: "new-refresh" } });
|
||||
expect((await secretService(db).getById(secretId))?.description).toBe("Changed during refresh");
|
||||
} finally { releaseExchange.resolve(); await runtime?.cleanup(); request.mockRestore(); lock?.mockRestore(); }
|
||||
}, 10_000);
|
||||
|
||||
it("serializes quota credential refresh, persists rotation, and defers while the provider is active", async () => {
|
||||
const owner = "quota-refresh-owner";
|
||||
await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
|
||||
const original = JSON.stringify({ tokens: { access_token: "expired-access", refresh_token: "single-use", id_token: "identity", account_id: null } });
|
||||
const account = await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: "Refresh account", loginSessionId: "fixture", agentIds: [], allAgents: true }, original);
|
||||
const [row] = await service.quotaAccounts(companyId, owner);
|
||||
const request = vi.spyOn(globalThis, "fetch").mockImplementation(async () => new Response(JSON.stringify({ access_token: "fresh-access", refresh_token: "rotated-once" })));
|
||||
const signal = new AbortController().signal;
|
||||
try {
|
||||
// A pre-existing transaction may own a secret row before asking for the
|
||||
// company file lock. Refresh must defer before exchanging its token.
|
||||
await db.transaction(async tx => {
|
||||
const secretId = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!.secretId;
|
||||
await tx.select().from(companySecrets).where(eq(companySecrets.id, secretId)).for("update");
|
||||
await expect(service.refreshQuotaCredential(row, original, signal)).rejects.toThrow();
|
||||
expect(request).not.toHaveBeenCalled();
|
||||
});
|
||||
const activeId = randomUUID();
|
||||
await db.insert(heartbeatRuns).values({ id: activeId, companyId, agentId, invocationSource: "on_demand", status: "running", contextSnapshot: { aiConnection: { grantId: account.grantId } } });
|
||||
await expect(service.refreshQuotaCredential(row, original, signal)).rejects.toThrow("provider_unavailable");
|
||||
expect(request).not.toHaveBeenCalled();
|
||||
await db.update(heartbeatRuns).set({ status: "succeeded" }).where(eq(heartbeatRuns.id, activeId));
|
||||
const [first, second] = await Promise.all([
|
||||
service.refreshQuotaCredential(row, original, signal),
|
||||
service.refreshQuotaCredential(row, original, signal),
|
||||
]);
|
||||
expect(first).toBe(second);
|
||||
expect(request).toHaveBeenCalledTimes(1);
|
||||
expect(JSON.parse(await service.credential(row))).toMatchObject({ tokens: { access_token: "fresh-access", refresh_token: "rotated-once", account_id: null } });
|
||||
await db.update(connectionGrants).set({ status: "revoked" }).where(eq(connectionGrants.id, account.grantId));
|
||||
await expect(service.refreshQuotaCredential(row, first, signal)).rejects.toThrow("credentials_unavailable");
|
||||
expect(request).toHaveBeenCalledTimes(1);
|
||||
} finally { request.mockRestore(); }
|
||||
});
|
||||
|
||||
it.each([false, true])("reports the authoritative connection-manager capability for custom grants (manager: %s)", async (manager) => {
|
||||
const userId = `custom-manager-${manager}`;
|
||||
await db.insert(companyMemberships).values({ companyId, principalType: "user", principalId: userId, status: "active", membershipRole: "member" });
|
||||
@@ -188,6 +700,18 @@ describe("managed AI connections", () => {
|
||||
await db.insert(heartbeatRuns).values({ companyId, agentId: id, status: "succeeded", responsibleUserId: "alice", contextSnapshot: { issueId }, createdAt: new Date(Date.now() + 1000) });
|
||||
expect(await intents.requestForRunAuthFailure(runId)).toBeNull();
|
||||
});
|
||||
|
||||
it("lists quotas only for the current member's accessible subscription accounts", async () => {
|
||||
const user = `quota-${randomUUID()}`;
|
||||
await db.insert(companyMemberships).values({ companyId, principalId: user, principalType: "user", status: "active", membershipRole: "member" });
|
||||
const subscription = await service.save(companyId, user, { provider: "openai", method: "subscription", ownership: "personal", name: "Private quota", loginSessionId: "fixture", allAgents: true, agentIds: [] }, JSON.stringify({ tokens: { access_token: "fixture", account_id: "fixture", refresh_token: "fixture" } }));
|
||||
expect((await service.quotaAccounts(companyId, user)).some(row => row.connection.id === subscription.connectionId)).toBe(true);
|
||||
expect((await service.quotaAccounts(companyId, "bob")).some(row => row.connection.id === subscription.connectionId)).toBe(false);
|
||||
expect(await service.quotaAccounts(otherCompanyId, user)).toEqual([]);
|
||||
await db.update(companyMemberships).set({ status: "inactive" }).where(and(eq(companyMemberships.companyId, companyId), eq(companyMemberships.principalId, user)));
|
||||
expect(await service.quotaAccounts(companyId, user)).toEqual([]);
|
||||
});
|
||||
|
||||
it.each([
|
||||
["anthropic", "claude_local", "CLAUDE_CODE_OAUTH_TOKEN", "ANTHROPIC_API_KEY"],
|
||||
["openai", "codex_local", "CODEX_HOME", "OPENAI_API_KEY"],
|
||||
@@ -992,3 +1516,9 @@ describe("AI connection recovery delivery", () => {
|
||||
}, 30000,
|
||||
);
|
||||
});
|
||||
|
||||
function deferredSignal() {
|
||||
let resolve!: () => void;
|
||||
const promise = new Promise<void>(done => { resolve = done; });
|
||||
return { promise, resolve };
|
||||
}
|
||||
@@ -0,0 +1,256 @@
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import type { Db } from "@paperclipai/db";
|
||||
const mocks = vi.hoisted(() => ({
|
||||
accounts: vi.fn(),
|
||||
credential: vi.fn(),
|
||||
refresh: vi.fn(),
|
||||
codex: vi.fn(),
|
||||
claude: vi.fn(),
|
||||
adapters: vi.fn(),
|
||||
}));
|
||||
vi.mock("../services/ai-connections.js", () => ({
|
||||
aiConnectionService: () => ({
|
||||
quotaAccounts: mocks.accounts,
|
||||
credential: mocks.credential,
|
||||
refreshQuotaCredential: mocks.refresh,
|
||||
}),
|
||||
}));
|
||||
vi.mock("../adapters/registry.js", () => ({
|
||||
listServerAdapters: mocks.adapters,
|
||||
}));
|
||||
vi.mock("@paperclipai/adapter-codex-local/server", () => ({
|
||||
fetchCodexQuota: mocks.codex,
|
||||
}));
|
||||
vi.mock("@paperclipai/adapter-claude-local/server", () => ({
|
||||
fetchClaudeQuota: mocks.claude,
|
||||
}));
|
||||
import { fetchCompanyQuotaWindows } from "../services/quota-windows.js";
|
||||
|
||||
let revision = 1;
|
||||
const db = {
|
||||
select: () => ({
|
||||
from: () => ({
|
||||
leftJoin: () => ({
|
||||
leftJoin: () => ({
|
||||
where: async () => [
|
||||
{
|
||||
id: "secret",
|
||||
updatedAt: new Date(0),
|
||||
latestVersion: revision,
|
||||
status: "active",
|
||||
},
|
||||
],
|
||||
}),
|
||||
}),
|
||||
}),
|
||||
}),
|
||||
} as unknown as Db;
|
||||
const account = (id: string, provider = "openai", status = "connected") => ({
|
||||
connection: { id, updatedAt: new Date(0) },
|
||||
grant: {
|
||||
id,
|
||||
updatedAt: new Date(0),
|
||||
credentialSecretRefs: [{ secretId: "secret" }],
|
||||
},
|
||||
summary: { provider, name: id, status },
|
||||
});
|
||||
describe("connected account quotas", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
revision = 1;
|
||||
mocks.refresh.mockRejectedValue(new Error("authentication_required"));
|
||||
mocks.credential.mockResolvedValue(
|
||||
JSON.stringify({
|
||||
tokens: { access_token: "private-token", account_id: "actual-account" },
|
||||
}),
|
||||
);
|
||||
mocks.codex.mockResolvedValue([
|
||||
{ label: "5h", usedPercent: 0.5, resetsAt: null, valueLabel: null },
|
||||
]);
|
||||
});
|
||||
afterEach(() => vi.useRealTimers());
|
||||
it("uses the selected account without invoking a local CLI or fallback login", async () => {
|
||||
mocks.accounts.mockResolvedValue([account("remote-codex")]);
|
||||
const result = await fetchCompanyQuotaWindows(db, "company-1", "user-1");
|
||||
expect(mocks.accounts).toHaveBeenCalledWith("company-1", "user-1");
|
||||
expect(mocks.codex).toHaveBeenCalledWith("private-token", "actual-account", expect.any(AbortSignal));
|
||||
expect(mocks.adapters).not.toHaveBeenCalled();
|
||||
expect(result[0]).toMatchObject({
|
||||
ok: true,
|
||||
accountLabel: "remote-codex",
|
||||
source: "managed-connection",
|
||||
});
|
||||
expect(result[0].capturedAt).toBeDefined();
|
||||
expect(JSON.stringify(result)).not.toContain("private-token");
|
||||
});
|
||||
it("deduplicates account requests and invalidates on credential revision", async () => {
|
||||
mocks.accounts.mockResolvedValue([account("cached-codex")]);
|
||||
const [a, b] = await Promise.all([
|
||||
fetchCompanyQuotaWindows(db, "company-2", "user-2"),
|
||||
fetchCompanyQuotaWindows(db, "company-2", "user-2"),
|
||||
]);
|
||||
expect(mocks.codex).toHaveBeenCalledTimes(1);
|
||||
expect(a[0].capturedAt).toBe(b[0].capturedAt);
|
||||
revision++;
|
||||
const next = await fetchCompanyQuotaWindows(db, "company-2", "user-2");
|
||||
expect(mocks.codex).toHaveBeenCalledTimes(2);
|
||||
expect(next[0].accountKey).not.toBe(a[0].accountKey);
|
||||
});
|
||||
it("keeps two accounts for the same provider separate", async () => {
|
||||
mocks.accounts.mockResolvedValue([account("first"), account("second")]);
|
||||
const result = await fetchCompanyQuotaWindows(db, "company-3", "user-3");
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0].accountKey).not.toBe(result[1].accountKey);
|
||||
});
|
||||
it("does not expose failures and marks authentication failures for stale-window removal", async () => {
|
||||
mocks.accounts.mockResolvedValue([account("expired")]);
|
||||
mocks.codex.mockRejectedValue(
|
||||
new Error("401 private-token raw provider diagnostics"),
|
||||
);
|
||||
const [result] = await fetchCompanyQuotaWindows(db, "company-4", "user-4");
|
||||
expect(result.errorFamily).toBe("authentication_required");
|
||||
expect(result.windows).toEqual([]);
|
||||
expect(JSON.stringify(result)).not.toContain("private-token");
|
||||
});
|
||||
it("does not read revoked credentials or use the server login for an unconnected remote account", async () => {
|
||||
mocks.accounts.mockResolvedValue([account("revoked", "openai", "revoked")]);
|
||||
const [result] = await fetchCompanyQuotaWindows(db, "company-5", "user-5");
|
||||
expect(result.errorFamily).toBe("credentials_unavailable");
|
||||
expect(mocks.credential).not.toHaveBeenCalled();
|
||||
mocks.accounts.mockResolvedValue([]);
|
||||
const empty = await fetchCompanyQuotaWindows(db, "company-5", "user-5");
|
||||
expect(
|
||||
empty.every((row) => row.errorFamily === "credentials_unavailable"),
|
||||
).toBe(true);
|
||||
expect(mocks.adapters).not.toHaveBeenCalled();
|
||||
});
|
||||
it("uses a managed Claude OAuth credential and keeps failed reads private", async () => {
|
||||
mocks.accounts.mockResolvedValue([account("claude-account", "anthropic")]);
|
||||
mocks.credential.mockResolvedValue("claude-private-oauth");
|
||||
mocks.claude.mockResolvedValue([
|
||||
{ label: "Weekly", usedPercent: 12, resetsAt: null, valueLabel: null },
|
||||
]);
|
||||
const [result] = await fetchCompanyQuotaWindows(db, "company-6", "user-6");
|
||||
expect(result.ok).toBe(true);
|
||||
expect(mocks.claude).toHaveBeenCalledWith("claude-private-oauth", expect.any(AbortSignal));
|
||||
mocks.accounts.mockResolvedValue([account("secret-failure")]);
|
||||
mocks.credential.mockRejectedValue(
|
||||
new Error("private credential store failure"),
|
||||
);
|
||||
const [failed] = await fetchCompanyQuotaWindows(db, "company-6", "user-6");
|
||||
expect(failed.errorFamily).toBe("provider_unavailable");
|
||||
expect(JSON.stringify(failed)).not.toContain("private credential");
|
||||
});
|
||||
it("reads quota for a verified login without an optional account ID", async () => {
|
||||
mocks.accounts.mockResolvedValue([account("missing-identity")]);
|
||||
mocks.credential.mockResolvedValue(
|
||||
JSON.stringify({ accessToken: "private" }),
|
||||
);
|
||||
const [result] = await fetchCompanyQuotaWindows(db, "company-7", "user-7");
|
||||
expect(result.ok).toBe(true);
|
||||
expect(mocks.codex).toHaveBeenCalledWith("private", null, expect.any(AbortSignal));
|
||||
});
|
||||
it("distinguishes permission errors from revoked credentials", async () => {
|
||||
mocks.accounts.mockResolvedValue([account("permission-error")]);
|
||||
mocks.codex.mockRejectedValueOnce(new Error("chatgpt wham api returned 403"));
|
||||
const [result] = await fetchCompanyQuotaWindows(db, "permission-company", "user");
|
||||
expect(result.errorFamily).toBe("permission_denied");
|
||||
expect(mocks.refresh).not.toHaveBeenCalled();
|
||||
});
|
||||
it("refreshes an expired selected credential once before retrying quota", async () => {
|
||||
mocks.accounts.mockResolvedValue([account("refresh-account")]);
|
||||
mocks.codex.mockRejectedValueOnce(new Error("chatgpt wham api returned 401"));
|
||||
mocks.refresh.mockImplementationOnce(async () => {
|
||||
const value = JSON.stringify({ tokens: { access_token: "new-token", account_id: "same-account" } });
|
||||
revision++;
|
||||
mocks.credential.mockResolvedValue(value);
|
||||
return value;
|
||||
});
|
||||
const [result] = await fetchCompanyQuotaWindows(db, "refresh-company", "user");
|
||||
expect(result.ok).toBe(true);
|
||||
expect(mocks.refresh).toHaveBeenCalledTimes(1);
|
||||
expect(mocks.codex).toHaveBeenLastCalledWith("new-token", "same-account", expect.any(AbortSignal));
|
||||
expect(JSON.stringify(result)).not.toContain("new-token");
|
||||
const again = await fetchCompanyQuotaWindows(db, "refresh-company", "user");
|
||||
expect(again[0]).toEqual(result);
|
||||
expect(mocks.codex).toHaveBeenCalledTimes(2); // initial 401 and one refreshed read
|
||||
|
||||
});
|
||||
it("does not cache a quota read under a credential revision that changed during resolution", async () => {
|
||||
mocks.accounts.mockResolvedValue([account("rotating-during-refresh")]);
|
||||
mocks.codex.mockRejectedValueOnce(new Error("401"));
|
||||
mocks.refresh.mockResolvedValue(JSON.stringify({ accessToken: "refreshed" }));
|
||||
mocks.credential.mockResolvedValueOnce(JSON.stringify({ accessToken: "expired" })).mockImplementationOnce(async () => {
|
||||
revision++;
|
||||
return JSON.stringify({ accessToken: "changed-again" });
|
||||
});
|
||||
const [result] = await fetchCompanyQuotaWindows(db, "changing-revision", "user");
|
||||
expect(result).toMatchObject({ ok: false, errorFamily: "provider_unavailable" });
|
||||
expect(mocks.codex).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("bounds a stalled credential lookup and preserves its account identity", async () => {
|
||||
vi.useFakeTimers();
|
||||
mocks.accounts.mockResolvedValue([account("stalled")]);
|
||||
mocks.credential.mockReturnValue(new Promise(() => {}));
|
||||
const pending = fetchCompanyQuotaWindows(db, "company-8", "user-8");
|
||||
await vi.advanceTimersByTimeAsync(20_001);
|
||||
const [result] = await pending;
|
||||
expect(result).toMatchObject({
|
||||
ok: false,
|
||||
accountLabel: "stalled",
|
||||
windows: [],
|
||||
});
|
||||
expect(result.accountKey).toBeTruthy();
|
||||
expect(result.error).not.toContain("timed out");
|
||||
});
|
||||
it("does not start a provider request when a credential resolves after its deadline", async () => {
|
||||
vi.useFakeTimers();
|
||||
mocks.accounts.mockResolvedValue([account("late-credential")]);
|
||||
let resolveCredential!: (value: string) => void;
|
||||
mocks.credential.mockReturnValue(new Promise<string>((resolve) => { resolveCredential = resolve; }));
|
||||
const pending = fetchCompanyQuotaWindows(db, "late-company", "late-user");
|
||||
await vi.advanceTimersByTimeAsync(20_001);
|
||||
expect((await pending)[0].ok).toBe(false);
|
||||
resolveCredential(JSON.stringify({ accessToken: "private", accountId: "actual" }));
|
||||
await vi.advanceTimersByTimeAsync(0);
|
||||
expect(mocks.codex).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("cancels an active provider read when the account deadline expires", async () => {
|
||||
vi.useFakeTimers();
|
||||
mocks.accounts.mockResolvedValue([account("slow-provider")]);
|
||||
let signal!: AbortSignal;
|
||||
mocks.codex.mockImplementation((_token, _id, abortSignal) => {
|
||||
signal = abortSignal;
|
||||
return new Promise((_resolve, reject) => signal.addEventListener("abort", () => reject(signal.reason), { once: true }));
|
||||
});
|
||||
const pending = fetchCompanyQuotaWindows(db, "slow-company", "slow-user");
|
||||
await vi.advanceTimersByTimeAsync(20_001);
|
||||
expect((await pending)[0].ok).toBe(false);
|
||||
expect(signal.aborted).toBe(true);
|
||||
});
|
||||
|
||||
it("limits account probes to four concurrent requests and preserves account order", async () => {
|
||||
mocks.accounts.mockResolvedValue(
|
||||
Array.from({ length: 9 }, (_, i) => account(`parallel-${i}`)),
|
||||
);
|
||||
mocks.credential.mockResolvedValue(
|
||||
JSON.stringify({ accessToken: "private", accountId: "actual" }),
|
||||
);
|
||||
let active = 0,
|
||||
maximum = 0;
|
||||
mocks.codex.mockImplementation(async () => {
|
||||
active++;
|
||||
maximum = Math.max(maximum, active);
|
||||
await new Promise((resolve) => setTimeout(resolve, 5));
|
||||
active--;
|
||||
return [];
|
||||
});
|
||||
const results = await fetchCompanyQuotaWindows(db, "company-9", "user-9");
|
||||
expect(maximum).toBe(4);
|
||||
expect(results.map((row) => row.accountLabel)).toEqual(
|
||||
Array.from({ length: 9 }, (_, i) => `parallel-${i}`),
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,60 @@
|
||||
import { promises as fs } from "node:fs";
|
||||
import os from "node:os";
|
||||
import path from "node:path";
|
||||
import { afterEach, describe, expect, it, vi } from "vitest";
|
||||
import { ensureDurableDirectory } from "../lib/durable-directory.js";
|
||||
|
||||
const roots: string[] = [];
|
||||
afterEach(async () => {
|
||||
vi.restoreAllMocks();
|
||||
await Promise.all(roots.splice(0).map(root => fs.rm(root, { recursive: true, force: true })));
|
||||
});
|
||||
|
||||
describe.skipIf(process.platform === "win32")("durable directory creation", () => {
|
||||
async function fixture() {
|
||||
const root = await fs.mkdtemp(path.join(os.tmpdir(), "paperclip-durable-directory-"));
|
||||
roots.push(root);
|
||||
return { root, directory: path.join(root, "instance", "recovery", "company") };
|
||||
}
|
||||
|
||||
function observeSyncs(failOn?: string) {
|
||||
const synced: string[] = [];
|
||||
const open = fs.open.bind(fs);
|
||||
vi.spyOn(fs, "open").mockImplementation(async (...args: Parameters<typeof fs.open>) => {
|
||||
const handle = await open(...args);
|
||||
const sync = handle.sync.bind(handle);
|
||||
vi.spyOn(handle, "sync").mockImplementation(async () => {
|
||||
if (args[0] === failOn) throw new Error("Injected directory fsync failure");
|
||||
await sync();
|
||||
synced.push(String(args[0]));
|
||||
});
|
||||
return handle;
|
||||
});
|
||||
return synced;
|
||||
}
|
||||
|
||||
it("flushes every created path entry and avoids ancestor fsyncs on subsequent writes", async () => {
|
||||
const { root, directory } = await fixture();
|
||||
const synced = observeSyncs();
|
||||
await ensureDurableDirectory(directory);
|
||||
expect(synced.slice(0, 4)).toEqual([directory, path.dirname(directory), path.join(root, "instance"), root]);
|
||||
synced.length = 0;
|
||||
await ensureDurableDirectory(directory);
|
||||
expect(synced).toEqual([]);
|
||||
await fs.rm(directory, { recursive: true });
|
||||
await ensureDurableDirectory(directory);
|
||||
expect(synced.slice(0, 4)).toEqual([directory, path.dirname(directory), path.join(root, "instance"), root]);
|
||||
});
|
||||
|
||||
it("retries all ancestors after a failed flush even though mkdir already succeeded", async () => {
|
||||
const { root, directory } = await fixture();
|
||||
observeSyncs(root);
|
||||
await expect(ensureDurableDirectory(directory)).rejects.toThrow("Injected directory fsync failure");
|
||||
expect((await fs.stat(directory)).isDirectory()).toBe(true);
|
||||
vi.restoreAllMocks();
|
||||
const synced = observeSyncs();
|
||||
await ensureDurableDirectory(directory);
|
||||
expect(synced).toContain(root);
|
||||
expect(synced).toContain(path.dirname(root));
|
||||
});
|
||||
});
|
||||
@@ -3,9 +3,11 @@ import { describe, expect, it, vi, beforeEach, afterEach } from "vitest";
|
||||
vi.mock("../adapters/registry.js", () => ({
|
||||
listServerAdapters: vi.fn(),
|
||||
}));
|
||||
vi.mock("../middleware/logger.js", () => ({ logger: { warn: vi.fn() } }));
|
||||
|
||||
import { listServerAdapters } from "../adapters/registry.js";
|
||||
import { fetchAllQuotaWindows } from "../services/quota-windows.js";
|
||||
import { logger } from "../middleware/logger.js";
|
||||
|
||||
describe("fetchAllQuotaWindows", () => {
|
||||
beforeEach(() => {
|
||||
@@ -48,9 +50,47 @@ describe("fetchAllQuotaWindows", () => {
|
||||
{
|
||||
provider: "anthropic",
|
||||
ok: false,
|
||||
error: "quota polling timed out after 20s",
|
||||
error: "Subscription quota is currently unavailable. Check usage with your provider.",
|
||||
windows: [],
|
||||
},
|
||||
]);
|
||||
});
|
||||
|
||||
it("keeps command diagnostics out of API responses and redacts them in server logs", async () => {
|
||||
const diagnostic = 'Command failed: sh -c probe --token fixture-private-token';
|
||||
vi.mocked(listServerAdapters).mockReturnValue([
|
||||
{ type: "claude_local", getQuotaWindows: vi.fn().mockResolvedValue({
|
||||
provider: "anthropic", ok: false, windows: [], error: diagnostic,
|
||||
}) },
|
||||
{ type: "codex_local", getQuotaWindows: vi.fn().mockRejectedValue(new Error(diagnostic)) },
|
||||
] as never);
|
||||
|
||||
const results = await fetchAllQuotaWindows();
|
||||
expect(results).toHaveLength(2);
|
||||
for (const result of results) {
|
||||
expect(result.ok).toBe(false);
|
||||
expect(result.error).toBe("Subscription quota is currently unavailable. Check usage with your provider.");
|
||||
}
|
||||
expect(JSON.stringify(results)).not.toContain("Command failed");
|
||||
expect(logger.warn).toHaveBeenCalledWith(expect.objectContaining({
|
||||
adapterType: "claude_local", diagnostic: expect.stringContaining("Command failed"),
|
||||
}), "Provider subscription quota unavailable");
|
||||
expect(JSON.stringify(vi.mocked(logger.warn).mock.calls)).not.toContain("fixture-private-token");
|
||||
});
|
||||
|
||||
it("isolates synchronous probe failures and preserves structured auth failure information", async () => {
|
||||
vi.mocked(listServerAdapters).mockReturnValue([
|
||||
{ type: "claude_local", getQuotaWindows: () => { throw new Error("local command unavailable"); } },
|
||||
{ type: "codex_local", getQuotaWindows: vi.fn().mockResolvedValue({
|
||||
provider: "openai", source: "codex-rpc", ok: false,
|
||||
errorFamily: "refresh_token_expired", error: "private diagnostic", windows: [],
|
||||
}) },
|
||||
] as never);
|
||||
const results = await fetchAllQuotaWindows();
|
||||
expect(results[0]).toMatchObject({ provider: "anthropic", ok: false });
|
||||
expect(results[1]).toMatchObject({
|
||||
provider: "openai", source: "codex-rpc", ok: false, errorFamily: "refresh_token_expired",
|
||||
});
|
||||
expect(JSON.stringify(results)).not.toContain("private diagnostic");
|
||||
});
|
||||
});
|
||||
@@ -137,7 +137,7 @@ describe("WHAM used_percent normalization via fetchCodexQuota", () => {
|
||||
expect(windows[0]!.usedPercent).toBe(50);
|
||||
});
|
||||
|
||||
it("treats values < 1 as fraction and multiplies by 100 (0.5 → 50%)", async () => {
|
||||
it("preserves fractional percentage points (0.5 → 0.5%)", async () => {
|
||||
mockFetch({
|
||||
rate_limit: {
|
||||
primary_window: {
|
||||
@@ -148,11 +148,10 @@ describe("WHAM used_percent normalization via fetchCodexQuota", () => {
|
||||
},
|
||||
});
|
||||
const windows = await fetchCodexQuota("token", null);
|
||||
expect(windows[0]!.usedPercent).toBe(50);
|
||||
expect(windows[0]!.usedPercent).toBe(0.5);
|
||||
});
|
||||
|
||||
it("treats value exactly 1.0 as 1% (not 100%) — the < 1 heuristic boundary", async () => {
|
||||
// 1.0 is NOT < 1, so it is treated as already-percentage → 1%
|
||||
it("treats value exactly 1.0 as 1%", async () => {
|
||||
mockFetch({
|
||||
rate_limit: {
|
||||
primary_window: {
|
||||
@@ -166,6 +165,16 @@ describe("WHAM used_percent normalization via fetchCodexQuota", () => {
|
||||
expect(windows[0]!.usedPercent).toBe(1);
|
||||
});
|
||||
|
||||
it.each([
|
||||
[0.01, 0.01], [0.99, 0.99], [1.01, 1.01], [-1, 0],
|
||||
[Number.NaN, null], [Number.POSITIVE_INFINITY, null], [undefined, null],
|
||||
])("handles percentage boundary %s consistently in RPC and WHAM", async (raw, expected) => {
|
||||
mockFetch({ rate_limit: { primary_window: { used_percent: raw } } });
|
||||
const windows = await fetchCodexQuota("token", null);
|
||||
expect(windows[0].usedPercent).toBe(expected);
|
||||
expect(mapCodexRpcQuota({ rateLimits: { primary: { usedPercent: raw } } }).windows[0].usedPercent).toBe(expected);
|
||||
});
|
||||
|
||||
it("treats value 0 as 0%", async () => {
|
||||
mockFetch({
|
||||
rate_limit: {
|
||||
@@ -820,6 +829,29 @@ describe("mapCodexRpcQuota", () => {
|
||||
// fetchWithTimeout — abort on timeout
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
describe("managed quota cancellation", () => {
|
||||
afterEach(() => vi.unstubAllGlobals());
|
||||
|
||||
it.each([
|
||||
["Claude", (signal: AbortSignal) => fetchClaudeQuota("token", signal)],
|
||||
["Codex", (signal: AbortSignal) => fetchCodexQuota("token", "account", signal)],
|
||||
] as const)("forwards cancellation to the %s provider", async (_provider, read) => {
|
||||
const controller = new AbortController();
|
||||
let requestSignal!: AbortSignal;
|
||||
vi.stubGlobal("fetch", vi.fn(async (_url: string, init: RequestInit) => {
|
||||
requestSignal = init.signal!;
|
||||
return new Promise<Response>((_resolve, reject) => {
|
||||
requestSignal.addEventListener("abort", () => reject(requestSignal.reason), { once: true });
|
||||
});
|
||||
}));
|
||||
const pending = read(controller.signal);
|
||||
const rejected = expect(pending).rejects.toThrow("account deadline");
|
||||
controller.abort(new Error("account deadline"));
|
||||
await rejected;
|
||||
expect(requestSignal.aborted).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe("fetchWithTimeout", () => {
|
||||
afterEach(() => {
|
||||
vi.unstubAllGlobals();
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
import { promises as fs } from "node:fs";
|
||||
import path from "node:path";
|
||||
|
||||
// These application-owned directories are stable after creation. Check their
|
||||
// identity on reuse, so deleting/recreating one cannot reuse its durability proof.
|
||||
const durableDirectories = new Map<string, string>();
|
||||
|
||||
export async function syncDirectory(directory: string) {
|
||||
if (process.platform === "win32") return;
|
||||
const handle = await fs.open(directory, "r");
|
||||
try { await handle.sync(); } finally { await handle.close(); }
|
||||
}
|
||||
|
||||
/** Make the complete directory path durable before acknowledging a file write.
|
||||
* A failed or interrupted mkdir/fsync may leave existing but unflushed ancestors,
|
||||
* so the first successful check must include them, not just newly created paths. */
|
||||
export async function ensureDurableDirectory(directory: string) {
|
||||
const resolved = path.resolve(directory);
|
||||
const created = await fs.mkdir(resolved, { recursive: true, mode: 0o700 });
|
||||
if (process.platform === "win32") return;
|
||||
const stat = await fs.stat(resolved, { bigint: true });
|
||||
const identity = `${stat.dev}:${stat.ino}:${stat.birthtimeNs}`;
|
||||
if (!created && durableDirectories.get(resolved) === identity) return;
|
||||
for (let current = resolved; ; current = path.dirname(current)) {
|
||||
await syncDirectory(current);
|
||||
if (path.dirname(current) === current) break;
|
||||
}
|
||||
// Bound bookkeeping across companies and disposable/custom spool paths.
|
||||
if (durableDirectories.size >= 1024) durableDirectories.delete(durableDirectories.keys().next().value!);
|
||||
durableDirectories.set(resolved, identity);
|
||||
}
|
||||
@@ -5,7 +5,7 @@ import { agentFileStore, agentFileTokenFromHash } from "../services/agent-file-s
|
||||
import { pipeline } from "node:stream/promises";
|
||||
import { resolveAgentAppearance, agentAvatarUrl } from "@paperclipai/shared";
|
||||
import { listOpenRouterModels } from "../services/openrouter-models.js";
|
||||
import { prepareManagedAiRuntime, assertManagedAiProjectAuth, stripAiAuthBindings } from "../services/ai-connection-runtime.js";
|
||||
import { prepareManagedAiRuntime, withManagedAiProbe, assertManagedAiProjectAuth, stripAiAuthBindings } from "../services/ai-connection-runtime.js";
|
||||
import { ADAPTER_AUTH_MISSING_CHECK_CODE, AI_CONNECTION_CAPABILITIES, aiConnectionBindingSchema, type AiConnectionBinding, aiRuntimeConnectionBindingSchema, type AiRuntimeConnectionBinding } from "@paperclipai/shared";
|
||||
import { toolConnections } from "@paperclipai/db";
|
||||
import { aiConnectionService } from "../services/ai-connections.js";
|
||||
@@ -3503,17 +3503,17 @@ export function agentRoutes(
|
||||
const testEnvironmentId = await resolveAdapterTestEnvironmentId(companyId, environmentId);
|
||||
if (testEnvironmentId) await assertAdapterTestEnvironmentForCompany(companyId, testEnvironmentId);
|
||||
const target = await resolveAdapterTestExecutionContext({ companyId, adapterType, environmentId: testEnvironmentId });
|
||||
let managed: Awaited<ReturnType<typeof prepareManagedAiRuntime>> | undefined;
|
||||
try {
|
||||
if (!target.executionTarget && target.fallbackChecks.length > 0) throw unprocessable("The agent environment is not available for adoption");
|
||||
managed = await prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: userId, adapterType, binding, config, allowUninstalledPersonal: newAgent, allowUninstalledShared, allowLegacyValidation: true });
|
||||
await withManagedAiProbe(db, { companyId, agentId, responsibleUserId: userId, adapterType, binding, config, allowUninstalledPersonal: newAgent, allowUninstalledShared, allowLegacyValidation: true }, async managed => {
|
||||
const result = await testManagedEnvironment(adapterType, { companyId, adapterType, config: managed.config, executionTarget: target.executionTarget, environmentName: target.environmentName }, binding, managed, agentId);
|
||||
if (result.status === "fail" || result.checks.some(check => check.code === ADAPTER_AUTH_MISSING_CHECK_CODE)) throw unprocessable("The selected AI connection failed validation in this agent’s environment. Run the agent test to see the failing checks.", {
|
||||
code: "ai_connection_validation_failed",
|
||||
checks: result.checks.filter(check => check.level === "error" || check.code === ADAPTER_AUTH_MISSING_CHECK_CODE).map(check => ({ code: check.code, level: check.level })),
|
||||
});
|
||||
if (selection.connection.config.aiLegacyAdoption === true) await db.update(toolConnections).set({ healthStatus: "ok", config: { ...selection.connection.config, aiLegacyAdoption: false }, updatedAt: new Date() }).where(eq(toolConnections.id, selection.connection.id));
|
||||
} finally { try { await managed?.cleanup(); } finally { await target.release("released"); } }
|
||||
});
|
||||
} finally { await target.release("released"); }
|
||||
}
|
||||
return { connectionId: selection.connection.id };
|
||||
}
|
||||
@@ -3706,12 +3706,11 @@ export function agentRoutes(
|
||||
effectiveAdapterConfig.apiKey = req.body.testCredentials.API_SERVER_KEY;
|
||||
}
|
||||
}
|
||||
const managed = aiBinding ? await prepareManagedAiRuntime(db, { companyId, agentId: req.body.agentId ?? "", responsibleUserId: responsibleUserForAiRequest(req), adapterType: type, binding: aiBinding, config: effectiveAdapterConfig, allowUninstalledPersonal: !req.body.agentId, allowUninstalledShared: !req.body.agentId && await canInstallSharedAiConnectionForNewAgent(db, req, companyId, aiBinding) }) : null;
|
||||
let result;
|
||||
try {
|
||||
result = managed && aiBinding ? await testManagedEnvironment(type, { companyId, adapterType: type, config: managed.config, executionTarget, environmentName }, aiBinding, managed, savedAgentId ?? undefined) : await adapter.testEnvironment({ companyId, adapterType: type, config: effectiveAdapterConfig, executionTarget, environmentName });
|
||||
if (managed) result.checks.unshift({ code: "ai_connection_tested", level: "info", message: `Tested ${managed.accountName} — ${managed.accountOwnerUserId ? managed.accountOwnerUserId === responsibleUserForAiRequest(req) ? "your personal account" : "the owner’s account authorized for this agent" : "company-shared account"}. Responsible user: ${req.actor.type === "agent" ? responsibleUserForAiRequest(req) ?? "unavailable" : "the signed-in user"}.` });
|
||||
} finally { await managed?.cleanup(); }
|
||||
const result = aiBinding ? await withManagedAiProbe(db, { companyId, agentId: req.body.agentId ?? "", responsibleUserId: responsibleUserForAiRequest(req), adapterType: type, binding: aiBinding, config: effectiveAdapterConfig, allowUninstalledPersonal: !req.body.agentId, allowUninstalledShared: !req.body.agentId && await canInstallSharedAiConnectionForNewAgent(db, req, companyId, aiBinding) }, async managed => {
|
||||
const tested = await testManagedEnvironment(type, { companyId, adapterType: type, config: managed.config, executionTarget, environmentName }, aiBinding, managed, savedAgentId ?? undefined);
|
||||
tested.checks.unshift({ code: "ai_connection_tested", level: "info", message: `Tested ${managed.accountName} — ${managed.accountOwnerUserId ? managed.accountOwnerUserId === responsibleUserForAiRequest(req) ? "your personal account" : "the owner’s account authorized for this agent" : "company-shared account"}. Responsible user: ${req.actor.type === "agent" ? responsibleUserForAiRequest(req) ?? "unavailable" : "the signed-in user"}.` });
|
||||
return tested;
|
||||
}) : await adapter.testEnvironment({ companyId, adapterType: type, config: effectiveAdapterConfig, executionTarget, environmentName });
|
||||
|
||||
const prefixChecks = [
|
||||
...(sandboxIdentityCheck ? [sandboxIdentityCheck] : []),
|
||||
|
||||
@@ -21,7 +21,7 @@ import {
|
||||
logActivity,
|
||||
} from "../services/index.js";
|
||||
import { assertBoard, assertCompanyAccess, getAccessibleResource, getActorInfo } from "./authz.js";
|
||||
import { fetchAllQuotaWindows } from "../services/quota-windows.js";
|
||||
import { fetchCompanyQuotaWindows } from "../services/quota-windows.js";
|
||||
import { badRequest } from "../errors.js";
|
||||
import type { PluginWorkerManager } from "../services/plugin-worker-manager.js";
|
||||
|
||||
@@ -281,7 +281,7 @@ export function costRoutes(
|
||||
res.status(404).json({ error: "Company not found" });
|
||||
return;
|
||||
}
|
||||
const results = await fetchAllQuotaWindows();
|
||||
const results = await fetchCompanyQuotaWindows(db, companyId, getActorInfo(req).actorId);
|
||||
res.json(results);
|
||||
});
|
||||
|
||||
|
||||
@@ -11,7 +11,8 @@ import {
|
||||
} from "@paperclipai/shared";
|
||||
import { aiConnectionService } from "./ai-connections.js";
|
||||
import { secretService } from "./secrets.js";
|
||||
import { decideCodexAuthMerge } from "@paperclipai/adapter-codex-local/server";
|
||||
import { decideCodexAuthMerge, withAccountHomeSecretMutationLock } from "@paperclipai/adapter-codex-local/server";
|
||||
import { WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE } from "@paperclipai/adapter-utils/workspace-restore-merge";
|
||||
import type { AdapterExecutionTarget } from "@paperclipai/adapter-utils/execution-target";
|
||||
import { runAdapterExecutionTargetProcess } from "@paperclipai/adapter-utils/execution-target";
|
||||
import { decideGrokAuthMerge } from "@paperclipai/adapter-grok-local/server";
|
||||
@@ -259,12 +260,29 @@ export async function prepareManagedAiRuntime(
|
||||
);
|
||||
const credentialRef = selection.grant.credentialSecretRefs.find((ref) => ref.configPath === "ai.credential");
|
||||
if (!credentialRef) throw unprocessable("The selected AI credential is unavailable");
|
||||
const readFreshness = async () => (await db.select({ epoch: companySecrets.aiSessionEpoch, version: companySecrets.latestVersion })
|
||||
.from(companySecrets).where(and(eq(companySecrets.companyId, input.companyId), eq(companySecrets.id, credentialRef.secretId))).limit(1))[0];
|
||||
const freshness = await readFreshness();
|
||||
const value = await service.credential(selection);
|
||||
const afterRead = await readFreshness();
|
||||
if (!freshness || freshness.version !== afterRead?.version || freshness.epoch !== afterRead.epoch) throw unprocessable("The AI credential changed during preparation; retry this execution");
|
||||
const readFreshness = async () => {
|
||||
const [currentGrant] = await db.select({ refs: connectionGrants.credentialSecretRefs, status: connectionGrants.status })
|
||||
.from(connectionGrants).where(and(eq(connectionGrants.companyId, input.companyId), eq(connectionGrants.id, selection.grant.id)));
|
||||
const currentRef = currentGrant?.refs.find(ref => ref.configPath === "ai.credential");
|
||||
if (currentGrant?.status !== "active" || currentRef?.secretId !== credentialRef.secretId || currentRef.versionSelector !== credentialRef.versionSelector) return undefined;
|
||||
return (await db.select({ epoch: companySecrets.aiSessionEpoch, version: companySecrets.latestVersion })
|
||||
.from(companySecrets).where(and(eq(companySecrets.companyId, input.companyId), eq(companySecrets.id, credentialRef.secretId))).limit(1))[0];
|
||||
};
|
||||
const { value, freshness } = await (async () => {
|
||||
// Recovering a rotated quota token can advance the secret version during
|
||||
// the first read. Re-read the saved credential, retaining the epoch guard
|
||||
// against a concurrent reconnect or explicit rotation.
|
||||
for (let attempt = 0; attempt < 2; attempt += 1) {
|
||||
const before = await readFreshness();
|
||||
const value = await service.runtimeCredential(selection);
|
||||
const after = await readFreshness();
|
||||
if (before && after && before.version === after.version && before.epoch === after.epoch) {
|
||||
return { value, freshness: after };
|
||||
}
|
||||
if (!before || before.epoch !== after?.epoch) break;
|
||||
}
|
||||
throw unprocessable("The AI credential changed during preparation; retry this execution");
|
||||
})();
|
||||
home = await mkdtemp(
|
||||
path.join(
|
||||
os.tmpdir(),
|
||||
@@ -325,75 +343,86 @@ export async function prepareManagedAiRuntime(
|
||||
identity,
|
||||
home,
|
||||
cleanup: async () => {
|
||||
try {
|
||||
if (subscriptionFile) {
|
||||
const refreshed = await readFile(authFile, "utf8");
|
||||
if (refreshed !== value)
|
||||
await db.transaction(async (tx) => {
|
||||
const [grant] = await tx
|
||||
.select()
|
||||
.from(connectionGrants)
|
||||
.where(
|
||||
and(
|
||||
eq(connectionGrants.id, selection.grant.id),
|
||||
eq(connectionGrants.companyId, input.companyId),
|
||||
),
|
||||
)
|
||||
.for("update");
|
||||
// A missing or revoked grant blocks the write-back. Among
|
||||
// active copies, the merge decision below keeps the
|
||||
// credential with the newest provider freshness field.
|
||||
if (!grant || grant.status !== "active") return;
|
||||
const ref = grant.credentialSecretRefs.find(
|
||||
(r) => r.configPath === "ai.credential",
|
||||
);
|
||||
if (!ref) return;
|
||||
// Lock the referenced secret row for the rest of this
|
||||
// transaction. The grant-row lock above does not cover it,
|
||||
// so an authorized rotation of this secret could otherwise
|
||||
// land between the read and the write below and be
|
||||
// overwritten by this stale write-back.
|
||||
await tx
|
||||
.select({ id: companySecrets.id })
|
||||
.from(companySecrets)
|
||||
.where(
|
||||
and(
|
||||
eq(companySecrets.id, ref.secretId),
|
||||
eq(companySecrets.companyId, input.companyId),
|
||||
),
|
||||
)
|
||||
.for("update");
|
||||
const current = await aiConnectionService(
|
||||
tx as unknown as Db,
|
||||
).credential({ ...selection, grant });
|
||||
const destination = path.join(
|
||||
providerHome,
|
||||
"current-auth.json",
|
||||
);
|
||||
await writeFile(destination, current, { mode: 0o600 });
|
||||
const decision =
|
||||
input.binding.provider === "openai"
|
||||
? await decideCodexAuthMerge(authFile, destination, {
|
||||
errorLabel: "AI account refresh",
|
||||
})
|
||||
: await decideGrokAuthMerge(authFile, destination, {
|
||||
errorLabel: "AI account refresh",
|
||||
});
|
||||
if (decision !== 10) return;
|
||||
await secretService(tx).rotate(
|
||||
ref.secretId,
|
||||
{ value: refreshed, preserveAiSessionEpoch: true },
|
||||
{ userId: grant.subjectUserId },
|
||||
);
|
||||
await tx
|
||||
.update(connectionGrants)
|
||||
.set({ updatedAt: new Date() })
|
||||
.where(eq(connectionGrants.id, grant.id));
|
||||
});
|
||||
if (subscriptionFile) {
|
||||
const refreshed = await readFile(authFile, "utf8");
|
||||
if (refreshed !== value) {
|
||||
// A quota exchange can hold this company lock for 60 seconds.
|
||||
// Retry its 30-second acquisition timeout without discarding the
|
||||
// provider's only copy of a rotated, single-use refresh token.
|
||||
const writeBack = () => withAccountHomeSecretMutationLock(undefined, input.companyId, () => db.transaction(async (tx) => {
|
||||
const [grant] = await tx
|
||||
.select()
|
||||
.from(connectionGrants)
|
||||
.where(
|
||||
and(
|
||||
eq(connectionGrants.id, selection.grant.id),
|
||||
eq(connectionGrants.companyId, input.companyId),
|
||||
),
|
||||
)
|
||||
.for("update");
|
||||
// A missing or revoked grant blocks the write-back. Among
|
||||
// active copies, the merge decision below keeps the
|
||||
// credential with the newest provider freshness field.
|
||||
if (!grant || grant.status !== "active") return;
|
||||
const ref = grant.credentialSecretRefs.find(
|
||||
(r) => r.configPath === "ai.credential",
|
||||
);
|
||||
if (!ref) return;
|
||||
// Lock the referenced secret row for the rest of this
|
||||
// transaction. The grant-row lock above does not cover it,
|
||||
// so an authorized rotation of this secret could otherwise
|
||||
// land between the read and the write below and be
|
||||
// overwritten by this stale write-back.
|
||||
await tx
|
||||
.select({ id: companySecrets.id })
|
||||
.from(companySecrets)
|
||||
.where(
|
||||
and(
|
||||
eq(companySecrets.id, ref.secretId),
|
||||
eq(companySecrets.companyId, input.companyId),
|
||||
),
|
||||
)
|
||||
.for("update");
|
||||
const current = await aiConnectionService(
|
||||
tx as unknown as Db,
|
||||
).credential({ ...selection, grant });
|
||||
const destination = path.join(
|
||||
providerHome,
|
||||
"current-auth.json",
|
||||
);
|
||||
await writeFile(destination, current, { mode: 0o600 });
|
||||
const decision =
|
||||
input.binding.provider === "openai"
|
||||
? await decideCodexAuthMerge(authFile, destination, {
|
||||
errorLabel: "AI account refresh",
|
||||
})
|
||||
: await decideGrokAuthMerge(authFile, destination, {
|
||||
errorLabel: "AI account refresh",
|
||||
});
|
||||
if (decision !== 10) return;
|
||||
await secretService(tx).rotate(
|
||||
ref.secretId,
|
||||
{ value: refreshed, preserveAiSessionEpoch: true },
|
||||
{ userId: grant.subjectUserId },
|
||||
);
|
||||
await tx
|
||||
.update(connectionGrants)
|
||||
.set({ updatedAt: new Date() })
|
||||
.where(eq(connectionGrants.id, grant.id));
|
||||
}));
|
||||
for (let attempt = 0; ; attempt++) {
|
||||
try { await writeBack(); break; }
|
||||
catch (error) {
|
||||
if (attempt >= 2 || (error as NodeJS.ErrnoException)?.code !== WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE) throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
if (home) await rm(home, { recursive: true, force: true });
|
||||
}
|
||||
// Preserve the private home on any failed write-back, including an
|
||||
// exhausted lock retry or database failure. The caller can retry the
|
||||
// same cleanup; only a committed save or an intentional discard
|
||||
// (revoked grant, older/different identity) permits deletion.
|
||||
if (home) await rm(home, { recursive: true, force: true });
|
||||
},
|
||||
};
|
||||
} catch (error) {
|
||||
@@ -401,3 +430,26 @@ export async function prepareManagedAiRuntime(
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
/** Manual tests/adoption have no heartbeat row for quota's active-run guard.
|
||||
* Hold the shared credential lock from resolution through the provider probe
|
||||
* and write-back. Nested credential writes reuse this lock's async ownership. */
|
||||
export async function withManagedAiProbe<T>(
|
||||
db: Db,
|
||||
input: Parameters<typeof prepareManagedAiRuntime>[1],
|
||||
probe: (runtime: Awaited<ReturnType<typeof prepareManagedAiRuntime>>) => Promise<T>,
|
||||
): Promise<T> {
|
||||
const run = async () => {
|
||||
const runtime = await prepareManagedAiRuntime(db, input);
|
||||
try { return await probe(runtime); }
|
||||
finally { await runtime.cleanup(); }
|
||||
};
|
||||
return input.binding.provider === "openai"
|
||||
? withAccountHomeSecretMutationLock(undefined, input.companyId, run).catch(error => {
|
||||
if ((error as NodeJS.ErrnoException)?.code === WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE) {
|
||||
throw unprocessable("AI credentials are being updated. Retry shortly.", { code: "ai_connection_busy" });
|
||||
}
|
||||
throw error;
|
||||
})
|
||||
: run();
|
||||
}
|
||||
@@ -1,6 +1,8 @@
|
||||
import { WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE } from "@paperclipai/adapter-utils/workspace-restore-merge";
|
||||
import { withAccountHomeSecretMutationLock } from "@paperclipai/adapter-codex-local/server";
|
||||
import { syncConnectionCredentialBindings } from "./connection-credential-bindings.js";
|
||||
import { createHash, randomUUID } from "node:crypto";
|
||||
import { and, eq, inArray, or } from "drizzle-orm";
|
||||
import { and, eq, inArray, or, sql } from "drizzle-orm";
|
||||
import {
|
||||
type Db,
|
||||
authUsers,
|
||||
@@ -8,6 +10,7 @@ import {
|
||||
aiConnectionDefaults,
|
||||
aiProviderDefaults,
|
||||
agents,
|
||||
heartbeatRuns,
|
||||
companyMemberships,
|
||||
companySecrets,
|
||||
userSecretDefinitions,
|
||||
@@ -32,6 +35,8 @@ import {
|
||||
type AiConnectionLoginIntent,
|
||||
} from "@paperclipai/shared";
|
||||
import { forbidden, notFound, unprocessable } from "../errors.js";
|
||||
import { quotaCredentialRecovery, quotaCredentialHash } from "./quota-credential-recovery.js";
|
||||
import { logger } from "../middleware/logger.js";
|
||||
import { logActivity } from "./activity-log.js";
|
||||
import { secretService } from "./secrets.js";
|
||||
import { probeAiConnectionUsage } from "./ai-connection-usage.js";
|
||||
@@ -400,7 +405,7 @@ export function aiConnectionService(db: Db) {
|
||||
} satisfies AiConnectionAttribution,
|
||||
};
|
||||
}
|
||||
async function credential(row: Pick<Awaited<ReturnType<typeof select>>, "connection" | "grant">) {
|
||||
async function credential(row: Pick<Awaited<ReturnType<typeof select>>, "connection" | "grant">, retry = 0): Promise<string> {
|
||||
const ref = row.grant.credentialSecretRefs.find(
|
||||
(r) => r.configPath === "ai.credential",
|
||||
);
|
||||
@@ -430,6 +435,7 @@ export function aiConnectionService(db: Db) {
|
||||
responsibleUserId: row.grant.subjectUserId,
|
||||
actorType: "system" as const,
|
||||
};
|
||||
let value: string;
|
||||
if (secret.scope === "user") {
|
||||
if (
|
||||
secret.ownerUserId !== row.grant.subjectUserId ||
|
||||
@@ -447,14 +453,21 @@ export function aiConnectionService(db: Db) {
|
||||
context,
|
||||
);
|
||||
if (!result) throw unprocessable("Reconnect this AI account");
|
||||
return result.value;
|
||||
value = result.value;
|
||||
} else {
|
||||
value = await secrets.resolveSecretValue(row.connection.companyId, secret.id, "latest", context);
|
||||
}
|
||||
return secrets.resolveSecretValue(
|
||||
row.connection.companyId,
|
||||
secret.id,
|
||||
"latest",
|
||||
context,
|
||||
);
|
||||
// Resolution records lastResolvedAt and can wait behind a rotating writer.
|
||||
// Re-read after that wait instead of handing a run the pre-rotation token.
|
||||
const [latest] = await db.select().from(companySecrets).where(and(
|
||||
eq(companySecrets.id, secret.id), eq(companySecrets.companyId, row.connection.companyId),
|
||||
));
|
||||
if (!latest || latest.status !== "active") throw unprocessable("Reconnect this AI account");
|
||||
if (latest.latestVersion !== secret.latestVersion) {
|
||||
if (retry >= 2) throw unprocessable("AI credentials are changing. Retry this execution.", { code: "ai_connection_busy" });
|
||||
return credential(row, retry + 1);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
async function save(
|
||||
companyId: string,
|
||||
@@ -497,7 +510,7 @@ export function aiConnectionService(db: Db) {
|
||||
);
|
||||
const id = reconnect?.connection.id ?? randomUUID();
|
||||
const grantId = reconnect?.grant.id ?? randomUUID();
|
||||
return db.transaction(async (tx) => {
|
||||
return withAccountHomeSecretMutationLock(undefined, companyId, () => db.transaction(async (tx) => {
|
||||
const secrets = secretService(tx);
|
||||
if (sessionId) {
|
||||
const [session] = await tx
|
||||
@@ -802,8 +815,79 @@ export function aiConnectionService(db: Db) {
|
||||
details: { provider: input.provider, method: input.method, grantId },
|
||||
});
|
||||
return { connectionId: id, grantId };
|
||||
}));
|
||||
}
|
||||
async function recoverQuotaCredentialInTransaction(
|
||||
tx: Db, row: Parameters<typeof credential>[0], current: string, onRecoveryStarted?: () => void,
|
||||
) {
|
||||
const metadata = aiConnectionMetadataSchema.safeParse(row.connection.config.ai);
|
||||
if (!metadata.success || metadata.data.provider !== "openai" || metadata.data.method !== "subscription") return null;
|
||||
const recovery = quotaCredentialRecovery(row.connection.companyId, row.grant.id);
|
||||
const pending = await recovery.read();
|
||||
if (!pending) return null;
|
||||
const ref = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential");
|
||||
if (!ref || pending.secretId !== ref.secretId || pending.connectionId !== row.connection.id || pending.baseHash !== quotaCredentialHash(current)) {
|
||||
await recovery.clear();
|
||||
return null;
|
||||
}
|
||||
onRecoveryStarted?.();
|
||||
await secretService(tx).rotate(ref.secretId, { value: pending.value, preserveAiSessionEpoch: true }, { userId: row.grant.subjectUserId });
|
||||
await tx.update(connectionGrants).set({ updatedAt: new Date() }).where(eq(connectionGrants.id, row.grant.id));
|
||||
await logActivity(tx, { companyId: row.connection.companyId, actorType: "system", actorId: "quota_recovery",
|
||||
action: "ai_connection.credential_refreshed", entityType: "tool_connection", entityId: row.connection.id,
|
||||
details: { provider: "openai", grantId: row.grant.id } });
|
||||
return pending.value;
|
||||
}
|
||||
|
||||
/** Never hand a consumed token to a new provider while its replacement awaits commit. */
|
||||
async function runtimeCredential(row: Parameters<typeof credential>[0]) {
|
||||
const metadata = aiConnectionMetadataSchema.safeParse(row.connection.config.ai);
|
||||
if (!metadata.success || metadata.data.provider !== "openai" || metadata.data.method !== "subscription") return credential(row);
|
||||
let savingReplacement = false;
|
||||
return withAccountHomeSecretMutationLock(undefined, row.connection.companyId, async () => {
|
||||
let recovered = false;
|
||||
const value = await db.transaction(async tx => {
|
||||
const [grant] = await tx.select().from(connectionGrants).where(and(
|
||||
eq(connectionGrants.id, row.grant.id), eq(connectionGrants.companyId, row.connection.companyId),
|
||||
eq(connectionGrants.connectionId, row.connection.id),
|
||||
)).for("update", { noWait: true });
|
||||
if (!grant || grant.status !== "active") throw new Error("credentials_unavailable");
|
||||
const ref = grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential");
|
||||
if (!ref) throw new Error("credentials_unavailable");
|
||||
await tx.select({ id: companySecrets.id }).from(companySecrets).where(and(
|
||||
eq(companySecrets.id, ref.secretId), eq(companySecrets.companyId, row.connection.companyId),
|
||||
)).for("update", { noWait: true });
|
||||
const selected = { connection: row.connection, grant };
|
||||
const current = await aiConnectionService(tx as unknown as Db).credential(selected);
|
||||
const pending = await recoverQuotaCredentialInTransaction(tx as unknown as Db, selected, current, () => {
|
||||
savingReplacement = true;
|
||||
});
|
||||
recovered = pending !== null;
|
||||
return pending ?? current;
|
||||
});
|
||||
if (recovered) await quotaCredentialRecovery(row.connection.companyId, row.grant.id).clear()
|
||||
.catch(error => logger.warn({ err: error, grantId: row.grant.id }, "Saved runtime credential recovery cleanup pending"));
|
||||
return value;
|
||||
}).catch(error => {
|
||||
// Once a matching replacement is found, a failed write or commit must
|
||||
// defer execution. Keep its journal so the next attempt can save it.
|
||||
if (savingReplacement) {
|
||||
throw unprocessable("AI credentials are being updated. Retry shortly.", { code: "ai_connection_busy" });
|
||||
}
|
||||
// Drizzle wraps Postgres errors in `cause`. Contention is a temporary
|
||||
// pre-provider wait, not evidence that this account needs reconnecting.
|
||||
let current: unknown = error;
|
||||
for (let depth = 0; depth < 4 && current && typeof current === "object"; depth += 1) {
|
||||
const value = current as { code?: unknown; cause?: unknown };
|
||||
if (value.code === WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE || value.code === "55P03") {
|
||||
throw unprocessable("AI credentials are being updated. Retry shortly.", { code: "ai_connection_busy" });
|
||||
}
|
||||
current = value.cause;
|
||||
}
|
||||
throw error;
|
||||
});
|
||||
}
|
||||
|
||||
/** A late failure must never invalidate credentials that were refreshed or reconnected meanwhile. */
|
||||
async function markAuthenticationFailed(input: {
|
||||
companyId: string;
|
||||
@@ -812,41 +896,158 @@ export function aiConnectionService(db: Db) {
|
||||
runStartedAt: Date;
|
||||
attribution: AiConnectionAttribution & { identity: string };
|
||||
}) {
|
||||
return db.transaction(async (tx) => {
|
||||
const { attribution } = input;
|
||||
const [grant] = await tx.select().from(connectionGrants).where(and(
|
||||
eq(connectionGrants.companyId, input.companyId),
|
||||
eq(connectionGrants.id, attribution.grantId),
|
||||
eq(connectionGrants.connectionId, attribution.connectionId),
|
||||
)).for("update");
|
||||
if (!grant || grant.status !== "active" || grant.updatedAt > input.runStartedAt) return;
|
||||
const [connection] = await tx.select().from(toolConnections).where(and(
|
||||
eq(toolConnections.companyId, input.companyId),
|
||||
eq(toolConnections.id, attribution.connectionId),
|
||||
eq(toolConnections.connectionPurpose, "ai"),
|
||||
));
|
||||
if (!connection) return;
|
||||
const metadata = aiConnectionMetadataSchema.safeParse(connection.config.ai);
|
||||
if (!metadata.success || metadata.data.provider !== attribution.provider || metadata.data.method !== attribution.method) return;
|
||||
const ref = grant.credentialSecretRefs.find((candidate) => candidate.configPath === "ai.credential");
|
||||
if (!ref) return;
|
||||
await tx.select({ id: companySecrets.id }).from(companySecrets).where(and(
|
||||
eq(companySecrets.companyId, input.companyId), eq(companySecrets.id, ref.secretId),
|
||||
)).for("update");
|
||||
const value = await aiConnectionService(tx as unknown as Db).credential({ connection, grant });
|
||||
const generation = createHash("sha256").update(value).digest("hex").slice(0, 16);
|
||||
if (attribution.identity !== `${grant.id}:${attribution.responsibleUserId ?? "shared"}:${generation}`) return;
|
||||
await tx.update(connectionGrants).set({ status: "needs_reauthorization", updatedAt: new Date() })
|
||||
.where(eq(connectionGrants.id, grant.id));
|
||||
await tx.update(toolConnections).set({ healthStatus: "error", healthMessage: "Sign in again to restore this AI connection.", updatedAt: new Date() })
|
||||
.where(eq(toolConnections.id, connection.id));
|
||||
await logActivity(tx as unknown as Db, {
|
||||
companyId: input.companyId, actorType: "system", actorId: input.runId ? "heartbeat" : "adapter_test",
|
||||
agentId: input.agentId, runId: input.runId, action: "ai_connection.authentication_failed",
|
||||
entityType: "tool_connection", entityId: connection.id,
|
||||
details: { provider: attribution.provider, grantId: grant.id },
|
||||
return withAccountHomeSecretMutationLock(undefined, input.companyId, async () => {
|
||||
let recovered = false;
|
||||
await db.transaction(async (tx) => {
|
||||
const { attribution } = input;
|
||||
const [grant] = await tx.select().from(connectionGrants).where(and(
|
||||
eq(connectionGrants.companyId, input.companyId),
|
||||
eq(connectionGrants.id, attribution.grantId),
|
||||
eq(connectionGrants.connectionId, attribution.connectionId),
|
||||
)).for("update");
|
||||
if (!grant || grant.status !== "active" || grant.updatedAt > input.runStartedAt) return;
|
||||
const [connection] = await tx.select().from(toolConnections).where(and(
|
||||
eq(toolConnections.companyId, input.companyId),
|
||||
eq(toolConnections.id, attribution.connectionId),
|
||||
eq(toolConnections.connectionPurpose, "ai"),
|
||||
));
|
||||
if (!connection) return;
|
||||
const metadata = aiConnectionMetadataSchema.safeParse(connection.config.ai);
|
||||
if (!metadata.success || metadata.data.provider !== attribution.provider || metadata.data.method !== attribution.method) return;
|
||||
const ref = grant.credentialSecretRefs.find((candidate) => candidate.configPath === "ai.credential");
|
||||
if (!ref) return;
|
||||
await tx.select({ id: companySecrets.id }).from(companySecrets).where(and(
|
||||
eq(companySecrets.companyId, input.companyId), eq(companySecrets.id, ref.secretId),
|
||||
)).for("update");
|
||||
const value = await aiConnectionService(tx as unknown as Db).credential({ connection, grant });
|
||||
const generation = createHash("sha256").update(value).digest("hex").slice(0, 16);
|
||||
if (attribution.identity !== `${grant.id}:${attribution.responsibleUserId ?? "shared"}:${generation}`) return;
|
||||
// A quota exchange may have consumed this run's token before a failed
|
||||
// database commit. Save its working replacement instead of revoking it.
|
||||
recovered = await recoverQuotaCredentialInTransaction(tx as unknown as Db, { connection, grant }, value) !== null;
|
||||
if (recovered) return;
|
||||
await tx.update(connectionGrants).set({ status: "needs_reauthorization", updatedAt: new Date() })
|
||||
.where(eq(connectionGrants.id, grant.id));
|
||||
await tx.update(toolConnections).set({ healthStatus: "error", healthMessage: "Sign in again to restore this AI connection.", updatedAt: new Date() })
|
||||
.where(eq(toolConnections.id, connection.id));
|
||||
await logActivity(tx as unknown as Db, {
|
||||
companyId: input.companyId, actorType: "system", actorId: input.runId ? "heartbeat" : "adapter_test",
|
||||
agentId: input.agentId, runId: input.runId, action: "ai_connection.authentication_failed",
|
||||
entityType: "tool_connection", entityId: connection.id,
|
||||
details: { provider: attribution.provider, grantId: grant.id },
|
||||
});
|
||||
});
|
||||
if (recovered) await quotaCredentialRecovery(input.companyId, input.attribution.grantId).clear()
|
||||
.catch(error => logger.warn({ err: error, grantId: input.attribution.grantId }, "Saved failed-run credential recovery cleanup pending"));
|
||||
});
|
||||
}
|
||||
return { list, select, credential, probeUsage, save, setDefault, membership, markAuthenticationFailed };
|
||||
/** Refresh only the selected vaulted identity, serializing rotation and reconnect. */
|
||||
async function refreshQuotaCredential(row: Parameters<typeof credential>[0], failedValue: string, signal: AbortSignal) {
|
||||
return withAccountHomeSecretMutationLock(undefined, row.connection.companyId, async () => {
|
||||
const recovery = quotaCredentialRecovery(row.connection.companyId, row.grant.id);
|
||||
let savingReplacement = false;
|
||||
const save = () => db.transaction(async (tx) => {
|
||||
// Do not wait on a transaction that might itself be waiting for the file
|
||||
// lock. All database locks must be available before consuming OAuth tokens.
|
||||
const companyId = row.connection.companyId;
|
||||
const [grant] = await tx.select().from(connectionGrants).where(and(
|
||||
eq(connectionGrants.id, row.grant.id), eq(connectionGrants.companyId, companyId),
|
||||
eq(connectionGrants.connectionId, row.connection.id),
|
||||
)).for("update", { noWait: true });
|
||||
if (!grant || grant.status !== "active") throw new Error("credentials_unavailable");
|
||||
const ref = grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential");
|
||||
if (!ref) throw new Error("credentials_unavailable");
|
||||
await tx.select({ id: companySecrets.id }).from(companySecrets).where(and(
|
||||
eq(companySecrets.id, ref.secretId), eq(companySecrets.companyId, companyId),
|
||||
)).for("update", { noWait: true });
|
||||
const current = await aiConnectionService(tx as unknown as Db).credential({ connection: row.connection, grant });
|
||||
signal.throwIfAborted();
|
||||
// Another poll, an agent, or reconnect already supplied a newer credential.
|
||||
if (current !== failedValue) return current;
|
||||
let pending = await recovery.read();
|
||||
if (pending && (pending.secretId !== ref.secretId || pending.connectionId !== row.connection.id || pending.baseHash !== quotaCredentialHash(current))) {
|
||||
// Reconnect may reuse this grant. Once the current authorized token
|
||||
// itself expires, an older journal must not prevent refreshing it.
|
||||
await recovery.clear();
|
||||
pending = null;
|
||||
}
|
||||
// An active provider may rotate this single-use token itself. Defer polling
|
||||
// until its credential write-back instead of invalidating its live copy.
|
||||
const [active] = await tx.select({ id: heartbeatRuns.id }).from(heartbeatRuns)
|
||||
.innerJoin(agents, and(eq(agents.id, heartbeatRuns.agentId), eq(agents.companyId, companyId))).where(and(
|
||||
eq(heartbeatRuns.companyId, companyId), inArray(heartbeatRuns.status, ["queued", "running"]),
|
||||
or(sql`${heartbeatRuns.contextSnapshot}->'aiConnection'->>'grantId' = ${grant.id}`,
|
||||
sql`${agents.runtimeConfig}->'aiConnection'->>'provider' = 'openai'`, eq(agents.adapterType, "codex_local")),
|
||||
)).limit(1);
|
||||
if (active && !pending) throw new Error("provider_unavailable");
|
||||
let refreshed = pending?.value;
|
||||
if (!refreshed) {
|
||||
await recovery.prepare();
|
||||
const auth = JSON.parse(current);
|
||||
const refreshToken = auth.tokens?.refresh_token;
|
||||
if (typeof refreshToken !== "string" || !refreshToken) throw new Error("authentication_required");
|
||||
// Once exchange starts, its lifetime is independent of the dashboard's
|
||||
// shorter deadline. A consumed single-use token must still be read/saved.
|
||||
signal.throwIfAborted();
|
||||
const refreshSignal = AbortSignal.timeout(60_000);
|
||||
// Matches the Codex OAuth client (openai/codex, login/src/auth/manager.rs).
|
||||
const response = await fetch("https://auth.openai.com/oauth/token", {
|
||||
method: "POST", redirect: "error", signal: refreshSignal,
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ client_id: "app_EMoamEEZ73f0CkXaXp7hrann", grant_type: "refresh_token", refresh_token: refreshToken }),
|
||||
});
|
||||
if (!response.ok) {
|
||||
await response.body?.cancel();
|
||||
throw new Error(response.status === 400 || response.status === 401 ? "authentication_required" : "provider_unavailable");
|
||||
}
|
||||
const tokens = await response.json() as Record<string, unknown>;
|
||||
if (typeof tokens.access_token !== "string" || !tokens.access_token) throw new Error("provider_unavailable");
|
||||
refreshed = JSON.stringify({ ...auth, tokens: { ...auth.tokens,
|
||||
access_token: tokens.access_token,
|
||||
...(typeof tokens.refresh_token === "string" && tokens.refresh_token ? { refresh_token: tokens.refresh_token } : {}),
|
||||
...(typeof tokens.id_token === "string" && tokens.id_token ? { id_token: tokens.id_token } : {}),
|
||||
}, last_refresh: new Date().toISOString() });
|
||||
// This fsynced encrypted journal survives transaction rollback and host
|
||||
// restart. Never exchange the consumed token again after a failed save.
|
||||
await recovery.write({ companyId, grantId: grant.id, connectionId: row.connection.id,
|
||||
secretId: ref.secretId, baseHash: quotaCredentialHash(current), value: refreshed });
|
||||
}
|
||||
savingReplacement = true;
|
||||
// Persist the rotated token even if the quota deadline expires after the
|
||||
// successful exchange; losing it would strand all subsequent executions.
|
||||
await secretService(tx).rotate(ref.secretId, { value: refreshed, preserveAiSessionEpoch: true }, { userId: grant.subjectUserId });
|
||||
await tx.update(connectionGrants).set({ updatedAt: new Date() }).where(eq(connectionGrants.id, grant.id));
|
||||
await logActivity(tx as unknown as Db, { companyId, actorType: "system", actorId: "quota",
|
||||
action: "ai_connection.credential_refreshed", entityType: "tool_connection", entityId: row.connection.id,
|
||||
details: { provider: "openai", grantId: grant.id } });
|
||||
return refreshed;
|
||||
});
|
||||
for (let attempt = 0; ; attempt++) {
|
||||
try {
|
||||
const value = await save();
|
||||
// Clear only after commit. A committed-but-lost response is detected
|
||||
// by the current credential comparison on the next attempt.
|
||||
await recovery.clear().catch((error) => logger.warn({ err: error, grantId: row.grant.id }, "Saved quota credential recovery cleanup pending"));
|
||||
return value;
|
||||
} catch (error) {
|
||||
if (!savingReplacement || attempt >= 2) throw error;
|
||||
// Saving retries have no provider side effect. Persistent failures
|
||||
// retain the journal for the next quota poll, including after restart.
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
// A quota read uses the same credential audience as execution. Operator status
|
||||
// alone never grants access to another member's personal subscription.
|
||||
async function quotaAccounts(companyId: string, userId: string) {
|
||||
if (!(await membership(companyId, userId))) return [];
|
||||
const visible = await list(companyId, userId);
|
||||
const accounts = await rows(companyId);
|
||||
return accounts.flatMap(row => {
|
||||
const summary = visible.find(item => item.id === row.connection.id && item.grantId === row.grant.id);
|
||||
if (!summary || summary.method !== "subscription" || !["openai", "anthropic"].includes(summary.provider)) return [];
|
||||
return [{ ...row, summary }];
|
||||
});
|
||||
}
|
||||
return { list, quotaAccounts, refreshQuotaCredential, select, credential, runtimeCredential, probeUsage, save, setDefault, membership, markAuthenticationFailed };
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
import { createHash, randomUUID } from "node:crypto";
|
||||
import { promises as fs } from "node:fs";
|
||||
import path from "node:path";
|
||||
import { z } from "zod";
|
||||
import { resolvePaperclipInstanceRoot } from "../home-paths.js";
|
||||
import { localEncryptedProvider } from "../secrets/local-encrypted-provider.js";
|
||||
import { ensureDurableDirectory, syncDirectory } from "../lib/durable-directory.js";
|
||||
|
||||
const recoverySchema = z.object({
|
||||
companyId: z.string().uuid(), grantId: z.string().uuid(), connectionId: z.string().uuid(), secretId: z.string().uuid(),
|
||||
baseHash: z.string().length(64), value: z.string().min(1),
|
||||
});
|
||||
type Recovery = z.infer<typeof recoverySchema>;
|
||||
export const quotaCredentialHash = (value: string) => createHash("sha256").update(value).digest("hex");
|
||||
export const quotaCredentialRecoveryPath = (companyId: string, grantId: string) =>
|
||||
path.join(resolvePaperclipInstanceRoot(), "quota-credential-recovery", companyId, `${grantId}.json`);
|
||||
|
||||
/** Called only while holding the company credential mutation lock. The local
|
||||
* encrypted provider uses the instance secrets master key, including when the
|
||||
* destination vault is remote. No plaintext token enters this journal. */
|
||||
export function quotaCredentialRecovery(companyId: string, grantId: string) {
|
||||
const file = quotaCredentialRecoveryPath(companyId, grantId);
|
||||
const directory = path.dirname(file);
|
||||
async function writeEncrypted(target: string, value: string) {
|
||||
const encrypted = await localEncryptedProvider.createVersion({ value });
|
||||
await ensureDurableDirectory(directory);
|
||||
const temporary = `${target}.${randomUUID()}.tmp`;
|
||||
const handle = await fs.open(temporary, "wx", 0o600);
|
||||
try {
|
||||
try { await handle.writeFile(JSON.stringify(encrypted.material)); await handle.sync(); }
|
||||
finally { await handle.close(); }
|
||||
await fs.rename(temporary, target);
|
||||
await syncDirectory(directory);
|
||||
} catch (error) { await fs.rm(temporary, { force: true }); throw error; }
|
||||
}
|
||||
return {
|
||||
// Verify encryption, storage, and fsync before exchanging a single-use token.
|
||||
prepare: async () => {
|
||||
const probe = `${file}.${randomUUID()}.probe`;
|
||||
try { await writeEncrypted(probe, "credential-recovery-probe"); }
|
||||
finally { await fs.rm(probe, { force: true }); }
|
||||
},
|
||||
read: async (): Promise<Recovery | null> => {
|
||||
let material: string;
|
||||
try { material = await fs.readFile(file, "utf8"); }
|
||||
catch (error) { if ((error as NodeJS.ErrnoException).code === "ENOENT") return null; throw error; }
|
||||
const pending = recoverySchema.parse(JSON.parse(await localEncryptedProvider.resolveVersion({ material: JSON.parse(material), externalRef: null })));
|
||||
if (pending.companyId !== companyId || pending.grantId !== grantId) throw new Error("credentials_unavailable");
|
||||
return pending;
|
||||
},
|
||||
write: (input: Recovery) => writeEncrypted(file, JSON.stringify(recoverySchema.parse(input))),
|
||||
clear: async () => { await fs.rm(file, { force: true }); },
|
||||
};
|
||||
}
|
||||
@@ -1,5 +1,135 @@
|
||||
import type { ProviderQuotaResult } from "@paperclipai/shared";
|
||||
import { redactDiagnosticText } from "@paperclipai/adapter-utils";
|
||||
import { listServerAdapters } from "../adapters/registry.js";
|
||||
import { logger } from "../middleware/logger.js";
|
||||
import { createHash } from "node:crypto";
|
||||
import { eq, inArray, and } from "drizzle-orm";
|
||||
import { companySecrets, companySecretVersions, userSecretDefinitions, type Db } from "@paperclipai/db";
|
||||
import { fetchCodexQuota } from "@paperclipai/adapter-codex-local/server";
|
||||
import { fetchClaudeQuota } from "@paperclipai/adapter-claude-local/server";
|
||||
import { aiConnectionService } from "./ai-connections.js";
|
||||
|
||||
const accountRequests = new Map<string, { expires: number; result: Promise<ProviderQuotaResult> }>();
|
||||
const publicError = "Subscription quota is currently unavailable. Check usage with your provider.";
|
||||
|
||||
/** Managed connections use their own credentials regardless of where agents run.
|
||||
* Never silently substitute the control-plane host's login for a remote account. */
|
||||
export async function fetchCompanyQuotaWindows(db: Db, companyId: string, userId: string): Promise<ProviderQuotaResult[]> {
|
||||
const service = aiConnectionService(db);
|
||||
const accounts = await service.quotaAccounts(companyId, userId);
|
||||
if (!accounts.length) {
|
||||
return ["anthropic", "openai"].map(provider => ({ provider, ok: false,
|
||||
accountKey: `unconnected:${companyId}:${provider}`, source: "managed-connection",
|
||||
errorFamily: "credentials_unavailable", error: "Connect a subscription account in AI connections to view its quota.", windows: [] }));
|
||||
}
|
||||
// Resolve revision metadata before consulting the cache, so rotation and
|
||||
// revocation cannot serve windows associated with an old credential.
|
||||
const secretIds = accounts.flatMap(row => row.grant.credentialSecretRefs.map(ref => ref.secretId));
|
||||
const readRevisions = async (ids: string[]) => ids.length ? await db.select({ id: companySecrets.id, latestVersion: companySecrets.latestVersion, status: companySecrets.status, versionStatus: companySecretVersions.status, revokedAt: companySecretVersions.revokedAt, definitionId: companySecrets.userSecretDefinitionId, definitionStatus: userSecretDefinitions.status, definitionDeletedAt: userSecretDefinitions.deletedAt })
|
||||
.from(companySecrets).leftJoin(companySecretVersions, and(eq(companySecretVersions.secretId, companySecrets.id), eq(companySecretVersions.version, companySecrets.latestVersion))).leftJoin(userSecretDefinitions, and(eq(userSecretDefinitions.id, companySecrets.userSecretDefinitionId), eq(userSecretDefinitions.companyId, companyId))).where(and(eq(companySecrets.companyId, companyId), inArray(companySecrets.id, ids))) : [];
|
||||
const revisions = await readRevisions(secretIds);
|
||||
const identity = (row: (typeof accounts)[number], metadata: typeof revisions) => ({
|
||||
provider: row.summary.provider,
|
||||
accountKey: createHash("sha256").update(JSON.stringify([companyId, row.connection.id, row.grant.id,
|
||||
row.connection.updatedAt, row.grant.updatedAt,
|
||||
row.grant.credentialSecretRefs.map(ref => [ref.secretId, metadata.find(r => r.id === ref.secretId)]),
|
||||
])).digest("hex"),
|
||||
accountLabel: row.summary.name, source: "managed-connection",
|
||||
});
|
||||
async function readAccount(row: (typeof accounts)[number]): Promise<ProviderQuotaResult> {
|
||||
let base = identity(row, revisions);
|
||||
const { accountKey } = base;
|
||||
const definitionUnavailable = row.grant.credentialSecretRefs.some(ref => {
|
||||
const revision = revisions.find(value => value.id === ref.secretId);
|
||||
return revision && (revision.status !== "active" || revision.versionStatus === "disabled" || revision.revokedAt != null
|
||||
|| (revision.definitionId && (revision.definitionStatus !== "active" || revision.definitionDeletedAt !== null)));
|
||||
});
|
||||
if (row.summary.status !== "connected" || definitionUnavailable) {
|
||||
return { ...base, ok: false, errorFamily: "credentials_unavailable", error: publicError, windows: [] };
|
||||
}
|
||||
const key = `${userId}:${accountKey}`;
|
||||
let pending = accountRequests.get(key);
|
||||
if (!pending || pending.expires <= Date.now()) {
|
||||
const controller = new AbortController();
|
||||
const result = (async (): Promise<ProviderQuotaResult> => {
|
||||
let credentialResolved = false;
|
||||
try {
|
||||
const value = await service.credential(row);
|
||||
controller.signal.throwIfAborted();
|
||||
credentialResolved = true;
|
||||
let windows;
|
||||
if (base.provider === "openai") {
|
||||
let auth;
|
||||
try { auth = JSON.parse(value); } catch { throw new Error("credentials_unavailable"); }
|
||||
const token = auth.tokens?.access_token ?? auth.accessToken;
|
||||
if (typeof token !== "string" || !token) {
|
||||
return { ...base, ok: false, errorFamily: "credentials_unavailable", error: publicError, windows: [] };
|
||||
}
|
||||
const read = (auth: { tokens?: { access_token?: string; account_id?: string | null }; accessToken?: string; accountId?: string | null }) =>
|
||||
fetchCodexQuota(auth.tokens?.access_token ?? auth.accessToken!, auth.tokens?.account_id ?? auth.accountId ?? null, controller.signal);
|
||||
try { windows = await read(auth); }
|
||||
catch (error) {
|
||||
if (!(error instanceof Error) || !/\b401\b/.test(error.message)) throw error;
|
||||
auth = JSON.parse(await service.refreshQuotaCredential(row, value, controller.signal));
|
||||
// Rotation changes both the grant and secret revision. Return and
|
||||
// cache the observation under the identity the next poll will see.
|
||||
const refreshed = (await service.quotaAccounts(companyId, userId)).find(candidate =>
|
||||
candidate.connection.id === row.connection.id && candidate.grant.id === row.grant.id);
|
||||
if (!refreshed || refreshed.summary.status !== "connected") throw new Error("credentials_unavailable");
|
||||
base = identity(refreshed, await readRevisions(refreshed.grant.credentialSecretRefs.map(ref => ref.secretId)));
|
||||
// A reconnect can race the completed refresh. Bind the value and
|
||||
// cache identity to one revision, rather than label old quota with
|
||||
// a newer account's key. A concurrent change defers this poll.
|
||||
const refreshedValue = await service.credential(refreshed);
|
||||
const verified = (await service.quotaAccounts(companyId, userId)).find(candidate =>
|
||||
candidate.connection.id === row.connection.id && candidate.grant.id === row.grant.id);
|
||||
if (!verified || verified.summary.status !== "connected") throw new Error("credentials_unavailable");
|
||||
const verifiedIdentity = identity(verified, await readRevisions(verified.grant.credentialSecretRefs.map(ref => ref.secretId)));
|
||||
if (verifiedIdentity.accountKey !== base.accountKey) throw new Error("provider_unavailable");
|
||||
auth = JSON.parse(refreshedValue);
|
||||
controller.signal.throwIfAborted();
|
||||
windows = await read(auth);
|
||||
}
|
||||
} else {
|
||||
windows = await fetchClaudeQuota(value, controller.signal);
|
||||
}
|
||||
return { ...base, ok: true, windows, capturedAt: new Date().toISOString() };
|
||||
} catch (error) {
|
||||
const message = error instanceof Error ? error.message : "Quota unavailable";
|
||||
const missing = message === "credentials_unavailable";
|
||||
const invalid = message === "authentication_required" || (credentialResolved && /\b401\b|refresh_token_(?:reused|expired|invalidated)/i.test(message));
|
||||
const errorFamily = missing ? "credentials_unavailable" : invalid ? "authentication_required"
|
||||
: credentialResolved && /\b403\b/.test(message) ? "permission_denied" : "provider_unavailable";
|
||||
logger.warn({ companyId, accountKey, provider: base.provider, authenticationFailed: invalid }, "Connected account quota unavailable");
|
||||
return { ...base, ok: false, windows: [], error: publicError,
|
||||
errorFamily };
|
||||
}
|
||||
})();
|
||||
const bounded: Promise<ProviderQuotaResult> = withQuotaTimeout(base.provider, result, () => controller.abort()).then(quota => {
|
||||
const observation = { ...base, ...quota, ...(quota.ok ? {} : { error: publicError }) };
|
||||
if (observation.ok && observation.accountKey !== accountKey) {
|
||||
if (accountRequests.get(key)?.result === bounded) accountRequests.delete(key);
|
||||
if (accountRequests.size >= 500) accountRequests.delete(accountRequests.keys().next().value!);
|
||||
accountRequests.set(`${userId}:${observation.accountKey}`, { result: Promise.resolve(observation), expires: Date.now() + 30_000 });
|
||||
}
|
||||
return observation;
|
||||
});
|
||||
pending = { result: bounded, expires: Date.now() + 30_000 };
|
||||
if (accountRequests.size >= 500) accountRequests.delete(accountRequests.keys().next().value!);
|
||||
accountRequests.set(key, pending);
|
||||
}
|
||||
return pending.result;
|
||||
}
|
||||
const results: ProviderQuotaResult[] = new Array(accounts.length);
|
||||
let next = 0;
|
||||
await Promise.all(Array.from({ length: Math.min(4, accounts.length) }, async () => {
|
||||
while (next < accounts.length) {
|
||||
const index = next++;
|
||||
results[index] = await readAccount(accounts[index]);
|
||||
}
|
||||
}));
|
||||
return results;
|
||||
}
|
||||
|
||||
const QUOTA_PROVIDER_TIMEOUT_MS = 20_000;
|
||||
|
||||
@@ -24,24 +154,31 @@ export async function fetchAllQuotaWindows(): Promise<ProviderQuotaResult[]> {
|
||||
const adapters = listServerAdapters().filter((a) => a.getQuotaWindows != null);
|
||||
|
||||
const settled = await Promise.allSettled(
|
||||
adapters.map((adapter) => withQuotaTimeout(adapter.type, adapter.getQuotaWindows!())),
|
||||
adapters.map((adapter) => withQuotaTimeout(adapter.type, Promise.resolve().then(() => adapter.getQuotaWindows!()))),
|
||||
);
|
||||
|
||||
return settled.map((result, i) => {
|
||||
if (result.status === "fulfilled") return result.value;
|
||||
const adapterType = adapters[i]!.type;
|
||||
return {
|
||||
const quota: ProviderQuotaResult = result.status === "fulfilled" ? result.value : {
|
||||
provider: providerSlugForAdapterType(adapterType),
|
||||
ok: false,
|
||||
error: String(result.reason),
|
||||
windows: [],
|
||||
};
|
||||
if (quota.ok) return { ...quota, error: undefined };
|
||||
logger.warn({
|
||||
adapterType,
|
||||
errorFamily: quota.errorFamily,
|
||||
diagnostic: redactDiagnosticText(quota.error ?? "Quota probe failed").slice(0, 2_000),
|
||||
}, "Provider subscription quota unavailable");
|
||||
return { ...quota, error: "Subscription quota is currently unavailable. Check usage with your provider." };
|
||||
});
|
||||
}
|
||||
|
||||
async function withQuotaTimeout(
|
||||
adapterType: string,
|
||||
task: Promise<ProviderQuotaResult>,
|
||||
onTimeout?: () => void,
|
||||
): Promise<ProviderQuotaResult> {
|
||||
let timeoutId: NodeJS.Timeout | null = null;
|
||||
try {
|
||||
@@ -49,6 +186,7 @@ async function withQuotaTimeout(
|
||||
task,
|
||||
new Promise<ProviderQuotaResult>((resolve) => {
|
||||
timeoutId = setTimeout(() => {
|
||||
onTimeout?.();
|
||||
resolve({
|
||||
provider: providerSlugForAdapterType(adapterType),
|
||||
ok: false,
|
||||
|
||||
Reference in new issue
Block a user