diff --git a/cli/src/__tests__/worktree.test.ts b/cli/src/__tests__/worktree.test.ts index 0fb6b2aee32..a37449a774c 100644 --- a/cli/src/__tests__/worktree.test.ts +++ b/cli/src/__tests__/worktree.test.ts @@ -157,12 +157,12 @@ async function seedValidWorktreeSource( userId, role: "instance_admin", }); - await db.insert(companies).values({ - id: companyId, - name: "Seed Source", - issuePrefix: "SEED", - requireBoardApprovalForNewAgents: false, - }); + // Seed only columns present in the historical source schema. The current + // model includes accounting columns that the worktree migration adds later. + await db.$client` + insert into companies (id, name, issue_prefix, require_board_approval_for_new_agents) + values (${companyId}, 'Seed Source', 'SEED', false) + `; await db.insert(companyMemberships).values({ companyId, principalType: "user", diff --git a/doc/COST-ACCOUNTING.md b/doc/COST-ACCOUNTING.md index 085824748c9..172bfd9abde 100644 --- a/doc/COST-ACCOUNTING.md +++ b/doc/COST-ACCOUNTING.md @@ -153,3 +153,5 @@ Company quota requests resolve the caller's accessible subscription AI connectio Responses carry an opaque account identity and successful capture timestamp. The 30-second cache is scoped to company, user, connection, grant and credential revision; authorization and revision metadata are checked before reading cached results. Transient failures retain the last successful UI observation with its original timestamp. Rotation, revocation and explicit authentication rejection clear obsolete windows. Public messages never contain provider subprocess commands or raw diagnostics. Both Costs page variants retain loaded reports and expanded rows during background refresh failures, with a small stale-data notice. Initial failures remain visible. Company/date changes isolate cached reports. The run-count label is `runs: 0 api ยท 11 sub`: distinct recorded runs, not tasks or individual model requests. Input totals include cache reads, with the cached portion shown explicitly. + +Receipt identity and supplemental charges: `heartbeat:` idempotency keys are reserved for internally generated run receipts. Additional API-reported charges may link to the same run with their own keys. They contribute to monthly and acknowledged-run lifetime totals, while receipt-integrity checks compare only the original provider receipt. Recovery isolates each budget scope so a deleted target or failed scope does not prevent recovery elsewhere. diff --git a/server/src/__tests__/cost-accounting-operators.test.ts b/server/src/__tests__/cost-accounting-operators.test.ts index b3c32714706..e47565f1894 100644 --- a/server/src/__tests__/cost-accounting-operators.test.ts +++ b/server/src/__tests__/cost-accounting-operators.test.ts @@ -86,6 +86,40 @@ const support = await getEmbeddedPostgresTestSupport(); expect((await service.inspect(f.company.id)).findings).toEqual([]); }); + it.each(["before", "after", "concurrent"])("accounts separate run charges %s the provider receipt without false integrity findings", async (timing) => { + const f = await fixture(); const run = await runFor(f, "succeeded"); + await db.update(heartbeatRuns).set({ costAccountingPending: true, usageJson: { costUsdExact: "0.01", inputTokens: 7, cachedInputTokens: 3, outputTokens: 11 } }).where(eq(heartbeatRuns.id, run.id)); + const extra = { agentId: f.agent.id, heartbeatRunId: run.id, provider: "test", model: "tool", costCents: "0.1234567", inputTokens: 2, cachedInputTokens: 4, outputTokens: 6, occurredAt: new Date(), idempotencyKey: "extra-charge" }; + const costs = costService(db); + if (timing === "before") await costs.createEvent(f.company.id, extra); + if (timing === "concurrent") await Promise.all([costs.createEvent(f.company.id, extra), accountRunCost(db, run.id)]); + else await accountRunCost(db, run.id); + const event = await costs.createEvent(f.company.id, extra); + // Replaying either source cannot count it twice. + await costs.createEvent(f.company.id, extra); + expect(await accountRunCost(db, run.id)).toBe(false); + const [runtime] = await db.select().from(agentRuntimeState).where(eq(agentRuntimeState.agentId, f.agent.id)); + expect(runtime).toMatchObject({ totalCostCents: 1.1234567, totalInputTokens: 9, totalCachedInputTokens: 7, totalOutputTokens: 17 }); + const integrity = accountingIntegrityService(db); + expect((await integrity.inspect(f.company.id)).findings).toEqual([]); + await billingReconciliationService(db).adjust(f.company.id, event.id, { idempotencyKey: "extra-correction", expectedCents: "0.1234567", correctedCents: "0.2234567", reason: "Provider correction", pricing: { source: "provider_invoice", evidence: "external charge" } }, "board"); + expect((await costs.summary(f.company.id)).spendCentsExact).toBe("1.2234567"); + expect((await integrity.inspect(f.company.id)).findings).toEqual([]); + await db.update(agentRuntimeState).set({ totalCostCents: 999 }).where(eq(agentRuntimeState.agentId, f.agent.id)); + const review = await integrity.inspect(f.company.id); + expect(review.findings.map(f => f.kind)).toEqual(["runtime_projection"]); + expect((await integrity.repair(f.company.id, review.fingerprint, "Fix drift", "board")).findings).toEqual([]); + // An extra charge cannot conceal a missing original provider receipt. + await db.delete(costEvents).where(and(eq(costEvents.heartbeatRunId, run.id), eq(costEvents.idempotencyKey, `heartbeat:${run.id}:final`))); + expect((await integrity.inspect(f.company.id)).findings.some(f => f.kind === "missing_receipt")).toBe(true); + }); + + it("rejects public charges that impersonate the provider receipt namespace", async () => { + const f = await fixture(); + await expect(costService(db).createEvent(f.company.id, { agentId: f.agent.id, provider: "test", model: "test", costCents: 1, occurredAt: new Date(), idempotencyKey: "heartbeat:fake:final" })).rejects.toThrow("reserved"); + expect(await db.select().from(costEvents).where(eq(costEvents.companyId, f.company.id))).toEqual([]); + }); + it("imports invoices idempotently and appends reviewed corrections without losing the provider receipt", async () => { const f = await fixture(); const originalPricing = { source: "rate_card" as const, version: "2026-09-01", inputCentsPerMillion: "12.3456789" }; diff --git a/server/src/__tests__/cost-accounting-reliability.test.ts b/server/src/__tests__/cost-accounting-reliability.test.ts index e3a1153d422..022de3e1948 100644 --- a/server/src/__tests__/cost-accounting-reliability.test.ts +++ b/server/src/__tests__/cost-accounting-reliability.test.ts @@ -125,6 +125,35 @@ databaseDescribe("cost accounting reliability (PostgreSQL)", () => { expect(await budgetService(db).getInvocationBlock(f.company.id, f.agent.id)).toMatchObject({ scopeType: "agent" }); }); + it("recovers live scopes despite deleted agent and project policies and an isolated transaction failure", async () => { + const f = await fixture(); const other = await fixture(); const budgets = budgetService(db); + const [removedAgent] = await db.insert(agents).values({ companyId: f.company.id, name: "Removed", adapterType: "process" }).returning(); + const [removedProject] = await db.insert(projects).values({ companyId: f.company.id, name: "Removed" }).returning(); + await budgets.upsertPolicy(f.company.id, { scopeType: "agent", scopeId: removedAgent.id, amount: 100 }, "board"); + await budgets.upsertPolicy(f.company.id, { scopeType: "project", scopeId: removedProject.id, amount: 100 }, "board"); + await db.delete(agents).where(eq(agents.id, removedAgent.id)); + await db.delete(projects).where(eq(projects.id, removedProject.id)); + for (const live of [f, other]) { + await budgets.upsertPolicy(live.company.id, { scopeType: "agent", scopeId: live.agent.id, amount: 100 }, "board"); + await db.update(agents).set({ status: "paused", pauseReason: "budget" }).where(eq(agents.id, live.agent.id)); + } + // Restrict the scanned scope list so the injected first transaction fails + // on the stale scope, ahead of both live policies in this real database. + const query = vi.spyOn(db, "selectDistinct").mockReturnValueOnce({ from: () => ({ orderBy: async () => [ + { companyId: f.company.id, scopeType: "project", scopeId: removedProject.id }, + { companyId: f.company.id, scopeType: "agent", scopeId: removedAgent.id }, + ...[f, other].map(live => ({ companyId: live.company.id, scopeType: "agent", scopeId: live.agent.id })), + ] }) } as never); + const transaction = vi.spyOn(db, "transaction").mockRejectedValueOnce(new Error("Injected scope recovery failure")); + try { await budgets.reconcilePolicies(); } finally { query.mockRestore(); transaction.mockRestore(); } + for (const live of [f, other]) { + const [agent] = await db.select().from(agents).where(eq(agents.id, live.agent.id)); + expect(agent).toMatchObject({ status: "idle", pauseReason: null }); + } + // A subsequent sweep uses actual policy discovery and skips both deleted scopes. + await expect(budgets.reconcilePolicies()).resolves.toBeUndefined(); + }); + it("recovers terminal run accounting exactly once across simultaneous live and recovery workers", async () => { const f = await fixture(); await db.update(heartbeatRuns).set({ status: "timed_out", finishedAt: new Date(), costAccountingPending: true, diff --git a/server/src/services/accounting-integrity.ts b/server/src/services/accounting-integrity.ts index c5b1709e71e..2e03fdd8e53 100644 --- a/server/src/services/accounting-integrity.ts +++ b/server/src/services/accounting-integrity.ts @@ -38,7 +38,7 @@ async function inspect(tx: Db, companyId: string): Promise } const incomplete = await tx.execute<{ id: string; acknowledged: boolean; events: number }>(sql` select r.id, r.cost_accounted_at is not null as acknowledged, count(e.id)::int as events from heartbeat_runs r - left join cost_events e on e.heartbeat_run_id = r.id and e.company_id = r.company_id + left join cost_events e on e.heartbeat_run_id = r.id and e.company_id = r.company_id and e.idempotency_key like ('heartbeat:' || r.id::text || ':%') where r.company_id = ${companyId} and r.status in ('succeeded','failed','timed_out','cancelled') and (r.cost_accounting_pending or r.accounting_projection_version = 'v2') and r.result_json->'executionRecovery'->>'providerWorkStarted' is distinct from 'false' @@ -52,7 +52,7 @@ async function inspect(tx: Db, companyId: string): Promise const receipts = await tx.execute<{ id: string; usage: Record; actual: Record }>(sql` select r.id, r.usage_json as usage, jsonb_build_object('cents', sum(coalesce(e.reported_cost_cents,e.cost_cents))::text, 'input', sum(e.input_tokens)::text, 'cached', sum(e.cached_input_tokens)::text, 'output', sum(e.output_tokens)::text) as actual - from heartbeat_runs r join cost_events e on e.heartbeat_run_id = r.id and e.company_id = r.company_id + from heartbeat_runs r join cost_events e on e.heartbeat_run_id = r.id and e.company_id = r.company_id and e.idempotency_key like ('heartbeat:' || r.id::text || ':%') where r.company_id = ${companyId} and r.accounting_projection_version = 'v2' and r.cost_accounted_at is not null group by r.id order by r.id`); for (const row of receipts) { const usage = row.usage ?? {}; diff --git a/server/src/services/budgets.ts b/server/src/services/budgets.ts index 4767a99b519..b26ffe092e2 100644 --- a/server/src/services/budgets.ts +++ b/server/src/services/budgets.ts @@ -24,7 +24,7 @@ import type { BudgetThresholdType, BudgetWindowKind, } from "@paperclipai/shared"; -import { notFound, unprocessable } from "../errors.js"; +import { HttpError, notFound, unprocessable } from "../errors.js"; import { logger } from "../middleware/logger.js"; import { withAccountingTransaction } from "./accounting-transaction.js"; import { logActivity, type LogActivityInput, type ActivityPublication } from "./activity-log.js"; @@ -586,16 +586,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica } return { - reconcileCompanyScopes: async (companyId: string) => { - const policies = await listPolicyRows(companyId); - const seen = new Set(); - for (const policy of policies) { - const key = `${policy.scopeType}:${policy.scopeId}`; - if (seen.has(key)) continue; - seen.add(key); - await reconcileScope(companyId, policy.scopeType as BudgetScopeType, policy.scopeId); - } - }, + reconcileScope, listPolicies: async (companyId: string): Promise => { const rows = await listPolicyRows(companyId); return rows.map((row) => ({ @@ -982,8 +973,18 @@ export function budgetService(db: Db, hooks: BudgetServiceHooks = {}) { ...reads, deliverPendingEnforcement: (companyId?: string) => deliverBudgetEnforcement(db, hooks, companyId), reconcilePolicies: async () => { - const rows = await db.selectDistinct({ companyId: budgetPolicies.companyId }).from(budgetPolicies); - for (const row of rows) await mutate(row.companyId, (service) => service.reconcileCompanyScopes(row.companyId)); + const scopes = await db.selectDistinct({ companyId: budgetPolicies.companyId, scopeType: budgetPolicies.scopeType, scopeId: budgetPolicies.scopeId }) + .from(budgetPolicies).orderBy(budgetPolicies.companyId, budgetPolicies.scopeType, budgetPolicies.scopeId); + // Each scope commits independently: deleted targets and a failed scope + // must not roll back recovery or cancellation delivery for other work. + for (const scope of scopes) { + try { + await mutate(scope.companyId, (service) => service.reconcileScope(scope.companyId, scope.scopeType as BudgetScopeType, scope.scopeId)); + } catch (error) { + if (error instanceof HttpError && error.status === 404) continue; + logger.warn({ err: error, ...scope }, "Budget scope recovery pending; continuing remaining scopes"); + } + } }, upsertPolicy: (companyId: string, input: BudgetPolicyUpsertInput, actorUserId: string | null) => mutate(companyId, (service) => service.upsertPolicy(companyId, input, actorUserId)), diff --git a/server/src/services/costs.ts b/server/src/services/costs.ts index fe1a583900d..431dd5a136a 100644 --- a/server/src/services/costs.ts +++ b/server/src/services/costs.ts @@ -3,7 +3,7 @@ import { agentAvatarUrl, resolveAgentAppearance, createCostEventSchema, normaliz import { and, desc, eq, gte, inArray, isNotNull, isNull, lt, lte, sql } from "drizzle-orm"; import { alias } from "drizzle-orm/pg-core"; import type { Db } from "@paperclipai/db"; -import { activityLog, agents, companies, costEvents, heartbeatRuns, issues, projects, goals } from "@paperclipai/db"; +import { activityLog, agents, agentRuntimeState, companies, costEvents, heartbeatRuns, issues, projects, goals } from "@paperclipai/db"; import { notFound, unprocessable, conflict } from "../errors.js"; import { budgetService, budgetServiceInTransaction, type BudgetServiceHooks } from "./budgets.js"; import { logActivity, type LogActivityInput, type ActivityPublication } from "./activity-log.js"; @@ -119,6 +119,20 @@ export async function createCostEventInTransaction(db: Db, companyId: string, da } const [event] = await db.insert(costEvents).values({ ...values, costCents: sql`${values.costCents}::numeric`, reportedCostCents: values.costCents, id: data.id, companyId, receiptHash }).returning(); await updateMonthlySpendProjections(db, companyId, event.agentId, values.costCents, event.occurredAt); + // Separately reported charges linked to an already-accounted run contribute + // to lifetime totals too. Before acknowledgement, accountRunCost includes + // all of the run's events atomically when it initializes the projection. + if (event.heartbeatRunId) { + const [run] = await db.select({ acknowledged: heartbeatRuns.costAccountedAt, version: heartbeatRuns.accountingProjectionVersion }) + .from(heartbeatRuns).where(and(eq(heartbeatRuns.id, event.heartbeatRunId), eq(heartbeatRuns.companyId, companyId))); + if (run.acknowledged && run.version === "v2") await db.update(agentRuntimeState).set({ + totalCostCents: sql`${agentRuntimeState.totalCostCents} + ${values.costCents}::numeric`, + totalInputTokens: sql`${agentRuntimeState.totalInputTokens} + ${event.inputTokens}`, + totalCachedInputTokens: sql`${agentRuntimeState.totalCachedInputTokens} + ${event.cachedInputTokens}`, + totalOutputTokens: sql`${agentRuntimeState.totalOutputTokens} + ${event.outputTokens}`, + updatedAt: new Date(), + }).where(and(eq(agentRuntimeState.agentId, event.agentId), eq(agentRuntimeState.companyId, companyId))); + } await budgetServiceInTransaction(db, publications).evaluateCostEvent(event); await logActivity(db, { companyId, actorType: actor?.actorType ?? "system", actorId: actor?.actorId ?? "cost_accounting", agentId: actor?.agentId ?? event.agentId, @@ -132,6 +146,7 @@ export function costService(db: Db, budgetHooks: BudgetServiceHooks = {}) { const budgets = budgetService(db, budgetHooks); return { createEvent: async (companyId: string, data: Omit & { costCents: MoneyInput }, actor?: Pick) => { + if (data.idempotencyKey?.startsWith("heartbeat:")) throw unprocessable("The heartbeat idempotency namespace is reserved for run receipts"); const event = await withAccountingTransaction(db, companyId, (tx, publications) => createCostEventInTransaction(tx, companyId, data, publications, actor)); await budgets.deliverPendingEnforcement(companyId); return event; diff --git a/server/src/services/recovery/service.ts b/server/src/services/recovery/service.ts index 113ea5ce1cd..27bb40f1c24 100644 --- a/server/src/services/recovery/service.ts +++ b/server/src/services/recovery/service.ts @@ -3827,6 +3827,12 @@ export function recoveryService( current.companyId, current.id, ); + // Budget admission may wait behind another recovery worker. Read the + // execution paths after the persisted counter so a newly reserved + // successor cannot be mistaken for permission to schedule the next one. + const currentState = await collectDispositionRepairSourceState(db, { issue: current }); + if (currentState.hasActiveExecutionPath || currentState.hasDurableWaitingPath) return "skipped"; + if (!episode && currentState.fingerprint !== state.fingerprint) return "skipped"; const runAttempt = previousAttempt?.fingerprint === state.fingerprint ? previousAttempt.attempt diff --git a/server/src/services/run-cost-accounting.ts b/server/src/services/run-cost-accounting.ts index e4d643f92d5..84b7b2d2f05 100644 --- a/server/src/services/run-cost-accounting.ts +++ b/server/src/services/run-cost-accounting.ts @@ -1,6 +1,6 @@ -import { usdToUnits, unitsToCents, centsToUnits } from "@paperclipai/shared"; +import { usdToUnits, unitsToCents } from "@paperclipai/shared"; import { and, asc, eq, inArray, sql } from "drizzle-orm"; -import { accountingRuntimeBaselines, budgetReservations, agentRuntimeState, agents, heartbeatRuns, issues, projects, type Db } from "@paperclipai/db"; +import { accountingRuntimeBaselines, budgetReservations, agentRuntimeState, agents, costEvents, heartbeatRuns, 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"; @@ -69,7 +69,6 @@ export async function accountRunCost(db: Db, runId: string, hooks: BudgetService const validParts = parts.length > 0 && tokensMatch && costMatches && parts.every((part) => text(part.model) && amount(part.costUsd) !== null) && new Set(parts.map((part) => part.model)).size === parts.length; - const events = []; if (validParts) { let allocated = 0n; for (const [index, part] of parts.entries()) { @@ -79,19 +78,21 @@ export async function accountRunCost(db: Db, runId: string, hooks: BudgetService const partUnits = billingType === "subscription_included" ? 0n : index === parts.length - 1 ? remainder : reported < remainder ? reported : remainder; allocated += partUnits; const partCents = unitsToCents(partUnits); - events.push(await createCostEventInTransaction(tx, run.companyId, { + await createCostEventInTransaction(tx, run.companyId, { ...receipt, idempotencyKey: `heartbeat:${run.id}:model:${index}`, model: text(part.model)!, inputTokens: amount(partUsage.inputTokens) ?? 0, cachedInputTokens: amount(partUsage.cachedInputTokens) ?? 0, outputTokens: amount(partUsage.outputTokens) ?? 0, costCents: partCents, - }, publications)); + }, publications); } } else { - events.push(await createCostEventInTransaction(tx, run.companyId, receipt, publications)); + await createCostEventInTransaction(tx, run.companyId, receipt, publications); } - const event = events.reduce((total, part) => ({ - inputTokens: total.inputTokens + part.inputTokens, cachedInputTokens: total.cachedInputTokens + part.cachedInputTokens, - outputTokens: total.outputTokens + part.outputTokens, costCents: total.costCents + centsToUnits(part.costCentsExact), - }), { inputTokens: 0, cachedInputTokens: 0, outputTokens: 0, costCents: 0n }); + const [event] = await tx.select({ + inputTokens: sql`coalesce(sum(${costEvents.inputTokens}),0)::text`, + cachedInputTokens: sql`coalesce(sum(${costEvents.cachedInputTokens}),0)::text`, + outputTokens: sql`coalesce(sum(${costEvents.outputTokens}),0)::text`, + costCents: sql`coalesce(sum(${costEvents.costCents}),0)::text`, + }).from(costEvents).where(and(eq(costEvents.companyId, run.companyId), eq(costEvents.heartbeatRunId, run.id))); const [agent] = await tx.select({ adapterType: agents.adapterType }).from(agents).where(eq(agents.id, run.agentId)); await tx.insert(agentRuntimeState).values({ agentId: run.agentId, companyId: run.companyId, adapterType: agent.adapterType }).onConflictDoNothing(); await tx.execute(sql`insert into accounting_runtime_baselines (agent_id,company_id,cost_cents,input_tokens,cached_input_tokens,output_tokens) @@ -101,7 +102,7 @@ export async function accountRunCost(db: Db, runId: string, hooks: BudgetService totalInputTokens: sql`${agentRuntimeState.totalInputTokens} + ${event.inputTokens}`, totalCachedInputTokens: sql`${agentRuntimeState.totalCachedInputTokens} + ${event.cachedInputTokens}`, totalOutputTokens: sql`${agentRuntimeState.totalOutputTokens} + ${event.outputTokens}`, - totalCostCents: sql`${agentRuntimeState.totalCostCents} + ${unitsToCents(event.costCents)}::numeric`, + totalCostCents: sql`${agentRuntimeState.totalCostCents} + ${event.costCents}::numeric`, updatedAt: new Date(), }).where(eq(agentRuntimeState.agentId, run.agentId)); await tx.update(heartbeatRuns).set({ costAccountingPending: false, costAccountedAt: new Date(), accountingProjectionVersion: "v2", accountingLastError: null, accountingLastAttemptAt: new Date() }).where(eq(heartbeatRuns.id, run.id));