diff --git a/doc/COST-ACCOUNTING.md b/doc/COST-ACCOUNTING.md index f8b287d3715..9e13573a48a 100644 --- a/doc/COST-ACCOUNTING.md +++ b/doc/COST-ACCOUNTING.md @@ -177,7 +177,25 @@ The mutation gate retains partial results and per-run reports/logs under 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. +that usage snapshot; ledger acknowledgement still waits until the native coordinator +has a result or a terminal failure. A retryable failed heartbeat keeps its +reservation and can resume. Replacing its recorder restores the native run's +cumulative snapshot; recovery does not add the same tokens twice, and an empty +final result cannot erase earlier observed usage. Missing terminal evidence and +session-only totals remain pending. + +Budget admission completes before the final dispatch ownership check. A rejected +handoff releases a newly created hold using explicit evidence that the adapter +was never entered. Process adapter configuration and spawn failures likewise +return pre-provider proof; failures after a child exists cannot claim it. + +Scheduled retry gates acquire company, issue, then run locks, in the same order +as accounting and attribution writes. Delayed budget cancellation rechecks the +policy version and current blocking state when it writes durable stop intent, +under the company admission lock. External shutdown happens after that lock is +released so the provider can save its final receipt. Raising a budget before the +stop is claimed preserves newly admitted work, including previously queued rows; +a run with an already claimed stop cannot enter provider work after the grant. 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/doc/connections/AI-CONNECTIONS.md b/doc/connections/AI-CONNECTIONS.md index be4ac8df20c..3cad3975d09 100644 --- a/doc/connections/AI-CONNECTIONS.md +++ b/doc/connections/AI-CONNECTIONS.md @@ -540,3 +540,15 @@ the private runtime home for retry; cleanup deletes it only after a successful transaction or an intentional freshness/authorization discard. These retained homes are recovery evidence, not a background replay queue; persistent database or lock failures still require operator intervention. + +Quota OAuth replacement tokens are encrypted with the instance secrets master +key and fsynced under `/quota-credential-recovery//` +before vault, grant, or activity writes. Storage and encryption are checked before +exchanging a single-use refresh token. Failed saves retry without another OAuth +exchange; persistent failures keep the encrypted record for the next quota poll, +including after a restart. The journal is removed only after database commit (or +when an authorized newer credential makes it obsolete). Replay checks the grant, +connection, secret, and original credential fingerprint and cannot reactivate a +revoked grant or overwrite a reconnect. Back up this directory and the instance +secrets key with the instance data. Loss of durable storage while receiving a +provider token can still require reconnecting the account. diff --git a/server/src/__tests__/adapter-cost-failure-paths.test.ts b/server/src/__tests__/adapter-cost-failure-paths.test.ts index b831a1e0476..2b6e5cd24e0 100644 --- a/server/src/__tests__/adapter-cost-failure-paths.test.ts +++ b/server/src/__tests__/adapter-cost-failure-paths.test.ts @@ -176,3 +176,32 @@ describe("incremental protocol accounting", () => { expect(second.usage).toMatchObject(Object.fromEntries(Object.entries(fixture.tokens!).map(([key, count]) => [key, count * multiplier]))); }); }); + + +describe("process adapter bootstrap accounting proof", () => { + const context = (config: Record): AdapterExecutionContext => ({ + runId: `bootstrap-${crypto.randomUUID()}`, + agent: { id: "agent", companyId: "company", name: "Process", adapterType: "process", adapterConfig: {} }, + runtime: { sessionId: null, sessionParams: null, sessionDisplayId: null, taskKey: null }, + config, context: {}, onLog: async () => {}, + }); + it.each([{}, { command: "/nonexistent-paperclip-accounting-command" }, + { command: process.execPath, cwd: "/nonexistent-paperclip-accounting-directory" }])( + "proves no provider work for a failure before spawn: %j", async (config) => { + const { execute } = await import("../adapters/process/execute.js"); + const onSpawn = vi.fn(); + const result = await execute({ ...context(config), onSpawn }); + expect(result).toMatchObject({ exitCode: 1, executionRecovery: { kind: "bootstrap", providerWorkStarted: false } }); + expect(result.errorMessage).toBeTruthy(); + expect(onSpawn).not.toHaveBeenCalled(); + }, + ); + it("does not invent bootstrap proof when a spawned process fails", async () => { + const { execute } = await import("../adapters/process/execute.js"); + const onSpawn = vi.fn(async () => { throw new Error("Metadata persistence failed"); }); + const result = await execute({ ...context({ command: process.execPath, args: ["-e", "process.exit(7)"] }), onSpawn }); + expect(onSpawn).toHaveBeenCalledOnce(); + expect(result.exitCode).toBe(7); + expect(result.executionRecovery).toBeUndefined(); + }); +}); diff --git a/server/src/__tests__/ai-connections.test.ts b/server/src/__tests__/ai-connections.test.ts index 30fa4587adf..5dcdb5400f8 100644 --- a/server/src/__tests__/ai-connections.test.ts +++ b/server/src/__tests__/ai-connections.test.ts @@ -11,6 +11,7 @@ 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, 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 * as codexAdapter from "@paperclipai/adapter-codex-local/server"; import { WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE } from "@paperclipai/adapter-utils/workspace-restore-merge"; @@ -104,6 +105,92 @@ describe("managed AI connections", () => { } 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(["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", "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 { + 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(); + } + 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" }); diff --git a/server/src/__tests__/budgets-service.test.ts b/server/src/__tests__/budgets-service.test.ts index 52d69cecd05..3414b332513 100644 --- a/server/src/__tests__/budgets-service.test.ts +++ b/server/src/__tests__/budgets-service.test.ts @@ -188,7 +188,7 @@ describeEmbeddedPostgres("budgetService release gate enforcement", () => { expect(agentAfterHardStop).toMatchObject({ status: "paused", pauseReason: "budget" }); expect(agentAfterHardStop?.pausedAt).toBeInstanceOf(Date); expect(cancelWorkForScope).toHaveBeenCalledTimes(2); - expect(cancelWorkForScope).toHaveBeenCalledWith({ companyId, scopeType: "agent", scopeId: agentId, createdBefore: expect.any(Date) }); + expect(cancelWorkForScope).toHaveBeenCalledWith({ companyId, scopeType: "agent", scopeId: agentId, createdBefore: expect.any(Date), enforcement: { policyId: expect.any(String), version: expect.any(Number) } }); const block = await service.getInvocationBlock(companyId, agentId); expect(block).toEqual({ @@ -273,7 +273,7 @@ describeEmbeddedPostgres("budgetService release gate enforcement", () => { .from(projects); expect(projectAfterHardStop?.pauseReason).toBe("budget"); expect(projectAfterHardStop?.pausedAt).toBeInstanceOf(Date); - expect(cancelWorkForScope).toHaveBeenCalledWith({ companyId, scopeType: "project", scopeId: projectId, createdBefore: expect.any(Date) }); + expect(cancelWorkForScope).toHaveBeenCalledWith({ companyId, scopeType: "project", scopeId: projectId, createdBefore: expect.any(Date), enforcement: { policyId: expect.any(String), version: expect.any(Number) } }); const overviewWhileBlocked = await service.overview(companyId); expect(overviewWhileBlocked.pausedProjectCount).toBe(1); diff --git a/server/src/__tests__/cost-accounting-operations-edges.test.ts b/server/src/__tests__/cost-accounting-operations-edges.test.ts index 6a97dcc5978..7398a58c818 100644 --- a/server/src/__tests__/cost-accounting-operations-edges.test.ts +++ b/server/src/__tests__/cost-accounting-operations-edges.test.ts @@ -4,7 +4,7 @@ import os from "node:os"; import path from "node:path"; import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; import { and, eq, sql } from "drizzle-orm"; -import { accountingRuntimeBaselines, agents, agentRuntimeState, billingInvoiceLines, budgetPolicies, budgetReservations, companies, costEvents, createDb, heartbeatRuns, projects, runUsageReceipts } from "@paperclipai/db"; +import { accountingRuntimeBaselines, agents, agentRuntimeState, billingInvoiceLines, budgetPolicies, budgetReservations, companies, costEvents, createDb, heartbeatRuns, nativeRunFinalizations, issues, projects, runUsageReceipts } from "@paperclipai/db"; import { accountingIntegrityService } from "../services/accounting-integrity.js"; import { billingReconciliationService } from "../services/billing-reconciliation.js"; import { budgetService } from "../services/budgets.js"; @@ -67,6 +67,35 @@ const support = await getEmbeddedPostgresTestSupport(); expect((await db.select().from(budgetReservations).where(eq(budgetReservations.runId,f.run.id)))[0].state).toBe("released"); expect(await db.select().from(costEvents).where(eq(costEvents.companyId,f.company.id))).toEqual([]); }); + it.each(["retryable_failure", "observed"])("keeps native %s receipts open through same-run recovery", async (phase) => { + const f = await fixture(); + const [issue] = await db.insert(issues).values({ companyId: f.company.id, title: "Recoverable native work" }).returning(); + await db.update(heartbeatRuns).set({ runtimeMode: "native", nativeIssueId: issue.id }).where(eq(heartbeatRuns.id, f.run.id)); + await reserveRunBudget(db, f.company.id, f.run.id, null); + await db.insert(nativeRunFinalizations).values({ companyId: f.company.id, runId: f.run.id, issueId: issue.id, phase }); + const input = { companyId: f.company.id, runId: f.run.id, adapterType: "paperclip_runner" }; + const spool = path.join(directory, randomUUID()); + const recorder = await createRunUsageRecorder(db, input, spool); + await recorder.capture({ complete: true, usageBasis: "per_run", usage: { inputTokens: 100, outputTokens: 10 }, costUsd: 0.1 }); + 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(false); + await reconcileRunCosts(db); + expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.run.id)))[0].costAccountedAt).toBeNull(); + expect((await db.select().from(budgetReservations).where(eq(budgetReservations.runId, f.run.id)))[0].state).toBe("held"); + await db.update(heartbeatRuns).set({ status: "running" }).where(eq(heartbeatRuns.id, f.run.id)); + const resumed = await createRunUsageRecorder(db, input, spool); + expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.run.id)))[0].usageJson?.accountingReceiptReady).toBe(false); + const empty = await resumed.complete({ exitCode: 0, signal: null, timedOut: false }); + expect(empty).toMatchObject({ complete: false, usage: { inputTokens: 100, outputTokens: 10 } }); + // The provider's native delta is cumulative for this run across recovery. + await resumed.capture({ complete: true, usageBasis: "per_run", usage: { inputTokens: 150, outputTokens: 15 }, costUsd: 0.15 }); + await db.update(nativeRunFinalizations).set({ phase: "terminal_failure" }).where(eq(nativeRunFinalizations.runId, f.run.id)); + 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 db.select().from(costEvents).where(eq(costEvents.heartbeatRunId, f.run.id)))[0]).toMatchObject({ inputTokens: 150, outputTokens: 15, costCents: 15 }); + }); + it("reports over-limit advisory budgets without blocking, and tolerates legacy active zero limits", async () => { const f = await fixture(), budgets = budgetService(db); const policy = await budgets.upsertPolicy(f.company.id, { scopeType: "company", scopeId: f.company.id, amount: 1, hardStopEnabled: false }, "board"); diff --git a/server/src/__tests__/cost-accounting-operators.test.ts b/server/src/__tests__/cost-accounting-operators.test.ts index cd9605d1917..42d4b89bf18 100644 --- a/server/src/__tests__/cost-accounting-operators.test.ts +++ b/server/src/__tests__/cost-accounting-operators.test.ts @@ -75,7 +75,7 @@ const support = await getEmbeddedPostgresTestSupport(); const original = await reserveRunBudget(db, f.company.id, run.id, null); await expect(reserveRunBudget(db, f.company.id, run.id, null)).rejects.toThrow("already reserved"); await expect(reserveRunBudget(db, f.company.id, run.id, null, {}, "stale-owner")).rejects.toThrow("already reserved"); - expect(await reserveRunBudget(db, f.company.id, run.id, null, {}, "successor")).toEqual(original); + expect(await reserveRunBudget(db, f.company.id, run.id, null, {}, "successor")).toEqual({ ...original, reused: true }); expect((await accountingIntegrityService(db).health(f.company.id)).heldReservationCents).toBe("6.0000000"); await db.update(nativeRunFinalizations).set({ leaseExpiresAt: new Date(0) }).where(eq(nativeRunFinalizations.runId, run.id)); await expect(reserveRunBudget(db, f.company.id, run.id, null, {}, "successor")).rejects.toThrow("already reserved"); diff --git a/server/src/__tests__/cost-accounting-reliability.test.ts b/server/src/__tests__/cost-accounting-reliability.test.ts index 8e285800d0f..3b7be23c970 100644 --- a/server/src/__tests__/cost-accounting-reliability.test.ts +++ b/server/src/__tests__/cost-accounting-reliability.test.ts @@ -2,9 +2,11 @@ import { upsertBudgetPolicySchema } from "@paperclipai/shared"; import { randomUUID } from "node:crypto"; import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; import { and, eq, sql } from "drizzle-orm"; -import { activityLog, agentRuntimeState, agents, approvals, budgetIncidents, budgetPolicies, companies, costEvents, createDb, financeEvents, goals, heartbeatRuns, issues, projects } from "@paperclipai/db"; -import { costService } from "../services/costs.js"; -import { budgetService } from "../services/budgets.js"; +import { activityLog, agentRuntimeState, agentWakeupRequests, agents, approvals, budgetIncidents, budgetPolicies, companies, costEvents, createDb, financeEvents, goals, heartbeatRuns, issues, projects } from "@paperclipai/db"; +import { reserveRunBudget } from "../services/budget-reservations.js"; +import { createRunDispatch } from "../modules/run-dispatch/index.js"; +import { costService, createCostEventInTransaction } from "../services/costs.js"; +import { budgetService, withCurrentBudgetEnforcement, type BudgetEnforcementScope } from "../services/budgets.js"; import { accountRunCost, reconcileRunCosts } from "../services/run-cost-accounting.js"; import { financeService } from "../services/finance.js"; import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "./helpers/embedded-postgres.js"; @@ -57,6 +59,31 @@ databaseDescribe("cost accounting reliability (PostgreSQL)", () => { }); }, 5000); + it("takes the company lock before retry issue/run locks during an attributed cost write", async () => { + const f = await fixture(); + await db.update(issues).set({ status: "in_progress", assigneeAgentId: f.agent.id }).where(eq(issues.id, f.issue.id)); + await db.update(heartbeatRuns).set({ status: "scheduled_retry", scheduledRetryAt: new Date(0), scheduledRetryReason: "workspace_busy", + contextSnapshot: { issueId: f.issue.id, projectId: f.project.id }, + }).where(eq(heartbeatRuns.id, f.run.id)); + let promotion: ReturnType["promoteScheduledRetry"]> | undefined; + await db.transaction(async (tx) => { + await tx.select().from(companies).where(eq(companies.id, f.company.id)).for("no key update"); + promotion = createRunDispatch(db).promoteScheduledRetry({ companyId: f.company.id, runId: f.run.id }); + // Observe the actual blocked company-lock query before checking the issue; + // no timing guess or timeout is used as proof of lock ordering. + await vi.waitFor(async () => { + const result = await db.execute(sql`select pid from pg_stat_activity where datname = current_database() + and pid <> pg_backend_pid() and wait_event_type = 'Lock' and query like '%companies%for no key update%'`); + expect(result.length).toBeGreaterThan(0); + }, { timeout: 3000 }); + await tx.execute(sql`select id from issues where id = ${f.issue.id} for key share nowait`); + await createCostEventInTransaction(tx as unknown as ReturnType, f.company.id, + { ...f.receipt, issueId: f.issue.id, projectId: f.project.id }, []); + }); + expect((await promotion)?.outcome).toBe("promoted"); + expect((await costService(db).summary(f.company.id)).spendCents).toBe(100); + }, 10000); + it("deduplicates concurrent retries and rejects changed receipt content", async () => { const f = await fixture(); const receipt = { ...f.receipt, idempotencyKey: "provider-receipt-1" }; @@ -419,4 +446,57 @@ databaseDescribe("cost accounting reliability (PostgreSQL)", () => { expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, newRun.id)))[0].status).toBe("scheduled_retry"); }); + it.each(["agent", "company", "project"] as const)("rechecks delayed %s enforcement after a grant admits an old queued run", async (scopeType) => { + const f = await fixture(); + const { heartbeatService } = await import("../services/heartbeat.js"); + const heartbeat = heartbeatService(db); + const budgets = budgetService(db); + const scopeId = scopeType === "agent" ? f.agent.id : scopeType === "company" ? f.company.id : f.project.id; + const policyInput = { scopeType, scopeId, amount: 1 }; + await db.update(heartbeatRuns).set({ status: "queued", contextSnapshot: { projectId: f.project.id } }).where(eq(heartbeatRuns.id, f.run.id)); + const [wake] = await db.insert(agentWakeupRequests).values({ companyId: f.company.id, agentId: f.agent.id, source: "automation", payload: { projectId: f.project.id } }).returning(); + await budgets.upsertPolicy(f.company.id, policyInput, "board"); + await costService(db).createEvent(f.company.id, { ...f.receipt, projectId: f.project.id, costCents: 2 }); + let release!: () => void; + const paused = new Promise(resolve => { release = resolve; }); + let snapshotReady!: (scope: BudgetEnforcementScope) => void; + const snapshot = new Promise(resolve => { snapshotReady = resolve; }); + const delivery = budgetService(db, { cancelWorkForScope: async (scope) => { + snapshotReady(scope); await paused; await heartbeat.cancelBudgetScopeWork(scope); + } }).deliverPendingEnforcement(f.company.id); + try { + expect((await snapshot).enforcement).toMatchObject({ policyId: expect.any(String), version: expect.any(Number) }); + await budgets.upsertPolicy(f.company.id, { ...policyInput, amount: 100 }, "board"); + await db.update(heartbeatRuns).set({ status: "running" }).where(eq(heartbeatRuns.id, f.run.id)); + await reserveRunBudget(db, f.company.id, f.run.id, f.project.id); + } finally { release(); } + await delivery; + expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.run.id)))[0]).toMatchObject({ status: "running", resultJson: null }); + expect((await db.select().from(agentWakeupRequests).where(eq(agentWakeupRequests.id, wake.id)))[0].status).toBe("queued"); + }); + + it("keeps a durably claimed budget stop fenced after a later grant", async () => { + const f = await fixture(), budgets = budgetService(db); + await budgets.upsertPolicy(f.company.id, { scopeType: "company", scopeId: f.company.id, amount: 1 }, "board"); + await costService(db).createEvent(f.company.id, { ...f.receipt, costCents: 2 }); + const [policy] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.companyId, f.company.id)); + const scope: BudgetEnforcementScope = { companyId: f.company.id, scopeType: "company", scopeId: f.company.id, + enforcement: { policyId: policy.id, version: policy.enforcementVersion } }; + expect(await withCurrentBudgetEnforcement(db, scope, async (tx) => { + await tx.update(heartbeatRuns).set({ status: "running", resultJson: { cancellation: { reason: "Budget stop" } } }).where(eq(heartbeatRuns.id, f.run.id)); + return true; + })).toBe(true); + await budgets.upsertPolicy(f.company.id, { scopeType: "company", scopeId: f.company.id, amount: 100 }, "board"); + await expect(reserveRunBudget(db, f.company.id, f.run.id, null)).rejects.toThrow("cancellation already requested"); + const effect = vi.fn(async () => true); + for (const invalid of [ + scope, + { ...scope, enforcement: { ...scope.enforcement!, policyId: randomUUID() } }, + { ...scope, scopeId: randomUUID() }, + { ...scope, scopeType: "agent" as const }, + { ...scope, enforcement: { ...scope.enforcement!, version: policy.enforcementVersion + 1 } }, + ]) expect(await withCurrentBudgetEnforcement(db, invalid, effect)).toBeNull(); + expect(effect).not.toHaveBeenCalled(); + }); + }); diff --git a/server/src/__tests__/heartbeat-process-recovery.test.ts b/server/src/__tests__/heartbeat-process-recovery.test.ts index eeb8f031de4..9208a72c27b 100644 --- a/server/src/__tests__/heartbeat-process-recovery.test.ts +++ b/server/src/__tests__/heartbeat-process-recovery.test.ts @@ -7553,6 +7553,19 @@ describeEmbeddedPostgres("heartbeat orphaned process recovery", () => { }, ); + it.each([{}, { command: "/nonexistent-paperclip-bootstrap-command" }])("releases holds for real process bootstrap failures: %j", async (config) => { + const { execute } = await import("../adapters/process/execute.js"); + mockAdapterExecute.mockImplementationOnce(execute); + const { runId, agentId } = await seedRunFixture({ runtimeMode: "legacy", agentStatus: "idle", runStatus: "queued", includeIssue: false }); + await db.update(agents).set({ adapterType: "process", adapterConfig: config }).where(eq(agents.id, agentId)); + const heartbeat = heartbeatService(db); + await heartbeat.resumeQueuedRuns(); + await heartbeat.drainActiveRunExecutions(); + expect(await heartbeat.getRun(runId)).toMatchObject({ status: "failed", costAccountingPending: false, costAccountedAt: expect.any(Date) }); + expect((await db.select().from(budgetReservations).where(eq(budgetReservations.runId, runId)))[0]?.state).toBe("released"); + expect(await db.select().from(costEvents).where(eq(costEvents.heartbeatRunId, runId))).toEqual([]); + }); + it("releases the dispatch reservation after a proven pre-provider adapter failure", async () => { mockAdapterExecute.mockResolvedValueOnce({ exitCode: 1, signal: null, timedOut: false, diff --git a/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts b/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts index 719e4ea9778..00fcd37b246 100644 --- a/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts +++ b/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts @@ -6,6 +6,7 @@ import { agentWakeupRequests, companies, costEvents, + budgetReservations, createDb, documentRevisions, documents, @@ -847,6 +848,9 @@ describeEmbeddedPostgres("heartbeat stale queued-run invalidation", () => { afterContinuationDispatchCheck = async ({ runId: guardedRunId, issueId: guardedIssueId }) => { expect(guardedRunId).toBe(runId); expect(guardedIssueId).toBe(issueId); + // Budget admission must finish before this last ownership check. Putting + // an awaited reservation inside dispatch lets work escape the row lock. + expect((await db.select().from(budgetReservations).where(eq(budgetReservations.runId, runId)))[0]?.state).toBe("held"); ordering.push("validated"); parkPromise = Promise.resolve( db @@ -890,6 +894,8 @@ describeEmbeddedPostgres("heartbeat stale queued-run invalidation", () => { expect(issue?.status).toBe("backlog"); expect(ordering).toEqual(["validated", "parked"]); expect(countExecuteCallsForRun(runId)).toBe(0); + await heartbeat.drainActiveRunExecutions(); + expect((await db.select().from(budgetReservations).where(eq(budgetReservations.runId, runId)))[0]?.state).toBe("released"); }); it("rate-limits skipped generic timer wakes by advancing the timer baseline", async () => { diff --git a/server/src/__tests__/workspace-runtime.test.ts b/server/src/__tests__/workspace-runtime.test.ts index 187135d1959..b602d38b281 100644 --- a/server/src/__tests__/workspace-runtime.test.ts +++ b/server/src/__tests__/workspace-runtime.test.ts @@ -43,6 +43,7 @@ import { releaseRuntimeServicesForRun, UnresolvedWorkspaceBaseRefError, resetRuntimeServicesForTests, + resetRuntimeServicePortReservationsForTests, MANAGED_RUNTIME_PUBLIC_URL_ENV, resolveManagedPaperclipRuntimePublicOrigin, resolveRuntimeProvisionCommand, @@ -454,6 +455,9 @@ afterEach(async () => { delete process.env.PAPERCLIP_WORKTREES_DIR; delete process.env.DATABASE_URL; await resetRuntimeServicesForTests(); + // Registry reset does not clear the process-local allocation claims. A + // failed-start fixture must not reserve another test's ephemeral port. + resetRuntimeServicePortReservationsForTests(); }); describe("sanitizeRuntimeServiceBaseEnv", () => { @@ -4099,6 +4103,7 @@ describe("realizeExecutionWorkspace", () => { "utf8", ); process.env.PAPERCLIP_WORKTREES_DIR = worktreesDir; + const canonicalInstanceRoot = await fs.realpath(instanceRoot); await cleanupExecutionWorkspaceArtifacts({ workspace: { @@ -4133,7 +4138,7 @@ describe("realizeExecutionWorkspace", () => { expect(operations[0]?.command).toBe("printf 'cleanup ok\\n'"); expect(operations[1]?.metadata).toMatchObject({ cleanupAction: "remove_worktree_instance", - instanceRoot, + instanceRoot: canonicalInstanceRoot, }); expect(operations[2]?.metadata).toMatchObject({ cleanupAction: "worktree_remove", @@ -6759,7 +6764,8 @@ describeEmbeddedPostgres("workspace runtime service control persistence", () => status: "starting", healthStatus: "unknown", }); - expect(startingRow.providerRef).toMatch(/^\d+$/); + // Starting is persisted before the asynchronous spawn metadata write. + // The child marker alone does not prove that write has committed. expect(startingRow.port).toEqual(expect.any(Number)); const services = await startPromise; @@ -6772,6 +6778,7 @@ describeEmbeddedPostgres("workspace runtime service control persistence", () => const runningRow = await waitForPersistedStatus("running"); expect(runningRow.id).toBe(startingRow.id); + expect(runningRow.providerRef).toMatch(/^\d+$/); await expect(fetch(services[0]!.url!)).resolves.toMatchObject({ ok: true }); const runtimeProvisionOperations = await db .select() @@ -7643,16 +7650,25 @@ describeEmbeddedPostgres("workspace runtime startup reconciliation", () => { process.env.PAPERCLIP_HOME = paperclipHome; process.env.PAPERCLIP_INSTANCE_ID = `runtime-desired-reconcile-${randomUUID()}`; - const reservePort = async () => { - const probe = net.createServer(); - await new Promise((resolve) => probe.listen(0, "127.0.0.1", resolve)); - const address = probe.address(); - const port = typeof address === "object" && address ? address.port : null; - await new Promise((resolve, reject) => { - probe.close((error) => error ? reject(error) : resolve()); - }); - if (!port) throw new Error("Failed to reserve runtime reconciliation test port"); - return port; + const reservePorts = async () => { + const probes = [net.createServer(), net.createServer()]; + try { + // Keep both sockets open until both allocations finish. Closing the + // first probe early lets the kernel return that same port again. + for (const probe of probes) await new Promise((resolve, reject) => { + probe.once("error", reject); + probe.listen(0, "127.0.0.1", resolve); + }); + return probes.map(probe => { + const address = probe.address(); + if (!address || typeof address !== "object") throw new Error("Failed to reserve runtime reconciliation test port"); + return address.port; + }); + } finally { + await Promise.all(probes.filter(probe => probe.listening).map(probe => new Promise((resolve, reject) => { + probe.close(error => error ? reject(error) : resolve()); + }))); + } }; const isLoopbackPortFree = async (port: number) => { const probe = net.createServer(); @@ -7673,8 +7689,8 @@ describeEmbeddedPostgres("workspace runtime startup reconciliation", () => { } throw new Error(`Port ${port} did not become free in time`); }; - const stoppedPort = await reservePort(); - const livePort = await reservePort(); + const [stoppedPort, livePort] = await reservePorts(); + expect(stoppedPort).not.toBe(livePort); const companyId = randomUUID(); const projectId = randomUUID(); const projectWorkspaceId = randomUUID(); @@ -9728,7 +9744,7 @@ describe("realizeExecutionWorkspace with an exact existing branch", () => { const workspace = await realizeExistingBranch(repoRoot, "feature/legacy-checkout"); - expect(workspace.cwd).toBe(path.resolve(legacyPath)); + expect(workspace.cwd).toBe(await fs.realpath(legacyPath)); expect(workspace.branchName).toBe("feature/legacy-checkout"); expect(workspace.created).toBe(false); expect(await readGit(workspace.cwd, ["rev-parse", "HEAD"])).toBe(branchTip); diff --git a/server/src/adapters/process/execute.ts b/server/src/adapters/process/execute.ts index 16ef40457e4..1087f688418 100644 --- a/server/src/adapters/process/execute.ts +++ b/server/src/adapters/process/execute.ts @@ -15,6 +15,24 @@ import { } from "../utils.js"; export async function execute(ctx: AdapterExecutionContext): Promise { + let spawned = false; + try { + return await executeProcess({ ...ctx, onSpawn: async (meta) => { + // Set this before metadata persistence: a failed callback cannot prove + // that the already-created process did no billable work. + spawned = true; + await ctx.onSpawn?.(meta); + } }); + } catch (error) { + if (spawned) throw error; + return { exitCode: 1, signal: null, timedOut: false, + errorMessage: error instanceof Error ? error.message : String(error), + executionRecovery: { kind: "bootstrap", providerWorkStarted: false }, + }; + } +} + +async function executeProcess(ctx: AdapterExecutionContext): Promise { const { runId, agent, config, onLog, onMeta, authToken } = ctx; const command = asString(config.command, ""); if (!command) throw new Error("Process adapter missing command"); diff --git a/server/src/modules/run-dispatch/adapters/postgres.ts b/server/src/modules/run-dispatch/adapters/postgres.ts index b882b09baca..026715f95a9 100644 --- a/server/src/modules/run-dispatch/adapters/postgres.ts +++ b/server/src/modules/run-dispatch/adapters/postgres.ts @@ -15,9 +15,11 @@ import { issues, } from "@paperclipai/db"; import { ISSUE_DISPOSITION_REPAIR_RETRY_REASON } from "@paperclipai/shared"; +import { withAccountingTransaction } from "../../../services/accounting-transaction.js"; +import type { ActivityPublication } from "../../../services/activity-log.js"; import { parseObject } from "../../../adapters/utils.js"; import { evaluateAgentInvokabilityFromDb } from "../../../services/agent-invokability.js"; -import { budgetService } from "../../../services/budgets.js"; +import { budgetService, budgetServiceInTransaction } from "../../../services/budgets.js"; import { isCompletedOnboardingHandoffWake } from "../../../services/chat-completion-delivery.js"; import { isHeartbeatWakeOnDemandEnabled } from "../../../services/heartbeat-policy.js"; import { collectDispositionRepairSourceState } from "../../../services/recovery/disposition-repair.js"; @@ -190,7 +192,7 @@ export function createPostgresRunDispatchAdapter( async function withIssueThenRunLocks( input: { runId: string; companyId: string }, onMissing: () => T, - operation: (tx: Db, run: HeartbeatRun) => Promise, + operation: (tx: Db, run: HeartbeatRun, publications: ActivityPublication[]) => Promise, ): Promise { // Queue editing and claiming already use issue -> wake -> run. Read the // immutable issue reference without a lock first, then acquire issue -> @@ -206,8 +208,9 @@ export function createPostgresRunDispatchAdapter( const hintedIssueId = readNonEmptyString(parseObject(hint.contextSnapshot).issueId); const result: { kind: "missing" | "retry" } | { kind: "value"; value: T } = - await db.transaction(async (tx) => { - const typedTx = tx as unknown as Db; + await withAccountingTransaction(db, input.companyId, async (typedTx, publications) => { + // Budget gates mutate policy state. Acquire company before issue/run, + // matching receipt writers (including their attribution FK locks). if (hintedIssueId) { await typedTx .select({ id: issues.id }) @@ -231,7 +234,7 @@ export function createPostgresRunDispatchAdapter( const lockedIssueId = readNonEmptyString(parseObject(run.contextSnapshot).issueId); if (lockedIssueId !== hintedIssueId) return { kind: "retry" as const }; - return { kind: "value" as const, value: await operation(typedTx, run) }; + return { kind: "value" as const, value: await operation(typedTx, run, publications) }; }); if (result.kind === "missing") return onMissing(); @@ -247,12 +250,13 @@ export function createPostgresRunDispatchAdapter( input: LoadGateFactsInput, now: Date, tx?: unknown, + publications?: ActivityPublication[], ): Promise { // Semantic adapter operations pass their transaction here so the fact // read and the state transition share one unit of work. This helper is // deliberately not exposed through the module's public API. const dbOrTx = (tx as Db | undefined) ?? db; - const budgetsForRead = tx ? budgetService(dbOrTx) : budgets; + const budgetsForRead = tx ? budgetServiceInTransaction(dbOrTx, publications) : budgets; const treeControlForRead = tx ? issueTreeControlService(dbOrTx) : treeControlSvc; const issuesSvcForRead = tx ? issueService(dbOrTx) : issuesSvc; @@ -747,7 +751,7 @@ export function createPostgresRunDispatchAdapter( ): Promise { const now = input.now; - const promoteLockedRun = async (tx: Db, run: HeartbeatRun) => { + const promoteLockedRun = async (tx: Db, run: HeartbeatRun, publications: ActivityPublication[]) => { if ( run.status !== "scheduled_retry" || !run.scheduledRetryAt || @@ -768,6 +772,7 @@ export function createPostgresRunDispatchAdapter( }, now, tx, + publications, ); if (!factsResult.agentFound) { @@ -1010,7 +1015,7 @@ export function createPostgresRunDispatchAdapter( input: DispatchResolvedInteractionInput, ): Promise> { const dispatchLockedRun = async (tx: Db, run: HeartbeatRun) => { - if (run.status !== input.expectedStatus) { + if (run.status !== input.expectedStatus || run.resultJson?.cancellation || run.resultJson?.startupCancellation) { return { dispatched: false as const, cancellation: { outcome: "lost_race" as const } }; } const { issueId, decision: initialDecision } = await decideCurrentRunStaleness( diff --git a/server/src/services/ai-connections.ts b/server/src/services/ai-connections.ts index 89ee73554bf..e162d0f5926 100644 --- a/server/src/services/ai-connections.ts +++ b/server/src/services/ai-connections.ts @@ -34,6 +34,8 @@ import { type AiConnectionLoginIntent, } from "@paperclipai/shared"; import { forbidden, notFound, unprocessable } from "../errors.js"; +import { quotaCredentialRecovery, quotaCredentialHash } from "./quota-credential-recovery.js"; +import { logger } from "../middleware/logger.js"; import { logActivity } from "./activity-log.js"; import { secretService } from "./secrets.js"; import { probeAiConnectionUsage } from "./ai-connection-usage.js"; @@ -860,66 +862,96 @@ 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", { noWait: true }); - if (!grant || grant.status !== "active") throw new Error("credentials_unavailable"); - const ref = grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential"); - if (!ref) throw new Error("credentials_unavailable"); - await tx.select({ id: companySecrets.id }).from(companySecrets).where(and( - eq(companySecrets.id, ref.secretId), eq(companySecrets.companyId, companyId), - )).for("update", { noWait: true }); - const current = await aiConnectionService(tx as unknown as Db).credential({ connection: row.connection, grant }); - signal.throwIfAborted(); - // Another poll, an agent, or reconnect already supplied a newer credential. - if (current !== failedValue) return current; - // An active provider may rotate this single-use token itself. Defer polling - // until its credential write-back instead of invalidating its live copy. - const [active] = await tx.select({ id: heartbeatRuns.id }).from(heartbeatRuns) - .innerJoin(agents, and(eq(agents.id, heartbeatRuns.agentId), eq(agents.companyId, companyId))).where(and( - eq(heartbeatRuns.companyId, companyId), inArray(heartbeatRuns.status, ["queued", "running"]), - or(sql`${heartbeatRuns.contextSnapshot}->'aiConnection'->>'grantId' = ${grant.id}`, - sql`${agents.runtimeConfig}->'aiConnection'->>'provider' = 'openai'`, eq(agents.adapterType, "codex_local")), - )).limit(1); - if (active) throw new Error("provider_unavailable"); - const auth = JSON.parse(current); - const refreshToken = auth.tokens?.refresh_token; - if (typeof refreshToken !== "string" || !refreshToken) throw new Error("authentication_required"); - // Once exchange starts, its lifetime is independent of the dashboard's - // shorter deadline. A consumed single-use token must still be read/saved. - signal.throwIfAborted(); - const refreshSignal = AbortSignal.timeout(60_000); - // Matches the Codex OAuth client (openai/codex, login/src/auth/manager.rs). - const response = await fetch("https://auth.openai.com/oauth/token", { - method: "POST", redirect: "error", signal: refreshSignal, - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ client_id: "app_EMoamEEZ73f0CkXaXp7hrann", grant_type: "refresh_token", refresh_token: refreshToken }), + return withAccountHomeSecretMutationLock(undefined, row.connection.companyId, async () => { + const recovery = quotaCredentialRecovery(row.connection.companyId, row.grant.id); + let savingReplacement = false; + const save = () => db.transaction(async (tx) => { + // Do not wait on a transaction that might itself be waiting for the file + // lock. All database locks must be available before consuming OAuth tokens. + const companyId = row.connection.companyId; + const [grant] = await tx.select().from(connectionGrants).where(and( + eq(connectionGrants.id, row.grant.id), eq(connectionGrants.companyId, companyId), + eq(connectionGrants.connectionId, row.connection.id), + )).for("update", { noWait: true }); + if (!grant || grant.status !== "active") throw new Error("credentials_unavailable"); + const ref = grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential"); + if (!ref) throw new Error("credentials_unavailable"); + await tx.select({ id: companySecrets.id }).from(companySecrets).where(and( + eq(companySecrets.id, ref.secretId), eq(companySecrets.companyId, companyId), + )).for("update", { noWait: true }); + const current = await aiConnectionService(tx as unknown as Db).credential({ connection: row.connection, grant }); + signal.throwIfAborted(); + // Another poll, an agent, or reconnect already supplied a newer credential. + if (current !== failedValue) return current; + const pending = await recovery.read(); + if (pending && (pending.secretId !== ref.secretId || pending.connectionId !== row.connection.id || pending.baseHash !== quotaCredentialHash(current))) { + throw new Error("credentials_unavailable"); + } + // An active provider may rotate this single-use token itself. Defer polling + // until its credential write-back instead of invalidating its live copy. + const [active] = await tx.select({ id: heartbeatRuns.id }).from(heartbeatRuns) + .innerJoin(agents, and(eq(agents.id, heartbeatRuns.agentId), eq(agents.companyId, companyId))).where(and( + eq(heartbeatRuns.companyId, companyId), inArray(heartbeatRuns.status, ["queued", "running"]), + or(sql`${heartbeatRuns.contextSnapshot}->'aiConnection'->>'grantId' = ${grant.id}`, + sql`${agents.runtimeConfig}->'aiConnection'->>'provider' = 'openai'`, eq(agents.adapterType, "codex_local")), + )).limit(1); + if (active && !pending) throw new Error("provider_unavailable"); + let refreshed = pending?.value; + if (!refreshed) { + await recovery.prepare(); + const auth = JSON.parse(current); + const refreshToken = auth.tokens?.refresh_token; + if (typeof refreshToken !== "string" || !refreshToken) throw new Error("authentication_required"); + // Once exchange starts, its lifetime is independent of the dashboard's + // shorter deadline. A consumed single-use token must still be read/saved. + signal.throwIfAborted(); + const refreshSignal = AbortSignal.timeout(60_000); + // Matches the Codex OAuth client (openai/codex, login/src/auth/manager.rs). + const response = await fetch("https://auth.openai.com/oauth/token", { + method: "POST", redirect: "error", signal: refreshSignal, + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ client_id: "app_EMoamEEZ73f0CkXaXp7hrann", grant_type: "refresh_token", refresh_token: refreshToken }), + }); + if (!response.ok) { + await response.body?.cancel(); + throw new Error(response.status === 400 || response.status === 401 ? "authentication_required" : "provider_unavailable"); + } + const tokens = await response.json() as Record; + if (typeof tokens.access_token !== "string" || !tokens.access_token) throw new Error("provider_unavailable"); + refreshed = JSON.stringify({ ...auth, tokens: { ...auth.tokens, + access_token: tokens.access_token, + ...(typeof tokens.refresh_token === "string" && tokens.refresh_token ? { refresh_token: tokens.refresh_token } : {}), + ...(typeof tokens.id_token === "string" && tokens.id_token ? { id_token: tokens.id_token } : {}), + }, last_refresh: new Date().toISOString() }); + // This fsynced encrypted journal survives transaction rollback and host + // restart. Never exchange the consumed token again after a failed save. + await recovery.write({ companyId, grantId: grant.id, connectionId: row.connection.id, + secretId: ref.secretId, baseHash: quotaCredentialHash(current), value: refreshed }); + } + savingReplacement = true; + // Persist the rotated token even if the quota deadline expires after the + // successful exchange; losing it would strand all subsequent executions. + await secretService(tx).rotate(ref.secretId, { value: refreshed }, { userId: grant.subjectUserId }); + await tx.update(connectionGrants).set({ updatedAt: new Date() }).where(eq(connectionGrants.id, grant.id)); + await logActivity(tx as unknown as Db, { companyId, actorType: "system", actorId: "quota", + action: "ai_connection.credential_refreshed", entityType: "tool_connection", entityId: row.connection.id, + details: { provider: "openai", grantId: grant.id } }); + return refreshed; }); - if (!response.ok) { - await response.body?.cancel(); - throw new Error(response.status === 400 || response.status === 401 ? "authentication_required" : "provider_unavailable"); + for (let attempt = 0; ; attempt++) { + try { + const value = await save(); + // Clear only after commit. A committed-but-lost response is detected + // by the current credential comparison on the next attempt. + await recovery.clear().catch((error) => logger.warn({ err: error, grantId: row.grant.id }, "Saved quota credential recovery cleanup pending")); + return value; + } catch (error) { + if (!savingReplacement || attempt >= 2) throw error; + // Saving retries have no provider side effect. Persistent failures + // retain the journal for the next quota poll, including after restart. + } } - const tokens = await response.json() as Record; - if (typeof tokens.access_token !== "string" || !tokens.access_token) throw new Error("provider_unavailable"); - const refreshed = JSON.stringify({ ...auth, tokens: { ...auth.tokens, - access_token: tokens.access_token, - ...(typeof tokens.refresh_token === "string" && tokens.refresh_token ? { refresh_token: tokens.refresh_token } : {}), - ...(typeof tokens.id_token === "string" && tokens.id_token ? { id_token: tokens.id_token } : {}), - }, last_refresh: new Date().toISOString() }); - // Persist the rotated token even if the quota deadline expires after the - // successful exchange; losing it would strand all subsequent executions. - await secretService(tx).rotate(ref.secretId, { value: refreshed }, { userId: grant.subjectUserId }); - await tx.update(connectionGrants).set({ updatedAt: new Date() }).where(eq(connectionGrants.id, grant.id)); - await logActivity(tx as unknown as Db, { companyId, actorType: "system", actorId: "quota", - action: "ai_connection.credential_refreshed", entityType: "tool_connection", entityId: row.connection.id, - details: { provider: "openai", grantId: grant.id } }); - return refreshed; - })); + }); } // 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/server/src/services/budget-reservations.ts b/server/src/services/budget-reservations.ts index 134048f0e51..cefb4709da8 100644 --- a/server/src/services/budget-reservations.ts +++ b/server/src/services/budget-reservations.ts @@ -12,6 +12,7 @@ export async function reserveRunBudget(db: Db, companyId: string, runId: string, return withAccountingTransaction(db, companyId, async (tx, publications) => { const [run] = await tx.select().from(heartbeatRuns).where(and(eq(heartbeatRuns.id, runId), eq(heartbeatRuns.companyId, companyId))).for("update"); if (!run) throw notFound("Run not found"); + if (run.resultJson?.cancellation || run.resultJson?.startupCancellation) throw conflict("Run cancellation already requested"); if (run.costAccountedAt || !["queued", "running"].includes(run.status)) throw conflict("Run can no longer start provider work"); const [existing] = await tx.select().from(budgetReservations).where(and(eq(budgetReservations.companyId, companyId), eq(budgetReservations.runId, runId))); if (existing) { @@ -22,7 +23,7 @@ export async function reserveRunBudget(db: Db, companyId: string, runId: string, .from(nativeRunFinalizations).where(and(eq(nativeRunFinalizations.companyId, companyId), eq(nativeRunFinalizations.runId, runId), eq(nativeRunFinalizations.leaseOwner, recoveryLeaseOwner), sql`${nativeRunFinalizations.leaseExpiresAt} > now()`)).for("update") : []; - if (owner) return existing; + if (owner) return { ...existing, reused: true }; throw conflict("Provider dispatch has already reserved this run"); } const block = await budgetServiceInTransaction(tx, publications).getInvocationBlock(companyId, run.agentId, { projectId }); @@ -50,6 +51,6 @@ export async function reserveRunBudget(db: Db, companyId: string, runId: string, await tx.update(heartbeatRuns).set({ costAccountingPending: true, usageJson: sql`coalesce(${heartbeatRuns.usageJson}, '{}'::jsonb) || ${JSON.stringify({ accountingReceiptReady: false, ledgerScope: { ...(run.usageJson?.ledgerScope as object ?? {}), ...ledgerScope, projectId } })}::jsonb`, }).where(eq(heartbeatRuns.id, runId)); - return reservation; + return { ...reservation, reused: false }; }); } diff --git a/server/src/services/budgets.ts b/server/src/services/budgets.ts index 83e304364aa..2b940306201 100644 --- a/server/src/services/budgets.ts +++ b/server/src/services/budgets.ts @@ -45,6 +45,8 @@ export type BudgetEnforcementScope = { scopeId: string; /** Snapshot taken under the company lock, before a later budget grant can admit new work. */ createdBefore?: Date; + /** Revalidated at the durable cancellation write, under the admission lock. */ + enforcement?: { policyId: string; version: number }; }; export type BudgetServiceHooks = { @@ -939,6 +941,22 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica }; } +/** Fence the durable stop intent against policy edits and admission. External + * process shutdown must happen after this transaction, without holding locks + * that the executor needs to persist its final usage receipt. */ +export async function withCurrentBudgetEnforcement(db: Db, scope: BudgetEnforcementScope, work: (tx: Db) => Promise): Promise { + return withAccountingTransaction(db, scope.companyId, async (tx) => { + if (scope.enforcement) { + const [policy] = await tx.select().from(budgetPolicies).where(and( + eq(budgetPolicies.id, scope.enforcement.policyId), eq(budgetPolicies.companyId, scope.companyId), + )); + if (!policy || policy.scopeType !== scope.scopeType || policy.scopeId !== scope.scopeId + || policy.enforcementVersion !== scope.enforcement.version || !(await policyBlocks(tx, policy))) return null; + } + return work(tx); + }); +} + /** Cancellation is an at-least-once external effect. A failed delivery never * rolls back committed spend, and its version remains pending for recovery. */ export async function deliverBudgetEnforcement(db: Db, hooks: BudgetServiceHooks, companyId?: string) { @@ -952,7 +970,7 @@ export async function deliverBudgetEnforcement(db: Db, hooks: BudgetServiceHooks const cancellation = await withAccountingTransaction(db, policy.companyId, async (tx) => { const [current] = await tx.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.id)); if (!current || current.enforcementVersion !== policy.enforcementVersion || !(await policyBlocks(tx, current))) return null; - return { companyId: policy.companyId, scopeType: policy.scopeType as BudgetScopeType, scopeId: policy.scopeId, createdBefore: new Date() }; + return { companyId: policy.companyId, scopeType: policy.scopeType as BudgetScopeType, scopeId: policy.scopeId, createdBefore: new Date(), enforcement: { policyId: policy.id, version: policy.enforcementVersion } }; }); if (cancellation) await hooks.cancelWorkForScope(cancellation); await db.update(budgetPolicies).set({ enforcementDeliveredVersion: policy.enforcementVersion }) diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 2097195e599..a590b4ef3b0 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -340,7 +340,7 @@ import { import { reportRunFailure } from "./run-failure-report.js"; import { collectRunFailureSecretValues, type RunFailureReportOptions } from "./run-failure-diagnostics.js"; import { companySkillService } from "./company-skills.js"; -import { budgetService, type BudgetEnforcementScope } from "./budgets.js"; +import { budgetService, withCurrentBudgetEnforcement, type BudgetEnforcementScope } from "./budgets.js"; import { secretService, type MissingRuntimeBinding } from "./secrets.js"; import { resolveDefaultAgentWorkspaceDir, @@ -22688,7 +22688,7 @@ export function heartbeatService( }, ); } - const dispatchResolvedInteractionContinuationWithAtomicGate = async ( + const dispatchResolvedInteractionContinuationAfterAdmission = async ( dispatch: (markDispatchStarted: () => void) => Promise, ): Promise< { dispatched: true; resultPromise: Promise } | { dispatched: false } @@ -22754,6 +22754,30 @@ export function heartbeatService( } return { dispatched: false }; }; + const dispatchResolvedInteractionContinuationWithAtomicGate = async ( + dispatch: (markDispatchStarted: () => void) => Promise, + ) => { + // Admission can wait for accounting locks. Finish that wait before the + // final ownership gate, whose callback must enter the adapter directly. + const reservation = await reserveRunBudget(db, run.companyId, run.id, + readNonEmptyString(runLedgerScope.projectId), runLedgerScope, runOptions.nativeLeaseOwner); + let entered = false; + try { + return await dispatchResolvedInteractionContinuationAfterAdmission((markDispatchStarted) => { + entered = true; + return dispatch(markDispatchStarted); + }); + } finally { + if (!entered && !reservation.reused) { + await db.update(heartbeatRuns).set({ costAccountingPending: true, + usageJson: sql`coalesce(${heartbeatRuns.usageJson}, '{}'::jsonb) || '{"accountingProviderWorkStarted":false}'::jsonb`, + }).where(and(eq(heartbeatRuns.id, run.id), isNull(heartbeatRuns.costAccountedAt))); + await accountRunCost(db, run.id, budgetHooks).catch((err) => { + logger.warn({ err, runId: run.id }, "Undispatched reservation release remains pending"); + }); + } + } + }; if (!executionTarget || executionTarget.kind === "local") { try { runScratch = await prepareHeartbeatRunScratch({ @@ -24656,8 +24680,7 @@ export function heartbeatService( ); const guardedDispatch = await dispatchResolvedInteractionContinuationWithAtomicGate( - async (markDispatchStarted) => { - await reserveRunBudget(db, run.companyId, run.id, readNonEmptyString(runLedgerScope.projectId), runLedgerScope, runOptions.nativeLeaseOwner); + (markDispatchStarted) => { return executePaperclipNativeSession({ db, execution: nativeExecution, @@ -24884,8 +24907,7 @@ export function heartbeatService( } const guardedDispatch = await dispatchResolvedInteractionContinuationWithAtomicGate( - async (markDispatchStarted) => { - await reserveRunBudget(db, run.companyId, run.id, readNonEmptyString(runLedgerScope.projectId), runLedgerScope); + (markDispatchStarted) => { legacyAdapterEntered = true; return adapter.execute({ getFreshSessionHandoff, @@ -29331,6 +29353,7 @@ export function heartbeatService( async function listProjectScopedWakeupIds( companyId: string, projectId: string, + database: Db = db, ) { const wakeIssueId = sql< string | null @@ -29339,7 +29362,7 @@ export function heartbeatService( string | null >`coalesce(${agentWakeupRequests.payload} ->> 'projectId', ${issues.projectId}::text)`; - const rows = await db + const rows = await database .selectDistinctOn([agentWakeupRequests.id], { id: agentWakeupRequests.id, }) @@ -29369,68 +29392,72 @@ export function heartbeatService( async function cancelPendingWakeupsForBudgetScope( scope: BudgetEnforcementScope, ) { - const now = new Date(); - let wakeupIds: string[] = []; + return withCurrentBudgetEnforcement(db, scope, async (tx) => { + const now = new Date(); + let wakeupIds: string[] = []; - if (scope.scopeType === "company") { - wakeupIds = await db - .select({ id: agentWakeupRequests.id }) - .from(agentWakeupRequests) - .where( - and( - eq(agentWakeupRequests.companyId, scope.companyId), - inArray(agentWakeupRequests.status, [ - "queued", - "deferred_issue_execution", - ]), - sql`${agentWakeupRequests.runId} is null`, - ), - ) - .then((rows) => rows.map((row) => row.id)); - } else if (scope.scopeType === "agent") { - wakeupIds = await db - .select({ id: agentWakeupRequests.id }) - .from(agentWakeupRequests) - .where( - and( - eq(agentWakeupRequests.companyId, scope.companyId), - eq(agentWakeupRequests.agentId, scope.scopeId), - inArray(agentWakeupRequests.status, [ - "queued", - "deferred_issue_execution", - ]), - sql`${agentWakeupRequests.runId} is null`, - ), - ) - .then((rows) => rows.map((row) => row.id)); - } else { - wakeupIds = await listProjectScopedWakeupIds( - scope.companyId, - scope.scopeId, - ); - } + if (scope.scopeType === "company") { + wakeupIds = await tx + .select({ id: agentWakeupRequests.id }) + .from(agentWakeupRequests) + .where( + and( + eq(agentWakeupRequests.companyId, scope.companyId), + inArray(agentWakeupRequests.status, [ + "queued", + "deferred_issue_execution", + ]), + sql`${agentWakeupRequests.runId} is null`, + ), + ) + .then((rows) => rows.map((row) => row.id)); + } else if (scope.scopeType === "agent") { + wakeupIds = await tx + .select({ id: agentWakeupRequests.id }) + .from(agentWakeupRequests) + .where( + and( + eq(agentWakeupRequests.companyId, scope.companyId), + eq(agentWakeupRequests.agentId, scope.scopeId), + inArray(agentWakeupRequests.status, [ + "queued", + "deferred_issue_execution", + ]), + sql`${agentWakeupRequests.runId} is null`, + ), + ) + .then((rows) => rows.map((row) => row.id)); + } else { + wakeupIds = await listProjectScopedWakeupIds( + scope.companyId, + scope.scopeId, + tx, + ); + } - if (wakeupIds.length === 0) return 0; + if (wakeupIds.length === 0) return 0; - await db - .update(agentWakeupRequests) - .set({ - status: "cancelled", - finishedAt: now, - error: "Cancelled due to budget pause", - updatedAt: now, - }) - .where(and( - inArray(agentWakeupRequests.id, wakeupIds), - inArray(agentWakeupRequests.status, ["queued", "deferred_issue_execution"]), - isNull(agentWakeupRequests.runId), - scope.createdBefore ? lt(agentWakeupRequests.createdAt, scope.createdBefore) : undefined, - )); + await tx + .update(agentWakeupRequests) + .set({ + status: "cancelled", + finishedAt: now, + error: "Cancelled due to budget pause", + updatedAt: now, + }) + .where(and( + inArray(agentWakeupRequests.id, wakeupIds), + inArray(agentWakeupRequests.status, ["queued", "deferred_issue_execution"]), + isNull(agentWakeupRequests.runId), + scope.createdBefore ? lt(agentWakeupRequests.createdAt, scope.createdBefore) : undefined, + )); - return wakeupIds.length; + return wakeupIds.length; + }); } type CancelRunOptions = { + budgetEnforcement?: BudgetEnforcementScope; errorCode?: string; resultJson?: Record; eventMessage?: string; @@ -29505,25 +29532,40 @@ export function heartbeatService( // Established legacy processes must still be stopped if the database is // unavailable. Only native or not-yet-dispatched preparation needs this // additional durable fence before its existing cancellation path. - if (run.runtimeMode === "native" || (!run.runtimeModeResolvedAt && !running && !control)) { - const [fenced] = await db.update(heartbeatRuns).set({ - // Record handoff intent before native cancellation can finalize and release - // the run. Only its audited stop acknowledgement suppresses recovery. - resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) || - ${JSON.stringify(options.errorCode === "issue_reassigned" && options.resultJson?.reassignmentStopConfirmed === true - ? { reassignmentStopRequested: true } : {})}::jsonb || - ${JSON.stringify({ cancellation })}::jsonb || - jsonb_build_object('startupCancellation', jsonb_build_object( - 'requestedAt', ${new Date().toISOString()}::text, - 'beforeNativeSelection', ${heartbeatRuns.runtimeMode} = 'legacy' - and ${heartbeatRuns.runtimeModeResolvedAt} is null - and ${heartbeatRuns.executionStage} = 'preparing' - and coalesce(${heartbeatRuns.runnerProfileJson}->'adapterDispatch'->>'adapterType' = 'paperclip_runner', false) - ))`, - }).where(and(eq(heartbeatRuns.id, runId), inArray(heartbeatRuns.status, - pendingNativeRetry ? [...CANCELLABLE_HEARTBEAT_RUN_STATUSES, "failed"] : [...CANCELLABLE_HEARTBEAT_RUN_STATUSES], - ))).returning(); - if (!fenced) return getRun(runId); + if (options.budgetEnforcement || run.runtimeMode === "native" || (!run.runtimeModeResolvedAt && !running && !control)) { + const fence = async (tx: Db) => { + const [fenced] = await tx.update(heartbeatRuns).set({ + // Record handoff intent before native cancellation can finalize and release + // the run. Only its audited stop acknowledgement suppresses recovery. + resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) || + ${JSON.stringify(options.errorCode === "issue_reassigned" && options.resultJson?.reassignmentStopConfirmed === true + ? { reassignmentStopRequested: true } : {})}::jsonb || + ${JSON.stringify({ cancellation })}::jsonb || + jsonb_build_object('startupCancellation', jsonb_build_object( + 'requestedAt', ${new Date().toISOString()}::text, + 'beforeNativeSelection', ${heartbeatRuns.runtimeMode} = 'legacy' + and ${heartbeatRuns.runtimeModeResolvedAt} is null + and ${heartbeatRuns.executionStage} = 'preparing' + and coalesce(${heartbeatRuns.runnerProfileJson}->'adapterDispatch'->>'adapterType' = 'paperclip_runner', false) + ))`, + }).where(and(eq(heartbeatRuns.id, runId), inArray(heartbeatRuns.status, + pendingNativeRetry ? [...CANCELLABLE_HEARTBEAT_RUN_STATUSES, "failed"] : [...CANCELLABLE_HEARTBEAT_RUN_STATUSES], + ))).returning(); + return fenced ?? null; + }; + let fenced: typeof run | null; + try { + fenced = options.budgetEnforcement + ? await withCurrentBudgetEnforcement(db, options.budgetEnforcement, fence) + : await fence(db); + } catch (error) { + stopOwnership?.release(); + throw error; + } + if (!fenced) { + stopOwnership?.release(); + return getRun(runId); + } run = fenced; } const resultJson = { ...(agent @@ -29903,7 +29945,7 @@ export function heartbeatService( scope.createdBefore ? lt(heartbeatRuns.createdAt, scope.createdBefore) : undefined, inArray(heartbeatRuns.status, [...CANCELLABLE_HEARTBEAT_RUN_STATUSES]), )).then((rows) => rows.map((row) => row.id)); - for (const runId of runIds) await cancelRunInternal(runId, "Cancelled due to budget pause"); + for (const runId of runIds) await cancelRunInternal(runId, "Cancelled due to budget pause", { budgetEnforcement: scope }); await cancelPendingWakeupsForBudgetScope(scope); } diff --git a/server/src/services/quota-credential-recovery.ts b/server/src/services/quota-credential-recovery.ts new file mode 100644 index 00000000000..2eedf691b50 --- /dev/null +++ b/server/src/services/quota-credential-recovery.ts @@ -0,0 +1,59 @@ +import { createHash, randomUUID } from "node:crypto"; +import { promises as fs } from "node:fs"; +import path from "node:path"; +import { z } from "zod"; +import { resolvePaperclipInstanceRoot } from "../home-paths.js"; +import { localEncryptedProvider } from "../secrets/local-encrypted-provider.js"; + +const recoverySchema = z.object({ + companyId: z.string().uuid(), grantId: z.string().uuid(), connectionId: z.string().uuid(), secretId: z.string().uuid(), + baseHash: z.string().length(64), value: z.string().min(1), +}); +type Recovery = z.infer; +export const quotaCredentialHash = (value: string) => createHash("sha256").update(value).digest("hex"); +export const quotaCredentialRecoveryPath = (companyId: string, grantId: string) => + path.join(resolvePaperclipInstanceRoot(), "quota-credential-recovery", companyId, `${grantId}.json`); + +async function syncDirectory(directory: string) { + if (process.platform === "win32") return; + const handle = await fs.open(directory, "r"); + try { await handle.sync(); } finally { await handle.close(); } +} + +/** Called only while holding the company credential mutation lock. The local + * encrypted provider uses the instance secrets master key, including when the + * destination vault is remote. No plaintext token enters this journal. */ +export function quotaCredentialRecovery(companyId: string, grantId: string) { + const file = quotaCredentialRecoveryPath(companyId, grantId); + const directory = path.dirname(file); + async function writeEncrypted(target: string, value: string) { + const encrypted = await localEncryptedProvider.createVersion({ value }); + await fs.mkdir(directory, { recursive: true, mode: 0o700 }); + const temporary = `${target}.${randomUUID()}.tmp`; + const handle = await fs.open(temporary, "wx", 0o600); + try { + try { await handle.writeFile(JSON.stringify(encrypted.material)); await handle.sync(); } + finally { await handle.close(); } + await fs.rename(temporary, target); + await syncDirectory(directory); + } catch (error) { await fs.rm(temporary, { force: true }); throw error; } + } + return { + // Verify encryption, storage, and fsync before exchanging a single-use token. + prepare: async () => { + const probe = `${file}.${randomUUID()}.probe`; + try { await writeEncrypted(probe, "credential-recovery-probe"); } + finally { await fs.rm(probe, { force: true }); } + }, + read: async (): Promise => { + let material: string; + try { material = await fs.readFile(file, "utf8"); } + catch (error) { if ((error as NodeJS.ErrnoException).code === "ENOENT") return null; throw error; } + const pending = recoverySchema.parse(JSON.parse(await localEncryptedProvider.resolveVersion({ material: JSON.parse(material), externalRef: null }))); + if (pending.companyId !== companyId || pending.grantId !== grantId) throw new Error("credentials_unavailable"); + return pending; + }, + write: (input: Recovery) => writeEncrypted(file, JSON.stringify(recoverySchema.parse(input))), + clear: async () => { await fs.rm(file, { force: true }); }, + }; +} diff --git a/server/src/services/run-cost-accounting.ts b/server/src/services/run-cost-accounting.ts index b1126462a8b..5f05cfa4437 100644 --- a/server/src/services/run-cost-accounting.ts +++ b/server/src/services/run-cost-accounting.ts @@ -1,6 +1,6 @@ import { usdToUnits, unitsToCents } from "@paperclipai/shared"; import { and, asc, eq, inArray, sql } from "drizzle-orm"; -import { accountingRuntimeBaselines, budgetReservations, agentRuntimeState, agents, costEvents, heartbeatRuns, issues, projects, type Db } from "@paperclipai/db"; +import { accountingRuntimeBaselines, budgetReservations, agentRuntimeState, agents, costEvents, heartbeatRuns, nativeRunFinalizations, issues, projects, type Db } from "@paperclipai/db"; import { withAccountingTransaction } from "./accounting-transaction.js"; import { createCostEventInTransaction } from "./costs.js"; import { budgetService, deliverBudgetEnforcement, type BudgetServiceHooks } from "./budgets.js"; @@ -25,6 +25,13 @@ export async function accountRunCost(db: Db, runId: string, hooks: BudgetService } const preProviderFailure = object(object(run.resultJson).executionRecovery).providerWorkStarted === false || object(run.usageJson).accountingProviderWorkStarted === false; + // A failed native turn is not necessarily the end of this heartbeat: the + // coordinator can resume the same run. Only its closed provider boundary + // authorizes acknowledgement; finalization may still run after a result. + if (run.runtimeMode === "native" && !preProviderFailure) { + const [coordinator] = await tx.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, run.id)); + if (coordinator && !coordinator.resultId && coordinator.phase !== "terminal_failure") return false; + } // A stop can mark the run terminal before the provider's shutdown receipt // arrives. Keep the debt pending instead of acknowledging an invented zero // and permanently discarding that late receipt. @@ -118,6 +125,11 @@ export async function accountRunCost(db: Db, runId: string, hooks: BudgetService export async function reconcileRunCosts(db: Db, hooks: BudgetServiceHooks = {}) { const pending = await db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(and( eq(heartbeatRuns.costAccountingPending, true), inArray(heartbeatRuns.status, terminalStatuses), + sql`(${heartbeatRuns.resultJson}->'executionRecovery'->>'providerWorkStarted' = 'false' + or ${heartbeatRuns.usageJson}->>'accountingProviderWorkStarted' = 'false' + or not exists (select 1 from ${nativeRunFinalizations} where ${nativeRunFinalizations.runId} = ${heartbeatRuns.id} + and ${heartbeatRuns.runtimeMode} = 'native' and ${nativeRunFinalizations.resultId} is null + and ${nativeRunFinalizations.phase} <> 'terminal_failure'))`, sql`(${heartbeatRuns.usageJson}->>'accountingReceiptReady' is distinct from 'false' or ${heartbeatRuns.resultJson}->'executionRecovery'->>'providerWorkStarted' = 'false' or ${heartbeatRuns.usageJson}->>'accountingProviderWorkStarted' = 'false')`, )).orderBy(asc(heartbeatRuns.updatedAt), asc(heartbeatRuns.id)).limit(100); let accounted = 0; diff --git a/server/src/services/usage-receipts.ts b/server/src/services/usage-receipts.ts index 0421712dda4..26f4547e25c 100644 --- a/server/src/services/usage-receipts.ts +++ b/server/src/services/usage-receipts.ts @@ -121,12 +121,22 @@ export async function createRunUsageRecorder(db: Db, input: { companyId: string; const probe = await fs.open(path.join(directory, `${sourceId}.probe`), "wx", 0o600); try { await probe.sync(); } finally { await probe.close(); await fs.rm(path.join(directory, `${sourceId}.probe`)); } await syncDirectory(directory); - const bound = await db.update(heartbeatRuns).set({ usageJson: sql`coalesce(${heartbeatRuns.usageJson}, '{}'::jsonb) || ${JSON.stringify({ accountingReceiptSourceId: sourceId, accountingReceiptSequence: 0 })}::jsonb` }) - .where(and(eq(heartbeatRuns.id, input.runId), eq(heartbeatRuns.companyId, input.companyId), isNull(heartbeatRuns.costAccountedAt))).returning({ id: heartbeatRuns.id }); + const bound = await db.update(heartbeatRuns).set({ usageJson: sql`coalesce(${heartbeatRuns.usageJson}, '{}'::jsonb) || ${JSON.stringify({ accountingReceiptSourceId: sourceId, accountingReceiptSequence: 0, accountingReceiptReady: false })}::jsonb` }) + .where(and(eq(heartbeatRuns.id, input.runId), eq(heartbeatRuns.companyId, input.companyId), isNull(heartbeatRuns.costAccountedAt))).returning({ id: heartbeatRuns.id, runtimeMode: heartbeatRuns.runtimeMode, usageJson: heartbeatRuns.usageJson }); if (bound.length !== 1) throw conflict("Run cannot accept a new usage recorder"); let sequence = 0; let currentAttempt = "default"; const attempts = new Map(); + // Native run deltas retain their baseline across same-run recovery. Restore + // the previous snapshot as the same attempt, never add it a second time. + // An empty final result must not erase the last durable observation. + const priorReceiptId = bound[0].usageJson?.accountingReceiptId; + if (bound[0].runtimeMode === "native" && typeof priorReceiptId === "string") { + const [previous] = await db.select().from(runUsageReceipts).where(and( + eq(runUsageReceipts.id, priorReceiptId), eq(runUsageReceipts.runId, input.runId), eq(runUsageReceipts.companyId, input.companyId), + )); + if (previous) attempts.set("default", envelopeSchema.parse(previous.receiptJson).receipt as AdapterUsageCheckpoint); + } let chain = Promise.resolve(); let captureFailed = false; async function capture(raw: AdapterUsageCheckpoint) {