From 7cf8df553d8324efdd7d4a28ce0bf22a0d78f8ed Mon Sep 17 00:00:00 2001 From: Devin Foley Date: Sat, 3 Oct 2026 21:30:51 -0700 Subject: [PATCH 1/6] fix: serialize managed credential refresh with reads and secret edits Co-Authored-By: Paperclip --- .../src/server/codex-auth-cache.test.ts | 36 +++++++++++++++++- .../src/server/codex-auth-cache.ts | 19 +++++++++- server/src/__tests__/ai-connections.test.ts | 37 +++++++++++++++++++ server/src/services/ai-connection-runtime.ts | 6 +-- server/src/services/ai-connections.ts | 16 +++++--- ui/src/components/AccountQuotaPanels.test.tsx | 3 +- ui/src/components/AccountQuotaPanels.tsx | 9 +++-- 7 files changed, 109 insertions(+), 17 deletions(-) diff --git a/packages/adapters/codex-local/src/server/codex-auth-cache.test.ts b/packages/adapters/codex-local/src/server/codex-auth-cache.test.ts index 0ca1bee28ca..2cd6ba7c173 100644 --- a/packages/adapters/codex-local/src/server/codex-auth-cache.test.ts +++ b/packages/adapters/codex-local/src/server/codex-auth-cache.test.ts @@ -463,13 +463,15 @@ describe("codex auth cache store", () => { releaseFirstCaller = resolve; }); + const entered = deferredSignal(); const firstCall = withAccountHomeSecretMutationLock(env, "company-shared", async () => { events.push("first-enter"); + entered.resolve(); await firstCallerGate; events.push("first-exit"); return "first"; }); - await new Promise((resolve) => setTimeout(resolve, 20)); + await entered.promise; const secondCall = withAccountHomeSecretMutationLock(env, "company-shared", async () => { events.push("second-enter"); events.push("second-exit"); @@ -497,12 +499,14 @@ describe("codex auth cache store", () => { releaseCompanyA = resolve; }); + const entered = deferredSignal(); const companyACall = withAccountHomeSecretMutationLock(env, "company-a", async () => { events.push("a-enter"); + entered.resolve(); await companyAGate; events.push("a-exit"); }); - await new Promise((resolve) => setTimeout(resolve, 20)); + await entered.promise; const companyBCall = withAccountHomeSecretMutationLock(env, "company-b", async () => { events.push("b-enter"); events.push("b-exit"); @@ -536,6 +540,28 @@ describe("codex auth cache store", () => { expect(events).toEqual(["promotion-enter", "mutation-enter", "mutation-exit", "promotion-exit"]); }); + it("allows awaited nested writes but expires inherited ownership after release", async () => { + const env = envFor(await makeInstanceRoot()); + const detachedGate = deferredSignal(); + let detached!: Promise; + let enteredDetached = false; + await withAccountHomeSecretMutationLock(env, "nested", async () => { + await expect(withAccountHomeSecretMutationLock(env, "nested", async () => "nested write")).resolves.toBe("nested write"); + detached = detachedGate.promise.then(() => withAccountHomeSecretMutationLock(env, "nested", async () => { enteredDetached = true; })); + }); + const holding = deferredSignal(); + const release = deferredSignal(); + const other = withAccountHomeSecretMutationLock(env, "nested", async () => { holding.resolve(); await release.promise; }); + await holding.promise; + detachedGate.resolve(); + try { + await new Promise(resolve => setTimeout(resolve, 30)); + expect(enteredDetached).toBe(false); + } finally { release.resolve(); } + await Promise.all([other, detached]); + expect(enteredDetached).toBe(true); + }); + }); describe("Phase 8: account-home directory validity before a secret write commits", () => { @@ -581,3 +607,9 @@ describe("codex auth cache store", () => { }); }); }); + +function deferredSignal() { + let resolve!: () => void; + const promise = new Promise(done => { resolve = done; }); + return { promise, resolve }; +} diff --git a/packages/adapters/codex-local/src/server/codex-auth-cache.ts b/packages/adapters/codex-local/src/server/codex-auth-cache.ts index c51921fff6b..d6dd757313a 100644 --- a/packages/adapters/codex-local/src/server/codex-auth-cache.ts +++ b/packages/adapters/codex-local/src/server/codex-auth-cache.ts @@ -1,4 +1,5 @@ -import { lstat, mkdir, open, readFile, rename, rm } from "node:fs/promises"; +import { AsyncLocalStorage } from "node:async_hooks"; +import { lstat, mkdir, open, readFile, realpath, rename, rm } from "node:fs/promises"; import path from "node:path"; import { randomUUID } from "node:crypto"; import { resolvePaperclipInstanceRootForAdapter } from "@paperclipai/adapter-utils/server-utils"; @@ -334,6 +335,11 @@ export async function withCodexAccountHomePromotionLock( * create) can still call into a secrets-service write that takes this lock * without deadlocking on itself. */ +// An outer credential transaction acquires this lock before database row locks. +// Nested secret-service writes must share that ownership, not acquire it again. +// Expire ownership on return so detached work cannot inherit a released lock. +const secretMutationOwnership = new AsyncLocalStorage>(); + export async function withAccountHomeSecretMutationLock( env: NodeJS.ProcessEnv = process.env, companyId: string, @@ -341,7 +347,16 @@ export async function withAccountHomeSecretMutationLock( ): Promise { const lockDir = resolveCodexAuthCacheNamedLockDir(env, companyId, "codex-account-home-secret-mutation-lock"); await ensurePrivateDir(lockDir); - return withDirectoryMergeLock(lockDir, fn, env); + const key = await realpath(lockDir); + const inherited = secretMutationOwnership.getStore(); + if (inherited?.get(key)?.active) return fn(); + return withDirectoryMergeLock(lockDir, async () => { + const ownership = { active: true }; + const scope = new Map(inherited); + scope.set(key, ownership); + try { return await secretMutationOwnership.run(scope, fn); } + finally { ownership.active = false; } + }, env); } /** diff --git a/server/src/__tests__/ai-connections.test.ts b/server/src/__tests__/ai-connections.test.ts index 3d8632114a4..52d54a3bfb4 100644 --- a/server/src/__tests__/ai-connections.test.ts +++ b/server/src/__tests__/ai-connections.test.ts @@ -84,6 +84,37 @@ describe("managed AI connections", () => { } finally { fetchSpy.mockRestore(); } }); + it("serializes quota exchange with a concurrent secret edit and runtime credential read", async () => { + const owner = "quota-race-owner"; + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const original = JSON.stringify({ tokens: { access_token: "old-access", refresh_token: "single-use", id_token: "identity", account_id: "race-account" } }); + const account = await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: "Racing refresh", loginSessionId: "fixture", agentIds: [], allAgents: true }, original); + const [row] = await service.quotaAccounts(companyId, owner); + const secretId = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!.secretId; + const exchangeStarted = deferredSignal(); + const releaseExchange = deferredSignal(); + const request = vi.spyOn(globalThis, "fetch").mockImplementation(async () => { + exchangeStarted.resolve(); await releaseExchange.promise; + return Response.json({ access_token: "new-access", refresh_token: "new-refresh" }); + }); + let runtime: Awaited> | undefined; + const refreshing = service.refreshQuotaCredential(row, original, new AbortController().signal); + try { + await exchangeStarted.promise; + const edit = secretService(db).update(secretId, { description: "Changed during refresh" }); + const preparing = prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local", binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} }); + // Keep the provider exchange pending while both contenders reach their + // lock/read paths. Neither may hold a conflicting row lock or vend old auth. + await new Promise(resolve => setTimeout(resolve, 50)); + 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(); } + }, 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" }); @@ -1069,3 +1100,9 @@ describe("AI connection recovery delivery", () => { }, 30000, ); }); + +function deferredSignal() { + let resolve!: () => void; + const promise = new Promise(done => { resolve = done; }); + return { promise, resolve }; +} diff --git a/server/src/services/ai-connection-runtime.ts b/server/src/services/ai-connection-runtime.ts index 534f5cf023a..b199dbc2bcf 100644 --- a/server/src/services/ai-connection-runtime.ts +++ b/server/src/services/ai-connection-runtime.ts @@ -11,7 +11,7 @@ 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 type { AdapterExecutionTarget } from "@paperclipai/adapter-utils/execution-target"; import { runAdapterExecutionTargetProcess } from "@paperclipai/adapter-utils/execution-target"; import { decideGrokAuthMerge } from "@paperclipai/adapter-grok-local/server"; @@ -314,7 +314,7 @@ export async function prepareManagedAiRuntime( if (subscriptionFile) { const refreshed = await readFile(authFile, "utf8"); if (refreshed !== value) - await db.transaction(async (tx) => { + await withAccountHomeSecretMutationLock(undefined, input.companyId, () => db.transaction(async (tx) => { const [grant] = await tx .select() .from(connectionGrants) @@ -374,7 +374,7 @@ export async function prepareManagedAiRuntime( .update(connectionGrants) .set({ updatedAt: new Date() }) .where(eq(connectionGrants.id, grant.id)); - }); + })); } } finally { if (home) await rm(home, { recursive: true, force: true }); diff --git a/server/src/services/ai-connections.ts b/server/src/services/ai-connections.ts index 767ae7ecb0c..770feba2419 100644 --- a/server/src/services/ai-connections.ts +++ b/server/src/services/ai-connections.ts @@ -1,3 +1,4 @@ +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, sql } from "drizzle-orm"; @@ -402,6 +403,9 @@ export function aiConnectionService(db: Db) { }; } async function credential(row: Pick>, "connection" | "grant">) { + return withAccountHomeSecretMutationLock(undefined, row.connection.companyId, () => credentialUnlocked(row)); + } + async function credentialUnlocked(row: Pick>, "connection" | "grant">) { const ref = row.grant.credentialSecretRefs.find( (r) => r.configPath === "ai.credential", ); @@ -498,7 +502,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 @@ -803,7 +807,7 @@ export function aiConnectionService(db: Db) { details: { provider: input.provider, method: input.method, grantId }, }); return { connectionId: id, grantId }; - }); + })); } /** A late failure must never invalidate credentials that were refreshed or reconnected meanwhile. */ async function markAuthenticationFailed(input: { @@ -813,7 +817,7 @@ export function aiConnectionService(db: Db) { runStartedAt: Date; attribution: AiConnectionAttribution & { identity: string }; }) { - return db.transaction(async (tx) => { + return withAccountHomeSecretMutationLock(undefined, input.companyId, () => db.transaction(async (tx) => { const { attribution } = input; const [grant] = await tx.select().from(connectionGrants).where(and( eq(connectionGrants.companyId, input.companyId), @@ -847,11 +851,11 @@ export function aiConnectionService(db: Db) { entityType: "tool_connection", entityId: connection.id, details: { provider: attribution.provider, grantId: grant.id }, }); - }); + })); } /** Refresh only the selected vaulted identity, serializing rotation and reconnect. */ async function refreshQuotaCredential(row: Parameters[0], failedValue: string, signal: AbortSignal) { - return db.transaction(async (tx) => { + return withAccountHomeSecretMutationLock(undefined, row.connection.companyId, () => db.transaction(async (tx) => { const companyId = row.connection.companyId; const [grant] = await tx.select().from(connectionGrants).where(and( eq(connectionGrants.id, row.grant.id), eq(connectionGrants.companyId, companyId), @@ -904,7 +908,7 @@ export function aiConnectionService(db: Db) { action: "ai_connection.credential_refreshed", entityType: "tool_connection", entityId: row.connection.id, details: { provider: "openai", grantId: grant.id } }); return refreshed; - }); + })); } // A quota read uses the same credential audience as execution. Operator status // alone never grants access to another member's personal subscription. diff --git a/ui/src/components/AccountQuotaPanels.test.tsx b/ui/src/components/AccountQuotaPanels.test.tsx index c87a755b924..e5fcef037ae 100644 --- a/ui/src/components/AccountQuotaPanels.test.tsx +++ b/ui/src/components/AccountQuotaPanels.test.tsx @@ -15,9 +15,10 @@ describe("account quota panels", () => { }); it("displays unknown utilization without inventing zero percent", () => { const html = renderToStaticMarkup(); expect(html).toContain("$5 remaining"); + expect(html).toContain(`Resets ${new Date("2026-10-05T00:00:00Z").toLocaleString()}`); expect(html).not.toContain("0%"); expect(html).not.toContain('role="progressbar"'); }); diff --git a/ui/src/components/AccountQuotaPanels.tsx b/ui/src/components/AccountQuotaPanels.tsx index 31cd43d05d7..ea9b84c9230 100644 --- a/ui/src/components/AccountQuotaPanels.tsx +++ b/ui/src/components/AccountQuotaPanels.tsx @@ -27,9 +27,12 @@ export function AccountQuotaPanels({

)} {account.windows.map((window, index) => window.usedPercent == null ? ( -
- {window.label} - {window.valueLabel ?? "Usage not reported"} +
+
+ {window.label} + {window.valueLabel ?? "Usage not reported"} +
+ {window.resetsAt &&

Resets {new Date(window.resetsAt).toLocaleString()}

}
) : ( Date: Sat, 3 Oct 2026 21:33:54 -0700 Subject: [PATCH 2/6] fix: retain full-stream usage and scope biller inference to its provider Co-Authored-By: Paperclip --- .../opencode-local/src/server/execute.ts | 20 +++++++---- .../adapters/pi-local/src/server/execute.ts | 20 +++++++---- .../adapter-cost-failure-paths.test.ts | 33 +++++++++++++++++++ 3 files changed, 59 insertions(+), 14 deletions(-) diff --git a/packages/adapters/opencode-local/src/server/execute.ts b/packages/adapters/opencode-local/src/server/execute.ts index 72dd350f6b4..2fbbf9ea6f2 100644 --- a/packages/adapters/opencode-local/src/server/execute.ts +++ b/packages/adapters/opencode-local/src/server/execute.ts @@ -84,7 +84,7 @@ function parseModelProvider(model: string | null): string | null { } function resolveOpenCodeBiller(env: Record, provider: string | null): string { - return inferOpenAiCompatibleBiller(env, null) ?? provider ?? "unknown"; + return provider === "openai" ? inferOpenAiCompatibleBiller(env, "openai") ?? "unknown" : provider ?? "unknown"; } const REMOTE_OPENCODE_MODELS_PROBE_DEFAULT_TIMEOUT_SEC = 20; @@ -647,7 +647,9 @@ export async function execute(ctx: AdapterExecutionContext): Promise { + let hasAccounting = false; + const accountingLog = createUsageCheckpointLog(onLog, ctx.onUsage ?? (async () => {}), stdout => { + hasAccounting = true; const parsed = consumeAccounting(stdout); const provider = parseModelProvider(model || null); return { usage: parsed.usage, costUsd: parsed.costUsd, usageBasis: "per_run", provider, biller: resolveOpenCodeBiller(runtimeEnv, provider), billingType: "unknown", model, complete: false }; @@ -666,11 +668,15 @@ export async function execute(ctx: AdapterExecutionContext): Promise): Promise, provider: string | null): string { - return inferOpenAiCompatibleBiller(env, null) ?? provider ?? "unknown"; + return provider === "openai" ? inferOpenAiCompatibleBiller(env, "openai") ?? "unknown" : provider ?? "unknown"; } async function ensureSessionsDir(): Promise { @@ -738,7 +738,9 @@ export async function execute(ctx: AdapterExecutionContext): Promise { + let hasAccounting = false; + const accountingLog = createUsageCheckpointLog(bufferedOnLog, ctx.onUsage ?? (async () => {}), stdout => { + hasAccounting = true; const parsed = consumeAccounting(stdout); return { usage: parsed.usage, costUsd: parsed.usage.costUsd, usageBasis: "per_run", provider, biller: resolvePiBiller(runtimeEnv, provider), billingType: "unknown", model, complete: parsed.sawAgentEnd }; }); @@ -761,11 +763,15 @@ export async function execute(ctx: AdapterExecutionContext): Promise { onUsage.mockRejectedValue(new Error("Receipt persistence failed")); await expect(execute(context)).rejects.toThrow("Receipt persistence failed"); }); + it.each(cases.filter(f => ["pi", "opencode"].includes(f.name)))("$name preserves full-stream accounting after stdout truncation", async (fixture) => { + const dir = await mkdtemp(join(tmpdir(), "paperclip-accounting-tail-")); directories.push(dir); + const command = join(dir, "runtime"); await writeFile(command, "#!/bin/sh\nprintf 'anthropic test\\n'\n", { mode: 0o755 }); + const execute = await fixture.execute(); + for (const checkpoint of [true, false]) for (const timedOut of [false, true]) for (const missingPrice of [false, true]) { + const early = structuredClone(fixture.event) as any; + if (missingPrice) { + if (fixture.name === "pi") delete early.message.usage.cost; + else delete early.part.cost; + } + const end = fixture.name === "pi" && !timedOut ? '\n{"type":"agent_end","messages":[]}' : ""; + const first = JSON.stringify(early) + "\n"; + const noise = JSON.stringify({ type: "text", part: { text: "x".repeat(5 * 1024 * 1024) } }) + "\n"; + const last = JSON.stringify(fixture.event) + end; + processResult.mockImplementation(async (_run, _target, _command, _args, options) => { + for (const chunk of [first, noise, last]) await options.onLog("stdout", chunk); + return { exitCode: timedOut ? null : 0, signal: timedOut ? "SIGTERM" : null, timedOut, stdout: (first + noise + last).slice(-4 * 1024 * 1024), stderr: "", pid: 123, startedAt: new Date().toISOString() }; + }); + const onUsage = checkpoint ? vi.fn() : undefined; + const result = await execute({ + runId: "tail-run", agent: { id: "test-agent", companyId: "test-company", name: "Accounting", adapterType: `${fixture.name}_local`, adapterConfig: {} }, + runtime: { sessionId: null, sessionParams: null, sessionDisplayId: null, taskKey: null }, + config: { engine: "cli", command, cwd: dir, model: "anthropic/test", paperclipRuntimeSkills: [], env: { ANTHROPIC_API_KEY: "fixture", OPENAI_BASE_URL: "https://unrelated-proxy.example/v1", OPENROUTER_API_KEY: "unrelated-key", OPENCODE_ALLOW_ALL_MODELS: "1" } }, + context: {}, onLog: async () => {}, onUsage, + }); + expect(result.usage).toMatchObject({ inputTokens: 40, outputTokens: 20, cachedInputTokens: 200 }); + expect(result.costUsd).toBe(missingPrice ? null : 0.008); + expect(result.biller).toBe("anthropic"); + expect(result.usageComplete).toBe(!timedOut); + if (onUsage) expect(onUsage.mock.calls.at(-1)![0]).toMatchObject({ usage: result.usage, costUsd: result.costUsd }); + } + }); + it.each(cases)("$name retains observed accounting when the process times out", async (fixture) => { const dir = await mkdtemp(join(tmpdir(), "paperclip-accounting-adapter-")); directories.push(dir); const command = join(dir, "runtime"); await writeFile(command, "#!/bin/sh\nprintf 'openai test\\n'\n", { mode: 0o755 }); From dfd62bd7e2f1189cb2a417f4b2c96b63a95e167f Mon Sep 17 00:00:00 2001 From: Devin Foley Date: Sat, 3 Oct 2026 21:38:46 -0700 Subject: [PATCH 3/6] fix: recheck credential versions after waits and defer conflicting refreshes Co-Authored-By: Paperclip --- doc/connections/AI-CONNECTIONS.md | 11 +++++++ server/src/__tests__/ai-connections.test.ts | 17 +++++++++-- server/src/services/ai-connections.ts | 33 +++++++++++++-------- 3 files changed, 45 insertions(+), 16 deletions(-) diff --git a/doc/connections/AI-CONNECTIONS.md b/doc/connections/AI-CONNECTIONS.md index d31ce5fea97..c0ca25fa2c6 100644 --- a/doc/connections/AI-CONNECTIONS.md +++ b/doc/connections/AI-CONNECTIONS.md @@ -186,6 +186,17 @@ and contain neither credentials nor raw provider errors. Expired credentials report `authentication_required`; the probe does not exchange refresh tokens or change connection health. +The Costs dashboard uses the company quota endpoint. It reads each authorized +managed subscription separately, retaining the last successful observation on +transient errors. For managed Codex accounts, a 401 can trigger one persisted +OAuth refresh; it defers while potentially competing OpenAI work is active. +Refresh takes the company secret-mutation lock before database row locks and +requires those row locks immediately, before exchanging a single-use token. +Reconnect and runtime credential write-back follow the same lock order. +Credential resolution checks its version after any database wait and reloads a +rotated value before giving it to a new run. Unknown utilization has no progress +bar, but any provider-reported reset time remains visible. + Verification: ```sh diff --git a/server/src/__tests__/ai-connections.test.ts b/server/src/__tests__/ai-connections.test.ts index 52d54a3bfb4..dbd467740bf 100644 --- a/server/src/__tests__/ai-connections.test.ts +++ b/server/src/__tests__/ai-connections.test.ts @@ -103,9 +103,12 @@ describe("managed AI connections", () => { await exchangeStarted.promise; const edit = secretService(db).update(secretId, { description: "Changed during refresh" }); const preparing = prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local", binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} }); - // Keep the provider exchange pending while both contenders reach their - // lock/read paths. Neither may hold a conflicting row lock or vend old auth. - await new Promise(resolve => setTimeout(resolve, 50)); + // Wait for the runtime's old-token read to reach its bookkeeping update. + // After the refresh commits, resolution must detect the changed version. + await vi.waitFor(async () => { + const waiting = await db.execute(sql`select 1 from pg_stat_activity where datname = current_database() and wait_event_type = 'Lock' and query like '%last_resolved_at%'`); + expect(waiting.length).toBeGreaterThan(0); + }); releaseExchange.resolve(); const results = await Promise.all([refreshing, edit, preparing]); runtime = results[2]; @@ -124,6 +127,14 @@ describe("managed AI connections", () => { 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"); diff --git a/server/src/services/ai-connections.ts b/server/src/services/ai-connections.ts index 770feba2419..3b452933784 100644 --- a/server/src/services/ai-connections.ts +++ b/server/src/services/ai-connections.ts @@ -402,10 +402,7 @@ export function aiConnectionService(db: Db) { } satisfies AiConnectionAttribution, }; } - async function credential(row: Pick>, "connection" | "grant">) { - return withAccountHomeSecretMutationLock(undefined, row.connection.companyId, () => credentialUnlocked(row)); - } - async function credentialUnlocked(row: Pick>, "connection" | "grant">) { + async function credential(row: Pick>, "connection" | "grant">, retry = 0): Promise { const ref = row.grant.credentialSecretRefs.find( (r) => r.configPath === "ai.credential", ); @@ -435,6 +432,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 || @@ -452,14 +450,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, @@ -856,17 +861,19 @@ export function aiConnectionService(db: Db) { /** Refresh only the selected vaulted identity, serializing rotation and reconnect. */ async function refreshQuotaCredential(row: Parameters[0], failedValue: string, signal: AbortSignal) { return withAccountHomeSecretMutationLock(undefined, row.connection.companyId, () => 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"); + )).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"); + )).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. From acba463d569586e611a59e03a6493bc455a61e0f Mon Sep 17 00:00:00 2001 From: Devin Foley Date: Sat, 3 Oct 2026 21:43:43 -0700 Subject: [PATCH 4/6] test: drain heartbeat finalizers before fixture teardown Co-Authored-By: Paperclip --- .../src/__tests__/heartbeat-stale-queue-invalidation.test.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts b/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts index 882a81e5d2b..719e4ea9778 100644 --- a/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts +++ b/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts @@ -196,7 +196,9 @@ describeEmbeddedPostgres("heartbeat stale queued-run invalidation", () => { } await new Promise((resolve) => setTimeout(resolve, 50)); } - await new Promise((resolve) => setTimeout(resolve, 50)); + // Terminal run status precedes accounting and comment finalization. Drain + // the registered executions before taking TRUNCATE locks during teardown. + await heartbeat.drainActiveRunExecutions(); await cleanupHeartbeatInvalidationFixture(db); }); From fb8b03a083c6e7aecde27d1736a4737221ecdea3 Mon Sep 17 00:00:00 2001 From: Devin Foley Date: Sat, 3 Oct 2026 21:45:23 -0700 Subject: [PATCH 5/6] fix: bound accounting strings before buffering provider records Co-Authored-By: Paperclip --- .../src/usage-checkpoint.test.ts | 23 +++++++++++ .../adapter-utils/src/usage-checkpoint.ts | 38 ++++++++++++++++++- .../adapter-cost-failure-paths.test.ts | 6 ++- 3 files changed, 64 insertions(+), 3 deletions(-) diff --git a/packages/adapter-utils/src/usage-checkpoint.test.ts b/packages/adapter-utils/src/usage-checkpoint.test.ts index 2b622ef60f0..a158f587c15 100644 --- a/packages/adapter-utils/src/usage-checkpoint.test.ts +++ b/packages/adapter-utils/src/usage-checkpoint.test.ts @@ -29,6 +29,29 @@ describe("usage checkpoint stream", () => { expect(parser.mock.calls[0][0]).not.toContain("private content"); }); + it("bounds large content strings before buffering and preserves usage on either side", async () => { + const saved = vi.fn(); const parser = vi.fn(createParser()); const output = vi.fn(); + const log = createUsageCheckpointLog(output, saved, parser); + const line = JSON.stringify({ type: "result", usage: { inputTokens: 2, costUsd: 0.25 }, messages: [{ content: '"\\\n'.repeat(2 * 1024 * 1024) }], model: "after-large-content" }); + for (let offset = 0; offset < line.length; offset += 4093) await log("stdout", line.slice(offset, offset + 4093)); + await log.flush(); + expect(saved).toHaveBeenLastCalledWith(expect.objectContaining({ costUsd: 0.25, complete: true })); + expect(parser.mock.calls[0][0]).toContain("after-large-content"); + expect(parser.mock.calls[0][0].length).toBeLessThan(1000); + expect(output.mock.calls.map(call => call[1]).join("")).toBe(line); + }); + + it("preserves Unicode and escaped quotes across the string limit and chunk boundaries", async () => { + for (const suffix of ['\\u1234', '\\"', '\\\\']) { + const saved = vi.fn(); const parser = vi.fn(createParser()); + const log = createUsageCheckpointLog(vi.fn(), saved, parser); + const line = '{"type":"result","text":"' + 'a'.repeat(249) + suffix + 'discarded","usage":{"costUsd":0.25}}\n'; + for (const character of line) await log("stdout", character); + await log.flush(); + expect(saved).toHaveBeenCalledWith(expect.objectContaining({ costUsd: 0.25 })); + } + }); + it("surfaces a persistence failure outside the local process's best-effort log path", async () => { const failure = new Error("Receipt storage unavailable"); const output = vi.fn(); const saved = vi.fn().mockRejectedValue(failure); diff --git a/packages/adapter-utils/src/usage-checkpoint.ts b/packages/adapter-utils/src/usage-checkpoint.ts index 1e0c44bd86e..123ae66615e 100644 --- a/packages/adapter-utils/src/usage-checkpoint.ts +++ b/packages/adapter-utils/src/usage-checkpoint.ts @@ -24,6 +24,41 @@ function project(value: unknown, depth = 0, field = ""): unknown { : project(entry, depth + 1, key)])); } +/** Bound strings before buffering whole JSON lines. Provider records can carry + * many MiB of conversation text alongside a small usage object. Preserve JSON + * escapes across chunks so discarded text cannot corrupt the accounting fields. */ +function createStringCompactor() { + let inString = false, length = 0, escapeRemaining = 0; + let escapeCode = false, keepEscape = false; + return (chunk: string) => { + let output = ""; + for (const character of chunk) { + if (character === "\n") { + inString = false; escapeRemaining = 0; escapeCode = false; + output += character; + } else if (!inString) { + output += character; + if (character === '"') { inString = true; length = 0; } + } else if (escapeRemaining > 0) { + if (keepEscape) output += character; + escapeRemaining = escapeCode && character === "u" ? 4 : escapeRemaining - 1; + escapeCode = false; + length++; + } else if (character === '"') { + output += character; inString = false; + } else if (character === "\\") { + keepEscape = length < 250; + if (keepEscape) output += character; + escapeRemaining = 1; escapeCode = true; length++; + } else { + if (length < 250) output += character; + length++; + } + } + return output; + }; +} + /** The consumer incrementally parses only new records (never the full history). * A fresh instance belongs to one CLI attempt. Call flush() after its process * exits, before retrying or final result handling. Checkpoint failures are @@ -34,6 +69,7 @@ export function createUsageCheckpointLog( consume: (records: string) => AdapterUsageCheckpoint | null, ) { const attemptId = randomUUID(); + const compactStrings = createStringCompactor(); let remainder = "", accounting = "", previous = ""; let snapshot: AdapterUsageCheckpoint | null = null; let failure: unknown; @@ -63,7 +99,7 @@ export function createUsageCheckpointLog( const log = async (stream: "stdout" | "stderr", chunk: string) => { if (onUsage && !failed && stream === "stdout") { try { - remainder += chunk; + remainder += compactStrings(chunk); const lines = remainder.split("\n"); remainder = lines.pop() ?? ""; if (remainder.length > 8 * 1024 * 1024) throw new Error("Accounting protocol line exceeds 8 MiB"); for (const line of lines) retain(line); diff --git a/server/src/__tests__/adapter-cost-failure-paths.test.ts b/server/src/__tests__/adapter-cost-failure-paths.test.ts index bde2346f785..b831a1e0476 100644 --- a/server/src/__tests__/adapter-cost-failure-paths.test.ts +++ b/server/src/__tests__/adapter-cost-failure-paths.test.ts @@ -98,10 +98,12 @@ describe("CLI adapter accounting on timeout", () => { } const end = fixture.name === "pi" && !timedOut ? '\n{"type":"agent_end","messages":[]}' : ""; const first = JSON.stringify(early) + "\n"; - const noise = JSON.stringify({ type: "text", part: { text: "x".repeat(5 * 1024 * 1024) } }) + "\n"; + const noise = JSON.stringify({ type: "text", part: { text: "x".repeat(9 * 1024 * 1024) } }) + "\n"; const last = JSON.stringify(fixture.event) + end; processResult.mockImplementation(async (_run, _target, _command, _args, options) => { - for (const chunk of [first, noise, last]) await options.onLog("stdout", chunk); + for (const record of [first, noise, last]) { + for (let offset = 0; offset < record.length; offset += 65521) await options.onLog("stdout", record.slice(offset, offset + 65521)); + } return { exitCode: timedOut ? null : 0, signal: timedOut ? "SIGTERM" : null, timedOut, stdout: (first + noise + last).slice(-4 * 1024 * 1024), stderr: "", pid: 123, startedAt: new Date().toISOString() }; }); const onUsage = checkpoint ? vi.fn() : undefined; From b2815ee3cdf10d363faec5978e4da2ce9cd7733d Mon Sep 17 00:00:00 2001 From: Devin Foley Date: Sat, 3 Oct 2026 21:51:37 -0700 Subject: [PATCH 6/6] fix: retain failed native receipts and known attempted-run charges Co-Authored-By: Paperclip --- doc/COST-ACCOUNTING.md | 7 +++ .../cost-accounting-operations-edges.test.ts | 21 +++++++- .../native-session-executor.test.ts | 38 +++++++++++++- .../native-runtime/native-session-executor.ts | 49 ++++++++++++++----- server/src/services/run-cost-accounting.ts | 2 +- server/src/services/usage-receipts.ts | 10 ++-- 6 files changed, 107 insertions(+), 20 deletions(-) diff --git a/doc/COST-ACCOUNTING.md b/doc/COST-ACCOUNTING.md index 4363f279cdf..54d0472e16a 100644 --- a/doc/COST-ACCOUNTING.md +++ b/doc/COST-ACCOUNTING.md @@ -174,3 +174,10 @@ edits deliver hard-stop cancellation after their transaction commits. Status-card update costs retain the ledger's fractional-cent storage precision. The mutation gate retains partial results and per-run reports/logs under `coverage/accounting/` when its baseline fails or a mutation survives. + +Failed native turns persist observed run-delta usage before semantic-result +finalization. A matching terminal failure, cancellation, or interruption closes +that receipt; missing terminal evidence and session-only totals remain pending. +For multiple provider attempts, known prices remain in the exact spend total +when another attempt is unpriced. The aggregate stays marked unpriced, so the +configured unknown-price policy still controls whether new work may start. diff --git a/server/src/__tests__/cost-accounting-operations-edges.test.ts b/server/src/__tests__/cost-accounting-operations-edges.test.ts index 156d3092c0d..6a97dcc5978 100644 --- a/server/src/__tests__/cost-accounting-operations-edges.test.ts +++ b/server/src/__tests__/cost-accounting-operations-edges.test.ts @@ -175,11 +175,30 @@ const support = await getEmbeddedPostgresTestSupport(); expect(result).toMatchObject({costUsdExact:"0.040000001",complete:true,usage:{inputTokens:14,cachedInputTokens:22,outputTokens:6}}); await recorder.capture({attemptId:randomUUID(),usage,complete:false,billingType:"subscription"}); const uncertain = await recorder.complete({exitCode:1,signal:null,timedOut:false,usage}); - expect(uncertain).toMatchObject({costUsd:null,costUsdExact:null,billingType:"unknown"}); + expect(uncertain).toMatchObject({costUsd:0.040000001,costUsdExact:"0.040000001",costStatus:"unpriced",billingType:"unknown"}); const partial = await recorder.complete({exitCode:1,signal:null,timedOut:false}); expect(partial.complete).toBe(false); expect(partial.usage?.inputTokens).toBe(21); }); + it.each(["block", "allow"] as const)("keeps known attempted-run spend when unknown prices %s new work", async (unpricedUsagePolicy) => { + const f = await fixture(), budgets = budgetService(db); + await budgets.upsertPolicy(f.company.id, { scopeType: "company", scopeId: f.company.id, amount: 100, unpricedUsagePolicy }, "board"); + await reserveRunBudget(db, f.company.id, f.run.id, null); + const recorder = await createRunUsageRecorder(db, { companyId: f.company.id, runId: f.run.id, adapterType: "process" }, path.join(directory, randomUUID())); + await recorder.capture({ attemptId: randomUUID(), complete: true, billingType: "metered_api", usage, costUsdExact: "0.010000001" }); + const aggregate = await recorder.capture({ attemptId: randomUUID(), complete: true, billingType: "metered_api", usage, costUsd: null, costStatus: "unpriced" }); + expect(aggregate).toMatchObject({ costUsdExact: "0.010000001", costStatus: "unpriced", complete: true }); + await db.update(heartbeatRuns).set({ status: "failed", finishedAt: new Date() }).where(eq(heartbeatRuns.id, f.run.id)); + expect(await accountRunCost(db, f.run.id)).toBe(true); + expect(await accountRunCost(db, f.run.id)).toBe(false); + expect(await costService(db).summary(f.company.id)).toMatchObject({ spendCentsExact: "1.0000001", unpricedEventCount: 1, pricingComplete: false, pendingRunCount: 0 }); + const [runtime] = await db.select().from(agentRuntimeState).where(eq(agentRuntimeState.agentId, f.agent.id)); + expect(runtime.totalCostCents).toBe(1.0000001); + const [reservation] = await db.select().from(budgetReservations).where(eq(budgetReservations.runId, f.run.id)); + expect(reservation.state).toBe("settled"); + expect((await budgets.getInvocationBlock(f.company.id, f.agent.id)) !== null).toBe(unpricedUsagePolicy === "block"); + }); + it("refuses invalid checkpoints and does not certify a run after capture failure", async () => { const f = await fixture(); const recorder = await createRunUsageRecorder(db,{companyId:f.company.id,runId:f.run.id,adapterType:"process"},path.join(directory,randomUUID())); diff --git a/server/src/services/native-runtime/native-session-executor.test.ts b/server/src/services/native-runtime/native-session-executor.test.ts index 185f3416533..d07fecace03 100644 --- a/server/src/services/native-runtime/native-session-executor.test.ts +++ b/server/src/services/native-runtime/native-session-executor.test.ts @@ -191,6 +191,17 @@ const state = vi.hoisted(() => ({ release: null as null | (() => void), })); +const accountingEvents = vi.hoisted(() => ({ committed: undefined as undefined | ((event: PrpEvent) => Promise) })); +vi.mock("./paperclip-control-plane-port.js", async importOriginal => { + const original = await importOriginal(); + return { ...original, PaperclipControlPlanePort: class extends original.PaperclipControlPlanePort { + constructor(...args: ConstructorParameters) { + super(...args); + accountingEvents.committed = args[2]?.onCommittedEvent; + } + } }; +}); + const grokCopyBack = vi.hoisted(() => vi.fn(async (_input: { readSandboxAuth: () => Promise; hostHomeDir: string }) => undefined)); vi.mock("@paperclipai/adapter-grok-local/server", async importOriginal => ({ ...await importOriginal(), @@ -4873,7 +4884,7 @@ function leaseDb( where: () => query, orderBy: () => query, for: () => query, - limit: () => Promise.resolve(rows), + limit: () => query, }; return query; }, @@ -4881,7 +4892,9 @@ function leaseDb( const insert = (table: unknown) => ({ values: (values: Record) => { updates.push({ table, values }); - return { returning: async () => [values] }; + const query = { returning: async () => [values], onConflictDoUpdate: () => query, onConflictDoNothing: () => query, + then: Promise.resolve([values]).then.bind(Promise.resolve([values])) }; + return query; }, }); const tx = { @@ -5134,6 +5147,27 @@ describe("native startup restart detachment", () => { }); }); +describe("native failed-turn accounting", () => { + beforeEach(() => vi.spyOn(issueServiceModule, "issueService").mockReturnValue({ update: vi.fn(async () => ({ status: "blocked", statusVersion: 1 })) } as unknown as ReturnType)); + afterEach(() => vi.restoreAllMocks()); + it.each(["turn.failed", "turn.cancelled", "turn.interrupted", "missing_terminal", "session_total", "empty_usage", "different_turn"])("preserves run usage for %s without certifying incomplete evidence", async (scenario) => { + const onUsage = vi.fn(async (_receipt: import("@paperclipai/adapter-utils").AdapterUsageCheckpoint) => {}); + const failure = new Error("provider_transport_failed: no semantic result"); + state.execute.mockReset().mockImplementationOnce(async () => { + const usage = { inputTokens: 120, cacheReadTokens: 100, outputTokens: 10, providerCostUsd: 0.25 }; + await accountingEvents.committed!({ eventType: "item.completed", turnId: "turn", emittedAt: new Date().toISOString(), payload: { kind: "usage", usage: scenario === "session_total" ? { total: usage } : { runDelta: scenario === "empty_usage" ? { requests: 1 } : usage } } } as unknown as PrpEvent); + if (scenario !== "missing_terminal") await accountingEvents.committed!({ eventType: scenario.startsWith("turn.") ? scenario : "turn.failed", turnId: scenario === "different_turn" ? "another-turn" : "turn", emittedAt: new Date().toISOString(), payload: { status: "failed" } } as unknown as PrpEvent); + throw failure; + }); + await expect(executePaperclipNativeSession({ db: leaseDb(), execution: { ...execution, provider: { kind: "codex", model: "gpt-6-astra" } } as NativeExecutionInput, runnerInstanceId: "runner", onUsage })).rejects.toBe(failure); + if (["session_total", "empty_usage"].includes(scenario)) expect(onUsage).not.toHaveBeenCalled(); + else { + expect(onUsage).toHaveBeenLastCalledWith(expect.objectContaining({ complete: scenario.startsWith("turn."), costUsd: 0.25, usageBasis: "per_run", usage: { inputTokens: 20, cachedInputTokens: 100, outputTokens: 10 } })); + expect(onUsage.mock.calls[0]?.[0]).toMatchObject({ complete: false }); + } + }); +}); + describe("native resumed preparation timing", () => { it("keeps answered-question ingress at the run root rather than charging it to preparation", async () => { const answeredAtMs = Date.now(); diff --git a/server/src/services/native-runtime/native-session-executor.ts b/server/src/services/native-runtime/native-session-executor.ts index 4ba162a2e29..4b7a7851f8e 100644 --- a/server/src/services/native-runtime/native-session-executor.ts +++ b/server/src/services/native-runtime/native-session-executor.ts @@ -7838,6 +7838,22 @@ async function executePaperclipNativeSessionWithinScope( binding: { ...input.execution.binding, normalizedSessionId: nativeSessionKey(input.execution), runnerSourceInstanceId: effectiveRunnerInstanceId }, resolve: resolveNativeRuntimeRequest, }); + let observedAccountingUsage: Record | null = null; + let observedAccountingTurn: string | undefined; + const persistAccountingUsage = async (usage: Record | null, complete: boolean) => { + const accountingUsage = normalizeNativeUsage(usage, { inputIncludesCacheReads: input.execution.provider.kind === "codex" }); + const accountingCost = nativeUsageCostUsd(usage, input.execution.provider); + const billing = resolveNativeBilling(input.execution.provider, input.runnerEnvironment, input.billingIdentity); + await input.onUsage?.({ usage: accountingUsage, ...billing, complete, usageBasis: "per_run", + model: input.execution.provider.model ?? "unknown", costUsd: accountingCost ?? null }); + if (!input.onUsage) await input.db.update(heartbeatRuns).set({ + costAccountingPending: true, + usageJson: sql`coalesce(${heartbeatRuns.usageJson}, '{}'::jsonb) || ${JSON.stringify({ + ...accountingUsage, ...billing, accountingReceiptReady: complete, model: input.execution.provider.model ?? "unknown", + costUsd: accountingCost ?? null, usageSource: "per_run", + })}::jsonb`, + }).where(and(eq(heartbeatRuns.id, input.execution.binding.runId), eq(heartbeatRuns.companyId, input.execution.binding.companyId), isNull(heartbeatRuns.costAccountedAt))); + }; let completedConversationReply: PrpEvent | null = null; const controlPlane = new PaperclipControlPlanePort( input.db, @@ -7854,6 +7870,25 @@ async function executePaperclipNativeSessionWithinScope( }, { onCommittedEvent: async (event) => { + const payload = record(event.payload); + const usage = record(payload.usage); + const delta = record(usage.runDelta); + if (event.eventType === "turn.started" && observedAccountingUsage) { + observedAccountingTurn = undefined; + await persistAccountingUsage(observedAccountingUsage, false); + } + // Only an explicit run delta is safe here: session totals can include + // earlier work. Save provisional evidence before ancillary finalizers. + if (payload.kind === "usage" && numericUsageField(delta, ["inputTokens"]) !== undefined + && numericUsageField(delta, ["outputTokens"]) !== undefined) { + observedAccountingUsage = { runDelta: usage.runDelta }; + observedAccountingTurn = event.turnId; + await persistAccountingUsage(observedAccountingUsage, false); + } + if (["turn.failed", "turn.cancelled", "turn.interrupted"].includes(event.eventType) + && observedAccountingUsage && event.turnId === observedAccountingTurn) { + await persistAccountingUsage(observedAccountingUsage, true); + } await toolTrace.observe(event); if (event.eventType === "item.completed" && record(event.payload).kind === "agentMessage" && @@ -8602,18 +8637,8 @@ async function executePaperclipNativeSessionWithinScope( ); // Persist provider accounting before any workspace/issue finalization. A // detached controller or failed finalizer must not lose a completed turn. - const accountingUsage = normalizeNativeUsage(native.usage, { inputIncludesCacheReads: input.execution.provider.kind === "codex" }); - const accountingCost = nativeUsageCostUsd(native.usage, input.execution.provider); - const billing = resolveNativeBilling(input.execution.provider, input.runnerEnvironment, input.billingIdentity); - await input.onUsage?.({ usage: accountingUsage, ...billing, complete: true, usageBasis: "per_run", - model: input.execution.provider.model ?? "unknown", costUsd: accountingCost ?? null }); - if (!input.onUsage) await input.db.update(heartbeatRuns).set({ - costAccountingPending: true, - usageJson: sql`coalesce(${heartbeatRuns.usageJson}, '{}'::jsonb) || ${JSON.stringify({ - ...accountingUsage, ...billing, accountingReceiptReady: true, model: input.execution.provider.model ?? "unknown", - costUsd: accountingCost ?? null, usageSource: "per_run", - })}::jsonb`, - }).where(and(eq(heartbeatRuns.id, input.execution.binding.runId), eq(heartbeatRuns.companyId, input.execution.binding.companyId), isNull(heartbeatRuns.costAccountedAt))); + native = { ...native, usage: native.usage ?? observedAccountingUsage }; + await persistAccountingUsage(native.usage, true); try { await completeManagedNativeCredentialTurn(managedCredentialSession); } catch { diff --git a/server/src/services/run-cost-accounting.ts b/server/src/services/run-cost-accounting.ts index 895e2443acb..b1126462a8b 100644 --- a/server/src/services/run-cost-accounting.ts +++ b/server/src/services/run-cost-accounting.ts @@ -62,7 +62,7 @@ export async function accountRunCost(db: Db, runId: string, hooks: BudgetService providerRequestId: text(usage.providerRequestId), pricingProvenance: object(usage.pricingProvenance), model: text(usage.model) ?? "unknown", billingType, - costStatus: costUsd === null && billingType !== "subscription_included" && hasBillableEvidence ? "unpriced" : usage.costStatus === "estimated" ? "estimated" : "reported", + costStatus: (costUsd === null || usage.costStatus === "unpriced") && billingType !== "subscription_included" && hasBillableEvidence ? "unpriced" : usage.costStatus === "estimated" ? "estimated" : "reported", inputTokens: amount(usage.inputTokens) ?? 0, cachedInputTokens: amount(usage.cachedInputTokens) ?? 0, outputTokens: amount(usage.outputTokens) ?? 0, costCents, occurredAt: run.finishedAt ?? run.createdAt, }; diff --git a/server/src/services/usage-receipts.ts b/server/src/services/usage-receipts.ts index c629b41c567..0421712dda4 100644 --- a/server/src/services/usage-receipts.ts +++ b/server/src/services/usage-receipts.ts @@ -152,13 +152,15 @@ export async function createRunUsageRecorder(db: Db, input: { companyId: string; if (parts.length > 1) { let total = 0n, priced = true; for (const part of parts) { - const cost = part.cacheAdjustedCostUsd ?? part.costUsdExact ?? part.costUsd; - if (cost == null) priced = false; else total += usdToUnits(cost); + const subscription = ["subscription", "subscription_included"].includes(part.billingType ?? ""); + const cost = subscription ? 0 : part.cacheAdjustedCostUsd ?? part.costUsdExact ?? part.costUsd; + if (cost == null || (!subscription && part.costStatus === "unpriced")) priced = false; + if (cost != null) total += usdToUnits(cost); } if (new Set(parts.map(part => part.billingType ?? "unknown")).size > 1) receipt.billingType = "unknown"; const decimal = `${total / 1_000_000_000n}.${String(total % 1_000_000_000n).padStart(9, "0")}`; - receipt.costUsdExact = priced ? decimal : null; - receipt.costUsd = priced ? Number(decimal) : null; + receipt.costUsdExact = decimal; + receipt.costUsd = Number(decimal); receipt.cacheAdjustedCostUsd = null; receipt.usageByModel = undefined; receipt.costStatus = priced && parts.some(part => part.costStatus === "estimated") ? "estimated" : priced ? "reported" : "unpriced";