mirror of
https://github.com/paperclipai/paperclip.git
synced 2026-10-23 17:33:51 +02:00
fix: isolate budget recovery and account supplemental run charges
Co-Authored-By: Paperclip <noreply@paperclip.ing>
This commit is contained in:
1 parent
5997745d77
commit
3ac5772c66
9 files changed
+121
-33
No files matched your search
@@ -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",
|
||||
|
||||
@@ -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.
|
||||
@@ -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" };
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -38,7 +38,7 @@ async function inspect(tx: Db, companyId: string): Promise<AccountingInspection>
|
||||
}
|
||||
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<AccountingInspection>
|
||||
const receipts = await tx.execute<{ id: string; usage: Record<string, unknown>; actual: Record<string, string> }>(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 ?? {};
|
||||
|
||||
@@ -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<string>();
|
||||
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<BudgetPolicy[]> => {
|
||||
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)),
|
||||
|
||||
@@ -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<typeof costEvents.$inferInsert, "companyId" | "receiptHash" | "costCents"> & { costCents: MoneyInput }, actor?: Pick<LogActivityInput, "actorType" | "actorId" | "agentId">) => {
|
||||
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;
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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<string>`coalesce(sum(${costEvents.inputTokens}),0)::text`,
|
||||
cachedInputTokens: sql<string>`coalesce(sum(${costEvents.cachedInputTokens}),0)::text`,
|
||||
outputTokens: sql<string>`coalesce(sum(${costEvents.outputTokens}),0)::text`,
|
||||
costCents: sql<string>`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));
|
||||
|
||||
Reference in new issue
Block a user