mirror of
https://github.com/paperclipai/paperclip.git
synced 2026-11-07 16:35:48 +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();
|
||||
|
||||
Reference in new issue
Block a user