perf(budgets): scope accounting and deliver enforcement asynchronously (#17046)

## Thinking Path

> - Paperclip manages AI agents and their paid work.
> - Budgets must prevent new work and stop reported overruns.
> - Cost settlement also ran serial lifecycle reconciliation for every
eligible company agent.
> - That made request latency grow with unrelated idle agents and
delayed results on stop delivery.
> - This change separates atomic admission and accounting from durable
asynchronous enforcement.
> - Scoped live usage detects overruns without writing provisional
ledger charges.
> - The result is constant query scope for ordinary settlement and
recoverable stop delivery.

## Linked Issues or Issue Description

**What happened?**

Settlement and budget checks entered company-wide lifecycle
reconciliation. A company with 1,000 unrelated idle agents performed
1,001 accounting transactions for one ordinary cost evaluation. User
requests could also wait for provider termination. Running usage
checkpoints did not enforce reported monetary exposure.

**Expected behavior**

Keep the final reservation gate authoritative. Commit charges, scope
blocks, and enforcement intent before returning. Deliver stop requests
after commit. Read only applicable policies and execution scopes for
ordinary charges. Stop long-running work when supported reported or
priced usage establishes an overrun.

**Steps to reproduce**

1. Create a disposable company with one company budget and idle agents.
2. Evaluate a charge below the budget with 0, 17, and 1,000 unrelated
idle agents.
3. Compare query counts, company accounting transactions, and duration.
4. Use the controlled partial-usage and delayed-stop fixtures to check
enforcement and request responsiveness.

**Paperclip version or commit**

The measured baseline is `6c9b87d0922f1ebfc43e8b0c36889ac5ba7b9837` from
`master`.

Related public work: Refs #16976, #2006, #3539, #13763, #14296, #16853,
#13075. These cover the affected service charge, cancellation, and
accounting paths. This change removes their shared reconciliation
overhead.

## What Changed

- Batch applicable policy exposure under the existing company accounting
lock. Preserve exact fractional money and zero-reservation dispatch
fences.
- Commit a durable final budget dispatch claim before provider I/O.
Retain a separate final issue/run ownership transaction through
synchronous handoff, without holding the company accounting lock or
awaiting provider work. Preserve native recovery lease checks and
duplicate handoff fencing. Browser handoff rechecks current policy; it
reuses a consumed rate slot only for the identical policy generation.
- Separate ledger spend, live observed spend, and effective held
capacity. Use the maximum of reservation and live observation for each
unfinished operation.
- Deliver budget cancellation through the existing coordinator.
Revalidate generations and persist queued/running stop fences before
external termination. Bound stop concurrency to four.
- Apply agent holds only to the affected agent. Project budgets leave
other projects running. Project real company transitions in bounded
lifecycle batches.
- Replace periodic broad reconciliation with durable pending work,
startup repair, receipt replay, and concrete retry/window deadlines.
Retain drain and idle-safety checks.
- Enforce accepted cumulative usage checkpoints without provisional
charges. Final accounting atomically replaces live exposure and settles
reservations.
- Add operator-gated budget spans, regression coverage, map updates, and
a real-PostgreSQL comparison harness. No schema or first-party telemetry
change.

## Verification

- Full workspace `pnpm -r typecheck` passes on the final ownership fix,
rebased onto `bbcf5dfe98`. Production `pnpm build` passes on final head
`4c5f4de8b1`. Map, module-boundary, token, and Storybook catalog checks
pass.
- Focused real-PostgreSQL coverage includes the 292-case rebased
accounting/dispatch run, 63 browser cases, 52 run-control/cancellation
cases, 42 policy-generation cases, 77 lifecycle/commentary cases, and
the exact Muse confirmation regression. Native executor unit coverage
passes all 768 cases. Final ownership coverage passes: 74
adapter/use-case cases, nine public regression cases, and the additional
deleted-project race. These exercise competing ownership, budget changes
and charges between phases, zero/nonzero duplicate claims, rollback,
recovery contention, read-only capacity rechecks, and controlled HTTP
startup after commit. Another 74 pure/startup/tracing cases pass. The
complete CI matrix and both 5/5 reviews passed on `1ac9b58`; its
external security scan identified a stale-owner handoff race. The final
fix restores synchronous issue/run ownership fencing after the
accounting transaction commits. Final head `8e84f55ce8` has fresh
Greptile 5/5 with zero new comments, Paperclip review 5/5, and a passing
security scan. The complete latest-head [CI
matrix](https://github.com/paperclipai/paperclip/actions/runs/39535189808)
passes, including budget, browser, Runner, build, and canary jobs. The
unchanged Runner close fixture passed on its sole same-head job retry.
- Broad local `pnpm test:run` invocations were stopped before required
source updates (rebase and review fixes). They are not claimed as full
passes. Final full-suite coverage comes from the complete latest-head CI
matrix; local focused coverage and the incomplete broad invocations are
reported separately.
- Reproduce performance with `pnpm --filter @paperclipai/server exec tsx
../scripts/benchmark-budget-enforcement.ts
--baseline-ref=6c9b87d0922f1ebfc43e8b0c36889ac5ba7b9837`; repeat with
`--reverse`. The harness uses disposable real PostgreSQL, 20 warm
samples per variant/count, and records raw timings, SQL/lock counts,
agent reads, unchanged agent state, and query plans. Final
production-source measurement revision:
`eff713a4f72c7f4b7a4002a701cb481d4bfccd95`; the extended harness is
committed in `4c5f4de8b1`.

| Unrelated idle agents | Baseline median / p95 | New median / p95 |
Baseline SQL / locks | New SQL / locks | New agent reads |
| --- | ---: | ---: | ---: | ---: | ---: |
| 0 | 1.80 / 2.33 ms | 2.99 / 3.88 ms | 5 / 1 | 7 / 1 | 0 |
| 17 | 15.53 / 19.55 ms | 2.24 / 2.60 ms | 90 / 18 | 7 / 1 | 0 |
| 1,000 | 978.35 / 1,790.43 ms | 2.16 / 2.52 ms | 5,005 / 1,001 | 7 / 1
| 0 |

- These measurements use local PostgreSQL 18.1 and Node 26.4.0 arm64.
They isolate the cost-evaluation core and complete admission path,
without hosted-region or provider latency. Prior-revision reversed-order
controls also preserved the constant query counts; the final
measurements above use the final production source.
- Complete reserve → final ownership gate → immediate offline provider
result is measured separately, with one responsible agent and one active
operation at each unrelated-agent count. Admission p95 baseline/new: 0
agents **11.54 / 29.77 ms**; 17 agents **9.31 / 27.62 ms**; 1,000 agents
**24.51 / 74.05 ms**. Baseline admission was already scoped; the
stronger final fence adds bounded overhead. New counts stay constant at
53 Drizzle statements, three transactions, two company-accounting locks,
and five agent-reading statements bound to the responsible agent. Zero
unrelated-agent reads, exactly one provider callback per sample, and no
leftover runs or holds. All admission scenarios stay below 250 ms. No
external provider calls are included.
- The partial-usage/delayed-termination fixture passed 20/20: receipt
persistence p95 13 ms, receipt return to observed durable Stop fence p95
28 ms, complete receipt-start-to-fence p95 50 ms, unrelated database
request p95 1 ms while termination was blocked. After releasing the
controlled termination fixture, acknowledgement p95 was 14 ms. These are
separately measured phases; fixture acknowledgement is not live-provider
termination latency.
- Models and provider deadlines are unchanged. No deployment, live
inference, or staging acceptance is claimed. A quiet 30-minute follow-up
waits for the merged revision on staging, then repeats classifier and
controlled budget-stop acceptance in staging, recording hosted
bookkeeping, queue delay, and actual provider acknowledgement
separately.
- Development failures retained: finalizer notification regressions
(fixed; 27/27 affected cases passed), invalid browser admission fixture
assumptions (fixed; browser coverage now passes 63/63), and first-head
lifecycle/dispatch CI regressions (reproduced and repaired). A
prior-head Cases E2E test hit its unchanged announcement-dismiss handler
race while opening a filter; the same journey passed on `1ac9b58`. The
external security review then found a stale-owner race between
accounting commit and provider entry; final verification covers a
competing run in that gap while retaining provider I/O outside the
accounting lock. Early PostgreSQL startup attempts hit local semaphore
pressure; serial reruns succeeded. The final-head Runner lane had one
failure among 1,823 tests: the unchanged `timeout-attachment` fixture
failed at its first process close after 6.28 seconds. The sole same-head
job retry passed all 1,819 active tests (four skipped), including that
fixture in 1.85 seconds. Three isolated local probes could not run
correctly because the Rust fake-provider fixture binary was absent;
cleanup masked the original startup failure. CI builds that fixture
through its normal lane. These probes are inconclusive and are not
counted as passes.

## Risks

- This changes accounting, lifecycle, and cancellation ordering. Durable
generations, source/sequence checks, ownership leases, and company
locking remain required for correctness.
- Provider stop acknowledgement can be slow or fail. Stop intent remains
pending until existing termination proof succeeds; cancelled row status
alone does not acknowledge delivery.
- Mid-run limits depend on provider usage reports or supported price
calculations. Providers that report only at completion cannot provide
exact mid-call detection.
- Signals are process-local wake hints for the supported
single-active-writer model. Durable state and startup recovery remain
authoritative.
- There are no migrations, new queue tables, or cached balances. Repeat
the classifier and controlled budget-stop acceptance after normal
rollout.

## Model Used

OpenAI Codex, based on GPT-6, with reasoning, tool use, code execution,
and parallel coding agents. The host does not expose the exact backend
model revision or context-window size, so those values are not asserted.

## Checklist

- [x] I have included a thinking path that traces from project context
to this change
- [x] I have specified the model used (with version and capability
details)
- [x] I have checked ROADMAP.md and confirmed this PR does not duplicate
planned core work
- [x] I have searched GitHub for duplicate or related PRs and linked
them above
- [x] I have either (a) linked existing issues with `Fixes: #` / `Closes
#` / `Refs #` OR (b) described the issue in-PR following the relevant
issue template
- [x] I have not referenced internal/instance-local Paperclip issues or
links (only public GitHub `#NNN` / `github.com/paperclipai/paperclip`
URLs)
- [x] My branch name describes the change (e.g. `docs/...`, `fix/...`)
and contains no internal Paperclip ticket id or instance-derived details
- [x] I have run tests locally and they pass
- [x] I have added or updated tests where applicable
- [x] I have updated relevant documentation to reflect my changes
- [x] I have considered and documented any risks above
- [x] All Paperclip CI gates are green
- [x] Greptile is 5/5 with no open P2s, recommendations, or follow-ups
- [x] I will address all Greptile and reviewer comments before
requesting merge

---------

Co-authored-by: Paperclip <noreply@paperclip.ing>
This commit is contained in:
DottaandPaperclip authored and GitHub committed 2026-10-21 07:23:00 -05:00
1 parent dff63551fe
commit 30b753f614
85 files changed
+5204 -1111

No files matched your search

+41 -1
View File
@@ -3084,7 +3084,7 @@ current task binding. They do not inventory peer workspaces. Existing
`isolatedRemoteSource` receipts remain restorable, but new local sources are not
replaced with remote clones because of historical file ownership.
The disposable performance harness is:
The disposable workspace-preparation performance harness is:
```sh
pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-workspace-preparation-history.ts --baseline-ref=df80616d98119cbf5abd40549794d9980a5fabf6
@@ -3093,6 +3093,46 @@ pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-workspace-prepar
It compares actual preparation guards with synthetic history. Local timings are
separate from Git materialization, environment acquisition, and provider startup.
### Budget accounting performance
Budget admission and settlement share the company accounting lock. They evaluate
only the company, responsible agent, and applicable project. Running usage is an
internal exposure projection, not a provisional ledger charge. Cancellation and
lifecycle effects run through the delivery coordinator after accounting commits.
The final provider gate commits a dispatch claim before entering provider code.
The reservation and native recovery lease generation fence duplicate handoffs;
neither provider startup nor termination holds the company accounting lock.
After accounting commits, a short ownership transaction rechecks the current
task/run fences and covers the synchronous adapter handoff. It commits without
awaiting provider bootstrap, results, or finalization.
Hosted browser operations retain their own hold and admitted project in existing
run metadata. Unknown provider outcomes retain that capacity until accounting
establishes the charge. Browser holds are separate from their parent run's usage.
Unfinished uncapped browsers and legacy rows without a reliable exposure record
remain unknown usage when a new budget is introduced. The existing unpriced-usage
policy controls admission until their terminal charge is known.
Use a disposable real PostgreSQL database to compare cost evaluation and complete admission:
```sh
pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-budget-enforcement.ts --baseline-ref=6c9b87d0922f1ebfc43e8b0c36889ac5ba7b9837 > /tmp/budget-forward.json
pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-budget-enforcement.ts --baseline-ref=6c9b87d0922f1ebfc43e8b0c36889ac5ba7b9837 --reverse > /tmp/budget-reverse.json
```
The cost comparison keeps one company policy and no active runs while adding 0,
17, and 1,000 unrelated idle agents. It records 20 warm samples, SQL statement counts,
company accounting locks, agent reads, and an `EXPLAIN (ANALYZE, BUFFERS)` plan.
It asserts unchanged agent state and no agent reads in the new evaluation path.
The harness also measures reservation, the final ownership gate, and an immediate
offline provider result with one responsible agent and one active operation. It
records transaction counts and asserts no reads of unrelated agents. Baseline
service, lifecycle, reservation, and dispatch modules come from the supplied Git
revision.
Provider inference, network latency, stop acknowledgement, and hosted database
latency need separate measurements. Do not run this benchmark beside test suites
or builds when reporting timing results.
### Skill-source background synchronization
Skill imports and automatic source refreshes use the delivery coordinator. Workspace
+16
View File
@@ -1649,6 +1649,22 @@ Scheduler must skip invocation when:
- an existing run is active
- hard budget limit has been hit
Budget admission and accounting serialize under the company accounting lock.
The dispatch boundary rechecks current company, agent, and project policies and
reserves capacity before paid work starts. Accounting commits budget blocks and
durable stop intent without waiting for provider shutdown or reconciling unrelated
agents. An asynchronous worker cancels only work in the affected scope; committed
stop intent survives policy changes and restart until existing stop proof succeeds.
Accepted cumulative per-run usage checkpoints contribute observed running spend
before the final ledger charge. Admission counts each unfinished operation's
maximum of its reservation and observed spend; hard stops count actual ledger
and observed running spend, rather than reservation estimates. Final settlement
atomically replaces the observation with the ledger charge. Incomplete or unknown
accounting remains conservative. Mid-run enforcement depends on provider usage
reports or existing supported pricing; providers reporting only at completion
cannot offer exact mid-call spend detection.
Legacy execution records a renewable controller lease when claiming a queued run,
before provisioning. A live lease protects the run during overlapping service
deployments. An expired controller loses dispatch authority; a recovery worker
+35
View File
@@ -46,6 +46,41 @@ without installed handler tracking keeps the same conservative hold. Missing fil
continue through the normal fallback and error handlers; they do not exempt
other GET requests or callback-only application work.
## Budget and agent lifecycle work
Budget accounting commits its charge, admission block, and enforcement generation
before returning to its caller. Cancellation delivery runs separately through the
delivery coordinator, using the existing policy version pair and committed run
Stop records. It can finish already-admitted stops during drain, but cannot start
another paid turn. A provider still stopping, an unacknowledged Stop, a pending
policy generation, and unresolved accounting remain sleep blockers.
Agent lifecycle uses the same transaction-aware coordinator. Hire approval,
configuration invalidation, manual commands, and computed agent budget holds
register work before their transaction writes. Company status transitions register
inherited-hold projection; changing a company name, description, or logo does not.
Projection selects only agents whose company pause/archive hold differs from the
committed company state, in pages of 100 under one company accounting transaction
per page. Pause/archive transactions also persist pausing operations for agents
with cancellable runs, reading only those involved agents in batches. Immediate
reactivation can remove an inherited hold but cannot erase an owed Stop; new
provider work waits for lifecycle readiness. Archive activity counts describe
requested pauses and cancellations, not completed provider termination. Agent
and project budget changes never enumerate company peers.
Startup recovers pending lifecycle operations and inherited-hold mismatches.
Thereafter notifications and stored retry/lease deadlines drive work. Stable ready,
paused, rejected, and terminated history does not trigger recurring policy scans.
Pending phases and inherited-hold mismatches block sleep. Lifecycle execution can
verify configurations or resume queued work, so it waits while admission is
draining or the instance is in warm standby; its durable obligation is retained.
A future in-process retry timer does not grant permission to sleep.
These notifications cover the existing single-active-writer ownership model.
Independent SQL writers and additional processes need an explicit wake protocol;
process-local notifications do not provide cross-replica delivery. Worker quiet
alone is not evidence that the host or database actually suspended.
## Scheduled routine wake protocol
Historical failed ephemeral leases and terminal custom-image setup sessions can
+19
View File
@@ -35,6 +35,25 @@ maintains as a dependency; once you install the packages below, they become
normal dependencies of **your own** project, and your own Dependabot updates
them.
## Budget admission, accounting, and enforcement
The operator-gated `paperclip.budgets` tracer records
`budget.admission.evaluate`, `budget.admission.reserve`,
`budget.accounting.transaction`, `budget.accounting.settlement`, `budget.accounting.lock_wait`,
`budget.exposure.evaluate`, `budget.enforcement.deliver`,
`budget.enforcement.scope`, `budget.cancellation.deliver`, and `budget.accounting.recover`.
Transaction spans include commit time. Lock-wait spans cover acquisition of the
company accounting lock. Settlement spans cover the ledger and reservation
transaction. Enforcement spans measure asynchronous delivery; cancellation
delivery spans start after acquiring a concurrency slot and cover the external
Stop attempt. Both remain separate from provider inference.
The closed attribute list is `budget.policy_count` (number of applicable
policies) and `budget.enforcement.queue_delay_ms` (time since the pending policy's
last intent or retry timestamp). These spans contain no tool content, usage
receipt bodies, credentials, or provider error bodies. They do not change
Paperclip's separate first-party Telemetry contract.
## Slack receipt reactions
Slack stages its eyes reaction in the same transaction as the authorized task,
+12
View File
@@ -4,6 +4,18 @@ Product map: [Costs, budgets, and spending controls](../feature-map/budgets-cost
Feature ID: `budgets-costs`.
Budget admission commits under the company accounting lock. A separate ownership
transaction rechecks scope and capacity immediately before provider handoff;
provider work does not hold the company accounting lock. Company, applicable
agent, and applicable project limits include unsettled
observed run usage and reservations without counting either twice. A committed
hard stop prevents fresh provider work immediately; scoped cancellation delivery
runs asynchronously and preserves existing Stop acknowledgement and recovery
requirements. Project enforcement does not pause an agent's unrelated projects.
Hosted browser dispatch reserves a separate scoped hold, rechecks authority at
the paid POST, and retains unknown charges until terminal accounting. The parent
run reservation does not authorize or replace this separate browser cost cap.
Review date: 2026-10-12. Source commit: `f981fa3cc10cb8663a87dec60345fbff46d54d0a`.
This is a source-reviewed baseline of representative operations and known gaps.
+6
View File
@@ -61,6 +61,12 @@ or unavailable cost data without treating missing data as zero spend.
Automated: [budget service](../../server/src/__tests__/budgets-service.test.ts) exercises soft incidents, hard stops, scope, and valid budget raises.
Automated: [scoped budget accounting](../../server/src/__tests__/budget-enforcement-scoping.test.ts) exercises cumulative running usage, reservation capacity, final settlement, project isolation, and concurrent admission. [Budget cancellation](../../server/src/services/heartbeat/budget-cancellation.test.ts) exercises durable scoped cancellation, delayed providers, stale generations, and recovery. These controlled-provider tests establish enforcement semantics, not live provider termination latency.
Automated: [hosted browser accounting](../../server/src/__tests__/browser-use-connection.test.ts) covers separate browser holds, concurrent starts, final dispatch races, unknown POST outcomes, and exact terminal settlement. Browser caps reserve additional capacity alongside the parent run; internal holds are never returned as provider results.
Manual: With disposable budgets and a controlled streaming adapter, admit a run, report cumulative usage across its hard limit, and verify that new work is denied immediately while the existing run receives Stop asynchronously. Delay the Stop acknowledgement and verify that unrelated requests remain responsive. Use separate projects to verify scope isolation. Raise the budget before an unfenced cancellation attempt and verify that newly admitted work survives; retain an already committed Stop until its own acknowledgement. Use the explicit budget benchmark to compare 0, 17, and 1,000 unrelated idle agents.
Manual: Create a small fixture budget, cross its threshold with synthetic recorded spend, and verify the warning or hard-stop state. Attempt new work and confirm it is gated. Raise the correct budget through an authorized action and verify eligible continuation without clearing unrelated manual pauses.
### `subscription-accounts`
+179
View File
@@ -0,0 +1,179 @@
import assert from "node:assert/strict";
import { randomUUID } from "node:crypto";
import { execFileSync } from "node:child_process";
import { promises as fs } from "node:fs";
import { createRequire } from "node:module";
import path from "node:path";
import { fileURLToPath, pathToFileURL } from "node:url";
import { agents, budgetPolicies, companies, createDb, heartbeatRuns, issues, startEmbeddedPostgresTestDatabase, type Db } from "../packages/db/src/index.js";
const { eq } = createRequire(new URL("../server/package.json", import.meta.url))("drizzle-orm");
const ref = process.argv.find(arg => arg.startsWith("--baseline-ref="))?.slice(15);
assert.ok(ref, "Pass --baseline-ref=<verified pre-fix commit>");
const baselineRevision = execFileSync("git", ["rev-parse", `${ref}^{commit}`], { encoding: "utf8" }).trim();
const baselineOnly = process.argv.includes("--baseline-only");
const sampleCount = 20;
const root = fileURLToPath(new URL("../", import.meta.url));
const directory = await fs.mkdtemp(path.join(root, "server/scripts/.budget-baseline-"));
const files = ["server/src/services/budgets.ts", "server/src/services/budget-reservations.ts", "server/src/services/agent-lifecycle.ts", "server/src/services/accounting-transaction.ts",
...execFileSync("git", ["ls-tree", "-r", "--name-only", baselineRevision, "server/src/modules/agent-lifecycle", "server/src/modules/run-dispatch"], { encoding: "utf8" })
.trim().split("\n").filter(file => file.endsWith(".ts") && !file.endsWith(".test.ts") && !file.includes("/__tests__/"))];
const frozen = new Map(files.map(file => [path.join(root, file).replace(/\.ts$/, ".js"), path.join(directory, file)]));
for (const file of files) {
const origin = path.join(root, file);
const source = execFileSync("git", ["show", `${baselineRevision}:${file}`], { encoding: "utf8" });
const rewritten = source.replace(/(from\s+|import\()(["'])(\.{1,2}\/[^"']+)\2/g, (_match, prefix: string, quote: string, specifier: string) => {
const absolute = path.resolve(path.dirname(origin), specifier);
return `${prefix}${quote}${frozen.get(absolute) ?? absolute}${quote}`;
});
const target = path.join(directory, file);
await fs.mkdir(path.dirname(target), { recursive: true });
await fs.writeFile(target, rewritten);
}
let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>> | undefined;
let db: Db | undefined;
const keepAlive = setInterval(() => {}, 1000);
try {
const before = await import(pathToFileURL(path.join(directory, "server/src/services/budgets.ts")).href);
const after = baselineOnly ? null : await import("../server/src/services/budgets.js");
const beforeReservations = await import(pathToFileURL(path.join(directory, "server/src/services/budget-reservations.ts")).href);
const beforeDispatch = await import(pathToFileURL(path.join(directory, "server/src/modules/run-dispatch/adapters/postgres.ts")).href);
const afterReservations = baselineOnly ? null : await import("../server/src/services/budget-reservations.js");
const afterDispatch = baselineOnly ? null : await import("../server/src/modules/run-dispatch/adapters/postgres.js");
database = await startEmbeddedPostgresTestDatabase("budget-enforcement-benchmark-");
db = createDb(database.connectionString);
// Instrument this disposable client's Drizzle sessions. Transaction sessions
// inherit options.logger; no runtime client or production setting changes.
let captured: Array<{ sql: string; params: unknown[] }> | null = null;
let capturedTransactions = 0;
const transaction = db.transaction.bind(db);
db.transaction = ((...args: Parameters<Db["transaction"]>) => {
if (captured) capturedTransactions++;
return transaction(...args);
}) as Db["transaction"];
const session = Reflect.get(db, "session");
const logger = { logQuery(query: string, params: unknown[]) { captured?.push({ sql: query, params }); } };
session.logger = logger;
session.options.logger = logger;
const [version] = await db.$client.unsafe("select version()");
const results = [];
const admissionResults = [];
let expectedTreatmentCounts: string | undefined;
for (const count of [0, 17, 1000]) {
const companyId = randomUUID();
await db.insert(companies).values({ id: companyId, name: "Synthetic budget benchmark", issuePrefix: `B${companyId.slice(0, 7).replaceAll("-", "")}` });
await db.insert(budgetPolicies).values({ companyId, scopeType: "company", scopeId: companyId,
metric: "billed_cents", windowKind: "calendar_month_utc", amount: 1_000_000, notifyEnabled: false });
if (count) await db.$client.unsafe(`insert into agents (company_id,name,role,status,adapter_type,lifecycle_state)
select $1::uuid, 'Idle fixture ' || n, 'engineer', 'idle', 'process', 'ready' from generate_series(1,$2::int) n`, [companyId, count]);
await db.$client.unsafe("analyze agents");
const measured = [];
const entries = [["baseline", before], ["treatment", after]] as const;
for (const [name, service] of process.argv.includes("--reverse") ? [...entries].reverse() : entries) {
if (!service) continue;
const samples = [];
for (let i = 0; i < sampleCount + 2; i++) {
const start = performance.now();
// Real budget cost-evaluation service boundary: before this change,
// even a charge below its limit reconciled all idle agents.
await service.budgetService(db).evaluateCostEvent({ companyId, agentId: null, projectId: null });
const elapsed = performance.now() - start;
if (i >= 2) samples.push(elapsed);
}
samples.sort((a, b) => a - b);
captured = [];
await service.budgetService(db).evaluateCostEvent({ companyId, agentId: null, projectId: null });
const queries = captured;
captured = null;
const agentReads = queries.filter(query => /\bfrom "agents"/.test(query.sql)).length;
const accountingLocks = queries.filter(query => query.sql.includes("for no key update")).length;
if (name === "treatment") {
assert.equal(agentReads, 0, "Ordinary evaluation must not read any agents");
assert.equal(accountingLocks, 1, "Only the charge's accounting transaction may run");
}
const candidate = queries.find(query => name === "baseline" ? /\bfrom "agents"/.test(query.sql) : /\bfrom "budget_policies"/.test(query.sql));
const plan = candidate ? await db.$client.unsafe(`EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) ${candidate.sql}`, candidate.params as never[]) : null;
measured.push({ name, samplesMs: samples, medianMs: samples[Math.floor(sampleCount / 2)],
p95Ms: samples[Math.ceil(sampleCount * 0.95) - 1],
drizzleStatementCount: queries.length, accountingTransactions: accountingLocks, agentReads,
candidateQuery: candidate, candidatePlan: plan });
}
const persisted = await db.$client.unsafe("select lifecycle_state, lifecycle_holds, status from agents where company_id = $1::uuid", [companyId]);
assert.equal(persisted.length, count);
assert.ok(persisted.every(agent => agent.lifecycle_state === "ready" && agent.lifecycle_holds.length === 0 && agent.status === "idle"));
results.push({ unrelatedIdleAgents: count, contractPreserved: true, measured });
// Keep exactly one admitted operation and one responsible agent at every
// headcount. Fixture writes and cleanup are outside the timed boundary.
const agentId = randomUUID(), issueId = randomUUID();
await db.insert(agents).values({ id: agentId, companyId, name: "Responsible fixture", role: "engineer",
status: "idle", adapterType: "process", lifecycleState: "ready", runtimeConfig: { heartbeat: { wakeOnDemand: true } } });
await db.insert(issues).values({ id: issueId, companyId, title: "Offline dispatch fixture", status: "in_progress", assigneeAgentId: agentId });
await db.update(budgetPolicies).set({ reservationCents: "100" }).where(eq(budgetPolicies.companyId, companyId));
const admissionMeasured = [];
const admissionEntries = [["baseline", beforeReservations, beforeDispatch], ["treatment", afterReservations, afterDispatch]] as const;
for (const [name, reservations, dispatch] of process.argv.includes("--reverse") ? [...admissionEntries].reverse() : admissionEntries) {
if (!reservations || !dispatch) continue;
const adapter = dispatch.createPostgresRunDispatchAdapter(db);
const samples = [];
let queries: Array<{ sql: string; params: unknown[] }> = [], databaseTransactions = 0;
for (let i = 0; i < sampleCount + 3; i++) {
const runId = randomUUID();
await db.insert(heartbeatRuns).values({ id: runId, companyId, agentId, status: "running", contextSnapshot: { issueId } });
await db.update(issues).set({ executionRunId: runId }).where(eq(issues.id, issueId));
let callbacks = 0;
const probe = i === sampleCount + 2;
if (probe) { captured = []; capturedTransactions = 0; }
const start = performance.now();
try {
const reservation = await reservations.reserveRunBudget(db, companyId, runId, null);
const outcome = await adapter.dispatchResolvedInteractionIfCurrent({ companyId, runId,
reservationId: reservation.id, expectedStatus: "running", now: new Date(),
dispatch: async () => { callbacks++; return "offline-provider-result"; } });
assert.equal(outcome.dispatched, true, `${name}: current admitted run must dispatch`);
if (outcome.dispatched) assert.equal(await outcome.resultPromise, "offline-provider-result");
const elapsed = performance.now() - start;
assert.equal(callbacks, 1, "Each admitted fixture calls the provider exactly once");
if (probe) { queries = captured!; databaseTransactions = capturedTransactions; }
else if (i >= 2) samples.push(elapsed);
} finally {
captured = null;
await db.update(issues).set({ executionRunId: null }).where(eq(issues.id, issueId));
await db.delete(heartbeatRuns).where(eq(heartbeatRuns.id, runId));
}
}
samples.sort((a, b) => a - b);
const agentQueries = queries.filter(query => /\b(?:from|join) "agents"/.test(query.sql));
const unrelatedAgentReadStatements = agentQueries.filter(query => !query.params.includes(agentId)).length;
const accountingTransactions = queries.filter(query => query.sql.includes("for no key update")).length;
const p95Ms = samples[Math.ceil(sampleCount * 0.95) - 1];
if (name === "treatment") {
assert.equal(unrelatedAgentReadStatements, 0, "Agent reads must be bound to the responsible agent");
const counts = JSON.stringify([queries.length, databaseTransactions, accountingTransactions, agentQueries.length]);
expectedTreatmentCounts ??= counts;
assert.equal(counts, expectedTreatmentCounts, "Admission query and transaction counts must not grow with idle agents");
assert.ok(p95Ms < 250, `Warm admission p95 ${p95Ms}ms exceeds 250ms`);
}
admissionMeasured.push({ name, samplesMs: samples, medianMs: samples[Math.floor(sampleCount / 2)], p95Ms,
drizzleStatementCount: queries.length, databaseTransactions, accountingTransactions, agentReads: agentQueries.length,
unrelatedAgentReadStatements, providerCallsPerSample: 1 });
}
const [remaining] = await db.$client.unsafe(`select (select count(*) from heartbeat_runs where company_id=$1::uuid) runs,
(select count(*) from budget_reservations where company_id=$1::uuid) reservations`, [companyId]);
assert.equal(Number(remaining.runs), 0); assert.equal(Number(remaining.reservations), 0);
const unchanged = await db.$client.unsafe("select lifecycle_state, lifecycle_holds, status from agents where company_id=$1::uuid", [companyId]);
assert.equal(unchanged.length, count + 1);
assert.ok(unchanged.every(agent => agent.lifecycle_state === "ready" && agent.lifecycle_holds.length === 0 && agent.status === "idle"));
admissionResults.push({ unrelatedIdleAgents: count, responsibleAgents: 1, activeRuns: 1, contractPreserved: true, measured: admissionMeasured });
console.error(`Budget benchmark: ${count} idle agents, cost evaluation and admission complete`);
}
console.log(JSON.stringify({ baselineRevision, treatmentRevision: execFileSync("git", ["rev-parse", "HEAD"], { encoding: "utf8" }).trim(),
note: "Local service timings against disposable real PostgreSQL; no provider/network latency. Treatment may include uncommitted source changes.",
policiesPerCompany: 1, activeRuns: 0, node: process.version, nodeArchitecture: process.arch,
postgres: version, results, admissionResults }, null, 2));
} finally {
await db?.$client.end({ timeout: 0 });
await database?.cleanup();
await fs.rm(directory, { recursive: true, force: true });
clearInterval(keepAlive);
}
@@ -1,4 +1,4 @@
import { budgetService } from "../services/budgets.js";
import { budgetService, computeObservedSpend } from "../services/budgets.js";
import { buildPaperclipRuntimeMcpServers } from "../services/heartbeat.js";
import { resolveNativeRuntimeMcpSnapshot } from "../services/native-runtime/runtime-context.js";
import express from "express";
@@ -13,6 +13,7 @@ import {
agents,
approvals,
toolActionRequests,
toolInvocations,
toolPolicies,
companies,
companyMemberships,
@@ -20,17 +21,21 @@ import {
createDb,
heartbeatRuns,
issues,
projects,
toolConnections,
toolApplications,
toolCatalogEntries,
toolConnectionInstalls,
toolProfiles,
toolProfileBindings,
toolProfileEntries,
browserUseSessions,
browserUseRuns,
browserUseBrowsers,
costEvents,
financeEvents,
budgetPolicies,
budgetReservations,
} from "@paperclipai/db";
import { toolAccessService } from "../services/tool-access.js";
import { createToolGatewayService } from "../services/tool-gateway.js";
@@ -84,6 +89,8 @@ const actor = { actorType: "user" as const, actorId: "browser-reviewer" };
providerBrowser = randomUUID();
let providerRun = randomUUID();
let acceptedTask = "", loseReply = false, rejectCode = 0, recoveryPages = 0;
let createGate: (() => Promise<void>) | undefined;
let totalCostUsd = "0.15";
let status = "running",
stopped = false,
browserReady = false,
@@ -112,6 +119,7 @@ const actor = { actorType: "user" as const, actorId: "browser-reviewer" };
if (path === "/runs" && init.method === "POST" && failCreate)
throw new Error("lost create response");
if (path === "/runs" && init.method === "POST") {
await createGate?.();
providerRun = randomUUID();
status = "running";
acceptedTask = JSON.parse(String(init.body)).task;
@@ -156,7 +164,7 @@ const actor = { actorType: "user" as const, actorId: "browser-reviewer" };
return Response.json({
status,
result: "Finished https://live.browser-use.com/canary",
totalCostUsd: "0.15",
totalCostUsd,
model: "fixture",
});
if (path === "/browsers")
@@ -280,7 +288,7 @@ const actor = { actorType: "user" as const, actorId: "browser-reviewer" };
.update(browserUseSessions)
.set({ nextPollAt: new Date(0) })
.where(eq(browserUseSessions.companyId, company.id));
await service.sweep();
await service.sweep(connection.connectionId);
};
return {
company,
@@ -311,6 +319,8 @@ const actor = { actorType: "user" as const, actorId: "browser-reviewer" };
loseReply: () => { loseReply = true; },
reject: (code: number) => { rejectCode = code; },
recoveryPages: (count: number) => { recoveryPages = count; },
delayCreate: (gate: () => Promise<void>) => { createGate = gate; },
usage: (value: string) => { totalCostUsd = value; },
paginate: () => {
paginatedEvents = true;
},
@@ -1389,6 +1399,219 @@ const actor = { actorType: "user" as const, actorId: "browser-reviewer" };
)[0].status,
).toBe("closed");
});
it("serializes browser holds with parent reservations while provider work is delayed", async () => {
const f = await fixture();
const policy = await budgetService(db).upsertPolicy(f.company.id,
{ scopeType: "company", scopeId: f.company.id, amount: 100, hardStopEnabled: true }, actor.actorId);
await db.insert(budgetReservations).values({ companyId: f.company.id, runId: f.run.id, agentId: f.agent.id, amountCents: "40" });
let release!: () => void, started!: () => void;
const providerEntered = new Promise<void>(resolve => { started = resolve; });
const providerRelease = new Promise<void>(resolve => { release = resolve; });
f.delayCreate(async () => { started(); await providerRelease; });
const first = f.service.execute(f.binding, f.grant, randomUUID(), "browser_start", { task: "Read", maxCostUsd: 0.8 });
// Attach the rejection handler before waiting for the controlled POST.
await Promise.race([providerEntered, first.then(() => undefined)]);
try {
await expect(f.service.execute(f.binding, f.grant, randomUUID(), "browser_start", { task: "Concurrent" })).rejects.toThrow("reserved");
const [row] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.policyId));
expect(await computeObservedSpend(db, row)).toMatchObject({ ledgerCents: "0.0000000", effectiveHeldCents: "100.0000000" });
expect(f.request.mock.calls.filter(([url, init]) => url.endsWith("/runs") && init.method === "POST")).toHaveLength(1);
const call = f.request.mock.calls.find(([url, init]) => url.endsWith("/runs") && init.method === "POST")!;
expect(JSON.parse(String(call[1].body)).maxCostUsd).toBe(0.6);
// Accounting lock is not held while the hosted browser is waiting.
expect(await db.select({ id: companies.id }).from(companies).where(eq(companies.id, f.company.id))).toHaveLength(1);
} finally { release(); await first; }
});
it.each(["reduction", "new_policy", "config", "charge", "cancel", "session_end", "policy_deny"])("rechecks %s at the actual provider handoff", async race => {
const f = await fixture();
if (race === "reduction" || race === "charge") await budgetService(db).upsertPolicy(f.company.id,
{ scopeType: "company", scopeId: f.company.id, amount: 100, hardStopEnabled: true }, actor.actorId);
const invocationId = randomUUID();
const transaction = db.transaction.bind(db);
let injected = false;
const barrier = vi.spyOn(db, "transaction").mockImplementation(async (callback, config) => {
const result = await transaction(callback, config);
const [run] = await db.select().from(browserUseRuns).where(eq(browserUseRuns.invocationId, invocationId));
if (!injected && run?.status === "creating" && !(run.result as { paperclipBudgetHold?: { dispatchCommittedAt?: string } })?.paperclipBudgetHold?.dispatchCommittedAt) {
injected = true;
if (race === "reduction" || race === "new_policy") await budgetService(db).upsertPolicy(f.company.id,
{ scopeType: "company", scopeId: f.company.id, amount: race === "new_policy" ? 100 : 10, hardStopEnabled: true }, actor.actorId);
else if (race === "config") await f.service.saveSettings(f.company.id, f.grant.id, actor.actorId, { allowedProfileIds: [], maxCostUsd: 0.1 });
else if (race === "charge") {
const { costService } = await import("../services/costs.js");
await costService(db).createEvent(f.company.id, { agentId: f.agent.id, provider: "fixture", model: "fixture", costCents: 20, occurredAt: new Date() });
} else if (race === "policy_deny") await db.insert(toolPolicies).values({ companyId: f.company.id,
name: "Deny before browser handoff", policyType: "block", selectors: { connectionId: f.connection.connectionId }, priority: 0 });
else if (race === "session_end") await db.update(browserUseSessions).set({ status: "stopping", stopRequested: "end" }).where(eq(browserUseSessions.id, run.sessionId));
else await db.update(heartbeatRuns).set({ resultJson: { startupCancellation: { requestedAt: new Date().toISOString() } } }).where(eq(heartbeatRuns.id, f.run.id));
}
return result;
});
try {
await expect(f.service.execute(f.binding, f.grant, invocationId, "browser_start", { task: "Must not start" })).rejects.toThrow();
} finally { barrier.mockRestore(); }
expect(injected).toBe(true);
expect(f.request.mock.calls.filter(([url, init]) => url.endsWith("/runs") && init.method === "POST")).toHaveLength(0);
const [run] = await db.select().from(browserUseRuns).where(eq(browserUseRuns.invocationId, invocationId));
expect(run).toMatchObject({ status: "failed", eventsDrained: 1, result: { paperclipBudgetHold: { state: "released" } } });
});
it.each(["default_allow", "tool_name"])("retains valid %s profile authorization at browser admission", async selector => {
const f = await fixture();
await db.update(toolProfiles).set({ status: "inactive" }).where(eq(toolProfiles.companyId, f.company.id));
const [profile] = await db.insert(toolProfiles).values({ companyId: f.company.id, name: "Scoped browser access",
profileKey: randomUUID(), defaultAction: selector === "default_allow" ? "allow" : "deny" }).returning();
await db.insert(toolProfileBindings).values({ companyId: f.company.id, profileId: profile.id, targetType: "agent", targetId: f.agent.id });
if (selector === "tool_name") await db.insert(toolProfileEntries).values({ companyId: f.company.id, profileId: profile.id,
selectorType: "tool_name", toolName: "browser_start", effect: "include" });
await expect(f.service.execute(f.binding, f.grant, randomUUID(), "browser_start", { task: "Allowed by profile" })).resolves.toMatchObject({ status: "running" });
expect(f.request.mock.calls.filter(([url, init]) => url.endsWith("/runs") && init.method === "POST")).toHaveLength(1);
});
it("checks the gateway invocation's namespaced tool identity for explicit denial", async () => {
const f = await fixture();
const invocationId = randomUUID(), toolName = "app.browser-fixture:browser-start";
const [entry] = await db.select().from(toolCatalogEntries).where(and(eq(toolCatalogEntries.connectionId, f.connection.connectionId), eq(toolCatalogEntries.toolName, "browser_start")));
await db.insert(toolInvocations).values({ id: invocationId, companyId: f.company.id, agentId: f.agent.id,
runId: f.run.id, issueId: f.issue.id, actorType: "agent", actorId: f.agent.id,
connectionId: f.connection.connectionId, applicationId: entry.applicationId, catalogEntryId: entry.id,
toolName, upstreamToolName: "browser_start", policyDecision: "allow", status: "executing" });
await db.insert(toolPolicies).values({ companyId: f.company.id, name: "Block namespaced invocation",
policyType: "block", selectors: { toolName }, priority: 0 });
await expect(f.service.execute(f.binding, f.grant, invocationId, "browser_start", { task: "Denied" })).rejects.toThrow("blocked");
expect(f.request.mock.calls.filter(([url, init]) => url.endsWith("/runs") && init.method === "POST")).toHaveLength(0);
});
it("releases a crashed reservation that never reached provider dispatch authority", async () => {
const f = await fixture();
const policy = await budgetService(db).upsertPolicy(f.company.id,
{ scopeType: "agent", scopeId: f.agent.id, amount: 100, hardStopEnabled: true }, actor.actorId);
const [session] = await db.insert(browserUseSessions).values({ companyId: f.company.id, issueId: f.issue.id,
agentId: f.agent.id, connectionId: f.connection.connectionId, grantId: f.grant.id, status: "starting" }).returning();
const invocationId = randomUUID();
const [run] = await db.insert(browserUseRuns).values({ companyId: f.company.id, sessionId: session.id,
heartbeatRunId: f.run.id, invocationId, createdAt: new Date(Date.now() - 61_000),
result: { paperclipBudgetHold: { invocationId, amountCents: "50.0000000", state: "held", projectId: null } } }).returning();
f.request.mockClear();
await f.tick();
const [policyRow] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.policyId));
expect((await computeObservedSpend(db, policyRow)).effectiveHeldCents).toBe("0.0000000");
expect((await db.select().from(browserUseRuns).where(eq(browserUseRuns.id, run.id)))[0]).toMatchObject({ status: "failed", eventsDrained: 1, result: { paperclipBudgetHold: { state: "released" } } });
expect((await db.select().from(browserUseSessions).where(eq(browserUseSessions.id, session.id)))[0].status).toBe("failed");
expect(f.request.mock.calls.some(([url]) => new URL(url).pathname === "/api/v4/runs")).toBe(false);
});
it("releases a definitively rejected browser hold without inventing a charge", async () => {
const f = await fixture();
const policy = await budgetService(db).upsertPolicy(f.company.id,
{ scopeType: "agent", scopeId: f.agent.id, amount: 100, hardStopEnabled: true }, actor.actorId);
f.reject(409);
await expect(f.service.execute(f.binding, f.grant, randomUUID(), "browser_start", { task: "Rejected", maxCostUsd: 0.5 })).rejects.toThrow("busy");
const [row] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.policyId));
expect(await computeObservedSpend(db, row)).toMatchObject({ ledgerCents: "0.0000000", effectiveHeldCents: "0.0000000" });
expect(await db.select().from(costEvents).where(eq(costEvents.companyId, f.company.id))).toHaveLength(0);
await db.update(browserUseRuns).set({ result: null }).where(eq(browserUseRuns.companyId, f.company.id));
expect((await computeObservedSpend(db, row)).unknownBrowserRunCount).toBe(0);
});
it.each(["company", "agent", "project"] as const)("treats a running uncapped browser as unknown when a %s budget is introduced", async scopeType => {
for (const legacy of [false, true]) {
const f = await fixture();
const [project] = await db.insert(projects).values({ companyId: f.company.id, name: "Browser project" }).returning();
const [otherProject] = await db.insert(projects).values({ companyId: f.company.id, name: "Unrelated project" }).returning();
await db.update(issues).set({ projectId: project.id }).where(eq(issues.id, f.issue.id));
const invocation = randomUUID();
await expect(f.service.execute(f.binding, f.grant, invocation, "browser_start", { task: "Uncapped running work" })).resolves.toMatchObject({ status: "running" });
const [run] = await db.select().from(browserUseRuns).where(eq(browserUseRuns.invocationId, invocation));
expect(run.result).toMatchObject({ paperclipBudgetHold: { amountCents: "0.0000000", uncapped: true } });
expect(JSON.parse(String(f.request.mock.calls.find(([url, init]) => url.endsWith("/runs") && init.method === "POST")![1].body))).not.toHaveProperty("maxCostUsd");
if (legacy) await db.update(browserUseRuns).set({ result: null }).where(eq(browserUseRuns.id, run.id));
const scopeId = scopeType === "company" ? f.company.id : scopeType === "agent" ? f.agent.id : project.id;
const policy = await budgetService(db).upsertPolicy(f.company.id,
{ scopeType, scopeId, amount: 100, hardStopEnabled: true }, actor.actorId);
const [row] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.policyId));
expect(await computeObservedSpend(db, row)).toMatchObject({ ledgerCents: "0.0000000", effectiveHeldCents: "0.0000000", unknownBrowserRunCount: 1 });
expect((await computeObservedSpend(db, { ...row, scopeType: "project", scopeId: otherProject.id })).unknownBrowserRunCount).toBe(0);
expect(await budgetService(db).getInvocationBlock(f.company.id, f.agent.id, { projectId: project.id })).not.toBeNull();
f.request.mockClear();
await expect(f.service.execute(f.binding, f.grant, randomUUID(), "browser_start", { task: "Must wait for priced work" })).rejects.toThrow();
expect(f.request.mock.calls.filter(([url, init]) => url.endsWith("/runs") && init.method === "POST")).toHaveLength(0);
await budgetService(db).upsertPolicy(f.company.id,
{ scopeType, scopeId, amount: 100, hardStopEnabled: true, unpricedUsagePolicy: "allow" }, actor.actorId);
expect(await budgetService(db).getInvocationBlock(f.company.id, f.agent.id, { projectId: project.id })).toBeNull();
f.complete(); await f.tick();
expect(await computeObservedSpend(db, row)).toMatchObject({ ledgerCents: "15.0000000", effectiveHeldCents: "0.0000000", unknownBrowserRunCount: 0 });
}
});
it.each([
{ label: "explicit capped zero", uncapped: false, state: "held", dispatched: true, status: "running", expected: 0 },
{ label: "legacy zero cap without a flag", uncapped: undefined, state: "held", dispatched: true, status: "running", expected: 1 },
{ label: "uncapped before dispatch", uncapped: true, state: "held", dispatched: false, status: "creating", expected: 0 },
{ label: "failed before dispatch", uncapped: true, state: "released", dispatched: false, status: "failed", expected: 0 },
])("distinguishes $label from unfinished uncapped spending", async item => {
const f = await fixture();
const [session] = await db.insert(browserUseSessions).values({ companyId: f.company.id, issueId: f.issue.id,
agentId: f.agent.id, connectionId: f.connection.connectionId, grantId: f.grant.id, status: "running" }).returning();
const invocationId = randomUUID();
await db.insert(browserUseRuns).values({ companyId: f.company.id, sessionId: session.id, heartbeatRunId: f.run.id,
invocationId, status: item.status, providerRunId: item.dispatched ? randomUUID() : null, eventsDrained: item.status === "failed" ? 1 : 0,
result: { paperclipBudgetHold: { invocationId, amountCents: "0.0000000", projectId: null, state: item.state,
...(item.uncapped === undefined ? {} : { uncapped: item.uncapped }),
...(item.dispatched ? { dispatchCommittedAt: new Date().toISOString() } : {}) } } });
const policy = await budgetService(db).upsertPolicy(f.company.id,
{ scopeType: "company", scopeId: f.company.id, amount: 100, hardStopEnabled: true }, actor.actorId);
const [row] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.policyId));
expect(await computeObservedSpend(db, row)).toMatchObject({ ledgerCents: "0.0000000", effectiveHeldCents: "0.0000000", unknownBrowserRunCount: item.expected });
const block = await budgetService(db).getInvocationBlock(f.company.id, f.agent.id);
expect(Boolean(block)).toBe(item.expected > 0);
});
it("blocks new budgeted work after an uncapped browser POST becomes uncertain", async () => {
const f = await fixture();
f.loseReply();
const uncertain = randomUUID();
await expect(f.service.execute(f.binding, f.grant, uncertain, "browser_start", { task: "Uncapped" })).rejects.toThrow("could not be confirmed");
const policy = await budgetService(db).upsertPolicy(f.company.id,
{ scopeType: "company", scopeId: f.company.id, amount: 100, hardStopEnabled: true }, actor.actorId);
const [row] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.policyId));
expect(await computeObservedSpend(db, row)).toMatchObject({ ledgerCents: "0.0000000", effectiveHeldCents: "0.0000000", unknownBrowserRunCount: 1 });
f.request.mockClear();
await expect(f.service.execute(f.binding, f.grant, randomUUID(), "browser_start", { task: "Must wait" })).rejects.toThrow();
expect(f.request.mock.calls.filter(([url, init]) => url.endsWith("/runs") && init.method === "POST")).toHaveLength(0);
// The existing unpriced-usage policy remains the explicit operator override.
await budgetService(db).upsertPolicy(f.company.id,
{ scopeType: "company", scopeId: f.company.id, amount: 100, hardStopEnabled: true, unpricedUsagePolicy: "allow" }, actor.actorId);
expect(await budgetService(db).getInvocationBlock(f.company.id, f.agent.id)).toBeNull();
f.usage("0.15");
await f.tick();
expect(await computeObservedSpend(db, row)).toMatchObject({ ledgerCents: "15.0000000", unknownBrowserRunCount: 0 });
});
it("retains unknown browser holds and settles exact terminal usage without exposing internal metadata", async () => {
const f = await fixture();
const policy = await budgetService(db).upsertPolicy(f.company.id,
{ scopeType: "company", scopeId: f.company.id, amount: 100, hardStopEnabled: true }, actor.actorId);
f.loseReply();
const invocation = randomUUID();
await expect(f.service.execute(f.binding, f.grant, invocation, "browser_start", { task: "Read", maxCostUsd: 0.5 })).rejects.toThrow("could not be confirmed");
const [row] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.policyId));
expect((await computeObservedSpend(db, row)).effectiveHeldCents).toBe("50.0000000");
f.usage("0.150000001");
await f.tick();
expect(await computeObservedSpend(db, row)).toMatchObject({ ledgerCents: "15.0000001", effectiveHeldCents: "0.0000000" });
const [run] = await db.select().from(browserUseRuns).where(eq(browserUseRuns.invocationId, invocation));
expect(run.result).toMatchObject({ paperclipBudgetHold: { state: "settled", accountedCentsExact: "15.0000001" } });
const replay = await f.service.execute(f.binding, f.grant, invocation, "browser_start", { task: "Read" });
expect(JSON.stringify(replay)).not.toContain("paperclipBudgetHold");
expect(await db.select().from(costEvents).where(eq(costEvents.companyId, f.company.id))).toHaveLength(1);
});
it("retains exact cumulative accounting when an older browser row has no hold projection", async () => {
const f = await fixture();
const invocation = randomUUID();
await f.service.execute(f.binding, f.grant, invocation, "browser_start", { task: "Legacy result" });
const [run] = await db.select().from(browserUseRuns).where(eq(browserUseRuns.invocationId, invocation));
await db.update(browserUseRuns).set({ result: null }).where(eq(browserUseRuns.id, run.id));
f.usage("0.150000001"); f.complete(); await f.tick();
await db.update(browserUseRuns).set({ eventsDrained: 0 }).where(eq(browserUseRuns.id, run.id));
f.usage("0.150000002"); await f.tick();
const [total] = await db.select({ cents: sql<string>`sum(${costEvents.costCents})::text` }).from(costEvents).where(eq(costEvents.companyId, f.company.id));
expect(total.cents).toBe("15.0000002");
expect((await db.select().from(browserUseRuns).where(eq(browserUseRuns.id, run.id)))[0].result)
.toMatchObject({ paperclipBudgetHold: { state: "settled", accountedCentsExact: "15.0000002" } });
});
it("caps paid work at the remaining hard budget and enforces reported spend", async () => {
const f = await fixture();
await budgetService(db).upsertPolicy(
+8
View File
@@ -39,6 +39,14 @@ describe("Browser Use v4 transport", () => {
).rejects.toThrow("will not be retried");
expect(request).toHaveBeenCalledTimes(1);
});
it("runs the authoritative dispatch gate before sending a paid request", async () => {
const request = vi.fn();
const gate = vi.fn(async () => { throw new Error("Budget changed before dispatch"); });
await expect(browserUseClient({}, request, randomUUID()).start({ task: "Do not dispatch" }, gate))
.rejects.toThrow("Budget changed before dispatch");
expect(gate).toHaveBeenCalledTimes(1);
expect(request).not.toHaveBeenCalled();
});
it("sanitizes nested and serialized credential addresses and sensitive fields", () => {
const url = "https://live.browser-use.com/watch?token=private";
const data = {
@@ -0,0 +1,251 @@
import { randomUUID } from "node:crypto";
import { and, eq, sql } from "drizzle-orm";
import { afterAll, afterEach, beforeAll, describe, expect, it } from "vitest";
import { agents, budgetPolicies, budgetReservations, companies, costEvents, createDb, heartbeatRuns, projects, type Db } from "@paperclipai/db";
import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "./helpers/embedded-postgres.js";
import { budgetService, computeObservedSpend } from "../services/budgets.js";
import { reserveRunBudget } from "../services/budget-reservations.js";
import { accountRunCost } from "../services/run-cost-accounting.js";
import { persistUsageReceipt, type UsageReceiptEnvelope } from "../services/usage-receipts.js";
import { configureAgentLifecycle } from "../services/agent-lifecycle.js";
const support = await getEmbeddedPostgresTestSupport();
(support.supported ? describe : describe.skip)("scoped live budget enforcement", () => {
let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>>;
let db: Db;
beforeAll(async () => {
database = await startEmbeddedPostgresTestDatabase("budget-live-scoping-");
db = createDb(database.connectionString);
}, 60_000);
afterEach(async () => { await db.transaction(async tx => {
await tx.execute(sql`SET LOCAL client_min_messages = warning`);
await tx.execute(sql`TRUNCATE companies CASCADE`);
}); });
afterAll(async () => { await database?.cleanup(); });
async function fixture(options: { scope?: "company" | "agent" | "project"; amount?: number; reservationCents?: string } = {}) {
const companyId = randomUUID(), agentId = randomUUID(), otherAgentId = randomUUID();
const projectId = randomUUID(), runId = randomUUID(), sourceId = randomUUID();
await db.insert(companies).values({ id: companyId, name: "Budget live", issuePrefix: "BLIVE" });
await db.insert(agents).values([
{ id: agentId, companyId, name: "Running", role: "engineer", status: "active", adapterType: "process" },
{ id: otherAgentId, companyId, name: "Unrelated idle", role: "engineer", status: "idle", adapterType: "process" },
]);
await db.insert(projects).values({ id: projectId, companyId, name: "Scoped project", status: "in_progress" });
await db.insert(heartbeatRuns).values({ id: runId, companyId, agentId, invocationSource: "on_demand", status: "running",
costAccountingPending: true, usageJson: { accountingReceiptSourceId: sourceId, accountingReceiptSequence: 0,
accountingReceiptReady: false, ledgerScope: { projectId } } });
const scopeType = options.scope ?? "agent";
const [policy] = await db.insert(budgetPolicies).values({ companyId, scopeType,
scopeId: scopeType === "company" ? companyId : scopeType === "project" ? projectId : agentId,
metric: "billed_cents", windowKind: "lifetime", amount: options.amount ?? 100,
reservationCents: options.reservationCents ?? "0", notifyEnabled: false }).returning();
const readRun = async () => (await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId)))[0];
const exposure = () => computeObservedSpend(db, policy);
const receipt = (sequence: number, costUsd?: number, extra: Partial<UsageReceiptEnvelope["receipt"]> = {}): UsageReceiptEnvelope => ({
schema: "paperclip/accounting-receipt/v1", id: randomUUID(), companyId, runId, sourceId, sequence,
receivedAt: new Date().toISOString(), adapterType: "process", receipt: {
complete: false, usageBasis: "per_run", provider: "test", biller: "test", model: "budget-fixture",
billingType: "metered_api", costStatus: "unpriced", ...(costUsd === undefined ? {} : { costUsd }), ...extra,
},
});
const queuedRun = async (id = randomUUID()) => {
await db.insert(heartbeatRuns).values({ id, companyId, agentId, invocationSource: "on_demand", status: "queued" });
return id;
};
return { companyId, agentId, otherAgentId, projectId, runId, sourceId, policy, readRun, exposure, receipt, queuedRun };
}
it("blocks a known streaming overrun without charging checkpoints or touching unrelated agents", async () => {
const f = await fixture();
const [before] = await db.select().from(agents).where(eq(agents.id, f.otherAgentId));
await persistUsageReceipt(db, f.receipt(1, 0.6));
expect((await f.exposure()).totalExact).toBe("60.0000000");
expect(await budgetService(db).getInvocationBlock(f.companyId, f.agentId, { projectId: f.projectId })).toBeNull();
await persistUsageReceipt(db, f.receipt(2, 1.1));
const [agent] = await db.select().from(agents).where(eq(agents.id, f.agentId));
const [policy] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, f.policy.id));
const [other] = await db.select().from(agents).where(eq(agents.id, f.otherAgentId));
expect(agent.lifecycleHolds).toContain("budget");
expect(policy.enforcementVersion).toBeGreaterThan(policy.enforcementDeliveredVersion);
expect(other).toEqual(before);
expect(await db.select().from(costEvents)).toHaveLength(0);
expect(await budgetService(db).getInvocationBlock(f.companyId, f.agentId, { projectId: f.projectId })).not.toBeNull();
expect((await f.exposure()).liveObservedCents).toBe("110.0000000");
});
it("replaces cumulative observations and preserves known exposure across missing, stale, and lower partial costs", async () => {
const f = await fixture({ amount: 1000 });
const first = f.receipt(1, 0.4);
await persistUsageReceipt(db, first);
await persistUsageReceipt(db, f.receipt(2, 0.7));
await persistUsageReceipt(db, first);
await persistUsageReceipt(db, f.receipt(3));
await persistUsageReceipt(db, { ...f.receipt(1, 20), sourceId: randomUUID() });
await persistUsageReceipt(db, f.receipt(4, 0.5));
expect((await f.exposure()).liveObservedCents).toBe("70.0000000");
expect((await f.readRun()).usageJson?.budgetExposure).toMatchObject({ sourceId: f.sourceId, sequence: 4, costCents: "70.0000000" });
expect(await db.select().from(costEvents)).toHaveLength(0);
});
it("counts the maximum of a reservation and live spend once at final dispatch", async () => {
const f = await fixture({ reservationCents: "40" });
await db.insert(budgetReservations).values({ companyId: f.companyId, runId: f.runId, agentId: f.agentId,
projectId: f.projectId, amountCents: "50", providerStartedAt: new Date() });
await persistUsageReceipt(db, f.receipt(1, 0.6));
expect((await f.exposure()).effectiveHeldCents).toBe("60.0000000");
const nextRunId = await f.queuedRun();
expect((await reserveRunBudget(db, f.companyId, nextRunId, f.projectId)).amountCents).toBe("40.0000000");
expect((await f.exposure()).effectiveHeldCents).toBe("100.0000000");
expect(await budgetService(db).getInvocationBlock(f.companyId, f.agentId, { projectId: f.projectId })).toBeNull();
await expect(reserveRunBudget(db, f.companyId, await f.queuedRun(), f.projectId)).rejects.toThrow("reserved");
});
it("serializes competing dispatch reservations without overspending available capacity", async () => {
const f = await fixture({ reservationCents: "60" });
const first = await f.queuedRun(), second = await f.queuedRun();
const outcomes = await Promise.allSettled([
reserveRunBudget(db, f.companyId, first, f.projectId), reserveRunBudget(db, f.companyId, second, f.projectId),
]);
expect(outcomes.filter(result => result.status === "fulfilled")).toHaveLength(1);
expect(outcomes.filter(result => result.status === "rejected")).toHaveLength(1);
expect((await f.exposure()).effectiveHeldCents).toBe("60.0000000");
});
it("commits a budget release while dispatch waits for asynchronous lifecycle readiness", async () => {
const f = await fixture();
const worker = configureAgentLifecycle(db, {
requiredPluginIds: async () => [], runHost: async () => "complete", runPlugin: async () => "complete",
});
try {
await persistUsageReceipt(db, f.receipt(1, 1.1));
await worker.process(f.agentId);
expect((await db.select().from(agents).where(eq(agents.id, f.agentId)))[0].lifecycleState).toBe("paused");
await db.update(budgetPolicies).set({ amount: 200 }).where(eq(budgetPolicies.id, f.policy.id));
const nextRunId = await f.queuedRun();
await expect(reserveRunBudget(db, f.companyId, nextRunId, f.projectId)).rejects.toThrow("lifecycle");
const [resuming] = await db.select().from(agents).where(eq(agents.id, f.agentId));
expect(resuming.lifecycleHolds).not.toContain("budget");
expect(resuming.lifecycleState).toBe("resuming");
expect(await db.select().from(budgetReservations).where(eq(budgetReservations.runId, nextRunId))).toHaveLength(0);
await worker.process(f.agentId);
expect((await db.select().from(agents).where(eq(agents.id, f.agentId)))[0].lifecycleState).toBe("ready");
await expect(reserveRunBudget(db, f.companyId, nextRunId, f.projectId)).resolves.toMatchObject({ reused: false });
} finally { await worker.stop(); }
});
it("atomically replaces live exposure and its hold with one final ledger charge", async () => {
const f = await fixture({ amount: 1000 });
await db.insert(budgetReservations).values({ companyId: f.companyId, runId: f.runId, agentId: f.agentId,
projectId: f.projectId, amountCents: "50", providerStartedAt: new Date() });
await persistUsageReceipt(db, f.receipt(1, 0.6));
await persistUsageReceipt(db, f.receipt(2, 0.8, { complete: true, costStatus: "reported" }));
await db.update(heartbeatRuns).set({ status: "succeeded", finishedAt: new Date() }).where(eq(heartbeatRuns.id, f.runId));
expect(await accountRunCost(db, f.runId)).toBe(true);
expect(await accountRunCost(db, f.runId)).toBe(false);
expect(await f.exposure()).toMatchObject({ ledgerCents: "80.0000000", liveObservedCents: "0.0000000",
effectiveHeldCents: "0.0000000", totalExact: "80.0000000", pendingTerminalRunCount: 0 });
expect(await db.select().from(costEvents).where(eq(costEvents.companyId, f.companyId))).toHaveLength(1);
const [hold] = await db.select().from(budgetReservations).where(eq(budgetReservations.runId, f.runId));
expect(hold.state).toBe("settled");
});
it("keeps an over-budget hold stable while replacing live usage with a multipart final receipt", async () => {
const f = await fixture();
await persistUsageReceipt(db, f.receipt(1, 1.1));
await persistUsageReceipt(db, f.receipt(2, 1.2, { complete: true, costStatus: "reported", usage: { inputTokens: 2, outputTokens: 2 },
usageByModel: [
{ model: "first", costUsd: 0.3, usage: { inputTokens: 1, outputTokens: 1 } },
{ model: "second", costUsd: 0.9, usage: { inputTokens: 1, outputTokens: 1 } },
] }));
await db.update(heartbeatRuns).set({ status: "succeeded", finishedAt: new Date() }).where(eq(heartbeatRuns.id, f.runId));
const [before] = await db.select().from(agents).where(eq(agents.id, f.agentId));
expect(before.lifecycleHolds).toContain("budget");
await accountRunCost(db, f.runId);
const [after] = await db.select().from(agents).where(eq(agents.id, f.agentId));
expect(after.lifecycleHolds).toContain("budget");
expect(after.lifecycleVersion).toBe(before.lifecycleVersion);
expect(after.lifecycleOperation).toEqual(before.lifecycleOperation);
expect(await f.exposure()).toMatchObject({ ledgerCents: "120.0000000", liveObservedCents: "0.0000000" });
expect(await db.select().from(costEvents).where(eq(costEvents.companyId, f.companyId))).toHaveLength(2);
});
it("recognizes legacy active-run usage without a projection and excludes subscription charges", async () => {
const f = await fixture();
const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, f.runId));
await db.update(heartbeatRuns).set({ usageJson: { ...run.usageJson, costUsd: 0.75, usageSource: "per_run", billingType: "metered_api" } })
.where(eq(heartbeatRuns.id, f.runId));
expect((await f.exposure()).liveObservedCents).toBe("75.0000000");
await persistUsageReceipt(db, f.receipt(1, 10, { billingType: "subscription_included" }));
expect((await f.exposure()).liveObservedCents).toBe("0.0000000");
expect(await budgetService(db).getInvocationBlock(f.companyId, f.agentId, { projectId: f.projectId })).toBeNull();
});
it("retains established per-run exposure when a later receipt has an unreconciled session basis", async () => {
const f = await fixture();
await persistUsageReceipt(db, f.receipt(1, 0.7));
await persistUsageReceipt(db, f.receipt(2, 99, { usageBasis: "session_cumulative" }));
expect((await f.exposure()).liveObservedCents).toBe("70.0000000");
expect((await f.readRun()).usageJson?.budgetExposure).toMatchObject({ costCents: "70.0000000" });
});
it("ignores out-of-range legacy numeric text without losing a usable fallback price", async () => {
const f = await fixture();
const run = await f.readRun();
await db.update(heartbeatRuns).set({ usageJson: { ...run.usageJson, usageSource: "per_run", billingType: "metered_api",
cacheAdjustedCostUsd: "1e999", costUsdExact: "1e999", costUsd: 0.25 } }).where(eq(heartbeatRuns.id, f.runId));
expect((await f.exposure()).liveObservedCents).toBe("25.0000000");
});
it("rounds each legacy run to nanodollars before aggregating its exposure", async () => {
const f = await fixture();
const secondRunId = randomUUID();
const usage = { usageSource: "per_run", billingType: "metered_api", costUsdExact: "0.0000000004" };
await db.update(heartbeatRuns).set({ usageJson: usage }).where(eq(heartbeatRuns.id, f.runId));
await db.insert(heartbeatRuns).values({ id: secondRunId, companyId: f.companyId, agentId: f.agentId,
invocationSource: "on_demand", status: "running", costAccountingPending: true, usageJson: usage });
expect((await f.exposure()).liveObservedCents).toBe("0.0000000");
await db.update(heartbeatRuns).set({ usageJson: { ...usage, costUsdExact: "0.0000000005" } })
.where(eq(heartbeatRuns.companyId, f.companyId));
expect((await f.exposure()).liveObservedCents).toBe("0.0000002");
});
it("limits project enforcement to that project while the same agent can serve another", async () => {
const f = await fixture({ scope: "project" });
const [otherProject] = await db.insert(projects).values({ companyId: f.companyId, name: "Other project", status: "in_progress" }).returning();
await persistUsageReceipt(db, f.receipt(1, 1.1));
const [agent] = await db.select().from(agents).where(eq(agents.id, f.agentId));
const [project] = await db.select().from(projects).where(eq(projects.id, f.projectId));
expect(project.pauseReason).toBe("budget");
expect(agent.lifecycleHolds).not.toContain("budget");
expect(await budgetService(db).getInvocationBlock(f.companyId, f.agentId, { projectId: f.projectId })).not.toBeNull();
expect(await budgetService(db).getInvocationBlock(f.companyId, f.agentId, { projectId: otherProject.id })).toBeNull();
await expect(reserveRunBudget(db, f.companyId, await f.queuedRun(), otherProject.id)).resolves.toMatchObject({ reused: false });
});
it("keeps missing-price terminal work held until accounting evidence arrives", async () => {
const f = await fixture();
await db.update(heartbeatRuns).set({ status: "cancelled", finishedAt: new Date() }).where(eq(heartbeatRuns.id, f.runId));
expect(await budgetService(db).getInvocationBlock(f.companyId, f.agentId, { projectId: f.projectId })).not.toBeNull();
expect(await f.exposure()).toMatchObject({ pendingTerminalRunCount: 1, ledgerCents: "0.0000000", liveObservedCents: "0.0000000" });
expect(await accountRunCost(db, f.runId)).toBe(false);
expect((await f.readRun()).costAccountingPending).toBe(true);
expect(await db.select().from(costEvents).where(and(eq(costEvents.companyId, f.companyId), eq(costEvents.heartbeatRunId, f.runId)))).toHaveLength(0);
});
it.each(["pre-provider release", "stale pending marker"])("releases accounting-induced holds after %s", async kind => {
const f = await fixture();
await db.update(heartbeatRuns).set({ status: "cancelled", finishedAt: new Date() }).where(eq(heartbeatRuns.id, f.runId));
expect(await budgetService(db).getInvocationBlock(f.companyId, f.agentId, { projectId: f.projectId })).not.toBeNull();
const [held] = await db.select().from(agents).where(eq(agents.id, f.agentId));
expect(held.lifecycleHolds).toContain("budget");
await db.update(heartbeatRuns).set(kind === "pre-provider release"
? { resultJson: { executionRecovery: { providerWorkStarted: false } } }
: { costAccountedAt: new Date() }).where(eq(heartbeatRuns.id, f.runId));
expect(await accountRunCost(db, f.runId)).toBe(kind === "pre-provider release");
const [released] = await db.select().from(agents).where(eq(agents.id, f.agentId));
expect(released.lifecycleHolds).not.toContain("budget");
expect((await f.readRun()).costAccountingPending).toBe(false);
expect(await db.select().from(costEvents).where(eq(costEvents.companyId, f.companyId))).toHaveLength(0);
});
});
@@ -0,0 +1,140 @@
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import { installDatabaseWorkSignals } from "../../../packages/db/src/work-signals.js";
import type { Db } from "@paperclipai/db";
import { createDeliveryWorkCoordinator } from "../services/delivery-work-coordinator.js";
import { DELIVERY_QUEUES, notifyDeliveryWork } from "../services/delivery-work-notifications.js";
import { publishLiveEvent } from "../services/live-events.js";
import { registerBudgetEnforcementWork, registerCostAccountingWork, wakeCostAccountingWork } from "../services/budget-enforcement-worker.js";
const work = vi.hoisted(() => ({
deliver: vi.fn(), pending: vi.fn(), next: vi.fn(), reconcile: vi.fn(), accounting: vi.fn(),
}));
vi.mock("../services/budgets.js", () => ({
deliverBudgetEnforcement: work.deliver,
hasPendingBudgetEnforcement: work.pending,
nextBudgetReconciliationAt: work.next,
reconcileDueBudgetPolicies: work.reconcile,
}));
vi.mock("../services/run-cost-accounting.js", () => ({ readCostAccountingWork: work.accounting }));
const coordinators: ReturnType<typeof createDeliveryWorkCoordinator>[] = [];
const subscriptions: (() => void)[] = [];
beforeEach(() => {
vi.useFakeTimers();
vi.resetAllMocks();
work.deliver.mockResolvedValue(undefined);
work.pending.mockResolvedValue(false);
work.next.mockResolvedValue(null);
work.reconcile.mockResolvedValue({ nextCursor: null });
work.accounting.mockResolvedValue({ pending: false, nextRunAt: null });
});
afterEach(async () => {
for (const stop of subscriptions.splice(0)) stop();
for (const c of coordinators.splice(0)) await c.stop();
vi.useRealTimers();
});
function setup() {
const execute = vi.fn(async () => [{ xid: "42", status: "committed" }]);
const owner = installDatabaseWorkSignals({ execute, async transaction(callback: (tx: unknown) => Promise<unknown>) {
return callback({ execute, transaction: owner.transaction });
} });
const onError = vi.fn();
const coordinator = createDeliveryWorkCoordinator({ owner, canRun: () => true, onError });
coordinators.push(coordinator);
return { owner, db: owner as unknown as Db, coordinator, onError };
}
describe("budget accounting delivery", () => {
it("recovers at startup, then leaves unchanged budget history completely quiet", async () => {
const s = setup();
await registerBudgetEnforcementWork(s.db, s.coordinator, {}).ready;
expect(work.reconcile).toHaveBeenCalledWith(s.db, { startup: true });
expect(work.deliver).toHaveBeenCalledOnce();
expect(s.coordinator.nextWakeAt()).toBeNull();
await vi.advanceTimersByTimeAsync(3_600_000);
expect(work.pending).toHaveBeenCalledOnce();
expect(work.reconcile).toHaveBeenCalledOnce();
await s.owner.transaction(tx => notifyDeliveryWork(tx as object, DELIVERY_QUEUES.budgetEnforcement));
await vi.advanceTimersByTimeAsync(1);
expect(work.reconcile).toHaveBeenCalledOnce();
expect(work.deliver).toHaveBeenCalledTimes(2);
expect(s.coordinator.nextWakeAt()).toBeNull();
});
it("recovers a concrete policy-window deadline without polling before it", async () => {
const s = setup();
const due = Date.now() + 60_000;
work.next.mockResolvedValue(new Date(due));
work.reconcile.mockImplementation(async () => {
if (Date.now() >= due) work.next.mockResolvedValue(null);
return { nextCursor: null };
});
await registerBudgetEnforcementWork(s.db, s.coordinator, {}).ready;
await vi.advanceTimersByTimeAsync(59_999);
expect(work.reconcile).toHaveBeenCalledOnce();
expect(s.coordinator.nextWakeAt()).toBe(due);
await vi.advanceTimersByTimeAsync(1);
expect(work.reconcile).toHaveBeenCalledTimes(2);
expect(work.reconcile).toHaveBeenLastCalledWith(s.db, { startup: false });
expect(s.coordinator.nextWakeAt()).toBeNull();
});
it("retries interrupted delivery and disarms after acknowledgement", async () => {
const s = setup();
work.deliver.mockRejectedValueOnce(new Error("Provider stop unavailable"));
await registerBudgetEnforcementWork(s.db, s.coordinator, {}).ready;
expect(s.onError).toHaveBeenCalledOnce();
await vi.advanceTimersByTimeAsync(1000);
expect(work.deliver).toHaveBeenCalledTimes(2);
expect(s.coordinator.nextWakeAt()).toBeNull();
});
it("continues bounded startup pages without forgetting the cursor after notifications", async () => {
const s = setup();
work.reconcile.mockResolvedValueOnce({ nextCursor: "company:agent:scope" });
await registerBudgetEnforcementWork(s.db, s.coordinator, {}).ready;
expect(work.reconcile).toHaveBeenLastCalledWith(s.db, { startup: true });
await s.owner.transaction(tx => notifyDeliveryWork(tx as object, DELIVERY_QUEUES.budgetEnforcement));
await vi.advanceTimersByTimeAsync(1);
expect(work.reconcile).toHaveBeenLastCalledWith(s.db, { startup: true, cursor: "company:agent:scope" });
expect(s.coordinator.nextWakeAt()).toBeNull();
});
it("waits on an interrupted-service deadline and wakes terminal accounting from committed events", async () => {
const s = setup();
const recover = vi.fn(async () => {});
const due = Date.now() + 120_001;
work.accounting.mockResolvedValue({ pending: false, nextRunAt: due });
const worker = registerCostAccountingWork(s.db, s.coordinator, recover);
subscriptions.push(worker.stop);
await worker.ready;
await vi.advanceTimersByTimeAsync(60_000);
expect(recover).toHaveBeenCalledOnce();
publishLiveEvent({ companyId: "company", type: "heartbeat.run.status", payload: { status: "succeeded" } });
await vi.advanceTimersByTimeAsync(1);
expect(recover).toHaveBeenCalledTimes(2);
work.accounting.mockResolvedValue({ pending: false, nextRunAt: null });
await vi.advanceTimersByTimeAsync(due - Date.now());
expect(recover).toHaveBeenCalledTimes(3);
expect(s.coordinator.nextWakeAt()).toBeNull();
await vi.advanceTimersByTimeAsync(3_600_000);
expect(work.accounting).toHaveBeenCalledTimes(3);
});
it("replays durable spool work after a pre-transaction failure and releases its local wake on stop", async () => {
const s = setup();
const recover = vi.fn(async () => {});
const worker = registerCostAccountingWork(s.db, s.coordinator, recover);
subscriptions.push(worker.stop);
await worker.ready;
wakeCostAccountingWork(s.db);
wakeCostAccountingWork(s.db);
await vi.advanceTimersByTimeAsync(1);
expect(recover).toHaveBeenCalledTimes(2);
worker.stop();
wakeCostAccountingWork(s.db);
await vi.advanceTimersByTimeAsync(1);
expect(recover).toHaveBeenCalledTimes(2);
expect(s.coordinator.nextWakeAt()).toBeNull();
});
});
@@ -0,0 +1,121 @@
import { randomUUID } from "node:crypto";
import { promises as fs } from "node:fs";
import os from "node:os";
import path from "node:path";
import { eq, sql } from "drizzle-orm";
import { afterAll, afterEach, beforeAll, describe, expect, it } from "vitest";
import { agents, budgetPolicies, budgetReservations, companies, costEvents, createDb, heartbeatRuns, type Db } from "@paperclipai/db";
import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "./helpers/embedded-postgres.js";
import { computeObservedSpend } from "../services/budgets.js";
import { accountRunCost } from "../services/run-cost-accounting.js";
import { createRunUsageRecorder, replayUsageReceipts, type UsageReceiptEnvelope } from "../services/usage-receipts.js";
const support = await getEmbeddedPostgresTestSupport();
(support.supported ? describe : describe.skip)("live budget exposure across recorder attempts", () => {
let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>>;
let db: Db;
const directories: string[] = [];
beforeAll(async () => {
database = await startEmbeddedPostgresTestDatabase("budget-live-attempts-");
db = createDb(database.connectionString);
}, 60_000);
afterEach(async () => {
await db.execute(sql`TRUNCATE companies CASCADE`);
for (const directory of directories.splice(0)) await fs.rm(directory, { force: true, recursive: true });
});
afterAll(async () => { await database?.cleanup(); });
async function fixture(rejectedZero = true) {
const companyId = randomUUID(), agentId = randomUUID(), runId = randomUUID();
await db.insert(companies).values({ id: companyId, name: "Live attempts", issuePrefix: "LAT" });
await db.insert(agents).values({ id: agentId, companyId, name: "Provider", adapterType: "codex_local" });
await db.insert(heartbeatRuns).values({ id: runId, companyId, agentId, status: "running" });
const [policy] = await db.insert(budgetPolicies).values({ companyId, scopeType: "agent", scopeId: agentId,
metric: "billed_cents", windowKind: "lifetime", amount: 100, notifyEnabled: false }).returning();
const directory = await fs.mkdtemp(path.join(os.tmpdir(), "budget-live-attempts-"));
directories.push(directory);
let offline = false;
const proxy = new Proxy(db, { get(target, property) {
if (property === "transaction" && offline) return () => Promise.reject(new Error("Database unavailable"));
const value = Reflect.get(target, property);
return typeof value === "function" ? value.bind(target) : value;
} }) as Db;
const recorder = await createRunUsageRecorder(proxy, { companyId, runId, adapterType: "codex_local" }, directory);
const identity = { usageBasis: "per_run" as const, provider: "openai", biller: "openai",
billingType: "metered_api" as const, model: "gpt-6.1-sol" };
if (rejectedZero) await recorder.capture({ ...identity, attemptId: randomUUID(), complete: true,
usage: { inputTokens: 0, outputTokens: 0 }, costUsd: 0, costStatus: "reported" });
const paidAttemptId = randomUUID();
const partial = { ...identity, attemptId: paidAttemptId, complete: false,
usage: { inputTokens: 1000, outputTokens: 100000 }, costStatus: "unpriced" as const };
const readRun = async () => (await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId)))[0];
return { companyId, agentId, runId, policy, directory, proxy, recorder, identity, partial, readRun,
offline: (value: boolean) => { offline = value; } };
}
it("enforces a token-priced retry while preserving final receipt uncertainty and one ledger charge", async () => {
const f = await fixture();
await f.recorder.capture(f.partial);
expect(await computeObservedSpend(db, f.policy)).toMatchObject({ liveObservedCents: "100.2000000" });
const partial = await f.readRun();
expect(partial.usageJson).toMatchObject({ costUsdExact: "0.000000000", costStatus: "unpriced",
accountingReceiptReady: false, budgetExposure: { costCents: "100.2000000" } });
expect(await db.select().from(costEvents)).toHaveLength(0);
expect((await db.select().from(agents).where(eq(agents.id, f.agentId)))[0].lifecycleHolds).toContain("budget");
await f.recorder.capture({ ...f.partial, complete: true, costUsd: 0.8, costStatus: "reported" });
expect((await f.readRun()).usageJson?.accountingReceiptReady).toBe(true);
await db.update(heartbeatRuns).set({ status: "succeeded", finishedAt: new Date() }).where(eq(heartbeatRuns.id, f.runId));
expect(await accountRunCost(db, f.runId)).toBe(true);
expect(await computeObservedSpend(db, f.policy)).toMatchObject({ liveObservedCents: "0.0000000", ledgerCents: "80.0000000" });
expect(await db.select().from(costEvents)).toHaveLength(1);
});
it("replays the fsynced per-attempt exposure after a failed database write", async () => {
const f = await fixture();
f.offline(true);
await f.recorder.capture(f.partial);
const files = (await fs.readdir(f.directory)).filter(file => file.endsWith(".json"));
expect(files).toHaveLength(1);
const saved = JSON.parse(await fs.readFile(path.join(f.directory, files[0]), "utf8")) as UsageReceiptEnvelope;
expect(saved).toMatchObject({ liveCostCents: "100.2000000", receipt: { complete: false, costStatus: "unpriced", costUsdExact: "0.000000000" } });
expect((await computeObservedSpend(db, f.policy)).liveObservedCents).toBe("0.0000000");
f.offline(false);
await replayUsageReceipts(f.proxy, f.directory);
expect((await computeObservedSpend(db, f.policy)).liveObservedCents).toBe("100.2000000");
expect((await db.select().from(agents).where(eq(agents.id, f.agentId)))[0].lifecycleHolds).toContain("budget");
expect((await fs.readdir(f.directory)).filter(file => file.endsWith(".json"))).toEqual([]);
});
it("uses normalized per-run exposure instead of cumulative USD retained in aggregate compatibility fields", async () => {
const f = await fixture();
await f.recorder.capture({ ...f.identity, attemptId: randomUUID(), usageBasis: "session_cumulative",
complete: true, costUsd: 100, usage: { inputTokens: 100, outputTokens: 100 } });
await f.recorder.capture({ ...f.identity, attemptId: randomUUID(), complete: false,
costUsd: 0.2, usage: { inputTokens: 20, outputTokens: 20 } });
expect((await f.readRun()).usageJson).toMatchObject({ costUsdExact: "100.200000000",
usageSource: "per_run", budgetExposure: { costCents: "20.0000000" } });
expect((await computeObservedSpend(db, f.policy)).liveObservedCents).toBe("20.0000000");
expect((await db.select().from(agents).where(eq(agents.id, f.agentId)))[0].lifecycleHolds).not.toContain("budget");
});
it("does not resurrect cumulative USD when the normalized per-run attempt remains unpriced", async () => {
const f = await fixture(false);
await db.insert(budgetReservations).values({ companyId: f.companyId, agentId: f.agentId, runId: f.runId, amountCents: "10" });
await f.recorder.capture({ ...f.identity, attemptId: randomUUID(), usageBasis: "session_cumulative",
complete: true, costUsd: 100, usage: { inputTokens: 100, outputTokens: 100 } });
const unknownAttemptId = randomUUID();
await f.recorder.capture({ ...f.identity, attemptId: unknownAttemptId, model: "unsupported", complete: false,
usage: { inputTokens: 20, outputTokens: 20 } });
await f.recorder.capture({ ...f.identity, attemptId: unknownAttemptId, model: "unsupported", complete: false,
usage: { inputTokens: 30, outputTokens: 30 } });
expect((await f.readRun()).usageJson).toMatchObject({ costUsdExact: "100.000000000",
costStatus: "unpriced", accountingReceiptReady: false, usageSource: "per_run", budgetExposure: { costCents: null } });
expect(await computeObservedSpend(db, f.policy)).toMatchObject({ liveObservedCents: "0.0000000", effectiveHeldCents: "10.0000000" });
await db.update(heartbeatRuns).set({ status: "cancelled", finishedAt: new Date() }).where(eq(heartbeatRuns.id, f.runId));
expect(await accountRunCost(db, f.runId)).toBe(false);
expect((await db.select().from(budgetReservations).where(eq(budgetReservations.runId, f.runId)))[0].state).toBe("held");
expect((await f.readRun()).costAccountingPending).toBe(true);
expect(await db.select().from(costEvents)).toEqual([]);
});
});
+60 -8
View File
@@ -11,7 +11,7 @@ import {
createDb,
projects,
} from "@paperclipai/db";
import { budgetService } from "../services/budgets.ts";
import { budgetService, deliverBudgetEnforcement, hasPendingBudgetEnforcement } from "../services/budgets.ts";
import {
getEmbeddedPostgresTestSupport,
startEmbeddedPostgresTestDatabase,
@@ -113,7 +113,7 @@ describeEmbeddedPostgres("budgetService release gate enforcement", () => {
return event!;
}
it("raises one soft incident per window before hard-stopping and safely logging agent telemetry", async () => {
it("raises one soft incident per window before hard-stopping and preserving safe local audit details", async () => {
const { companyId, agentId } = await createBudgetFixture();
const cancelWorkForScope = vi.fn().mockResolvedValue(undefined);
const service = budgetService(db, { cancelWorkForScope });
@@ -188,7 +188,12 @@ describeEmbeddedPostgres("budgetService release gate enforcement", () => {
.from(agents);
expect(agentAfterHardStop).toMatchObject({ status: "paused", pauseReason: "budget" });
expect(agentAfterHardStop?.pausedAt).toBeInstanceOf(Date);
expect(cancelWorkForScope).toHaveBeenCalledTimes(2);
expect(cancelWorkForScope).not.toHaveBeenCalled();
expect(await hasPendingBudgetEnforcement(db)).toBe(true);
const [pendingPolicy] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy!.id));
expect(pendingPolicy.enforcementVersion).toBe(1);
await deliverBudgetEnforcement(db, { cancelWorkForScope }, companyId);
expect(cancelWorkForScope).toHaveBeenCalledTimes(1);
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);
@@ -199,10 +204,14 @@ describeEmbeddedPostgres("budgetService release gate enforcement", () => {
reason: "Agent is paused because its budget hard-stop was reached.",
});
const telemetryCalls = mockLogActivity.mock.calls.map(([, input]) => input);
expect(telemetryCalls.filter((call) => call.action === "budget.soft_threshold_crossed")).toHaveLength(1);
expect(telemetryCalls.filter((call) => call.action === "budget.hard_threshold_crossed")).toHaveLength(1);
expect(telemetryCalls).toEqual(
const [afterAdmission] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy!.id));
expect(afterAdmission.enforcementVersion).toBe(pendingPolicy.enforcementVersion);
await deliverBudgetEnforcement(db, { cancelWorkForScope }, companyId);
expect(cancelWorkForScope).toHaveBeenCalledTimes(1);
const activityCalls = mockLogActivity.mock.calls.map(([, input]) => input);
expect(activityCalls.filter((call) => call.action === "budget.soft_threshold_crossed")).toHaveLength(1);
expect(activityCalls.filter((call) => call.action === "budget.hard_threshold_crossed")).toHaveLength(1);
expect(activityCalls).toEqual(
expect.arrayContaining([
expect.objectContaining({
action: "budget.soft_threshold_crossed",
@@ -227,7 +236,7 @@ describeEmbeddedPostgres("budgetService release gate enforcement", () => {
}),
]),
);
for (const call of telemetryCalls) {
for (const call of activityCalls) {
expect(JSON.stringify(call.details)).not.toContain("SECRET_TOKEN_SHOULD_NOT_LEAK");
expect(call.details).not.toHaveProperty("prompt");
expect(call.details).not.toHaveProperty("message");
@@ -291,6 +300,11 @@ describeEmbeddedPostgres("budgetService release gate enforcement", () => {
.from(projects);
expect(projectAfterHardStop?.pauseReason).toBe("budget");
expect(projectAfterHardStop?.pausedAt).toBeInstanceOf(Date);
expect(cancelWorkForScope).not.toHaveBeenCalled();
expect((await db.select().from(agents).where(eq(agents.id, agentId)))[0]).toMatchObject({ status: "active", lifecycleHolds: [] });
expect(await service.getInvocationBlock(companyId, agentId)).toBeNull();
await deliverBudgetEnforcement(db, { cancelWorkForScope }, companyId);
expect(cancelWorkForScope).toHaveBeenCalledTimes(1);
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);
@@ -355,4 +369,42 @@ describeEmbeddedPostgres("budgetService release gate enforcement", () => {
});
expect(overviewAfterResume.activeIncidents).toHaveLength(0);
});
it("discards stale undelivered enforcement when a budget is raised before dispatch", async () => {
const { companyId, agentId, projectId } = await createBudgetFixture();
const cancelWorkForScope = vi.fn().mockResolvedValue(undefined);
const service = budgetService(db, { cancelWorkForScope });
await service.upsertPolicy(companyId, { scopeType: "project", scopeId: projectId,
amount: 100, windowKind: "lifetime" }, "board");
await service.evaluateCostEvent(await insertCostEvent({ companyId, agentId, projectId, costCents: 125 }));
expect(await hasPendingBudgetEnforcement(db)).toBe(true);
expect(cancelWorkForScope).not.toHaveBeenCalled();
await service.upsertPolicy(companyId, { scopeType: "project", scopeId: projectId,
amount: 200, windowKind: "lifetime" }, "board");
await deliverBudgetEnforcement(db, { cancelWorkForScope }, companyId);
expect(cancelWorkForScope).not.toHaveBeenCalled();
expect(await hasPendingBudgetEnforcement(db)).toBe(false);
expect(await service.getInvocationBlock(companyId, agentId, { projectId })).toBeNull();
});
it("retains failed cancellation delivery and retries without replaying accounting", async () => {
const { companyId, agentId, projectId } = await createBudgetFixture();
const service = budgetService(db);
await service.upsertPolicy(companyId, { scopeType: "project", scopeId: projectId,
amount: 100, windowKind: "lifetime" }, "board");
const event = await insertCostEvent({ companyId, agentId, projectId, costCents: 125 });
await service.evaluateCostEvent(event);
const cancelWorkForScope = vi.fn().mockRejectedValueOnce(new Error("Controlled stop delivery failure")).mockResolvedValue(undefined);
await deliverBudgetEnforcement(db, { cancelWorkForScope }, companyId);
expect(await hasPendingBudgetEnforcement(db)).toBe(true);
expect((await db.select().from(projects).where(eq(projects.id, projectId)))[0].pauseReason).toBe("budget");
const [afterFailure] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.companyId, companyId));
expect(afterFailure.enforcementVersion).toBeGreaterThan(afterFailure.enforcementDeliveredVersion);
await deliverBudgetEnforcement(db, { cancelWorkForScope }, companyId);
expect(cancelWorkForScope).toHaveBeenCalledTimes(2);
expect(await hasPendingBudgetEnforcement(db)).toBe(false);
expect(await db.select().from(costEvents).where(eq(costEvents.companyId, companyId))).toHaveLength(1);
expect((await db.select().from(budgetIncidents)).filter(incident => incident.thresholdType === "hard")).toHaveLength(1);
});
});
+53 -9
View File
@@ -33,7 +33,9 @@ vi.mock("../services/cloud-lifecycle-sync.js", () => ({
notifyCloudOfPrimaryCompanyLifecycleChange: notifyCloudSpy,
}));
import { configureAgentLifecycle } from "../services/agent-lifecycle.js";
import { startAgentLifecycle } from "../services/agent-lifecycle.js";
import { createLifecycleDriver } from "../services/agent-lifecycle-driver.js";
import { reserveRunBudget } from "../services/budget-reservations.js";
import { companyService } from "../services/companies.js";
import { deriveIssuePrefixBase } from "../services/issue-prefix.js";
import { readBuiltInAgentMarker } from "../lib/built-in-agent-metadata.js";
@@ -50,8 +52,11 @@ if (!embeddedPostgresSupport.supported) {
describeEmbeddedPostgres("companyService", () => {
let db!: ReturnType<typeof createDb>;
let lifecycleWorker: ReturnType<typeof configureAgentLifecycle>;
beforeEach(() => { lifecycleWorker = configureAgentLifecycle(db, { requiredPluginIds: async () => [], runPlugin: async () => "complete", runHost: async () => "complete" }); });
let lifecycleWorker: ReturnType<typeof startAgentLifecycle>;
beforeEach(() => {
lifecycleWorker = startAgentLifecycle(db, { ...createLifecycleDriver(db, {} as never),
requiredPluginIds: async () => [], runPlugin: async () => "complete" }, () => true);
});
async function insertAgents(input: typeof agents.$inferInsert | Array<typeof agents.$inferInsert>) {
const rows = Array.isArray(input) ? input : [input];
await db.insert(agents).values(rows.map(row => ({ ...row,
@@ -107,6 +112,17 @@ describeEmbeddedPostgres("companyService", () => {
expect(rows.map((row) => row.issuePrefix).sort()).toEqual(["ARO", "AROA"]);
});
it("leaves unrelated company-hold projection out of name edits", async () => {
const [company] = await db.insert(companies).values({ name: "Rename only", issuePrefix: "REN" }).returning();
const [agent] = await db.insert(agents).values({ companyId: company.id, name: "Unrelated agent",
status: "paused", lifecycleState: "paused", lifecycleHolds: ["manual", "company_paused"], pauseReason: "manual" }).returning();
await companyService(db).update(company.id, { description: "Updated description" });
expect((await db.select().from(agents).where(eq(agents.id, agent.id)))[0].lifecycleHolds)
.toEqual(["manual", "company_paused"]);
await lifecycleWorker.sweep();
expect((await db.select().from(agents).where(eq(agents.id, agent.id)))[0].lifecycleHolds).toEqual(["manual"]);
});
it("does not auto-provision bundled built-in agents for a freshly created company", async () => {
const created = await companyService(db).create({
name: "Fresh Company",
@@ -273,6 +289,10 @@ describeEmbeddedPostgres("companyService", () => {
});
expect(archived?.status).toBe("archived");
// The request commits the admission block; stopping providers is owed work.
expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId)))[0].status).toBe("running");
expect(await lifecycleWorker.hasPending()).toBe(true);
await lifecycleWorker.sweep();
const archiveActivity = await db
.select({
@@ -289,7 +309,7 @@ describeEmbeddedPostgres("companyService", () => {
expect(archiveActivity[0]).toMatchObject({
actorType: "user",
actorId: "test-user",
details: { agentsPaused: 3, runsCancelled: 1 },
details: { agentsPauseRequested: 3, runsCancellationRequested: 1 },
});
const rows = await db
@@ -333,6 +353,27 @@ describeEmbeddedPostgres("companyService", () => {
});
});
it.each(["archive", "pause"] as const)("preserves owed stops across immediate %s and reactivation", async action => {
const companyId = randomUUID();
const agentId = randomUUID();
await db.insert(companies).values({ id: companyId, name: "Rapid lifecycle", issuePrefix: `R${companyId.slice(0, 6)}` });
await insertAgents({ id: agentId, companyId, name: "Running target", adapterType: "process", status: "running" });
const [original] = await db.insert(heartbeatRuns).values({ companyId, agentId, status: "running", invocationSource: "on_demand" }).returning();
const service = companyService(db);
if (action === "archive") await service.archive(companyId);
else await service.update(companyId, { status: "paused" });
expect((await db.select().from(agents).where(eq(agents.id, agentId)))[0]).toMatchObject({ lifecycleState: "pausing", status: "paused" });
expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, original.id)))[0].status).toBe("running");
await service.update(companyId, { status: "active" });
const [waiting] = await db.insert(heartbeatRuns).values({ companyId, agentId, status: "queued", invocationSource: "on_demand" }).returning();
await expect(reserveRunBudget(db, companyId, waiting.id, null)).rejects.toThrow(/paused|not ready/);
await lifecycleWorker.sweep();
expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, original.id)))[0].status).toBe("cancelled");
expect((await db.select().from(agents).where(eq(agents.id, agentId)))[0]).toMatchObject({ lifecycleState: "ready", status: "idle" });
const [fresh] = await db.insert(heartbeatRuns).values({ companyId, agentId, status: "queued", invocationSource: "on_demand" }).returning();
await expect(reserveRunBudget(db, companyId, fresh.id, null)).resolves.toMatchObject({ runId: fresh.id, reused: false });
});
it("reactivates only agents paused because the company was archived", async () => {
const companyId = randomUUID();
const archivedPausedAgentId = randomUUID();
@@ -394,7 +435,7 @@ describeEmbeddedPostgres("companyService", () => {
);
expect(reactivated?.status).toBe("active");
await lifecycleWorker.process(archivedPausedAgentId);
await lifecycleWorker.sweep();
const reactivateActivity = await db
.select({
@@ -514,6 +555,7 @@ describeEmbeddedPostgres("companyService", () => {
);
expect(archived?.status).toBe("archived");
await lifecycleWorker.sweep();
const rows = await db
.select({ id: agents.id, status: agents.status, pauseReason: agents.pauseReason })
@@ -547,7 +589,7 @@ describeEmbeddedPostgres("companyService", () => {
expect(archiveActivity[0]).toMatchObject({
actorType: "user",
actorId: "test-user",
details: { agentsPaused: 2, runsCancelled: 1 },
details: { agentsPauseRequested: 2, runsCancellationRequested: 1 },
});
});
@@ -600,7 +642,7 @@ describeEmbeddedPostgres("companyService", () => {
);
expect(reactivated?.status).toBe("active");
await lifecycleWorker.process(archivedPausedAgentId);
await lifecycleWorker.sweep();
const rows = await db
.select({ id: agents.id, status: agents.status, pauseReason: agents.pauseReason })
@@ -772,6 +814,7 @@ describeEmbeddedPostgres("companyService", () => {
runId: null,
});
expect(archived?.status).toBe("archived");
await lifecycleWorker.sweep();
const wakeups = await db
.select({
@@ -829,7 +872,7 @@ describeEmbeddedPostgres("companyService", () => {
eq(activityLog.action, "company.archived"),
));
expect(archiveActivity).toHaveLength(1);
expect(archiveActivity[0]).toMatchObject({ details: { agentsPaused: 1, runsCancelled: 0 } });
expect(archiveActivity[0]).toMatchObject({ details: { agentsPauseRequested: 1, runsCancellationRequested: 0 } });
});
it("rings the Cloud lifecycle doorbell exactly on archived-boundary transitions", async () => {
@@ -904,6 +947,7 @@ describeEmbeddedPostgres("companyService", () => {
);
expect(archived?.status).toBe("archived");
await lifecycleWorker.sweep();
const agent = await db
.select({ status: agents.status, pauseReason: agents.pauseReason })
@@ -926,7 +970,7 @@ describeEmbeddedPostgres("companyService", () => {
));
expect(archiveActivity).toHaveLength(1);
expect(archiveActivity[0]).toMatchObject({
details: { agentsPaused: 1, runsCancelled: 1 },
details: { agentsPauseRequested: 1, runsCancellationRequested: 1 },
});
});
@@ -22,6 +22,7 @@ import express from "express";
import request from "supertest";
import { decisionModelRoutes } from "../routes/decision-models.js";
import { errorHandler } from "../middleware/index.js";
import { deliverBudgetEnforcement } from "../services/budgets.js";
let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>>;
let db: ReturnType<typeof createDb>;
@@ -53,6 +54,36 @@ async function fixture() {
return { companyId, agentId, binding, provider, service, context, config, connections };
}
describe("company decision service", () => {
it("returns answers and a committed charge before slow budget stop delivery", async () => {
const f = await fixture();
await db.insert(budgetPolicies).values({ companyId: f.companyId, scopeType: "company", scopeId: f.companyId,
windowKind: "calendar_month_utc", amount: 1, notifyEnabled: false });
f.provider.mockResolvedValue({ ...structuredClone(outcome), receipt: { ...outcome.receipt, costCents: "2.0000000" } });
let release!: () => void;
const gate = new Promise<void>(resolve => { release = resolve; });
const cancelWorkForScope = vi.fn(() => gate);
const service = decisionModelService(db, { provider: f.provider, budgetHooks: { cancelWorkForScope } });
const pending = service.decide(f.context, DECISION_TEST_REQUEST);
let delivery: Promise<void> | undefined;
let finished = false;
void pending.then(() => { finished = true; }, () => { finished = true; });
try {
await vi.waitFor(() => expect(finished).toBe(true), { timeout: 3000 });
expect(await pending).toMatchObject({ status: "succeeded", answers: outcome.answers });
expect(f.provider).toHaveBeenCalledOnce();
expect(cancelWorkForScope).not.toHaveBeenCalled();
const events = await db.select().from(costEvents).where(eq(costEvents.companyId, f.companyId));
expect(events).toHaveLength(1);
expect(events[0]).toMatchObject({ usageKind: "decision", costCents: 2 });
const [policy] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.companyId, f.companyId));
expect(policy.enforcementVersion).toBeGreaterThan(policy.enforcementDeliveredVersion);
delivery = deliverBudgetEnforcement(db, { cancelWorkForScope }, f.companyId);
await vi.waitFor(() => expect(cancelWorkForScope).toHaveBeenCalledOnce());
expect((await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.id)))[0].enforcementDeliveredVersion).toBe(policy.enforcementDeliveredVersion);
release(); await delivery;
expect((await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.id)))[0].enforcementDeliveredVersion).toBe(policy.enforcementVersion);
} finally { release(); await pending.catch(() => undefined); await delivery?.catch(() => undefined); }
});
it("starts unconfigured with background sponsorship on and makes availability free", async () => {
const f = await fixture();
await db.delete(companyDecisionModels).where(eq(companyDecisionModels.companyId, f.companyId));
@@ -214,6 +214,24 @@ describe("delivery work coordinator", () => {
expect(t.run).toHaveBeenCalledTimes(2);
expect(idleWorkSnapshot().active).toBe(0);
});
it("delivers already-admitted stop work during drain while respecting warm standby", async () => {
const s = setup(), t = task();
await s.coordinator.register(DELIVERY_QUEUES.budgetEnforcement, { ...t, runDuringDrain: true }).ready;
s.canRun.mockReturnValue(false);
s.canReconcile.mockReturnValue(true);
await s.enqueue(DELIVERY_QUEUES.budgetEnforcement);
await vi.advanceTimersByTimeAsync(1);
expect(t.run).toHaveBeenCalledTimes(2);
expect(s.coordinator.nextWakeAt()).toBeNull();
s.canReconcile.mockReturnValue(false);
await s.enqueue(DELIVERY_QUEUES.budgetEnforcement);
await vi.advanceTimersByTimeAsync(5001);
expect(t.run).toHaveBeenCalledTimes(2);
s.canReconcile.mockReturnValue(true);
await vi.advanceTimersByTimeAsync(5000);
expect(t.run).toHaveBeenCalledTimes(3);
expect(idleWorkSnapshot().active).toBe(0);
});
it("bounds an expired deadline retry during idle reconciliation and resumes after drain", async () => {
const s = setup(), t = task();
const due = Date.now() + 10_000;
@@ -36,6 +36,7 @@ import {
import { instanceSettingsService } from "../services/instance-settings.js";
import { settleConversationTurn } from "../services/agent-conversations.js";
import { buildLowTrustSourceTrust, LOW_TRUST_QUARANTINED_BODY } from "../services/source-trust.js";
import { deliverBudgetEnforcement } from "../services/budgets.js";
let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>>,
db: ReturnType<typeof createDb>,
home: string;
@@ -162,6 +163,36 @@ async function fixture() {
};
}
describe("fast response accepted turns", () => {
it("returns acknowledgement text and a committed charge before slow budget stop delivery", async () => {
const f = await fixture();
await db.insert(budgetPolicies).values({ companyId: f.companyId, scopeType: "company", scopeId: f.companyId,
windowKind: "calendar_month_utc", amount: 1, notifyEnabled: false });
f.provider.mockResolvedValue({ ...structuredClone(outcome), receipt: { ...outcome.receipt, costCents: "2.0000000" } });
let release!: () => void;
const gate = new Promise<void>(resolve => { release = resolve; });
const cancelWorkForScope = vi.fn(() => gate);
const service = fastResponseService(db, { provider: f.provider, budgetHooks: { cancelWorkForScope } });
const pending = service.test(f.companyId, "alice");
let delivery: Promise<void> | undefined;
let finished = false;
void pending.then(() => { finished = true; }, () => { finished = true; });
try {
await vi.waitFor(() => expect(finished).toBe(true), { timeout: 3000 });
expect(await pending).toMatchObject({ status: "succeeded", text: outcome.text });
expect(f.provider).toHaveBeenCalledOnce();
expect(cancelWorkForScope).not.toHaveBeenCalled();
const events = await db.select().from(costEvents).where(eq(costEvents.companyId, f.companyId));
expect(events).toHaveLength(1);
expect(events[0]).toMatchObject({ usageKind: "fast_response", costCents: 2 });
const [policy] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.companyId, f.companyId));
expect(policy.enforcementVersion).toBeGreaterThan(policy.enforcementDeliveredVersion);
delivery = deliverBudgetEnforcement(db, { cancelWorkForScope }, f.companyId);
await vi.waitFor(() => expect(cancelWorkForScope).toHaveBeenCalledOnce());
expect((await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.id)))[0].enforcementDeliveredVersion).toBe(policy.enforcementDeliveredVersion);
release(); await delivery;
expect((await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.id)))[0].enforcementDeliveredVersion).toBe(policy.enforcementVersion);
} finally { release(); await pending.catch(() => undefined); await delivery?.catch(() => undefined); }
});
it("does not enqueue or infer while the experimental setting is disabled", async () => {
const f = await fixture();
await instanceSettingsService(db).updateExperimental({ enableFastResponses: false });
+31 -1
View File
@@ -3,7 +3,7 @@ import { randomUUID } from "node:crypto";
import { eq, sql } from "drizzle-orm";
import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest";
import {
agentInstructionWorkingCopies, agents, agentApiKeys, agentWakeupRequests, companies, companySecretProposals, createDb,
agentInstructionWorkingCopies, agents, agentApiKeys, agentWakeupRequests, budgetPolicies, companies, companySecretProposals, createDb,
adapterAuthSessions, environments, environmentLeases, executionWorkspaces,
feedbackExports, feedbackVotes,
heartbeatRuns, issues, issueWatchdogs, issueThreadInteractions, issueQuestionResponseDeliveries,
@@ -187,6 +187,36 @@ if (!support.supported) console.warn(`Skipping idle sleep Postgres tests: ${supp
return { companyId, agentId, runId };
}
it("retains committed budget stop intent until its exact generation is delivered", async () => {
const f = await seed();
const [policy] = await db.insert(budgetPolicies).values({ companyId: f.companyId, scopeType: "company", scopeId: f.companyId,
windowKind: "calendar_month_utc", enforcementVersion: 2, enforcementDeliveredVersion: 1 }).returning();
expect(await read()).toEqual(present);
await db.update(budgetPolicies).set({ enforcementDeliveredVersion: 2 }).where(eq(budgetPolicies.id, policy.id));
expect(await read()).toEqual(none);
});
it("retains company hold projection and lifecycle operations through admission drain", async () => {
const f = await seed();
await db.update(companies).set({ status: "paused" }).where(eq(companies.id, f.companyId));
expect(await read()).toEqual(present);
await db.update(agents).set({ status: "paused", lifecycleState: "pausing", lifecycleHolds: ["company_paused"] }).where(eq(agents.id, f.agentId));
expect(await read()).toEqual(present);
await db.update(agents).set({ lifecycleState: "paused" }).where(eq(agents.id, f.agentId));
expect(await read()).toEqual(none);
});
it("retains a committed budget Stop after terminal status until delivery is acknowledged", async () => {
const f = await seed();
const intent = { companyId: f.companyId, scopeType: "agent", scopeId: f.agentId, policyId: randomUUID(), version: 1 };
await db.update(heartbeatRuns).set({ status: "cancelled", resultJson: { startupCancellation: { budgetEnforcement: intent } } })
.where(eq(heartbeatRuns.id, f.runId));
expect(await read()).toEqual(present);
await db.update(heartbeatRuns).set({ resultJson: { startupCancellation: { budgetEnforcement: { ...intent, deliveredAt: new Date(now).toISOString() } } } })
.where(eq(heartbeatRuns.id, f.runId));
expect(await read()).toEqual(none);
});
async function feedbackHistory(status: string, change: Partial<typeof feedbackExports.$inferInsert> = {}) {
const f = await seed();
const [issue] = await db.insert(issues).values({ companyId: f.companyId, title: "Finished fixture", status: "done" }).returning();
+10 -2
View File
@@ -40,6 +40,7 @@ import { assertNoExternalOverlap, withExternalAdmissionGuard } from "../modules/
import { digestPaperclipSemanticContent, externalOperationDigest, parseNativeExecutionInput, canonicalNativeRuntimeContextDigest, nativeRuntimePromptDigest, NATIVE_RUNTIME_ASSET_SCHEMA, PAPERCLIP_EXECUTION_PROMPT, PAPERCLIP_EXECUTION_PROMPT_REVISION, type ExternalProviderOperation } from "../vendor/paperclip-runner/index.js";
import { actorMiddleware } from "../middleware/auth.js";
import { createLocalAgentJwt } from "../agent-auth-jwt.js";
import { createDeliveryWorkCoordinator } from "../services/delivery-work-coordinator.js";
/** Real database/transport authority tests; synthetic provider receipts do not
* count as live Muse/native qualification evidence. */
@@ -2388,14 +2389,21 @@ describe("personal Muse control plane",()=>{
await db.update(agents).set({status:"paused",lifecycleState:"verifying",lifecycleOperation:{id:randomUUID(),hostComplete:false,completedPluginIds:[],attempts:0,responsibleUserId:f.operatorId},lifecycleRequiredPluginIds:[]}).where(eq(agents.id,f.agent.id));
const harness=agentHarnessVerificationService(db,{} as PluginWorkerManager),onReady=vi.fn();
const lifecycle=startAgentLifecycle(db,{requiredPluginIds:async()=>[],runPlugin:async()=>"complete",runHost:snapshot=>harness.verify(snapshot),onReady},()=>true);
const deliveryErrors=vi.fn(),coordinator=createDeliveryWorkCoordinator({owner:db,canRun:()=>true,onError:deliveryErrors});
const delivery=coordinator.register(DELIVERY_QUEUES.agentLifecycle,{retryMs:1000,run:()=>lifecycle.sweep(),hasPending:()=>lifecycle.hasPending(),nextRunAt:()=>lifecycle.nextRunAt()});
lifecycle.setWake(delivery.wake);
try {
await lifecycle.sweep();expect((await db.select().from(agents).where(eq(agents.id,f.agent.id)))[0].lifecycleState).toBe("verifying");
await delivery.ready;expect((await db.select().from(agents).where(eq(agents.id,f.agent.id)))[0].lifecycleState).toBe("verifying");
// The authenticated confirmation must retry this hire, not wait for a
// periodic scan or a previously persisted long backoff to expire.
await db.update(agents).set({lifecycleOperation:sql`${agents.lifecycleOperation} || jsonb_build_object('retryAt', ${new Date(Date.now()+60_000).toISOString()}::text)`}).where(eq(agents.id,f.agent.id));
await museReceiver(db).signal(f.binding.id);await museReceiver(db).flush();
const items=await f.broker.inspect(f.subject,{version:1,query:"mailbox",after:0}) as {items:Array<{references:{nonce:string}}>};
await f.broker.act(f.subject,{version:1,command:"challenge.confirm",nonce:items.items[0].references.nonce,requestId:randomUUID()});
await expect.poll(async()=> (await db.select().from(agents).where(eq(agents.id,f.agent.id)))[0].lifecycleState,{timeout:8000,interval:100}).toBe("ready");
expect((await db.select().from(agents).where(eq(agents.id,f.agent.id)))[0].status).toBe("idle");expect(onReady).toHaveBeenCalledOnce();
}finally{await lifecycle.stop();}
expect(deliveryErrors).not.toHaveBeenCalled();
}finally{await coordinator.stop();await lifecycle.stop();}
});
it("serves opaque public wake digests while worker and cleanup APIs still require private bearer authority",async()=>{
@@ -24,6 +24,7 @@ import { errorHandler } from "../middleware/index.js";
import { localEncryptedProvider } from "../secrets/local-encrypted-provider.js";
import { SecretProviderClientError } from "../secrets/types.js";
import { validateToolContent } from "../services/tool-content-guards.js";
import { deliverBudgetEnforcement } from "../services/budgets.js";
let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>>, db: ReturnType<typeof createDb>, home: string;
beforeAll(async () => {
home = await realpath(await mkdtemp(path.join(os.tmpdir(), "paperclip-injection-")));
@@ -49,6 +50,36 @@ async function fixture() {
return { companyId, agentId, connections, binding, provider, service, context, config };
}
describe("prompt injection service", () => {
it("returns the model verdict and committed charge before slow budget stop delivery", async () => {
const f = await fixture();
await db.insert(budgetPolicies).values({ companyId: f.companyId, scopeType: "company", scopeId: f.companyId,
windowKind: "calendar_month_utc", amount: 1, notifyEnabled: false });
f.provider.mockResolvedValue({ ...structuredClone(outcome), receipt: { ...outcome.receipt, costCents: "2.0000000" } });
let release!: () => void;
const gate = new Promise<void>(resolve => { release = resolve; });
const cancelWorkForScope = vi.fn(() => gate);
const service = promptInjectionPreventionService(db, { provider: f.provider, budgetHooks: { cancelWorkForScope } });
const pending = service.classify(f.context, "Benign source code");
let delivery: Promise<void> | undefined;
let finished = false;
void pending.then(() => { finished = true; }, () => { finished = true; });
try {
await vi.waitFor(() => expect(finished).toBe(true), { timeout: 3000 });
expect(await pending).toMatchObject({ status: "classified", violation: false });
expect(f.provider).toHaveBeenCalledOnce();
expect(cancelWorkForScope).not.toHaveBeenCalled();
const events = await db.select().from(costEvents).where(eq(costEvents.companyId, f.companyId));
expect(events).toHaveLength(1);
expect(events[0]).toMatchObject({ usageKind: "prompt_injection", costCents: 2 });
const [policy] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.companyId, f.companyId));
expect(policy.enforcementVersion).toBeGreaterThan(policy.enforcementDeliveredVersion);
delivery = deliverBudgetEnforcement(db, { cancelWorkForScope }, f.companyId);
await vi.waitFor(() => expect(cancelWorkForScope).toHaveBeenCalledOnce());
expect((await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.id)))[0].enforcementDeliveredVersion).toBe(policy.enforcementDeliveredVersion);
release(); await delivery;
expect((await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.id)))[0].enforcementDeliveredVersion).toBe(policy.enforcementVersion);
} finally { release(); await pending.catch(() => undefined); await delivery?.catch(() => undefined); }
});
it.each(["missing_key", "wrong_key", "remote_vault_outage"])("uses full regex fallback for %s without dispatch or charge", async failure => {
const f = await fixture();
const previousKey = process.env.PAPERCLIP_SECRETS_MASTER_KEY;
@@ -17,6 +17,7 @@ import { remoteTerminationReceipt } from "../services/remote-execution-terminati
import { createDeliveryWorkCoordinator } from "../services/delivery-work-coordinator.js";
import { idleWorkSnapshot } from "../services/task-admission.js";
import { readIdleSleepSafety, type IdleSleepDrainStatus } from "../services/idle-sleep-safety.js";
import { reserveRunBudget } from "../services/budget-reservations.js";
const support = await getEmbeddedPostgresTestSupport();
const describePostgres = support.supported ? describe : describe.skip;
@@ -290,8 +291,9 @@ describePostgres("already available connection requests", () => {
expectedStatus: "queued", now: new Date() })).toMatchObject({ outcome: "cancelled", errorCode: "issue_waiting_for_response" });
} else {
const invoke = vi.fn(async () => null);
const reservation = await reserveRunBudget(db, f.companyId, run.id, null);
expect(await dispatch.dispatchResolvedInteractionIfCurrent({ companyId: f.companyId, runId: run.id,
expectedStatus: "running", now: new Date(), dispatch: invoke })).toMatchObject({ dispatched: false,
reservationId: reservation.id, expectedStatus: "running", now: new Date(), dispatch: invoke })).toMatchObject({ dispatched: false,
cancellation: { outcome: "cancelled", errorCode: "issue_waiting_for_response" } });
expect(invoke).not.toHaveBeenCalled();
}
@@ -57,7 +57,7 @@ const {
retryMs?: number;
run: (signal?: AbortSignal) => Promise<unknown>;
hasPending?: () => Promise<boolean>;
}) => ({ ready: task.run().catch(() => undefined), wake: vi.fn() }));
}) => ({ ready: task.run(new AbortController().signal).catch(() => undefined), wake: vi.fn() }));
const createAppMock = vi.fn(async () => Object.assign((_: unknown, __: unknown) => {}, {
locals: {
fastResponses: { hasPending: vi.fn(async () => false), sweepPending: vi.fn(async () => undefined) },
@@ -249,7 +249,8 @@ vi.mock("detect-port", () => ({
default: detectPortMock,
}));
vi.mock("@paperclipai/db", () => ({
vi.mock("@paperclipai/db", async (importOriginal) => ({
...await importOriginal<typeof import("@paperclipai/db")>(),
createDb: createDbMock,
ensurePostgresDatabase: vi.fn(),
getPostgresDataDirectory: vi.fn(),
@@ -259,16 +260,24 @@ vi.mock("@paperclipai/db", () => ({
formatDatabaseBackupResult: vi.fn(() => "ok"),
runDatabaseBackup: vi.fn(),
closeLivenessDb: vi.fn(async () => {}),
authUsers: {},
companies: {},
companyMemberships: {},
instanceUserRoles: {},
}));
vi.mock("../app.js", () => ({
createApp: createAppMock,
}));
// Startup owns registration. Durable budget/accounting state and coordinator
// retry delivery are exercised against PostgreSQL in the worker suites.
vi.mock("../services/budgets.js", () => ({
deliverBudgetEnforcement: vi.fn(async () => undefined),
hasPendingBudgetEnforcement: vi.fn(async () => false),
nextBudgetReconciliationAt: vi.fn(async () => null),
reconcileDueBudgetPolicies: vi.fn(async () => ({ nextCursor: null })),
}));
vi.mock("../services/run-cost-accounting.js", () => ({
readCostAccountingWork: vi.fn(async () => ({ pending: false, nextRunAt: null })),
}));
// This suite exercises startup coordination; source persistence and authority
// are covered by the database-backed source worker tests.
vi.mock("../services/skill-source-sync.js", () => ({
@@ -773,7 +782,7 @@ describe("startServer feedback export wiring", () => {
intervalCallback?.();
await new Promise<void>((resolve) => setImmediate(resolve));
expect(heartbeatServiceMock.sweepStaleIssueLocks).toHaveBeenCalledTimes(2);
expect(heartbeatServiceMock.reconcileCostAccounting).toHaveBeenCalledTimes(2);
expect(heartbeatServiceMock.reconcileCostAccounting).toHaveBeenCalledTimes(1);
expect(retiredDetector).not.toHaveBeenCalled();
} finally {
delete (runtime as Partial<typeof runtime>).reconcileProductivityReviews;
@@ -879,14 +888,14 @@ describe("startServer feedback export wiring", () => {
expect(externalObjectsServiceMock.refreshDueObjectsForActiveCompanies).toHaveBeenCalledTimes(1);
expect(routineServiceMock.tickScheduledTriggers).not.toHaveBeenCalled();
expect(heartbeatServiceMock.reconcileCostAccounting).toHaveBeenCalledTimes(2);
expect(heartbeatServiceMock.reconcileCostAccounting).toHaveBeenCalledTimes(1);
expect(environmentCustomImagesServiceMock.cleanupExpiredSetupSessions).not.toHaveBeenCalled();
} finally {
setIntervalSpy.mockRestore();
}
});
it.each([false, true])("retries failed accounting recovery with heartbeat scheduling enabled: %s", async (enabled) => {
it.each([false, true])("registers failed accounting recovery for delivery, independent of heartbeat scheduling: %s", async (enabled) => {
loadConfigMock.mockReturnValue(buildTestConfig({ heartbeatSchedulerEnabled: enabled }));
heartbeatServiceMock.reconcileCostAccounting.mockRejectedValueOnce(new Error("temporary ledger outage"));
let tick: (() => void) | undefined;
@@ -901,6 +910,10 @@ describe("startServer feedback export wiring", () => {
expect(tick).toBeDefined();
tick?.();
await new Promise<void>((resolve) => setImmediate(resolve));
expect(heartbeatServiceMock.reconcileCostAccounting).toHaveBeenCalledTimes(1);
const registration = deliveryWorkRegisterMock.mock.calls.find(([queue]) => queue === DELIVERY_QUEUES.accounting);
expect(registration?.[1]).toMatchObject({ retryMs: 5_000, runDuringDrain: true });
await registration?.[1].run(new AbortController().signal);
expect(heartbeatServiceMock.reconcileCostAccounting).toHaveBeenCalledTimes(2);
} finally {
timer.mockRestore();
@@ -18,6 +18,7 @@ import {
toolCallEvents,
toolConnections,
toolInvocations,
toolMcpGateways,
toolPolicies,
toolProfileBindings,
toolProfileEntries,
@@ -217,6 +218,7 @@ describeEmbeddedPostgres("tool access policy service", () => {
await db.delete(toolPolicies);
await db.delete(toolProfileEntries);
await db.delete(toolProfileBindings);
await db.delete(toolMcpGateways);
await db.delete(toolProfiles);
await db.delete(toolCatalogEntries);
await db.delete(toolConnections);
@@ -953,6 +955,113 @@ describeEmbeddedPostgres("tool access policy service", () => {
expect(decisions.filter((decision) => decision.reasonCode === "rate_limited")).toHaveLength(1);
});
async function admittedRateLimitedInvocation(limit = 1) {
const company = await createCompany(db);
const agent = await createAgent(db, company.id);
const issue = await createIssue(db, company.id);
const run = await createRun(db, company.id, agent.id);
const { connection, catalogEntry } = await createTool(db, company.id);
const [profile] = await db.insert(toolProfiles).values({ companyId: company.id,
profileKey: "admitted-gateway", name: "Admitted gateway", defaultAction: "deny" }).returning();
const [gateway] = await db.insert(toolMcpGateways).values({ companyId: company.id,
name: "Admitted gateway", slug: "admitted-gateway", profileId: profile.id }).returning();
await db.insert(principalPermissionGrants).values(sendEmailGrantRow(company, agent));
const [ratePolicy] = await db.insert(toolPolicies).values({ companyId: company.id, name: "One admitted send",
policyType: "rate_limit", selectors: { toolName: "send_email" }, priority: 1,
config: { limit, windowSeconds: 3600, keyBy: ["agent", "tool"] } }).returning();
const input = { ...sendEmailDecisionInput(company, agent, connection, catalogEntry),
runContext: { heartbeatRunId: run.id, issueId: issue.id, gatewayId: gateway.id }, consumeRateLimit: true,
request: { connectionId: connection.id, catalogEntryId: catalogEntry.id, toolName: "send_email", arguments: { to: "recipient@example.test" } } };
const svc = toolAccessPolicyService(db);
const initial = await svc.decide(input);
expect(initial.allowed).toBe(true);
const { invocation } = await svc.recordInvocation(input, initial);
expect(invocation.gatewayId).toBe(gateway.id);
await db.update(toolInvocations).set({ status: "executing" }).where(eq(toolInvocations.id, invocation.id));
return { company, agent, input: { ...input, consumeRateLimit: false }, invocation, ratePolicy, svc };
}
it.each(["allow", "approved"] as const)("retains the consumed final rate-limit slot for the same executing %s invocation", async authorization => {
const { input, invocation, ratePolicy, svc } = await admittedRateLimitedInvocation();
if (authorization === "approved") await db.update(toolInvocations).set({ policyDecision: "require_approval", approvalState: "approved" })
.where(eq(toolInvocations.id, invocation.id));
expect(await svc.decide(input)).toMatchObject({ decision: "rate_limited" });
expect(invocation.policyExplanation?.consumedRateAdmissions).toEqual([
{ policyId: ratePolicy.id, versionHash: expect.stringMatching(/^[a-f0-9]{64}$/) },
]);
const counters = await db.select().from(toolRateLimitCounters);
expect(await svc.decide(input, { admittedInvocationId: invocation.id })).toMatchObject({ allowed: true });
expect(await svc.decide(input, { admittedInvocationId: invocation.id })).toMatchObject({ allowed: true });
expect(await db.select().from(toolRateLimitCounters)).toEqual(counters);
expect((await db.select().from(toolRateLimitCounters))[0].remaining).toBe(0);
// Neither a new invocation nor an accidental consuming recheck can borrow
// another invocation's already-spent slot.
expect(await svc.decide(input, { admittedInvocationId: randomUUID() })).toMatchObject({ decision: "rate_limited" });
expect(await svc.decide({ ...input, consumeRateLimit: true }, { admittedInvocationId: invocation.id }))
.toMatchObject({ decision: "rate_limited" });
});
it("rechecks a newly added rate policy rather than borrowing an older policy's consumed slot", async () => {
const { company, input, invocation, svc } = await admittedRateLimitedInvocation();
await svc.createPolicy(company.id, { name: "New send limit", policyType: "rate_limit", priority: 0,
selectors: { toolName: "send_email" }, config: { limit: 1, windowSeconds: 3600, keyBy: ["agent", "tool"] } });
// Consume the new rule through normal admission; the exhausted original
// rule then denies this separate request before any provider dispatch.
expect(await svc.decide({ ...input, consumeRateLimit: true })).toMatchObject({ decision: "rate_limited" });
const counters = await db.select().from(toolRateLimitCounters);
expect(counters).toHaveLength(2);
expect(await svc.decide(input, { admittedInvocationId: invocation.id })).toMatchObject({ decision: "rate_limited" });
expect(await db.select().from(toolRateLimitCounters)).toEqual(counters);
});
it("rechecks an edited rate-policy generation without consuming another slot", async () => {
const { company, input, invocation, ratePolicy, svc } = await admittedRateLimitedInvocation(2);
expect(await svc.decide({ ...input, consumeRateLimit: true })).toMatchObject({ allowed: true });
await svc.updatePolicy({ companyId: company.id, policyId: ratePolicy.id,
body: { config: { limit: 1, windowSeconds: 3600, keyBy: ["agent", "tool"] } } });
const counters = await db.select().from(toolRateLimitCounters);
expect(await svc.decide(input, { admittedInvocationId: invocation.id })).toMatchObject({ decision: "rate_limited" });
expect(await db.select().from(toolRateLimitCounters)).toEqual(counters);
});
it("rechecks rate limits for historical executing invocations without durable consumption proof", async () => {
const { input, invocation, svc } = await admittedRateLimitedInvocation();
await db.update(toolInvocations).set({ policyExplanation: null }).where(eq(toolInvocations.id, invocation.id));
expect(await svc.decide(input, { admittedInvocationId: invocation.id })).toMatchObject({ decision: "rate_limited" });
});
it("reevaluates current blocks, grants, and approval policies for admitted work", async () => {
const { company, input, invocation, svc } = await admittedRateLimitedInvocation();
const options = { admittedInvocationId: invocation.id };
const [block] = await db.insert(toolPolicies).values({ companyId: company.id, name: "Block admitted sends",
policyType: "block", priority: 0, selectors: { toolName: "send_email" } }).returning();
expect(await svc.decide(input, options)).toMatchObject({ decision: "deny", reasonCode: "deny_policy_block" });
await db.delete(toolPolicies).where(eq(toolPolicies.id, block.id));
await db.delete(principalPermissionGrants).where(eq(principalPermissionGrants.companyId, company.id));
expect(await svc.decide(input, options)).toMatchObject({ decision: "deny", reasonCode: "deny_default" });
await db.insert(toolPolicies).values({ companyId: company.id, name: "Require new review",
policyType: "require_approval", priority: 2, selectors: { toolName: "send_email" } });
expect(await svc.decide(input, options)).toMatchObject({ decision: "require_approval" });
});
it("does not retain rate admission for a mismatched or no-longer-authorized invocation", async () => {
const { input, invocation, svc } = await admittedRateLimitedInvocation();
const options = { admittedInvocationId: invocation.id };
const mutations: Partial<typeof toolInvocations.$inferInsert>[] = [
{ actorType: "system" }, { actorId: randomUUID() }, { agentId: null }, { runId: null }, { issueId: null }, { gatewayId: null },
{ applicationId: null }, { connectionId: null }, { catalogEntryId: null }, { toolName: "different_tool" },
{ providerType: null }, { applicationKey: null }, { upstreamToolName: null }, { argumentsHash: "different" },
{ status: "authorized" }, { policyDecision: "deny" }, { policyDecision: "require_approval", approvalState: "pending" },
];
for (const mutation of mutations) {
await db.update(toolInvocations).set(mutation).where(eq(toolInvocations.id, invocation.id));
expect(await svc.decide(input, options), JSON.stringify(mutation)).toMatchObject({ decision: "rate_limited" });
await db.update(toolInvocations).set({ ...invocation, status: "executing" }).where(eq(toolInvocations.id, invocation.id));
}
expect(await svc.decide({ ...input, request: { ...input.request, arguments: { to: "another@example.test" } } }, options))
.toMatchObject({ decision: "rate_limited" });
});
it("rejects unsupported policy semantics at create and update time", async () => {
const company = await createCompany(db);
const svc = toolAccessPolicyService(db);
+8 -1
View File
@@ -1276,6 +1276,13 @@ export async function createApp(
onError: (err, queue) => logger.error({ err, queue }, "Delivery reconciliation failed"),
});
app.locals.deliveryWork = deliveryWork;
const agentLifecycleDelivery = deliveryWork.register(DELIVERY_QUEUES.agentLifecycle, {
retryMs: 1_000,
run: () => agentLifecycle.sweep(),
hasPending: () => agentLifecycle.hasPending(),
nextRunAt: () => agentLifecycle.nextRunAt(),
});
agentLifecycle.setWake(agentLifecycleDelivery.wake);
const stopMuseMaintenance = startMuseMaintenance(db, deliveryWork);
const resourceLifecycleDelivery = registerPluginLifecycleDelivery(deliveryWork, db, workerManager, lifecycle);
await resourceLifecycleDelivery.ready;
@@ -1439,7 +1446,7 @@ export async function createApp(
logger.error({ err }, "Failed to load ready plugins on startup");
});
app.locals.bundledPluginsStartup = trackIdleWork(bundledPluginsStartup);
void bundledPluginsStartup.then(() => { lifecyclePluginsReady = true; resourceLifecycleDelivery.wake(); return agentLifecycle.sweep(); }).catch(() => logger.warn("Agent lifecycle recovery failed; retrying."));
void bundledPluginsStartup.then(() => { lifecyclePluginsReady = true; resourceLifecycleDelivery.wake(); agentLifecycleDelivery.wake(); }).catch(() => logger.warn("Agent lifecycle recovery failed; retrying."));
// The shutdown hook runs at most once. It caches the in-flight promise, so a
// second caller (for example the `exit` handler) awaits the same completion
// instead of starting a second teardown.
+7 -6
View File
@@ -12,6 +12,7 @@ import {
} from "./services/native-runtime/native-harness-backup-expiration.js";
import { subscribeAllCompanyLiveEvents } from "./services/live-events.js";
import { DELIVERY_QUEUES } from "./services/delivery-work-notifications.js";
import { registerBudgetEnforcementWork, registerCostAccountingWork } from "./services/budget-enforcement-worker.js";
import { caseHandoffDeliveryService } from "./services/case-handoff-delivery.js";
import { registerScheduledWork } from "./services/scheduled-work.js";
import type { createDeliveryWorkCoordinator } from "./services/delivery-work-coordinator.js";
@@ -1345,6 +1346,11 @@ async function startServerWithDatabaseTeardown(
resolveNativeQuestion: (interaction) => deliverNativeQuestionResponse(db as any, interaction),
});
const deliveryWork = app.locals.deliveryWork as ReturnType<typeof createDeliveryWorkCoordinator>;
const budgetEnforcementWorker = registerBudgetEnforcementWork(db, deliveryWork, {
cancelWorkForScope: accountingHeartbeat.cancelBudgetScopeWork,
});
const accountingWorker = registerCostAccountingWork(db, deliveryWork, accountingHeartbeat.reconcileCostAccounting);
server.on("close", accountingWorker.stop);
skillSourceSync.start(deliveryWork);
const caseHandoffDeliveries = caseHandoffDeliveryService(db as any, { heartbeat: environmentLeaseCleanupHeartbeat });
deliveryWork.register(DELIVERY_QUEUES.caseHandoffs, {
@@ -1446,7 +1452,7 @@ async function startServerWithDatabaseTeardown(
await timeStartupRecovery("tool_reviews", () => app.locals.toolGateway.sweepActionReviews()).catch((err: unknown) => logger.error({ err }, "startup tool review recovery failed"));
await timeStartupRecovery("gateway_sessions", () => app.locals.toolGateway.cleanupExpiredSessions()).catch((err: unknown) => logger.error({ err }, "startup gateway token cleanup failed"));
await timeStartupRecovery("delivery_workers", () => Promise.all([fastResponseWorker, ...deliveryWorkers].map(worker => worker.ready)));
await timeStartupRecovery("delivery_workers", () => Promise.all([budgetEnforcementWorker, accountingWorker, fastResponseWorker, ...deliveryWorkers].map(worker => worker.ready)));
scheduleGitHubConnectionEventPoll();
scheduleGitHubConnectionContinuitySweep();
@@ -1589,7 +1595,6 @@ async function startServerWithDatabaseTeardown(
},
"worktree run-execution cutoff state",
);
await timeStartupRecovery("accounting", () => accountingHeartbeat.reconcileCostAccounting()).catch((err) => logger.error({ err }, "Cost accounting recovery failed; pending receipts will retry"));
const heartbeatSchedulingSuppression = await heartbeat.resolveSchedulingSuppression();
// Reap orphaned runs before timer ticks start so wakeups cannot coalesce
@@ -1883,8 +1888,6 @@ async function startServerWithDatabaseTeardown(
logger.error({ err }, "periodic secret proposal expiry sweep failed");
}));
trackHeartbeatSchedulerWork(accountingHeartbeat.reconcileCostAccounting().catch((err) => logger.error({ err }, "Cost accounting recovery failed")));
trackHeartbeatSchedulerWork(app.locals.toolGateway.sweepActionReviews().catch((err: unknown) => logger.error({ err }, "tool review recovery failed")));
trackHeartbeatSchedulerWork(app.locals.toolGateway.cleanupExpiredSessions().catch((err: unknown) => logger.error({ err }, "gateway token cleanup failed")));
@@ -1968,10 +1971,8 @@ async function startServerWithDatabaseTeardown(
// is still required. A failed acquire can leak a paid provider sandbox, so
// this path retries the teardown at startup and on the interval, exactly as
// the enabled path does.
await timeStartupRecovery("accounting", () => accountingHeartbeat.reconcileCostAccounting()).catch((err) => logger.error({ err }, "Cost accounting recovery failed; pending receipts will retry"));
await runEnvironmentLeaseCleanupSweep(0);
startHeartbeatSchedulerInterval(() => {
trackHeartbeatSchedulerWork(accountingHeartbeat.reconcileCostAccounting().catch((err) => logger.error({ err }, "Cost accounting recovery failed")));
scheduleExternalObjectRefreshSweep(new Date());
scheduleEnvironmentLeaseCleanupSweep();
scheduleGitHubConnectionEventPoll();
@@ -110,14 +110,30 @@ describe("lifecycle failure diagnostics", () => {
await worker.stop();
});
it("reports periodic policy scan failures", async () => {
vi.useFakeTimers();
it("bounds retries for due work while plugins are unavailable and stops reading after shutdown", async () => {
const where = vi.fn().mockResolvedValue([{ due: 0 }]);
const select = vi.fn(() => ({ from: () => ({ where }) }));
const db = { select } as unknown as Db;
const work = startAgentLifecycle(db, createAgentLifecycleEffects(), fixture().driver, () => false);
try {
expect(await work.hasPending()).toBe(true);
expect(work.nextRunAt()).toBeNull();
await work.sweep();
expect(select).toHaveBeenCalledTimes(1);
await work.stop();
expect(await work.hasPending()).toBe(false);
expect(work.nextRunAt()).toBeNull();
expect(select).toHaveBeenCalledTimes(1);
} finally { await work.stop(); }
});
it("reports inherited company projection failures", async () => {
const error = new Error("private database detail");
const db = { select: () => { throw error; } } as unknown as Db;
const effects = { ...createAgentLifecycleEffects(), reportFailure: vi.fn() };
const work = startAgentLifecycle(db, effects, fixture().driver, () => true);
try {
await vi.advanceTimersByTimeAsync(5_000);
await expect(work.sweep()).rejects.toThrow(error);
expect(effects.reportFailure).toHaveBeenCalledWith({ stage: "policy_scan" }, error);
} finally { await work.stop(); }
});
@@ -0,0 +1,18 @@
import type { Db } from "@paperclipai/db";
import { describe, expect, it, vi } from "vitest";
const stop = vi.hoisted(() => vi.fn(async () => true));
vi.mock("../../../services/heartbeat.js", () => ({ heartbeatService: () => ({ stopInvocationsForAgents: stop }) }));
vi.mock("../../../services/agent-harness-verification.js", () => ({ agentHarnessVerificationService: () => ({ verify: vi.fn() }) }));
import { createLifecycleDriver } from "../../../services/agent-lifecycle-driver.js";
describe("lifecycle host stop scope", () => {
it("pauses only the requested agent without loading its company peers", async () => {
const db = { select: vi.fn(() => { throw new Error("No peer lookup expected"); }) } as unknown as Db;
const driver = createLifecycleDriver(db, {} as never);
expect(await driver.runHost({ id: "target", companyId: "company", lifecycleState: "pausing",
lifecycleVersion: 2, lifecycleRequiredPluginIds: [], lifecycleError: null, lifecycleOperation: null })).toBe("complete");
expect(stop).toHaveBeenCalledWith(["target"], "Agent lifecycle stop requested");
expect(db.select).not.toHaveBeenCalled();
});
});
@@ -19,7 +19,8 @@ import { and, eq, sql } from "drizzle-orm";
import { agents, approvals, projects, issues, heartbeatRuns, nativeRunFinalizations, budgetPolicies, environmentLeases, agentApiKeys, activityLog, agentConfigRevisions, userCompanyPreferences, companies, companyMemberships, principalPermissionGrants, plugins, pluginCompanySettings, createDb, type Db } from "@paperclipai/db";
import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../../../__tests__/helpers/embedded-postgres.js";
import { createAgentLifecycle, configureAgentLifecycle } from "../../../services/agent-lifecycle.js";
import { createAgentLifecycle, configureAgentLifecycle, startAgentLifecycle } from "../../../services/agent-lifecycle.js";
import { projectCompanyHolds } from "../adapters/company-policy.js";
import { createLifecycleStore as createStore } from "../adapters/postgres.js";
import { transition } from "../domain/policy.js";
import type { LifecycleAgent } from "../application/ports.js";
@@ -394,6 +395,25 @@ const support = await getEmbeddedPostgresTestSupport();
expect(await db.select().from(companies).where(eq(companies.id, companyId))).toHaveLength(0);
});
it("awaits only the deleted company's agents when a coordinator owns lifecycle wakeups", async () => {
const agent = await hire();
const [otherCompany] = await db.insert(companies).values({ name: "Unaffected", issuePrefix: `U${companyId.slice(0, 7)}` }).returning();
const [otherAgent] = await db.insert(agents).values({ companyId: otherCompany.id, name: "Pending elsewhere",
lifecycleState: "preparing", lifecycleOperation: { id: randomUUID(), hostComplete: false, completedPluginIds: [], attempts: 0 } }).returning();
const phases: string[] = [];
const work = startAgentLifecycle(db, { requiredPluginIds: async () => [], runPlugin: async () => "complete",
runHost: async snapshot => { expect(snapshot.id).toBe(agent.id); phases.push(snapshot.lifecycleState); return "complete"; } }, () => true);
const wake = vi.fn();
work.setWake(wake);
try {
await deleteCompany(db, companyId);
expect(phases).toEqual(["terminating", "cleaning_up"]);
expect(wake).toHaveBeenCalled();
expect(await db.select().from(companies).where(eq(companies.id, companyId))).toHaveLength(0);
expect((await current(otherAgent.id)).lifecycleState).toBe("preparing");
} finally { await work.stop(); }
});
it("requires a transaction for the company deletion integration", async () => {
await expect(agentLifecycleCompanyDeletion.deleteCompanyData(db as never, companyId))
.rejects.toThrow("Company deletion requires a database transaction");
@@ -603,16 +623,17 @@ const support = await getEmbeddedPostgresTestSupport();
await db.update(companies).set({ status: "paused" }).where(eq(companies.id, companyId));
if (command === "updateConfiguration") await lifecycle.updateConfiguration(one.id, { budgetMonthlyCents: 1000 });
else await lifecycle.updateAndTransition(one.id, "pause", { budgetMonthlyCents: 1000 });
expect((await current(one.id)).lifecycleHolds).toContain("company_paused");
// Company state blocks admission; its inherited holds project separately.
expect((await current(one.id)).lifecycleHolds).not.toContain("company_paused");
expect((await current(two.id)).lifecycleHolds).toEqual([]);
});
it("limits invocation policy checks to the requested agent", async () => {
it("blocks invocation without synchronously projecting inherited company holds", async () => {
const one = await hire();
const two = await createAgentLifecycle(db).requestHire(companyId, { name: "Another agent" });
await db.update(companies).set({ status: "paused" }).where(eq(companies.id, companyId));
await budgetService(db).getInvocationBlock(companyId, one.id);
expect((await current(one.id)).lifecycleHolds).toContain("company_paused");
expect(await budgetService(db).getInvocationBlock(companyId, one.id)).not.toBeNull();
expect((await current(one.id)).lifecycleHolds).toEqual([]);
expect((await current(two.id)).lifecycleHolds).toEqual([]);
});
@@ -620,9 +641,9 @@ const support = await getEmbeddedPostgresTestSupport();
const work = worker(); const lifecycle = createAgentLifecycle(db); const agent = await hire(); await work.process(agent.id);
await lifecycle.pauseAgent(agent.id); await work.process(agent.id);
await db.update(companies).set({ status: "archived" }).where(eq(companies.id, companyId));
await createAgentLifecycle(db).reconcilePolicyHolds(companyId); await work.process(agent.id);
await projectCompanyHolds(db, createAgentLifecycleEffects(), companyId); await work.process(agent.id);
await db.update(companies).set({ status: "active" }).where(eq(companies.id, companyId));
await createAgentLifecycle(db).reconcilePolicyHolds(companyId); await work.process(agent.id);
await projectCompanyHolds(db, createAgentLifecycleEffects(), companyId); await work.process(agent.id);
expect(await current(agent.id)).toMatchObject({ lifecycleState: "paused", lifecycleHolds: ["manual"] });
await lifecycle.resumeAgent(agent.id); await work.process(agent.id);
expect((await current(agent.id)).lifecycleState).toBe("ready");
@@ -2,11 +2,15 @@ import { agents, companies, createDb, type Db } from "@paperclipai/db";
import { asc, eq } from "drizzle-orm";
import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
import { startAgentLifecycle } from "../index.js";
import { applyAgentPolicyHoldsInTransaction } from "../adapters/postgres.js";
import { createAgentLifecycleEffects } from "../../../services/agent-lifecycle.js";
import { createDeliveryWorkCoordinator } from "../../../services/delivery-work-coordinator.js";
import { DELIVERY_QUEUES } from "../../../services/delivery-work-notifications.js";
import { withAccountingTransaction } from "../../../services/accounting-transaction.js";
import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../../../__tests__/helpers/embedded-postgres.js";
const support = await getEmbeddedPostgresTestSupport();
(support.supported ? describe : describe.skip)("lifecycle policy sweep", () => {
(support.supported ? describe : describe.skip)("scoped lifecycle policy projection", () => {
let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>>;
let db: Db;
beforeAll(async () => {
@@ -15,49 +19,138 @@ const support = await getEmbeddedPostgresTestSupport();
}, 90_000);
afterAll(async () => { await database?.cleanup(); });
it("caps each policy page, skips ineligible agents, and returns to the first page", async () => {
it("projects only company mismatches in pages of 100 using one transaction per page", async () => {
const [company] = await db.insert(companies).values({ name: "Policy pages", issuePrefix: "POL" }).returning();
const eligible = await db.insert(agents).values(Array.from({ length: 101 }, (_, index) => ({
id: `00000000-0000-4000-8000-${String(index + 1).padStart(12, "0")}`,
companyId: company.id, name: `Agent ${index}`, status: "idle", lifecycleState: "ready" as const,
lifecycleHolds: ["company_paused"], pauseReason: "company_paused",
}))).returning();
const excluded = await db.insert(agents).values((["pending_approval", "terminating", "cleaning_up", "terminated", "rejected"] as const).map(lifecycleState => ({
companyId: company.id, name: lifecycleState, status: "terminated", lifecycleState,
lifecycleHolds: ["company_paused"], pauseReason: "company_paused",
// Ineligible lifecycle work is not due while this test projects policy.
lifecycleOperation: { id: "saved", hostComplete: false, completedPluginIds: [], attempts: 1, retryAt: new Date(Date.now() + 60_000).toISOString() },
}))).returning();
const effects = createAgentLifecycleEffects();
const transaction = vi.spyOn(effects, "transaction");
const now = Date.now();
const clock = vi.spyOn(Date, "now").mockReturnValue(now);
const timer = vi.spyOn(globalThis, "setInterval").mockImplementation(() => ({ unref() {} }) as ReturnType<typeof setInterval>);
const clearTimer = vi.spyOn(globalThis, "clearInterval").mockImplementation(() => {});
const policy = vi.spyOn(effects, "policyBlocks").mockRejectedValue(new Error("Projection must not scan budgets"));
const work = startAgentLifecycle(db, effects, {
requiredPluginIds: async () => [], runHost: async () => "pending", runPlugin: async () => "pending",
}, () => true);
const rows = () => db.select().from(agents).where(eq(agents.companyId, company.id)).orderBy(asc(agents.id));
try {
expect(await work.hasPending()).toBe(true);
await work.sweep();
expect(transaction).toHaveBeenCalledTimes(100);
expect(transaction).toHaveBeenCalledTimes(1);
expect((await rows()).filter(row => row.lifecycleHolds.length === 0)).toHaveLength(100);
expect(await work.hasPending()).toBe(true);
await work.sweep();
expect(transaction).toHaveBeenCalledTimes(100);
clock.mockReturnValue(now + 60_001);
await work.sweep();
expect(transaction).toHaveBeenCalledTimes(101);
expect(transaction).toHaveBeenCalledTimes(2);
expect((await rows()).filter(row => row.lifecycleHolds.length > 0).map(row => row.id).sort())
.toEqual(excluded.map(row => row.id).sort());
await db.update(agents).set({ lifecycleHolds: ["company_paused"] }).where(eq(agents.id, eligible[0].id));
clock.mockReturnValue(now + 120_002);
await work.sweep();
expect(transaction).toHaveBeenCalledTimes(201);
expect((await rows())[0].lifecycleHolds).toEqual([]);
} finally {
await work.stop();
transaction.mockRestore();
clock.mockRestore();
timer.mockRestore();
clearTimer.mockRestore();
}
expect(policy).not.toHaveBeenCalled();
expect((await rows()).filter(row => eligible.some(agent => agent.id === row.id))).toHaveLength(101);
expect(await work.hasPending()).toBe(false);
expect(work.nextRunAt()).toBeGreaterThan(Date.now());
} finally { await work.stop(); transaction.mockRestore(); policy.mockRestore(); }
});
it("repairs an inherited hold whose lifecycle phase was not projected before restart", async () => {
const [company] = await db.insert(companies).values({ name: "Interrupted projection", issuePrefix: "INT", status: "paused" }).returning();
const [agent] = await db.insert(agents).values({ companyId: company.id, name: "Interrupted agent",
status: "paused", lifecycleState: "preparing", lifecycleHolds: ["company_paused"], pauseReason: "company_paused" }).returning();
const host = vi.fn(async () => "complete" as const);
const work = startAgentLifecycle(db, createAgentLifecycleEffects(), {
requiredPluginIds: async () => [], runHost: host, runPlugin: host,
}, () => true);
try {
expect(await work.hasPending()).toBe(true);
await work.sweep();
expect((await db.select().from(agents).where(eq(agents.id, agent.id)))[0]).toMatchObject({
lifecycleState: "paused", lifecycleHolds: ["company_paused"],
});
expect(host).toHaveBeenCalledWith(expect.objectContaining({ id: agent.id, lifecycleState: "pausing" }));
expect(await work.hasPending()).toBe(false);
} finally { await work.stop(); }
});
it("does not install an idle policy polling timer", async () => {
const isolatedDatabase = { select: vi.fn() } as unknown as Db;
const timer = vi.spyOn(globalThis, "setInterval");
const work = startAgentLifecycle(isolatedDatabase, createAgentLifecycleEffects(), {
requiredPluginIds: async () => [], runHost: async () => "complete", runPlugin: async () => "complete",
}, () => true);
try {
expect(timer).not.toHaveBeenCalled();
expect(isolatedDatabase.select).not.toHaveBeenCalled();
} finally { await work.stop(); timer.mockRestore(); }
});
it("applies a supplied agent budget hold inside the caller transaction without reading policies", async () => {
const [company] = await db.insert(companies).values({ name: "Scoped hold", issuePrefix: "SCP" }).returning();
const [target, other] = await db.insert(agents).values(["Target", "Other"].map(name => ({
companyId: company.id, name, status: "paused", lifecycleState: "paused" as const,
lifecycleHolds: ["manual"], pauseReason: "manual",
}))).returning();
const effects = createAgentLifecycleEffects();
const transaction = vi.spyOn(effects, "transaction");
const policy = vi.spyOn(effects, "policyBlocks").mockRejectedValue(new Error("No policy lookup allowed"));
try {
await withAccountingTransaction(db, company.id, tx =>
applyAgentPolicyHoldsInTransaction(tx, effects, company.id, target.id, { budget: true }));
expect(transaction).not.toHaveBeenCalled();
expect(policy).not.toHaveBeenCalled();
const current = await db.select().from(agents).where(eq(agents.companyId, company.id));
expect(current.find(row => row.id === target.id)?.lifecycleHolds).toEqual(["manual", "budget"]);
expect(current.find(row => row.id === other.id)?.lifecycleHolds).toEqual(["manual"]);
await expect(withAccountingTransaction(db, company.id, async tx => {
await applyAgentPolicyHoldsInTransaction(tx, effects, company.id, target.id, { budget: false });
throw new Error("Rollback accounting");
})).rejects.toThrow("Rollback accounting");
expect((await db.select().from(agents).where(eq(agents.id, target.id)))[0].lifecycleHolds).toEqual(["manual", "budget"]);
await withAccountingTransaction(db, company.id, tx =>
applyAgentPolicyHoldsInTransaction(tx, effects, company.id, target.id, { budget: false }));
expect((await db.select().from(agents).where(eq(agents.id, target.id)))[0]).toMatchObject({
lifecycleState: "paused", lifecycleHolds: ["manual"], pauseReason: "manual",
});
} finally { transaction.mockRestore(); policy.mockRestore(); }
});
it("recovers committed holds on startup and wakes after new commits without idle polling", async () => {
const [company] = await db.insert(companies).values({ name: "Durable lifecycle", issuePrefix: "DUR" }).returning();
const [recover, later] = await db.insert(agents).values(["Recover", "Later"].map(name => ({
companyId: company.id, name, status: "idle", lifecycleState: "ready" as const,
}))).returning();
const effects = createAgentLifecycleEffects();
await withAccountingTransaction(db, company.id, tx =>
applyAgentPolicyHoldsInTransaction(tx, effects, company.id, recover.id, { budget: true }));
const host = vi.fn(async () => "complete" as const);
const work = startAgentLifecycle(db, effects, { requiredPluginIds: async () => [], runHost: host, runPlugin: host }, () => true);
const coordinator = createDeliveryWorkCoordinator({ owner: db, canRun: () => true, onError: error => { throw error; } });
const registered = coordinator.register(DELIVERY_QUEUES.agentLifecycle, {
retryMs: 10, run: () => work.sweep(), hasPending: () => work.hasPending(), nextRunAt: () => work.nextRunAt(),
});
work.setWake(registered.wake);
const current = async (id: string) => (await db.select().from(agents).where(eq(agents.id, id)))[0];
try {
await registered.ready;
expect(await current(recover.id)).toMatchObject({ lifecycleState: "paused", lifecycleHolds: ["budget"] });
await expect(withAccountingTransaction(db, company.id, async tx => {
await applyAgentPolicyHoldsInTransaction(tx, effects, company.id, later.id, { budget: true });
throw new Error("Rollback producer");
})).rejects.toThrow("Rollback producer");
expect(await current(later.id)).toMatchObject({ lifecycleState: "ready", lifecycleHolds: [] });
await withAccountingTransaction(db, company.id, tx =>
applyAgentPolicyHoldsInTransaction(tx, effects, company.id, later.id, { budget: true }));
await vi.waitFor(async () => expect(await current(later.id)).toMatchObject({ lifecycleState: "paused", lifecycleHolds: ["budget"] }));
await new Promise(resolve => setTimeout(resolve, 50));
const reads = vi.spyOn(db, "select");
try {
await new Promise(resolve => setTimeout(resolve, 100));
expect(reads).not.toHaveBeenCalled();
} finally { reads.mockRestore(); }
expect(host).toHaveBeenCalledTimes(2);
} finally { await coordinator.stop(); await work.stop(); }
});
});
@@ -0,0 +1,73 @@
import { agents, companies, heartbeatRuns, type Db } from "@paperclipai/db";
import { and, asc, eq, inArray, sql, type SQLWrapper } from "drizzle-orm";
import type { LifecycleEffects } from "./effects.js";
import { applyAgentPolicyHoldsInTransaction } from "./postgres.js";
const inheritedHoldStates = ["preparing", "verifying", "ready", "pausing", "paused", "resuming"] as const;
/** Durable mismatch, also recoverable after a lost notification or a restart. */
export function companyHoldMismatch(companyStatus: SQLWrapper = companies.status) {
return and(
inArray(agents.lifecycleState, inheritedHoldStates),
sql`((coalesce(${agents.lifecycleHolds}, '[]'::jsonb) @> '["company_paused"]'::jsonb) <> (${companyStatus} = 'paused')
or (coalesce(${agents.lifecycleHolds}, '[]'::jsonb) @> '["company_archived"]'::jsonb) <> (${companyStatus} = 'archived')
or (${companyStatus} in ('paused', 'archived') and ${agents.lifecycleState} in ('preparing', 'verifying', 'ready', 'resuming')))`,
);
}
export async function hasCompanyHoldMismatch(db: Db) {
const rows = await db.select({ id: agents.id }).from(agents).innerJoin(companies, eq(companies.id, agents.companyId))
.where(companyHoldMismatch()).limit(1);
return rows.length > 0;
}
/** One accounting transaction per changed company page, never one per member. */
export async function projectCompanyHolds(db: Db, effects: LifecycleEffects, companyId?: string) {
const [candidate] = await db.select({ companyId: agents.companyId }).from(agents)
.innerJoin(companies, eq(companies.id, agents.companyId))
.where(and(companyHoldMismatch(), companyId ? eq(agents.companyId, companyId) : undefined))
.orderBy(asc(agents.companyId)).limit(1);
if (!candidate) return 0;
return effects.transaction(db, candidate.companyId, async tx => {
const [company] = await tx.select({ status: companies.status }).from(companies).where(eq(companies.id, candidate.companyId));
if (!company) return 0;
// Company is locked by the accounting boundary; select only stale projections.
const rows = await tx.select().from(agents).where(and(
eq(agents.companyId, candidate.companyId),
companyHoldMismatch(sql`${company.status}::text`),
)).orderBy(asc(agents.id)).limit(100).for("update", { skipLocked: true });
const stopping = company.status === "paused" || company.status === "archived";
const running = stopping && rows.length ? await tx.selectDistinct({ agentId: heartbeatRuns.agentId }).from(heartbeatRuns).where(and(
eq(heartbeatRuns.companyId, candidate.companyId), inArray(heartbeatRuns.agentId, rows.map(row => row.id)),
inArray(heartbeatRuns.status, ["queued", "running", "scheduled_retry"]),
)) : [];
const activeAgents = new Set(running.map(row => row.agentId));
for (const row of rows) {
await applyAgentPolicyHoldsInTransaction(tx, effects, candidate.companyId, row.id,
{ companyStatus: company.status, restartPause: activeAgents.has(row.id) }, row);
}
return rows.length;
});
}
/** Persist stop ownership for involved agents before a company pause commits.
* Later reactivation may remove the inherited hold, but cannot erase pausing. */
export async function fenceActiveCompanyAgentStopsInTransaction(tx: Db, effects: LifecycleEffects,
companyId: string, status: "paused" | "archived") {
const active = await tx.selectDistinct({ id: heartbeatRuns.agentId }).from(heartbeatRuns).where(and(
eq(heartbeatRuns.companyId, companyId), inArray(heartbeatRuns.status, ["queued", "running", "scheduled_retry"]),
)).orderBy(asc(heartbeatRuns.agentId));
let fenced = 0;
for (let offset = 0; offset < active.length; offset += 100) {
const rows = await tx.select().from(agents).where(and(eq(agents.companyId, companyId),
inArray(agents.id, active.slice(offset, offset + 100).map(agent => agent.id)),
inArray(agents.lifecycleState, inheritedHoldStates),
)).orderBy(asc(agents.id)).for("update");
for (const row of rows) {
await applyAgentPolicyHoldsInTransaction(tx, effects, companyId, row.id,
{ companyStatus: status, restartPause: true }, row);
fenced += 1;
}
}
return fenced;
}
@@ -16,6 +16,7 @@ export async function updateAgentConfiguration(tx: Db, effects: LifecycleEffects
const changed = ["adapterType", "adapterConfig", "runtimeConfig", "defaultEnvironmentId"].some(key =>
Object.prototype.hasOwnProperty.call(patch, key) && !jsonEqual(patch[key as keyof typeof patch], current[key as keyof typeof current]));
const invalidate = changed && ["preparing", "verifying", "resuming"].includes(current.lifecycleState);
if (invalidate) await effects.notifyWork(tx);
const [updated] = await tx.update(agents).set({ ...patch, updatedAt: new Date(), ...(invalidate ? {
lifecycleVersion: current.lifecycleVersion + 1, lifecycleError: null,
lifecycleOperation: { ...current.lifecycleOperation!, id: randomUUID(), hostComplete: false,
@@ -6,6 +6,7 @@ import type { ActivityPublication } from "../../../types/activity-publication.js
/** Services supply integrations; lifecycle owns the transaction and state writes. */
export interface LifecycleEffects {
notifyWork(db: Db): Promise<void>;
reportFailure(context: LifecycleFailure, error: unknown): void;
deleteDependencies(db: Db, companyId: string, agentId: string): Promise<void>;
prepareConfiguration(db: Db, current: AgentConfigurationRecord, data: AgentConfigurationPatch, options?: UpdateAgentOptions): Promise<AgentConfigurationPatch>;
@@ -13,7 +14,6 @@ export interface LifecycleEffects {
transaction<T>(db: Db, companyId: string, work: (tx: Db, publications: ActivityPublication[]) => Promise<T>): Promise<T>;
policyBlocks(db: Db, policy: typeof budgetPolicies.$inferSelect): Promise<boolean>;
setAgentBudget(db: Db, publications: ActivityPublication[], companyId: string, agentId: string, amount: number, userId: string | null, isActive?: boolean): Promise<unknown>;
enforceBudget(db: Db, companyId: string, agentId: string): Promise<void>;
recordCreation(db: Db, companyId: string, agentId: string): Promise<void>;
recordStatus(db: Db, companyId: string, agentId: string, before: string, after: string): Promise<void>;
clearPrimary(db: Db, companyId: string, agentId: string): Promise<unknown>;
@@ -2,7 +2,7 @@ import { randomUUID } from "node:crypto";
import { and, asc, eq, inArray, sql } from "drizzle-orm";
import { agents, agentApiKeys, companies, budgetPolicies, type Db } from "@paperclipai/db";
import type { AgentLifecycleOperation } from "@paperclipai/shared";
import { AgentLifecycleConflict, compatibilityStatus, transition } from "../domain/policy.js";
import { AgentLifecycleConflict, compatibilityStatus, isLifecycleWorkPending, transition, type LifecycleCommand } from "../domain/policy.js";
import type { LifecycleEffects } from "./effects.js";
import type { LifecycleStore } from "../application/ports.js";
@@ -10,6 +10,109 @@ export function assertRootDatabase(db: Db) {
if ("nestedIndex" in db) throw new AgentLifecycleConflict("Lifecycle commands require the root database connection");
}
type LifecycleRow = typeof agents.$inferSelect;
type ChangeOptions = { version?: number; owner?: string; reason?: string };
/** The caller owns the accounting transaction and has locked this agent. */
async function changeLockedAgent(tx: Db, effects: LifecycleEffects, agent: LifecycleRow,
command: LifecycleCommand, options: ChangeOptions = {}, policyHolds?: string[], restartPause = false) {
if (options.version !== undefined && options.version !== agent.lifecycleVersion) return null;
if (command === "resume" && agent.lifecycleHolds.includes("setup")) throw new AgentLifecycleConflict("Finish agent setup before resuming");
if (command === "complete" && (!options.owner || !agent.lifecycleOperation?.leaseUntil || agent.lifecycleOperation?.leaseOwner !== options.owner ||
Date.parse(agent.lifecycleOperation?.leaseUntil ?? "") <= Date.now() ||
!agent.lifecycleOperation.hostComplete || agent.lifecycleRequiredPluginIds === null ||
agent.lifecycleRequiredPluginIds.some(id => !agent.lifecycleOperation!.completedPluginIds.includes(id)))) return null;
if (command === "retry") {
transition(agent.lifecycleState, command);
if (agent.lifecycleOperation?.leaseUntil && Date.parse(agent.lifecycleOperation.leaseUntil) > Date.now()) return agent;
await effects.notifyWork(tx);
const [updated] = await tx.update(agents).set({ lifecycleError: null,
lifecycleOperation: sql`${agents.lifecycleOperation} - 'retryAt' - 'leaseOwner' - 'leaseUntil'`, updatedAt: new Date() }).where(eq(agents.id, agent.id)).returning();
return updated;
}
if (command === "resume" && !["ready", "paused", "pausing", "resuming"].includes(agent.lifecycleState)) transition(agent.lifecycleState, command);
let holds = policyHolds ?? agent.lifecycleHolds;
const reason = options.reason ?? "manual";
if (command === "reconcile" && policyHolds === undefined) {
if (["pending_approval", "terminating", "cleaning_up", "terminated", "rejected"].includes(agent.lifecycleState)) return agent;
const [company] = await tx.select().from(companies).where(eq(companies.id, agent.companyId));
const policies = await tx.select().from(budgetPolicies).where(and(eq(budgetPolicies.companyId, agent.companyId),
eq(budgetPolicies.scopeType, "agent"), eq(budgetPolicies.scopeId, agent.id)));
let budgetBlocked = false;
for (const policy of policies) if (await effects.policyBlocks(tx, policy)) budgetBlocked = true;
holds = holds.filter(hold => hold !== "budget" && hold !== "company_archived" && hold !== "company_paused");
if (budgetBlocked) holds.push("budget");
if (company.status === "archived") holds.push("company_archived");
if (company.status === "paused") holds.push("company_paused");
}
if (command === "pause") holds = [...new Set([...holds, reason])];
if (command === "resume") holds = holds.filter(hold => reason === "user" ? ["budget", "company_archived", "company_paused"].includes(hold) : hold !== reason);
let next = command === "reconcile"
? holds.length && !["pausing", "paused"].includes(agent.lifecycleState) ? "pausing" as const
: !holds.length && agent.lifecycleState === "paused" ? "resuming" as const : agent.lifecycleState
: command === "resume" && (holds.length || agent.lifecycleState === "pausing") ? agent.lifecycleState
: transition(agent.lifecycleState, command, agent.lifecycleOperation?.resumeState);
if (restartPause && next === "paused") next = "pausing";
if (command === "complete" && next === "ready" && holds.length) next = "pausing";
if (command === "complete" && next === "paused" && !holds.length) next = "resuming";
if (next === agent.lifecycleState) {
if (JSON.stringify(holds) !== JSON.stringify(agent.lifecycleHolds) || agent.pauseReason !== (holds[0] ?? null)) {
const [updated] = await tx.update(agents).set({ lifecycleHolds: holds, pauseReason: holds[0] ?? null }).where(eq(agents.id, agent.id)).returning();
return updated;
}
return agent;
}
const operation: AgentLifecycleOperation = {
responsibleUserId: agent.lifecycleOperation?.responsibleUserId,
id: randomUUID(), hostComplete: false, completedPluginIds: [], attempts: 0,
resumeState: next === "pausing"
? agent.lifecycleState === "preparing" || agent.lifecycleState === "verifying" ? agent.lifecycleState : agent.lifecycleOperation?.resumeState ?? "ready"
: agent.lifecycleOperation?.resumeState,
};
if (next === "ready") operation.resumeState = "ready";
const status = compatibilityStatus(next, agent.status);
if (isLifecycleWorkPending(next)) await effects.notifyWork(tx);
const [updated] = await tx.update(agents).set({
lifecycleHolds: holds, lifecycleState: next, lifecycleVersion: agent.lifecycleVersion + 1, lifecycleError: null,
lifecycleOperation: operation, status,
pauseReason: holds[0] ?? null,
pausedAt: next === "pausing" ? new Date() : next === "ready" ? null : agent.pausedAt,
updatedAt: new Date(),
}).where(eq(agents.id, agent.id)).returning();
if (next === "terminating" || next === "rejected") {
await tx.update(agentApiKeys).set({ revokedAt: new Date() }).where(eq(agentApiKeys.agentId, agent.id));
await effects.clearPrimary(tx, agent.companyId, agent.id);
}
if (command === "approve") await effects.recordCreation(tx, agent.companyId, agent.id);
if (next === "pausing" || next === "resuming" || next === "terminating" || next === "rejected") {
await effects.recordStatus(tx, agent.companyId, agent.id,
next === "resuming" ? "paused" : "idle",
next === "resuming" ? "idle" : next === "pausing" ? "paused" : "terminated");
}
return updated;
}
/** Apply a computed hold without resolving policies or opening another transaction. */
export async function applyAgentPolicyHoldsInTransaction(tx: Db, effects: LifecycleEffects,
companyId: string, agentId: string, patch: { budget?: boolean; companyStatus?: string; restartPause?: boolean }, lockedAgent?: LifecycleRow) {
if (!("nestedIndex" in tx)) throw new AgentLifecycleConflict("Policy holds require the caller's accounting transaction");
const agent = lockedAgent ?? (await tx.select().from(agents).where(and(
eq(agents.id, agentId), eq(agents.companyId, companyId),
)).for("update"))[0];
if (!agent || agent.companyId !== companyId || ["pending_approval", "terminating", "cleaning_up", "terminated", "rejected"].includes(agent.lifecycleState)) return null;
let holds = [...agent.lifecycleHolds];
if (patch.budget !== undefined) {
holds = holds.filter(hold => hold !== "budget");
if (patch.budget) holds.push("budget");
}
if (patch.companyStatus !== undefined) {
holds = holds.filter(hold => hold !== "company_paused" && hold !== "company_archived");
if (patch.companyStatus === "paused") holds.push("company_paused");
if (patch.companyStatus === "archived") holds.push("company_archived");
}
return changeLockedAgent(tx, effects, agent, "reconcile", {}, holds, patch.restartPause);
}
export function createLifecycleStore(db: Db, effects: LifecycleEffects): LifecycleStore {
const get = (id: string) => db.select().from(agents).where(eq(agents.id, id)).then(rows => rows[0] ?? null);
return {
@@ -20,76 +123,7 @@ export function createLifecycleStore(db: Db, effects: LifecycleEffects): Lifecyc
return effects.transaction(db, existing.companyId, async tx => {
const [agent] = await tx.select().from(agents).where(eq(agents.id, id)).for("update", command === "reconcile" ? { skipLocked: true } : undefined);
if (!agent) return null;
if (options.version !== undefined && options.version !== agent.lifecycleVersion) return null;
if (command === "resume" && agent.lifecycleHolds.includes("setup")) throw new AgentLifecycleConflict("Finish agent setup before resuming");
if (command === "complete" && (!options.owner || !agent.lifecycleOperation?.leaseUntil || agent.lifecycleOperation?.leaseOwner !== options.owner ||
Date.parse(agent.lifecycleOperation?.leaseUntil ?? "") <= Date.now() ||
!agent.lifecycleOperation.hostComplete || agent.lifecycleRequiredPluginIds === null ||
agent.lifecycleRequiredPluginIds.some(id => !agent.lifecycleOperation!.completedPluginIds.includes(id)))) return null;
if (command === "retry") {
transition(agent.lifecycleState, command);
if (agent.lifecycleOperation?.leaseUntil && Date.parse(agent.lifecycleOperation.leaseUntil) > Date.now()) return agent;
const [updated] = await tx.update(agents).set({ lifecycleError: null,
lifecycleOperation: sql`${agents.lifecycleOperation} - 'retryAt' - 'leaseOwner' - 'leaseUntil'`, updatedAt: new Date() }).where(eq(agents.id, id)).returning();
return updated;
}
if (command === "resume" && !["ready", "paused", "pausing", "resuming"].includes(agent.lifecycleState)) transition(agent.lifecycleState, command);
let holds = agent.lifecycleHolds;
const reason = options.reason ?? "manual";
if (command === "reconcile") {
if (["pending_approval", "terminating", "cleaning_up", "terminated", "rejected"].includes(agent.lifecycleState)) return agent;
const [company] = await tx.select().from(companies).where(eq(companies.id, agent.companyId));
const policies = await tx.select().from(budgetPolicies).where(and(eq(budgetPolicies.companyId, agent.companyId),
eq(budgetPolicies.scopeType, "agent"), eq(budgetPolicies.scopeId, id)));
let budgetBlocked = false;
for (const policy of policies) if (await effects.policyBlocks(tx, policy)) budgetBlocked = true;
holds = holds.filter(hold => hold !== "budget" && hold !== "company_archived" && hold !== "company_paused");
if (budgetBlocked) holds.push("budget");
if (company.status === "archived") holds.push("company_archived");
if (company.status === "paused") holds.push("company_paused");
}
if (command === "pause") holds = [...new Set([...holds, reason])];
if (command === "resume") holds = holds.filter(hold => reason === "user" ? ["budget", "company_archived", "company_paused"].includes(hold) : hold !== reason);
let next = command === "reconcile"
? holds.length && !["pausing", "paused"].includes(agent.lifecycleState) ? "pausing" as const
: !holds.length && agent.lifecycleState === "paused" ? "resuming" as const : agent.lifecycleState
: command === "resume" && (holds.length || agent.lifecycleState === "pausing") ? agent.lifecycleState
: transition(agent.lifecycleState, command, agent.lifecycleOperation?.resumeState);
if (command === "complete" && next === "ready" && holds.length) next = "pausing";
if (command === "complete" && next === "paused" && !holds.length) next = "resuming";
if (next === agent.lifecycleState) {
if (JSON.stringify(holds) !== JSON.stringify(agent.lifecycleHolds) || agent.pauseReason !== (holds[0] ?? null)) {
await tx.update(agents).set({ lifecycleHolds: holds, pauseReason: holds[0] ?? null }).where(eq(agents.id, id));
}
return agent;
}
const operation: AgentLifecycleOperation = {
responsibleUserId: agent.lifecycleOperation?.responsibleUserId,
id: randomUUID(), hostComplete: false, completedPluginIds: [], attempts: 0,
resumeState: next === "pausing"
? agent.lifecycleState === "preparing" || agent.lifecycleState === "verifying" ? agent.lifecycleState : agent.lifecycleOperation?.resumeState ?? "ready"
: agent.lifecycleOperation?.resumeState,
};
if (next === "ready") operation.resumeState = "ready";
const status = compatibilityStatus(next, agent.status);
const [updated] = await tx.update(agents).set({
lifecycleHolds: holds, lifecycleState: next, lifecycleVersion: agent.lifecycleVersion + 1, lifecycleError: null,
lifecycleOperation: operation, status,
pauseReason: holds[0] ?? null,
pausedAt: next === "pausing" ? new Date() : next === "ready" ? null : agent.pausedAt,
updatedAt: new Date(),
}).where(eq(agents.id, id)).returning();
if (next === "terminating" || next === "rejected") {
await tx.update(agentApiKeys).set({ revokedAt: new Date() }).where(eq(agentApiKeys.agentId, id));
await effects.clearPrimary(tx, agent.companyId, id);
}
if (command === "approve") await effects.recordCreation(tx, agent.companyId, id);
if (next === "pausing" || next === "resuming" || next === "terminating" || next === "rejected") {
await effects.recordStatus(tx, agent.companyId, id,
next === "resuming" ? "paused" : "idle",
next === "resuming" ? "idle" : next === "pausing" ? "paused" : "terminated");
}
return updated;
return changeLockedAgent(tx, effects, agent, command, options);
});
},
async claim(id, owner, now) {
@@ -98,6 +132,8 @@ export function createLifecycleStore(db: Db, effects: LifecycleEffects): Lifecyc
|| jsonb_build_object('leaseOwner', ${owner}::text, 'leaseUntil', ${(new Date(now.getTime() + 120_000)).toISOString()}::text,
'attempts', coalesce((${agents.lifecycleOperation}->>'attempts')::integer, 0) + 1)`,
}).where(and(eq(agents.id, id), inArray(agents.lifecycleState, ["preparing", "verifying", "pausing", "resuming", "terminating", "cleaning_up"]),
sql`(${agents.lifecycleState} in ('pausing', 'terminating', 'cleaning_up') or exists (
select 1 from ${companies} where ${companies.id} = ${agents.companyId} and ${companies.status} = 'active'))`,
sql`coalesce((${agents.lifecycleOperation}->>'leaseUntil')::timestamptz, '-infinity') <= ${now.toISOString()}::timestamptz`,
sql`coalesce((${agents.lifecycleOperation}->>'retryAt')::timestamptz, '-infinity') <= ${now.toISOString()}::timestamptz`,
)).returning();
@@ -1,7 +1,7 @@
import type { LifecycleEffects } from "./effects.js";
import { randomUUID } from "node:crypto";
import { and, eq, sql } from "drizzle-orm";
import { agents, agentApiKeys, type Db } from "@paperclipai/db";
import { agents, agentApiKeys, companies, type Db } from "@paperclipai/db";
import { conflict, notFound } from "../../../errors.js";
import { agentRecordQueries, type CreateAgentOptions, type CreateAgentData } from "../../../lib/agent-records.js";
@@ -32,14 +32,20 @@ export function agentRecords(db: Db, effects: LifecycleEffects) {
const prepared = await effects.prepareHire(txDb, companyId, data, options);
const status = options?.setupOnly ? "paused" : data.status;
const pauseReason = options?.setupOnly ? "setup" : data.pauseReason;
const [company] = await tx.select({ status: companies.status }).from(companies).where(eq(companies.id, companyId));
const companyHolds = company?.status === "paused" ? ["company_paused"] : company?.status === "archived" ? ["company_archived"] : [];
const holds = [...(status === "paused" ? [pauseReason ?? "manual"] : []), ...companyHolds];
const lifecycleState = status === "pending_approval" ? "pending_approval" : status === "terminated" ? "terminated"
: status === "paused" ? "paused" : companyHolds.length ? "pausing" : "preparing";
if (lifecycleState === "preparing" || lifecycleState === "pausing") await effects.notifyWork(tx);
const values: typeof agents.$inferInsert = {
...prepared,
createdByUserId: setup?.createdByUserId ?? options?.createdByUserId ?? null,
lifecycleState: status === "pending_approval" ? "pending_approval" : status === "terminated" ? "terminated" : status === "paused" ? "paused" : "preparing",
lifecycleState,
lifecycleVersion: setup ? setup.lifecycleVersion + 1 : 1,
lifecycleHolds: status === "paused" ? [pauseReason ?? "manual"] : [],
pauseReason: status === "paused" ? pauseReason ?? "manual" : null,
pausedAt: status === "paused" ? new Date() : null,
lifecycleHolds: holds,
pauseReason: holds[0] ?? null,
pausedAt: holds.length ? new Date() : null,
lifecycleOperation: { id: randomUUID(), hostComplete: false, completedPluginIds: [], attempts: 0, resumeState: "preparing" as const, responsibleUserId: options?.responsibleUserId ?? options?.createdByUserId },
status: status === "pending_approval" ? "pending_approval" : status === "terminated" ? "terminated" : "paused",
companyId,
@@ -109,9 +115,14 @@ export function agentRecords(db: Db, effects: LifecycleEffects) {
const existing = await agentRecordQueries(txDb).getById(id);
if (!existing || existing.status !== "pending_approval") return null;
const patch = await effects.prepareHireApproval(txDb, existing, approvedPayload);
const [company] = await tx.select({ status: companies.status }).from(companies).where(eq(companies.id, existing.companyId));
const holds = existing.lifecycleHolds.filter(hold => hold !== "company_paused" && hold !== "company_archived");
if (company?.status === "paused") holds.push("company_paused");
if (company?.status === "archived") holds.push("company_archived");
await effects.notifyWork(txDb);
const updated = await tx
.update(agents)
.set({ ...patch, status: "paused", lifecycleState: "preparing", lifecycleVersion: existing.lifecycleVersion + 1, lifecycleError: null, lifecycleOperation: sql`jsonb_build_object('id', ${randomUUID()}::text, 'hostComplete', false, 'completedPluginIds', '[]'::jsonb, 'attempts', 0, 'responsibleUserId', coalesce(nullif(${agents.lifecycleOperation}->'responsibleUserId', 'null'::jsonb), to_jsonb(${requestedByUserId ?? null}::text)))`, updatedAt: new Date() })
.set({ ...patch, status: "paused", lifecycleState: holds.length ? "pausing" : "preparing", lifecycleHolds: holds, pauseReason: holds[0] ?? null, lifecycleVersion: existing.lifecycleVersion + 1, lifecycleError: null, lifecycleOperation: sql`jsonb_build_object('id', ${randomUUID()}::text, 'hostComplete', false, 'completedPluginIds', '[]'::jsonb, 'attempts', 0, 'responsibleUserId', coalesce(nullif(${agents.lifecycleOperation}->'responsibleUserId', 'null'::jsonb), to_jsonb(${requestedByUserId ?? null}::text)))`, updatedAt: new Date() })
.where(and(eq(agents.id, id), eq(agents.status, "pending_approval")))
.returning()
.then((rows) => rows[0] ?? null);
@@ -6,6 +6,7 @@ export interface LifecycleAgent {
companyId: string;
lifecycleState: AgentLifecycleState;
lifecycleVersion: number;
lifecycleHolds?: readonly string[];
lifecycleRequiredPluginIds: string[] | null;
lifecycleError: string | null;
lifecycleOperation: AgentLifecycleOperation | null;
+43 -45
View File
@@ -1,6 +1,6 @@
import { hireApprovalService as approvalRecords, type HireDecisionTarget } from "./adapters/approvals.js";
import { agents } from "@paperclipai/db";
import { and, asc, eq, gt, inArray } from "drizzle-orm";
import { inArray, sql } from "drizzle-orm";
import type { Db } from "@paperclipai/db";
import type { LifecycleEffects } from "./adapters/effects.js";
export type { LifecycleEffects } from "./adapters/effects.js";
@@ -9,16 +9,20 @@ import type { AgentConfigurationPatch, UpdateAgentOptions } from "../../lib/agen
import { purgeAgent } from "./adapters/deletion.js";
import { agentRecords } from "./adapters/records.js";
import { assertRootDatabase, createLifecycleStore } from "./adapters/postgres.js";
import { hasCompanyHoldMismatch, projectCompanyHolds } from "./adapters/company-policy.js";
import { createLifecycleWorker } from "./application/worker.js";
import type { LifecycleDriver } from "./application/ports.js";
export { applyAgentPolicyHoldsInTransaction } from "./adapters/postgres.js";
export { companyHoldMismatch, fenceActiveCompanyAgentStopsInTransaction } from "./adapters/company-policy.js";
export { AgentLifecycleConflict, canConfigureAgentConnection, isAgentAwaitingSetup } from "./domain/policy.js";
export type { HireDecisionTarget } from "./adapters/approvals.js";
export type { LifecycleDriver, LifecycleAgent, LifecycleFailure } from "./application/ports.js";
const policyStates = ["preparing", "verifying", "ready", "pausing", "paused", "resuming"] as const;
const policyPageSize = 100;
const pendingStates = ["preparing", "verifying", "pausing", "resuming", "terminating", "cleaning_up"] as const;
const wakes = new WeakMap<Db, () => void>();
const workers = new WeakMap<Db, ReturnType<typeof createLifecycleWorker>>();
export function configureAgentLifecycle(db: Db, effects: LifecycleEffects, driver: LifecycleDriver, canRun = () => true) {
@@ -28,10 +32,18 @@ export function configureAgentLifecycle(db: Db, effects: LifecycleEffects, drive
return worker;
}
export function scheduleAgentLifecycle(db: Db, id: string) {
// The periodic sweep retries work if the process stops before this call.
const wake = wakes.get(db);
if (wake) { wake(); return; }
// Standalone/test owners may explicitly configure a worker without a coordinator.
return workers.get(db)?.process(id).catch(() => {});
}
/** Destructive deletion must establish this agent's completed cleanup before
* purging dependent records. Ordinary producers only schedule asynchronous work. */
export function processAgentLifecycle(db: Db, id: string): Promise<void> {
return workers.get(db)?.process(id) ?? Promise.resolve();
}
export function createAgentLifecycle(db: Db, effects: LifecycleEffects) {
assertRootDatabase(db);
const records = agentRecords(db, effects);
@@ -66,7 +78,6 @@ export function createAgentLifecycle(db: Db, effects: LifecycleEffects) {
if (!existing) return null;
const agent = await effects.transaction(db, existing.companyId, (tx, publications) =>
updateAgentConfiguration(tx, effects, id, data, options, publications));
if (data.budgetMonthlyCents !== undefined) await effects.enforceBudget(db, existing.companyId, id);
scheduleAgentLifecycle(db, id);
return agent;
},
@@ -78,21 +89,9 @@ export function createAgentLifecycle(db: Db, effects: LifecycleEffects) {
await createLifecycleStore(tx, effects).change(id, command, { reason: command === "resume" ? "user" : "manual" });
return agentRecords(tx, effects).getById(id);
});
if (data.budgetMonthlyCents !== undefined) await effects.enforceBudget(db, existing.companyId, id);
scheduleAgentLifecycle(db, id);
return agent;
},
async reconcilePolicyHolds(companyId: string, agentId?: string | null) {
if (agentId === null) return;
const rows = await db.select({ id: agents.id }).from(agents).where(and(
eq(agents.companyId, companyId), agentId ? eq(agents.id, agentId) : undefined,
inArray(agents.lifecycleState, policyStates),
));
for (const row of rows) {
await store.change(row.id, "reconcile");
scheduleAgentLifecycle(db, row.id);
}
},
pauseAgent: (id: string, reason = "manual") => change(id, "pause", reason),
resumeAgent: (id: string, reason = "user") => change(id, "resume", reason),
terminateAgent: (id: string) => change(id, "terminate"),
@@ -106,37 +105,36 @@ export function startAgentLifecycle(db: Db, effects: LifecycleEffects, driver: L
const worker = configureAgentLifecycle(db, effects, driver, canRun);
let running: Promise<void> | undefined;
let stopped = false;
let nextPolicySweep = 0;
let policyCursor: string | undefined;
const store = createLifecycleStore(db, effects);
function sweep() {
if (stopped || !canRun()) return Promise.resolve();
return running ??= (async () => {
if (Date.now() >= nextPolicySweep) {
try {
const rows = await db.select({ id: agents.id }).from(agents).where(and(
inArray(agents.lifecycleState, policyStates), policyCursor ? gt(agents.id, policyCursor) : undefined,
)).orderBy(asc(agents.id)).limit(policyPageSize);
for (const row of rows) {
await store.change(row.id, "reconcile");
scheduleAgentLifecycle(db, row.id);
policyCursor = row.id;
}
if (rows.length < policyPageSize) policyCursor = undefined;
}
catch (error) { effects.reportFailure({ stage: "policy_scan" }, error); throw error; }
nextPolicySweep = Date.now() + 60_000;
}
await worker.sweep();
})().finally(() => { running = undefined; });
}
const timer = setInterval(() => { void sweep().catch(() => {}); }, 5_000);
timer.unref?.();
let nextDeadline: number | null = null;
// The coordinator supplies the wake after construction. No direct work escapes it.
wakes.set(db, () => {});
return {
sweep,
setWake(wake: () => void) { wakes.set(db, wake); },
nextRunAt: () => nextDeadline,
async hasPending() {
if (stopped) { nextDeadline = null; return false; }
const [pending] = await db.select({ due: sql<number | null>`min(greatest(
coalesce(extract(epoch from (${agents.lifecycleOperation}->>'leaseUntil')::timestamptz) * 1000, 0),
coalesce(extract(epoch from (${agents.lifecycleOperation}->>'retryAt')::timestamptz) * 1000, 0)))::double precision` })
.from(agents).where(inArray(agents.lifecycleState, pendingStates));
const due = pending?.due == null ? null : Number(pending.due);
const now = Date.now();
// Already-due work uses the coordinator's bounded retry. A zero-delay
// deadline would spin if plugins are not ready to process it yet.
nextDeadline = due !== null && due > now ? due : null;
return (due !== null && due <= now) || await hasCompanyHoldMismatch(db);
},
sweep() {
if (stopped || !canRun()) return Promise.resolve();
return running ??= (async () => {
try { await projectCompanyHolds(db, effects); }
catch (error) { effects.reportFailure({ stage: "policy_scan" }, error); throw error; }
await worker.sweep();
})().finally(() => { running = undefined; });
},
async stop() {
stopped = true;
clearInterval(timer);
wakes.delete(db);
workers.delete(db);
try { await running; } finally { await worker.stop(); }
},
+14 -2
View File
@@ -13,7 +13,8 @@ import { conflict, forbidden, notFound, unprocessable } from "../../errors.js";
import { museIdentity, museCredentialHash, assertMuseCurrentAuthority, assertMuseSubjectCredential, inspectMuseConnectionCredential, assertMuseCleanupCredential, museCleanupBoundaryIds, assertMuseDisabledStopControlCredential, museStopControlBoundaryIds } from "../../services/muse-identity.js";
import { instanceSettingsService } from "../../services/instance-settings.js";
import { updateAgentConfigurationInTransaction } from "../../services/agent-configuration-transaction.js";
import { canConfigureAgentConnection, scheduleAgentLifecycle } from "../agent-lifecycle/index.js";
import { AgentLifecycleConflict, canConfigureAgentConnection } from "../agent-lifecycle/index.js";
import { createAgentLifecycle } from "../../services/agent-lifecycle.js";
import { authorizationService, issueReadSqlCondition, type AuthorizationActor } from "../../services/authorization.js";
import { logActivity, publishActivity, type ActivityPublication } from "../../services/activity-log.js";
import { issueService } from "../../services/issues.js";
@@ -383,7 +384,18 @@ function createMuseBroker(db:Db, publicOrigin?:string) {
await logActivity(tx as unknown as Db,{companyId:b.companyId,actorType:"agent",actorId:b.agentId,agentId:b.agentId,responsibleUserIdOverride:b.operatorId,action:"muse.background_reply_verified",entityType:"agent",entityId:b.agentId,details:{bindingId:b.id,generation:b.generation,authorizingUserId:b.operatorId,requestId:c.requestId}},publications);
});
publications.forEach(publishActivity);
void scheduleAgentLifecycle(db,b.agentId);
// The completed challenge changes this hire's verification facts. Clear
// its deferred retry durably; a wake alone cannot bypass that deadline.
const lifecycle=createAgentLifecycle(db),agent=await lifecycle.get(b.agentId);
if(agent&&["preparing","verifying"].includes(agent.lifecycleState)) {
try { await lifecycle.retry(b.agentId); }
catch(error) {
// A concurrent lifecycle completion or Stop owns the newer phase.
// Confirmation remains accepted; it must not restart that phase.
const current=error instanceof AgentLifecycleConflict?await lifecycle.get(b.agentId):null;
if(!current||["preparing","verifying"].includes(current.lifecycleState))throw error;
}
}
return {status:"ready",bindingId:b.id,generation:b.generation};
}
if(c.command==="task.create"||c.command==="task.comment")return idleMutation(s,c);
@@ -1,10 +1,16 @@
import { randomUUID } from "node:crypto";
import { createServer } from "node:http";
import { eq, sql } from "drizzle-orm";
import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest";
import {
activityLog,
agents,
approvals,
budgetIncidents,
budgetPolicies,
budgetReservations,
companies,
costEvents,
createDb,
documentRevisions,
documents,
@@ -16,6 +22,8 @@ import {
issueRecoveryActions,
issueTreeHolds,
issues,
nativeRunFinalizations,
projects,
} from "@paperclipai/db";
import { ISSUE_CONTINUATION_SUMMARY_DOCUMENT_KEY } from "@paperclipai/shared";
import {
@@ -25,6 +33,9 @@ import {
import { createPostgresRunDispatchAdapter } from "./postgres.js";
import { settleUnrecoverableExecutions } from "../../../services/execution-recovery-resolution.js";
import { getExecutionBlocker } from "../../../services/execution-blocker.js";
import { reserveRunBudget } from "../../../services/budget-reservations.js";
import { budgetService, recheckAdmittedRunBudgetReadOnly } from "../../../services/budgets.js";
import { costService } from "../../../services/costs.js";
// Proves the DB-to-facts mapping this adapter owns for each state the two
// run-dispatch gates decide on. `application/use-cases.test.ts` and
@@ -52,7 +63,12 @@ describeEmbeddedPostgres("run-dispatch postgres adapter", () => {
}, 20_000);
afterEach(async () => {
await db.delete(budgetIncidents);
await db.delete(approvals);
await db.delete(budgetPolicies);
await db.delete(costEvents);
await db.delete(environmentLeases);
await db.delete(nativeRunFinalizations);
await db.delete(activityLog);
await db.delete(issueDocuments);
await db.delete(documentRevisions);
@@ -62,6 +78,7 @@ describeEmbeddedPostgres("run-dispatch postgres adapter", () => {
await db.delete(issues);
await db.delete(heartbeatRunEvents);
await db.delete(heartbeatRuns);
await db.delete(projects);
await db.delete(agents);
await db.delete(companies);
});
@@ -317,6 +334,7 @@ describeEmbeddedPostgres("run-dispatch postgres adapter", () => {
expect(await adapter.evaluateScheduledRetryGate({ companyId, runId: replacementId, retryReasonOverride: retryReason, now: new Date() }))
.toMatchObject({ allowed: false, errorCode: "issue_execution_lock_changed" });
await db.update(heartbeatRuns).set({ status: "queued" }).where(eq(heartbeatRuns.id, replacementId));
const reservation = await reserveRunBudget(db, companyId, replacementId, null);
expect(await adapter.cancelStaleQueuedRun({ companyId, runId: replacementId, expectedStatus: "queued", now: new Date() }))
.toMatchObject({ outcome: "cancelled", errorCode: "issue_execution_lock_changed" });
@@ -325,7 +343,7 @@ describeEmbeddedPostgres("run-dispatch postgres adapter", () => {
await db.update(heartbeatRuns).set({ status: "running" }).where(eq(heartbeatRuns.id, replacementId));
let dispatched = false;
const outcome = await adapter.dispatchResolvedInteractionIfCurrent({ companyId, runId: replacementId,
expectedStatus: "running", now: new Date(), dispatch: async () => { dispatched = true; } });
reservationId: reservation.id, expectedStatus: "running", now: new Date(), dispatch: async () => { dispatched = true; } });
expect(outcome).toMatchObject({ dispatched: false, cancellation: { outcome: "cancelled" } });
expect(dispatched).toBe(false);
expect((await db.select().from(issues).where(eq(issues.id, issueId)))[0]![lock]).toBe(competingId);
@@ -343,11 +361,12 @@ describeEmbeddedPostgres("run-dispatch postgres adapter", () => {
expect(await adapter.evaluateScheduledRetryGate({ companyId, runId: replacementId, retryReasonOverride: retryReason, now: new Date() }))
.toMatchObject({ allowed: false, errorCode: "issue_blocked" });
await db.update(heartbeatRuns).set({ status: "queued" }).where(eq(heartbeatRuns.id, replacementId));
const reservation = await reserveRunBudget(db, companyId, replacementId, null);
expect(await adapter.cancelStaleQueuedRun({ companyId, runId: replacementId, expectedStatus: "queued", now: new Date() }))
.toMatchObject({ outcome: "cancelled", errorCode: "issue_blocked" });
await db.update(heartbeatRuns).set({ status: "running" }).where(eq(heartbeatRuns.id, replacementId));
const dispatch = vi.fn(async () => undefined);
expect(await adapter.dispatchResolvedInteractionIfCurrent({ companyId, runId: replacementId, expectedStatus: "running", now: new Date(), dispatch }))
expect(await adapter.dispatchResolvedInteractionIfCurrent({ companyId, runId: replacementId, reservationId: reservation.id, expectedStatus: "running", now: new Date(), dispatch }))
.toMatchObject({ dispatched: false, cancellation: { outcome: "cancelled" } });
expect(dispatch).not.toHaveBeenCalled();
expect((await db.select().from(issues).where(eq(issues.id, issueId)))[0]!.status).toBe("blocked");
@@ -373,7 +392,8 @@ describeEmbeddedPostgres("run-dispatch postgres adapter", () => {
if (mode === "resolved") await db.update(issues).set({ status: "done" }).where(eq(issues.id, blockerId));
await db.update(issues).set({ executionRunId: replacementId }).where(eq(issues.id, issueId));
await db.update(heartbeatRuns).set({ status: "running" }).where(eq(heartbeatRuns.id, replacementId));
const result = await adapter.dispatchResolvedInteractionIfCurrent({ companyId, runId: replacementId, expectedStatus: "running", now: new Date(), dispatch });
const reservation = await reserveRunBudget(db, companyId, replacementId, null);
const result = await adapter.dispatchResolvedInteractionIfCurrent({ companyId, runId: replacementId, reservationId: reservation.id, expectedStatus: "running", now: new Date(), dispatch });
expect(result).toMatchObject(mode === "resolved" ? { dispatched: true } : {
dispatched: false, cancellation: { outcome: "cancelled", errorCode: "issue_dependencies_blocked" },
});
@@ -389,8 +409,9 @@ describeEmbeddedPostgres("run-dispatch postgres adapter", () => {
const runId = await seedRun({ companyId, agentId, status: "running", contextSnapshot: { issueId } });
await db.update(issues).set({ executionRunId: runId }).where(eq(issues.id, issueId));
const adapter = createPostgresRunDispatchAdapter(db);
const reservation = await reserveRunBudget(db, companyId, runId, null);
const gate = await adapter.dispatchResolvedInteractionIfCurrent({
companyId, runId, expectedStatus: "running", now: new Date(),
companyId, runId, reservationId: reservation.id, expectedStatus: "running", now: new Date(),
dispatch: async () => {
// The incident's third attempt failed here before onSpawn. A held
// admission lock makes this finalization fail with lock_timeout.
@@ -407,33 +428,491 @@ describeEmbeddedPostgres("run-dispatch postgres adapter", () => {
expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId)))[0]?.status).toBe("failed");
});
it("initiates dispatch before admission locks can be released to a competing owner", async () => {
it("commits dispatch authority before provider HTTP and permits accounting while that provider waits", async () => {
const { companyId, agentId } = await seedCompanyAndAgent();
const runId = await seedRun({ companyId, agentId, status: "running" });
await budgetService(db).upsertPolicy(companyId, { scopeType: "company", scopeId: companyId,
metric: "billed_cents", windowKind: "lifetime", amount: 100, notifyEnabled: false }, null);
const reservation = await reserveRunBudget(db, companyId, runId, null);
let readyToCommit!: () => void, commit!: () => void, providerArrived!: () => void, finishProvider!: () => void;
let readyToHandoff!: () => void, commitHandoff!: () => void;
const ready = new Promise<void>(resolve => { readyToCommit = resolve; });
const commitAllowed = new Promise<void>(resolve => { commit = resolve; });
const arrived = new Promise<void>(resolve => { providerArrived = resolve; });
const providerResponse = new Promise<void>(resolve => { finishProvider = resolve; });
const handoffReady = new Promise<void>(resolve => { readyToHandoff = resolve; });
const handoffCommit = new Promise<void>(resolve => { commitHandoff = resolve; });
let requests = 0;
const server = createServer((_request, response) => {
requests += 1;
providerArrived();
void providerResponse.then(() => response.end("provider result"));
});
await new Promise<void>(resolve => { server.listen(0, "127.0.0.1", resolve); });
const address = server.address();
if (!address || typeof address === "string") throw new Error("Expected a provider HTTP fixture port");
const dispatch = vi.fn(async () => (await fetch(`http://127.0.0.1:${address.port}`)).text());
const transaction = db.transaction.bind(db);
let transactionNumber = 0;
// Freeze the exact SQL commit boundary. A provider callback inside the
// accounting transaction would contact this server before we release it.
const transactionSpy = vi.spyOn(db, "transaction").mockImplementation(async (callback, config) => {
const phase = ++transactionNumber;
return transaction(async tx => {
const value = await callback(tx);
if (phase === 1) { readyToCommit(); await commitAllowed; }
if (phase === 2) { readyToHandoff(); await handoffCommit; }
return value;
}, config);
});
const admission = createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent({
companyId, runId, reservationId: reservation.id, expectedStatus: "running", now: new Date(), dispatch,
});
let result: Promise<string> | undefined;
let charge: Promise<unknown> | undefined;
try {
await ready;
expect(dispatch).not.toHaveBeenCalled();
expect(requests).toBe(0);
commit();
await Promise.all([arrived, handoffReady]);
const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId));
expect(run.usageJson?.budgetDispatch).toMatchObject({ reservationId: reservation.id, recoveryLeaseOwner: null });
// This charge serializes after dispatch admission. It creates stop intent
// for already-admitted work without waiting for the provider's response.
let charged = false;
charge = costService(db).createEvent(companyId, { agentId, provider: "fixture", model: "post-admission-charge",
costCents: 100, inputTokens: 1, outputTokens: 1, occurredAt: new Date() }).then(() => { charged = true; });
void charge.catch(() => {});
// Phase two still owns issue/run locks and the provider is waiting.
// Holding company accounting here would stall this independent charge.
await vi.waitFor(() => expect(charged).toBe(true), { timeout: 2_000 });
const [policy] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.companyId, companyId));
expect(policy.enforcementVersion).toBeGreaterThan(policy.enforcementDeliveredVersion);
expect(requests).toBe(1);
commitHandoff();
const gate = await admission;
transactionSpy.mockRestore();
expect(gate.dispatched).toBe(true);
if (!gate.dispatched) throw new Error("Expected committed provider admission");
result = gate.resultPromise;
finishProvider();
expect(await result).toBe("provider result");
} finally {
commit();
commitHandoff();
finishProvider();
transactionSpy.mockRestore();
await admission.catch(() => undefined);
await result?.catch(() => undefined);
await charge?.catch(() => undefined);
server.closeAllConnections();
await new Promise<void>(resolve => { server.close(() => resolve()); });
}
});
it("does not hand off when another run takes issue ownership after accounting admission commits", async () => {
const { companyId, agentId } = await seedCompanyAndAgent();
const issueId = randomUUID();
await seedIssue({ companyId, issueId, assigneeAgentId: agentId, status: "in_progress" });
await seedIssue({ companyId, issueId, status: "in_progress", assigneeAgentId: agentId });
const runId = await seedRun({ companyId, agentId, status: "running", contextSnapshot: { issueId } });
const competingId = await seedRun({ companyId, agentId, status: "running", contextSnapshot: { issueId } });
await db.update(issues).set({ executionRunId: runId }).where(eq(issues.id, issueId));
const reservation = await reserveRunBudget(db, companyId, runId, null);
const transaction = db.transaction.bind(db);
const ordering: string[] = [];
// Inject a competing claim immediately after commit, before control returns
// to the adapter. A callback outside the transaction would run too late.
const transactionSpy = vi.spyOn(db, "transaction").mockImplementation(async (callback, config) => {
let admitted = false;
const race = vi.spyOn(db, "transaction").mockImplementation(async (callback, config) => {
const value = await transaction(callback, config);
await db.update(issues).set({ executionRunId: competingId }).where(eq(issues.id, issueId));
ordering.push("competing-owner");
if (!admitted) { admitted = true; await db.update(issues).set({ executionRunId: competingId }).where(eq(issues.id, issueId)); }
return value;
});
const dispatch = vi.fn(async () => "must not run");
try {
expect(await createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent({ companyId, runId,
reservationId: reservation.id, expectedStatus: "running", now: new Date(), dispatch })).toMatchObject({ dispatched: false,
cancellation: { outcome: "cancelled", errorCode: "issue_execution_lock_changed" } });
expect(dispatch).not.toHaveBeenCalled();
expect((await db.select().from(issues).where(eq(issues.id, issueId)))[0].executionRunId).toBe(competingId);
} finally { race.mockRestore(); }
});
it.each(["charge", "budget reduction"] as const)("rejects a committed %s between accounting admission and final ownership handoff", async change => {
const { companyId, agentId } = await seedCompanyAndAgent();
const runId = await seedRun({ companyId, agentId, status: "running" });
const budgets = budgetService(db);
await budgets.upsertPolicy(companyId, { scopeType: "company", scopeId: companyId,
metric: "billed_cents", windowKind: "lifetime", amount: 100, reservationCents: "60", notifyEnabled: false }, null);
const reservation = await reserveRunBudget(db, companyId, runId, null);
const transaction = db.transaction.bind(db);
let admitted = false;
const race = vi.spyOn(db, "transaction").mockImplementation(async (callback, config) => {
const value = await transaction(callback, config);
if (!admitted) {
admitted = true;
if (change === "charge") await costService(db).createEvent(companyId, { agentId, provider: "fixture",
model: "between-phase-charge", inputTokens: 1, outputTokens: 1, costCents: 50, occurredAt: new Date() });
else await budgets.upsertPolicy(companyId, { scopeType: "company", scopeId: companyId,
metric: "billed_cents", windowKind: "lifetime", amount: 50 }, null);
}
return value;
});
const dispatch = vi.fn(async () => "must not run");
try {
expect(await createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent({ companyId, runId,
reservationId: reservation.id, expectedStatus: "running", now: new Date(), dispatch })).toMatchObject({ dispatched: false,
admissionBlockReason: "Available budget is reserved by unfinished runs" });
expect(dispatch).not.toHaveBeenCalled();
// Neither case has actually spent its hard-stop limit. The denial must
// observe capacity, rather than relying only on a projected pause flag.
expect((await db.select().from(companies).where(eq(companies.id, companyId)))[0].status).toBe("active");
} finally { race.mockRestore(); }
});
it("performs the post-admission capacity check in a PostgreSQL read-only transaction", async () => {
const { companyId, agentId } = await seedCompanyAndAgent();
const runId = await seedRun({ companyId, agentId, status: "running" });
await budgetService(db).upsertPolicy(companyId, { scopeType: "company", scopeId: companyId,
metric: "billed_cents", windowKind: "lifetime", amount: 100, reservationCents: "60", notifyEnabled: false }, null);
const reservation = await reserveRunBudget(db, companyId, runId, null);
await db.transaction(async tx => {
await tx.execute(sql`set transaction read only`);
// PostgreSQL rejects both writes and SELECT FOR UPDATE in this mode.
expect(await recheckAdmittedRunBudgetReadOnly(tx as unknown as typeof db,
{ companyId, agentId, runId, projectId: null, reservationId: reservation.id })).toBeNull();
});
});
it("does not dispatch into an archival project deleted after accounting admission", async () => {
const { companyId, agentId } = await seedCompanyAndAgent();
const runId = await seedRun({ companyId, agentId, status: "running" });
const [project] = await db.insert(projects).values({ companyId, name: "Dispatch project" }).returning();
const reservation = await reserveRunBudget(db, companyId, runId, project.id);
const transaction = db.transaction.bind(db);
let admitted = false;
const race = vi.spyOn(db, "transaction").mockImplementation(async (callback, config) => {
const value = await transaction(callback, config);
if (!admitted) { admitted = true; await db.delete(projects).where(eq(projects.id, project.id)); }
return value;
});
const dispatch = vi.fn(async () => "must not run");
try {
expect(await createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent({ companyId, runId,
reservationId: reservation.id, expectedStatus: "running", now: new Date(), dispatch })).toMatchObject({ dispatched: false,
admissionBlockReason: "Reserved project no longer exists" });
expect(dispatch).not.toHaveBeenCalled();
expect((await db.select().from(budgetReservations).where(eq(budgetReservations.id, reservation.id)))[0])
.toMatchObject({ state: "held", projectId: project.id });
} finally { race.mockRestore(); }
});
it("initiates the provider callback while the final issue ownership lock is held without awaiting its response", async () => {
const { companyId, agentId } = await seedCompanyAndAgent();
const issueId = randomUUID();
await seedIssue({ companyId, issueId, status: "in_progress", assigneeAgentId: agentId });
const runId = await seedRun({ companyId, agentId, status: "running", contextSnapshot: { issueId } });
const competingId = await seedRun({ companyId, agentId, status: "running", contextSnapshot: { issueId } });
await db.update(issues).set({ executionRunId: runId }).where(eq(issues.id, issueId));
const reservation = await reserveRunBudget(db, companyId, runId, null);
const transaction = db.transaction.bind(db);
let phase = 0, finishProvider!: () => void;
const response = new Promise<void>(resolve => { finishProvider = resolve; });
const ordering: string[] = [];
const race = vi.spyOn(db, "transaction").mockImplementation(async (callback, config) => {
const current = ++phase;
const value = await transaction(callback, config);
if (current === 2) {
await db.update(issues).set({ executionRunId: competingId }).where(eq(issues.id, issueId));
ordering.push("competing-owner");
}
return value;
});
try {
const gate = await createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent({
companyId, runId, expectedStatus: "running", now: new Date(),
dispatch: async () => { ordering.push("handoff"); return "started"; },
});
const gate = await createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent({ companyId, runId,
reservationId: reservation.id, expectedStatus: "running", now: new Date(),
dispatch: () => { ordering.push("handoff"); return response; } });
expect(gate.dispatched).toBe(true);
if (gate.dispatched) expect(await gate.resultPromise).toBe("started");
expect(ordering).toEqual(["handoff", "competing-owner"]);
finishProvider();
if (gate.dispatched) await gate.resultPromise;
} finally { finishProvider(); race.mockRestore(); }
});
it.each(["0", "25"])("fences duplicate handoff of a %s-cent reservation generation", async reservationCents => {
const { companyId, agentId } = await seedCompanyAndAgent();
const runId = await seedRun({ companyId, agentId, status: "running" });
await budgetService(db).upsertPolicy(companyId, { scopeType: "company", scopeId: companyId,
metric: "billed_cents", windowKind: "lifetime", amount: 100, reservationCents, notifyEnabled: false }, null);
const reservation = await reserveRunBudget(db, companyId, runId, null);
const dispatch = vi.fn(async () => "one handoff");
const input = { companyId, runId, reservationId: reservation.id, expectedStatus: "running" as const, now: new Date(), dispatch };
const adapter = createPostgresRunDispatchAdapter(db);
const first = await adapter.dispatchResolvedInteractionIfCurrent(input);
expect(first.dispatched).toBe(true);
if (first.dispatched) await first.resultPromise;
expect(await adapter.dispatchResolvedInteractionIfCurrent(input)).toMatchObject({ dispatched: false,
admissionBlockReason: "Provider dispatch already admitted this reservation generation" });
expect(dispatch).toHaveBeenCalledOnce();
});
it("does not enter a provider when its dispatch claim rolls back", async () => {
const { companyId, agentId } = await seedCompanyAndAgent();
const runId = await seedRun({ companyId, agentId, status: "running" });
const reservation = await reserveRunBudget(db, companyId, runId, null);
const transaction = db.transaction.bind(db);
const rollback = vi.spyOn(db, "transaction").mockImplementation(async (callback, config) => transaction(async tx => {
await callback(tx);
throw new Error("fixture commit rejected");
}, config));
const dispatch = vi.fn(async () => "must not run");
try {
await expect(createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent({ companyId, runId,
reservationId: reservation.id, expectedStatus: "running", now: new Date(), dispatch })).rejects.toThrow("fixture commit rejected");
expect(dispatch).not.toHaveBeenCalled();
} finally { rollback.mockRestore(); }
expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId)))[0].usageJson?.budgetDispatch).toBeUndefined();
});
it("retains uncertain dispatch authority when the ownership commit fails after callback entry", async () => {
const { companyId, agentId } = await seedCompanyAndAgent();
const runId = await seedRun({ companyId, agentId, status: "running" });
const reservation = await reserveRunBudget(db, companyId, runId, null);
const transaction = db.transaction.bind(db);
let phase = 0;
const rollback = vi.spyOn(db, "transaction").mockImplementation(async (callback, config) => {
const current = ++phase;
return transaction(async tx => {
const value = await callback(tx);
if (current === 2) throw new Error("fixture ownership commit rejected");
return value;
}, config);
});
const dispatch = vi.fn(async () => "provider already entered");
const input = { companyId, runId, reservationId: reservation.id, expectedStatus: "running" as const, now: new Date(), dispatch };
try {
await expect(createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent(input)).rejects.toThrow("fixture ownership commit rejected");
expect(dispatch).toHaveBeenCalledOnce();
} finally { rollback.mockRestore(); }
const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId));
expect(run).toMatchObject({ costAccountingPending: true, costAccountedAt: null,
usageJson: { budgetDispatch: { reservationId: reservation.id } } });
expect((await db.select().from(budgetReservations).where(eq(budgetReservations.id, reservation.id)))[0].state).toBe("held");
expect(await createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent(input))
.toMatchObject({ dispatched: false, admissionBlockReason: "Provider dispatch already admitted this reservation generation" });
expect(dispatch).toHaveBeenCalledOnce();
expect(await db.select().from(costEvents)).toHaveLength(0);
});
it("commits the ownership handoff when a provider callback throws synchronously", async () => {
const { companyId, agentId } = await seedCompanyAndAgent();
const runId = await seedRun({ companyId, agentId, status: "running" });
const reservation = await reserveRunBudget(db, companyId, runId, null);
const failure = new Error("synchronous provider bootstrap failure");
const dispatch = vi.fn(() => { throw failure; });
const input = { companyId, runId, reservationId: reservation.id, expectedStatus: "running" as const, now: new Date(), dispatch };
const gate = await createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent(input);
expect(gate.dispatched).toBe(true);
if (gate.dispatched) await expect(gate.resultPromise).rejects.toBe(failure);
expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId)))[0].usageJson?.budgetDispatch)
.toMatchObject({ reservationId: reservation.id, handedOffAt: expect.any(String) });
expect(await createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent(input))
.toMatchObject({ dispatched: false, admissionBlockReason: "Provider dispatch already admitted this reservation generation" });
expect(dispatch).toHaveBeenCalledOnce();
});
it.each(["charge", "budget reduction"] as const)("rechecks a late %s after ordinary run preparation before any provider handoff", async change => {
const { companyId, agentId } = await seedCompanyAndAgent();
const runId = await seedRun({ companyId, agentId, status: "running" });
const budgets = budgetService(db);
await budgets.upsertPolicy(companyId, { scopeType: "company", scopeId: companyId,
metric: "billed_cents", windowKind: "lifetime", amount: 100, reservationCents: "60", notifyEnabled: false }, null);
const reservation = await reserveRunBudget(db, companyId, runId, null);
let finishPreparation!: () => void;
const preparation = new Promise<void>(resolve => { finishPreparation = resolve; });
const dispatch = vi.fn(async () => "provider started");
const preparedDispatch = preparation.then(() => createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent({
companyId, runId, reservationId: reservation.id, expectedStatus: "running", now: new Date(), dispatch,
}));
try {
if (change === "charge") await costService(db).createEvent(companyId, {
agentId, provider: "test", model: "dispatch-race", costCents: 50,
inputTokens: 1, outputTokens: 1, occurredAt: new Date(),
});
else await budgets.upsertPolicy(companyId, { scopeType: "company", scopeId: companyId,
metric: "billed_cents", windowKind: "lifetime", amount: 50 }, null);
finishPreparation();
expect(await preparedDispatch).toMatchObject({ dispatched: false,
admissionBlockReason: "Available budget is reserved by unfinished runs" });
expect(dispatch).not.toHaveBeenCalled();
} finally { finishPreparation(); await preparedDispatch; }
});
it("does not count its own existing reservation twice at final handoff", async () => {
const { companyId, agentId } = await seedCompanyAndAgent();
const runId = await seedRun({ companyId, agentId, status: "running" });
await budgetService(db).upsertPolicy(companyId, { scopeType: "company", scopeId: companyId,
metric: "billed_cents", windowKind: "lifetime", amount: 100, reservationCents: "100", notifyEnabled: false }, null);
const reservation = await reserveRunBudget(db, companyId, runId, null);
const dispatch = vi.fn(async () => "started");
const gate = await createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent({
companyId, runId, reservationId: reservation.id, expectedStatus: "running", now: new Date(), dispatch,
});
expect(gate.dispatched).toBe(true);
expect(dispatch).toHaveBeenCalledTimes(1);
if (gate.dispatched) expect(await gate.resultPromise).toBe("started");
});
it("does not dispatch a zero reservation once other work has fully reserved capacity", async () => {
const { companyId, agentId } = await seedCompanyAndAgent();
const runId = await seedRun({ companyId, agentId, status: "running" });
const budgets = budgetService(db);
await budgets.upsertPolicy(companyId, { scopeType: "company", scopeId: companyId,
metric: "billed_cents", windowKind: "lifetime", amount: 100, reservationCents: "0", notifyEnabled: false }, null);
const reservation = await reserveRunBudget(db, companyId, runId, null);
await budgets.upsertPolicy(companyId, { scopeType: "company", scopeId: companyId,
metric: "billed_cents", windowKind: "lifetime", reservationCents: "100" }, null);
const otherRunId = await seedRun({ companyId, agentId, status: "running" });
await reserveRunBudget(db, companyId, otherRunId, null);
const dispatch = vi.fn(async () => undefined);
expect(await createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent({
companyId, runId, reservationId: reservation.id, expectedStatus: "running", now: new Date(), dispatch,
})).toMatchObject({ dispatched: false, admissionBlockReason: "Available budget is reserved by unfinished runs" });
expect(dispatch).not.toHaveBeenCalled();
});
it("commits a discovered budget stop even when final handoff is denied", async () => {
const { companyId, agentId } = await seedCompanyAndAgent();
const runId = await seedRun({ companyId, agentId, status: "running" });
await budgetService(db).upsertPolicy(companyId, { scopeType: "company", scopeId: companyId,
metric: "billed_cents", windowKind: "lifetime", amount: 100, notifyEnabled: false }, null);
const reservation = await reserveRunBudget(db, companyId, runId, null);
// Seed committed historical spend whose lifecycle projection has not yet
// been repaired. The gate must persist that repair rather than roll it back.
await db.insert(costEvents).values({ companyId, agentId, provider: "test", model: "unprojected-charge",
inputTokens: 1, outputTokens: 1, costCents: 110, occurredAt: new Date() });
const dispatch = vi.fn(async () => undefined);
expect(await createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent({
companyId, runId, reservationId: reservation.id, expectedStatus: "running", now: new Date(), dispatch,
})).toMatchObject({ dispatched: false, admissionBlockReason: expect.any(String) });
expect(dispatch).not.toHaveBeenCalled();
expect((await db.select().from(companies).where(eq(companies.id, companyId)))[0])
.toMatchObject({ status: "paused", pauseReason: "budget" });
const [policy] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.companyId, companyId));
expect(policy.enforcementVersion).toBeGreaterThan(policy.enforcementDeliveredVersion);
});
it.each([false, true])("preserves accepted non-assignee work while retaining its final budget gate (late reduction: %s)", async reduceBudget => {
const { companyId, agentId } = await seedCompanyAndAgent();
const assigneeId = randomUUID(), issueId = randomUUID();
await seedAgent({ id: assigneeId, companyId, name: "Assigned owner" });
await seedIssue({ companyId, issueId, assigneeAgentId: assigneeId, status: "in_progress" });
const runId = await seedRun({ companyId, agentId, status: "running",
contextSnapshot: { issueId, wakeReason: "issue_comment_mentioned", commentId: randomUUID() } });
const budgets = budgetService(db);
await budgets.upsertPolicy(companyId, { scopeType: "company", scopeId: companyId,
metric: "billed_cents", windowKind: "lifetime", amount: 100, reservationCents: "60", notifyEnabled: false }, null);
const reservation = await reserveRunBudget(db, companyId, runId, null);
if (reduceBudget) await budgets.upsertPolicy(companyId, { scopeType: "company", scopeId: companyId,
metric: "billed_cents", windowKind: "lifetime", amount: 50 }, null);
const dispatch = vi.fn(async () => "accepted mention delivered");
const gate = await createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent({ companyId, runId,
reservationId: reservation.id, expectedStatus: "running", enforceContinuationOwnership: false, now: new Date(), dispatch });
expect(gate.dispatched).toBe(!reduceBudget);
expect(dispatch).toHaveBeenCalledTimes(reduceBudget ? 0 : 1);
if (!gate.dispatched) expect(gate).toMatchObject({ admissionBlockReason: "Available budget is reserved by unfinished runs" });
const [issue] = await db.select().from(issues).where(eq(issues.id, issueId));
expect(issue).toMatchObject({ assigneeAgentId: assigneeId, executionRunId: null });
});
async function seedRecoveryAdmission() {
const { companyId, agentId } = await seedCompanyAndAgent();
const issueId = randomUUID(), owner = `recovery:${randomUUID()}`;
await seedIssue({ companyId, issueId, assigneeAgentId: agentId, status: "in_progress" });
const runId = await seedRun({ companyId, agentId, status: "running", contextSnapshot: { issueId } });
await db.update(heartbeatRuns).set({ runtimeMode: "native", nativeIssueId: issueId }).where(eq(heartbeatRuns.id, runId));
await db.update(issues).set({ executionRunId: runId }).where(eq(issues.id, issueId));
const reservation = await reserveRunBudget(db, companyId, runId, null);
await db.insert(nativeRunFinalizations).values({ companyId, issueId, runId, phase: "observed",
leaseOwner: owner, leaseExpiresAt: new Date(Date.now() + 60_000) });
return { companyId, agentId, issueId, runId, owner, reservation };
}
it.each(["reservation", "dispatch"] as const)("orders recovery %s behind a coordinator before taking finalization's issue and run locks", async boundary => {
const f = await seedRecoveryAdmission();
let signalLocked!: () => void, allowFinalization!: () => void;
const locked = new Promise<void>(resolve => { signalLocked = resolve; });
const proceed = new Promise<void>(resolve => { allowFinalization = resolve; });
const finalization = db.transaction(async tx => {
await tx.execute(sql`set local lock_timeout = '2000ms'`);
await tx.select({ id: nativeRunFinalizations.runId }).from(nativeRunFinalizations)
.where(eq(nativeRunFinalizations.runId, f.runId)).for("update");
signalLocked();
await proceed;
await tx.select({ id: issues.id }).from(issues).where(eq(issues.id, f.issueId)).for("update");
await tx.update(heartbeatRuns).set({ updatedAt: new Date() }).where(eq(heartbeatRuns.id, f.runId));
});
await Promise.race([locked, finalization]);
const dispatch = vi.fn(async () => "one provider handoff");
const gateInput = { companyId: f.companyId, runId: f.runId, reservationId: f.reservation.id,
expectedStatus: "running" as const, recoveryLeaseOwner: f.owner, now: new Date(), dispatch };
const recovery = boundary === "dispatch"
? createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent(gateInput)
: reserveRunBudget(db, f.companyId, f.runId, null, {}, f.owner).then(reservation => {
expect(reservation.reused).toBe(true);
return createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent({ ...gateInput, reservationId: reservation.id });
});
try {
expect(await waitForBlockedForUpdate("native_run_finalizations")).toBe(true);
allowFinalization();
await finalization;
const gate = await recovery;
expect(gate.dispatched).toBe(true);
if (gate.dispatched) expect(await gate.resultPromise).toBe("one provider handoff");
expect(await createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent(gateInput))
.toMatchObject({ dispatched: false, admissionBlockReason: "Provider dispatch already admitted this reservation generation" });
expect(dispatch).toHaveBeenCalledOnce();
} finally {
transactionSpy.mockRestore();
allowFinalization();
await Promise.allSettled([finalization, recovery]);
}
});
it.each(["expired", "replaced"] as const)("does not admit recovery under an %s coordinator owner", async ownership => {
const f = await seedRecoveryAdmission();
await db.update(nativeRunFinalizations).set(ownership === "expired"
? { leaseExpiresAt: new Date(Date.now() - 1) } : { leaseOwner: "another-owner" }).where(eq(nativeRunFinalizations.runId, f.runId));
await expect(reserveRunBudget(db, f.companyId, f.runId, null, {}, f.owner)).rejects.toThrow("Native recovery no longer owns provider dispatch");
const dispatch = vi.fn(async () => undefined);
await expect(createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent({ companyId: f.companyId, runId: f.runId,
reservationId: f.reservation.id, recoveryLeaseOwner: f.owner, expectedStatus: "running", now: new Date(), dispatch }))
.rejects.toThrow("Native recovery no longer owns provider dispatch");
expect(dispatch).not.toHaveBeenCalled();
});
it("rechecks recovery lease expiry after waiting for the issue lock", async () => {
const f = await seedRecoveryAdmission();
const expiresAt = new Date(Date.now() + 1_000);
await db.update(nativeRunFinalizations).set({ leaseExpiresAt: expiresAt }).where(eq(nativeRunFinalizations.runId, f.runId));
let signalLocked!: () => void, releaseIssue!: () => void;
const locked = new Promise<void>(resolve => { signalLocked = resolve; });
const release = new Promise<void>(resolve => { releaseIssue = resolve; });
const holder = db.transaction(async tx => {
await tx.select({ id: issues.id }).from(issues).where(eq(issues.id, f.issueId)).for("update");
signalLocked();
await release;
});
await Promise.race([locked, holder]);
const dispatch = vi.fn(async () => undefined);
const recovery = createPostgresRunDispatchAdapter(db).dispatchResolvedInteractionIfCurrent({ companyId: f.companyId,
runId: f.runId, reservationId: f.reservation.id, recoveryLeaseOwner: f.owner, expectedStatus: "running", now: new Date(), dispatch });
try {
expect(await waitForBlockedForUpdate("issues")).toBe(true);
await new Promise(resolve => setTimeout(resolve, Math.max(0, expiresAt.getTime() - Date.now() + 10)));
releaseIssue();
expect(await recovery).toMatchObject({ dispatched: false, admissionBlockReason: "Native recovery no longer owns provider dispatch" });
expect(dispatch).not.toHaveBeenCalled();
} finally {
releaseIssue();
await Promise.allSettled([holder, recovery]);
}
});
@@ -7,18 +7,23 @@ import type { Db } from "@paperclipai/db";
import {
agentWakeupRequests,
agents,
budgetReservations,
companies,
connectionIntentDeliveries,
heartbeatRuns,
issueRecoveryActions,
issueComments,
issues,
nativeRunFinalizations,
projects,
} 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, budgetServiceInTransaction } from "../../../services/budgets.js";
import { budgetService, budgetServiceInTransaction, recheckAdmittedRunBudgetReadOnly } from "../../../services/budgets.js";
import { lockRunBudgetRecoveryOwnerInTransaction, recheckRunBudgetReservationInTransaction } from "../../../services/budget-reservations.js";
import { isCompletedOnboardingHandoffWake } from "../../../services/chat-completion-delivery.js";
import { isHeartbeatWakeOnDemandEnabled } from "../../../services/heartbeat-policy.js";
import { DELIVERY_QUEUES, notifyDeliveryWork } from "../../../services/delivery-work-notifications.js";
@@ -193,10 +198,11 @@ export function createPostgresRunDispatchAdapter(
const treeControlSvc = issueTreeControlService(db);
const issuesSvc = issueService(db);
async function withIssueThenRunLocks<T>(
input: { runId: string; companyId: string },
async function withRunOwnershipLocks<T>(
input: { runId: string; companyId: string; recoveryLeaseOwner?: string; enforceContinuationOwnership?: boolean },
onMissing: () => T,
operation: (tx: Db, run: HeartbeatRun, publications: ActivityPublication[]) => Promise<T>,
mode: "accounting" | "ownership",
): Promise<T> {
// Queue editing and claiming already use issue -> wake -> run. Read the
// immutable issue reference without a lock first, then acquire issue ->
@@ -209,37 +215,45 @@ export function createPostgresRunDispatchAdapter(
.where(and(eq(heartbeatRuns.id, input.runId), eq(heartbeatRuns.companyId, input.companyId)))
.then((rows) => rows[0] ?? null);
if (!hint) return onMissing();
const hintedIssueId = readNonEmptyString(parseObject(hint.contextSnapshot).issueId);
// Already-admitted ordinary work owes the budget/run gate, not another
// continuation or chat-control issue lock. Its authority is captured on
// the run and reservation; protected recovery retains its admission.
const hintedIssueId = input.enforceContinuationOwnership === false ? null
: readNonEmptyString(parseObject(hint.contextSnapshot).issueId);
const result: { kind: "missing" | "retry" } | { kind: "value"; value: T } =
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 })
.from(issues)
.where(and(eq(issues.id, hintedIssueId), eq(issues.companyId, input.companyId)))
.for("update");
}
const locked = async (typedTx: Db, publications: ActivityPublication[]) => {
// Recovery always takes coordinator before issue/run. Only the
// accounting phase acquires the company lock and evaluates budgets.
await lockRunBudgetRecoveryOwnerInTransaction(typedTx, input.companyId, input.runId, input.recoveryLeaseOwner);
if (hintedIssueId) {
await typedTx
.select({ id: issues.id })
.from(issues)
.where(and(eq(issues.id, hintedIssueId), eq(issues.companyId, input.companyId)))
.for("update");
}
const run = await typedTx
.select()
.from(heartbeatRuns)
.where(
and(eq(heartbeatRuns.id, input.runId), eq(heartbeatRuns.companyId, input.companyId)),
)
// Keep the run status stable through the semantic decision and any
// resulting mutation and synchronous dispatch handoff. Never await
// adapter-owned work while this transaction holds the row locks.
.for("update")
.then((rows) => rows[0] ?? null);
if (!run) return { kind: "missing" as const };
const run = await typedTx
.select()
.from(heartbeatRuns)
.where(
and(eq(heartbeatRuns.id, input.runId), eq(heartbeatRuns.companyId, input.companyId)),
)
// Keep the run status stable through the semantic decision and any
// resulting mutation and synchronous dispatch handoff. Never await
// adapter-owned work while this transaction holds the row locks.
.for("update")
.then((rows) => rows[0] ?? null);
if (!run) return { kind: "missing" as const };
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, publications) };
});
const lockedIssueId = input.enforceContinuationOwnership === false ? null
: readNonEmptyString(parseObject(run.contextSnapshot).issueId);
if (lockedIssueId !== hintedIssueId) return { kind: "retry" as const };
return { kind: "value" as const, value: await operation(typedTx, run, publications) };
};
const result: { kind: "missing" | "retry" } | { kind: "value"; value: T } = mode === "accounting"
? await withAccountingTransaction(db, input.companyId, locked)
: await db.transaction(tx => locked(tx as unknown as Db, []));
if (result.kind === "missing") return onMissing();
if (result.kind === "value") return result.value;
@@ -250,6 +264,14 @@ export function createPostgresRunDispatchAdapter(
);
}
function withIssueThenRunLocks<T>(
input: { runId: string; companyId: string; recoveryLeaseOwner?: string; enforceContinuationOwnership?: boolean },
onMissing: () => T,
operation: (tx: Db, run: HeartbeatRun, publications: ActivityPublication[]) => Promise<T>,
) {
return withRunOwnershipLocks(input, onMissing, operation, "accounting");
}
async function loadGateFacts(
input: LoadGateFactsInput,
now: Date,
@@ -888,6 +910,7 @@ export function createPostgresRunDispatchAdapter(
now: Date,
providerDispatchNotStarted = false,
): Promise<CancelStaleQueuedRunOutcome> {
await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, run.companyId);
const [row] = await tx
.update(heartbeatRuns)
.set({
@@ -1134,60 +1157,75 @@ export function createPostgresRunDispatchAdapter(
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(
tx,
run,
input.now,
publications,
);
const decision =
!initialDecision.stale && issueId
? await tx
.select({ executionRunId: issues.executionRunId })
.from(issues)
.where(and(eq(issues.id, issueId), eq(issues.companyId, run.companyId)))
.then((rows) => rows[0] ?? null)
.then((issue) =>
issue?.executionRunId === run.id
? initialDecision
: {
stale: true as const,
errorCode: "issue_execution_lock_changed" as const,
reason:
"Cancelled because continuation no longer owns the issue execution lock before adapter dispatch",
details: {
issueId,
expectedExecutionRunId: run.id,
currentExecutionRunId: issue?.executionRunId ?? null,
},
},
)
: initialDecision;
if (decision.stale && issueId) {
const cancellation = await cancelStaleRunInTx(
if (input.enforceContinuationOwnership !== false) {
const { issueId, decision: initialDecision } = await decideCurrentRunStaleness(
tx,
run,
issueId,
decision,
input.expectedStatus,
input.now,
true, // This locked gate has not handed off to the provider callback.
publications,
);
return { dispatched: false as const, cancellation };
const decision =
!initialDecision.stale && issueId
? await tx
.select({ executionRunId: issues.executionRunId })
.from(issues)
.where(and(eq(issues.id, issueId), eq(issues.companyId, run.companyId)))
.then((rows) => rows[0] ?? null)
.then((issue) =>
issue?.executionRunId === run.id
? initialDecision
: {
stale: true as const,
errorCode: "issue_execution_lock_changed" as const,
reason:
"Cancelled because continuation no longer owns the issue execution lock before adapter dispatch",
details: {
issueId,
expectedExecutionRunId: run.id,
currentExecutionRunId: issue?.executionRunId ?? null,
},
},
)
: initialDecision;
if (decision.stale && issueId) {
const cancellation = await cancelStaleRunInTx(
tx,
run,
issueId,
decision,
input.expectedStatus,
input.now,
true, // This locked gate has not handed off to the provider callback.
);
return { dispatched: false as const, cancellation };
}
}
// Hand off while ownership is still locked, but do not await the provider
// promise. Bootstrap and failure finalization can update these same rows;
// the transaction must commit independently of either callback completing.
const resultPromise = input.dispatch(() => {});
// A synchronous rejection can precede the commit response. Observe it
// immediately while preserving the original promise for the caller.
void resultPromise.catch(() => {});
return { dispatched: true as const, resultPromise };
const admissionBlockReason = await recheckRunBudgetReservationInTransaction(
tx, publications, run, input.reservationId, input.recoveryLeaseOwner,
);
// Commit budget-owned hold/enforcement changes before the caller handles
// denial. Throwing here would roll back the very stop intent we found.
if (admissionBlockReason) return { dispatched: false as const, admissionBlockReason };
const priorDispatch = parseObject(parseObject(run.usageJson).budgetDispatch);
const recoveryLeaseOwner = input.recoveryLeaseOwner ?? null;
if (priorDispatch.reservationId === input.reservationId && (priorDispatch.recoveryLeaseOwner ?? null) === recoveryLeaseOwner)
return { dispatched: false as const, admissionBlockReason: "Provider dispatch already admitted this reservation generation" };
// The committed claim is the dispatch linearization point. A concurrent
// charge wins before it, or targets already-admitted work afterwards.
// Preserve usage/cancellation evidence and fence duplicate handoffs even
// for zero-valued holds. A newly proven native recovery owner has its own
// generation; its ordinary stopped-execution checks remain upstream.
const claim = { reservationId: input.reservationId, recoveryLeaseOwner, admittedAt: new Date().toISOString() };
await tx.update(heartbeatRuns).set({
usageJson: sql`coalesce(${heartbeatRuns.usageJson}, '{}'::jsonb) || ${JSON.stringify({ budgetDispatch: claim })}::jsonb`,
}).where(and(eq(heartbeatRuns.id, run.id), eq(heartbeatRuns.companyId, run.companyId)));
return { dispatched: true as const, claim };
};
return withIssueThenRunLocks(
const admission = await withIssueThenRunLocks(
input,
() => {
throw new RunDispatchApplicationError(
@@ -1197,7 +1235,73 @@ export function createPostgresRunDispatchAdapter(
},
dispatchLockedRun,
);
if (!admission.dispatched) return admission;
// Accounting has committed. Close the separate issue/run ownership race
// under short row locks, without reacquiring company accounting or awaiting
// provider work. Mutating budget evaluation never runs in this second phase.
return withRunOwnershipLocks<DispatchResolvedInteractionOutcome<T>>(input, () => ({ dispatched: false as const,
cancellation: { outcome: "lost_race" as const } }), async (tx, run) => {
const claim = parseObject(parseObject(run.usageJson).budgetDispatch);
if (run.status !== input.expectedStatus || run.costAccountedAt || run.resultJson?.cancellation || run.resultJson?.startupCancellation
|| claim.reservationId !== admission.claim.reservationId || claim.recoveryLeaseOwner !== admission.claim.recoveryLeaseOwner
|| claim.admittedAt !== admission.claim.admittedAt || claim.handedOffAt)
return { dispatched: false as const, cancellation: { outcome: "lost_race" as const } };
const [scope] = await tx.select({ companyStatus: companies.status, agentStatus: agents.status,
lifecycleState: agents.lifecycleState, reservationId: budgetReservations.id,
projectId: budgetReservations.projectId, currentProjectId: projects.id, projectPausedAt: projects.pausedAt })
.from(companies).innerJoin(agents, and(eq(agents.companyId, companies.id), eq(agents.id, run.agentId)))
.leftJoin(budgetReservations, and(eq(budgetReservations.companyId, companies.id), eq(budgetReservations.runId, run.id),
eq(budgetReservations.id, input.reservationId), eq(budgetReservations.agentId, run.agentId), eq(budgetReservations.state, "held")))
.leftJoin(projects, and(eq(projects.companyId, companies.id), eq(projects.id, budgetReservations.projectId)))
.where(eq(companies.id, run.companyId));
if (!scope || scope.reservationId !== input.reservationId)
return { dispatched: false as const, admissionBlockReason: "Provider dispatch no longer owns a held budget reservation" };
if (scope.projectId && !scope.currentProjectId)
return { dispatched: false as const, admissionBlockReason: "Reserved project no longer exists" };
if (scope.companyStatus !== "active" || scope.projectPausedAt || scope.lifecycleState !== "ready"
|| ["paused", "terminated", "pending_approval"].includes(scope.agentStatus))
return { dispatched: false as const, admissionBlockReason: "Execution scope is paused or no longer ready for work" };
if (input.enforceContinuationOwnership !== false) {
const contextSnapshot = parseObject(run.contextSnapshot);
const issueId = readNonEmptyString(contextSnapshot.issueId);
if (issueId) {
// These facts read current issue authority. The reconstruction/source
// coordinator validation remains in the accounting phase, avoiding a
// source coordinator lock after issue/run in this final phase.
const facts = await loadStalenessFacts({ runId: run.id, companyId: run.companyId, agentId: run.agentId,
responsibleUserId: run.responsibleUserId,
conversationContinuation: run.runtimeMode === "legacy" && hasConversationContinuationPolicy(run.resultJson),
issueId, contextSnapshot, scheduledRetryReason: run.scheduledRetryReason }, input.now, tx);
const semantic = decideQueuedRunStaleness(facts, input.now);
const decision = semantic.stale || facts.issueExecutionRunId === run.id ? semantic : {
stale: true as const, errorCode: "issue_execution_lock_changed" as const,
reason: "Cancelled because continuation no longer owns the issue execution lock before adapter dispatch",
details: { issueId, expectedExecutionRunId: run.id, currentExecutionRunId: facts.issueExecutionRunId },
};
if (decision.stale) return { dispatched: false as const,
cancellation: await cancelStaleRunInTx(tx, run, issueId, decision, input.expectedStatus, input.now, true) };
}
}
const budgetBlock = await recheckAdmittedRunBudgetReadOnly(tx, { companyId: run.companyId, agentId: run.agentId,
runId: run.id, projectId: scope.projectId, reservationId: input.reservationId });
if (budgetBlock) return { dispatched: false as const, admissionBlockReason: budgetBlock };
if (input.recoveryLeaseOwner) {
const [owner] = run.runtimeMode === "native" ? await tx.select({ id: nativeRunFinalizations.runId }).from(nativeRunFinalizations).where(and(
eq(nativeRunFinalizations.companyId, run.companyId), eq(nativeRunFinalizations.runId, run.id),
eq(nativeRunFinalizations.leaseOwner, input.recoveryLeaseOwner), sql`${nativeRunFinalizations.leaseExpiresAt} > clock_timestamp()`)) : [];
if (!owner) return { dispatched: false as const, admissionBlockReason: "Native recovery no longer owns provider dispatch" };
}
await tx.update(heartbeatRuns).set({ usageJson: sql`coalesce(${heartbeatRuns.usageJson}, '{}'::jsonb)
|| ${JSON.stringify({ budgetDispatch: { ...claim, handedOffAt: new Date().toISOString() } })}::jsonb` })
.where(and(eq(heartbeatRuns.id, run.id), eq(heartbeatRuns.companyId, run.companyId)));
let resultPromise: Promise<T>;
try { resultPromise = Promise.resolve(input.dispatch(() => {})); }
catch (error) { resultPromise = Promise.reject(error); }
void resultPromise.catch(() => {});
return { dispatched: true as const, resultPromise };
}, "ownership");
}
return {
@@ -42,11 +42,15 @@ export type CancelStaleQueuedRunInput = {
};
export type DispatchResolvedInteractionInput<T> = CancelStaleQueuedRunInput & {
reservationId: string;
recoveryLeaseOwner?: string;
enforceContinuationOwnership?: boolean;
dispatch: (markDispatchStarted: () => void) => Promise<T>;
};
export type DispatchResolvedInteractionOutcome<T> =
| { dispatched: true; resultPromise: Promise<T> }
| { dispatched: false; admissionBlockReason: string }
| { dispatched: false; cancellation: CancelStaleQueuedRunOutcome };
/** Semantic database operations; persistence rows and transaction handles stay inside the adapter. */
@@ -161,10 +161,13 @@ describe("createDispatchResolvedInteractionIfCurrent", () => {
runId: "run-1",
companyId: "company-1",
expectedStatus: "running",
reservationId: "reservation-1",
enforceContinuationOwnership: false,
dispatch,
});
expect(result.dispatched).toBe(true);
expect(writer.dispatchCalls).toHaveLength(1);
expect(writer.dispatchCalls[0]).toMatchObject({ enforceContinuationOwnership: false });
});
});
@@ -62,6 +62,9 @@ export function createDispatchResolvedInteractionIfCurrent(deps: { writer: RunDi
runId: string;
companyId: string;
expectedStatus: "queued" | "running";
reservationId: string;
recoveryLeaseOwner?: string;
enforceContinuationOwnership?: boolean;
dispatch: (markDispatchStarted: () => void) => Promise<T>;
now?: Date;
}) => deps.writer.dispatchResolvedInteractionIfCurrent({
@@ -3,19 +3,20 @@ import { companies, type Db } from "@paperclipai/db";
import { publishActivity, type ActivityPublication } from "./activity-log.js";
import { notFound } from "../errors.js";
import { logger } from "../middleware/logger.js";
import { withBudgetSpan } from "./budget-tracing.js";
/** Serialize ledger and budget mutations in one company. Always acquire this
* before locking a run or a policy, so concurrent writers use one lock order.
* NO KEY UPDATE permits FK KEY SHARE locks held by recovery callers. */
export async function withAccountingTransaction<T>(db: Db, companyId: string, work: (tx: Db, publications: ActivityPublication[]) => Promise<T>): Promise<T> {
const publications: ActivityPublication[] = [];
const result = await db.transaction(async (transaction) => {
const result = await withBudgetSpan("budget.accounting.transaction", () => db.transaction(async (transaction) => {
const tx = transaction as unknown as Db;
const [company] = await tx.select({ id: companies.id }).from(companies)
.where(eq(companies.id, companyId)).for("no key update");
const [company] = await withBudgetSpan("budget.accounting.lock_wait", () => tx.select({ id: companies.id }).from(companies)
.where(eq(companies.id, companyId)).for("no key update"));
if (!company) throw notFound("Company not found");
return work(tx, publications);
});
}));
publishAccountingActivities(companyId, publications);
return result;
}
@@ -45,10 +45,12 @@ export function createLifecycleDriver(db: Db, manager: PluginWorkerManager): Lif
async runHost(agent) {
if (agent.lifecycleState === "verifying") return verify(agent);
if (["pausing", "terminating", "cleaning_up"].includes(agent.lifecycleState)) {
const companyAgents = await db.select().from(agents).where(eq(agents.companyId, agent.companyId));
const ids = agent.lifecycleState === "pausing" ? [agent.id]
: [agent.id, ...listInvalidOrgChainDescendantIds(agent.id, companyAgents)];
if (!await heartbeat.stopInvocationsForAgents(ids, "Agent lifecycle stop requested")) return "pending";
: [agent.id, ...listInvalidOrgChainDescendantIds(agent.id,
await db.select().from(agents).where(eq(agents.companyId, agent.companyId)))];
const reason = agent.lifecycleHolds?.includes("company_archived")
? "Cancelled because the company was archived" : "Agent lifecycle stop requested";
if (!await heartbeat.stopInvocationsForAgents(ids, reason)) return "pending";
}
return "complete";
},
+16 -4
View File
@@ -2,19 +2,23 @@ import { logger } from "../middleware/logger.js";
import type { Db } from "@paperclipai/db";
import * as lifecycle from "../modules/agent-lifecycle/index.js";
import { withAccountingTransaction } from "./accounting-transaction.js";
import { budgetServiceInTransaction, deliverBudgetEnforcement, policyBlocks, type BudgetServiceHooks } from "./budgets.js";
import { budgetServiceInTransaction, policyBlocks, type BudgetServiceHooks } from "./budgets.js";
import { recordAgentStatusEvent, recordResourceCreationEvent } from "./resource-lifecycle-events.js";
import { clearPrimaryAgent } from "./primary-agent.js";
import { deleteAgentDependencies } from "./agent-deletion.js";
import { prepareAgentConfiguration, completeAgentConfiguration } from "./agent-configuration.js";
import { prepareAgentHire, initializeAgentHire, prepareAgentHireApproval, completeAgentHireApproval } from "./agent-hiring.js";
import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notifications.js";
import { fenceActiveCompanyAgentStopsInTransaction as fenceActiveCompanyStops } from "../modules/agent-lifecycle/index.js";
import { applyAgentPolicyHoldsInTransaction } from "../modules/agent-lifecycle/index.js";
import { trackIdleWork } from "./task-admission.js";
export { scheduleAgentLifecycle } from "../modules/agent-lifecycle/index.js";
export { scheduleAgentLifecycle, processAgentLifecycle } from "../modules/agent-lifecycle/index.js";
export type { LifecycleDriver, LifecycleAgent } from "../modules/agent-lifecycle/index.js";
export function createAgentLifecycleEffects(hooks: BudgetServiceHooks = {}): lifecycle.LifecycleEffects {
export function createAgentLifecycleEffects(_hooks: BudgetServiceHooks = {}): lifecycle.LifecycleEffects {
return {
notifyWork: tx => notifyDeliveryWork(tx, DELIVERY_QUEUES.agentLifecycle),
reportFailure(context, error) {
const code = lifecycle.lifecycleFailureCode(error);
logger.warn({ ...context, failureCode: code ?? "unclassified" }, "Agent lifecycle operation failed");
@@ -26,7 +30,6 @@ export function createAgentLifecycleEffects(hooks: BudgetServiceHooks = {}): lif
budgetServiceInTransaction(db, publications).upsertPolicy(companyId, {
scopeType: "agent", scopeId: agentId, amount, isActive, windowKind: "calendar_month_utc",
}, userId),
enforceBudget: (db, companyId, agentId) => deliverBudgetEnforcement(db, hooks, companyId, agentId),
recordCreation: (db, companyId, agentId) => recordResourceCreationEvent(db, companyId, "agent", agentId),
recordStatus: recordAgentStatusEvent,
clearPrimary: clearPrimaryAgent,
@@ -55,3 +58,12 @@ export function configureAgentLifecycle(db: Db, driver: lifecycle.LifecycleDrive
export function startAgentLifecycle(db: Db, driver: lifecycle.LifecycleDriver, canRun: () => boolean) {
return lifecycle.startAgentLifecycle(db, createAgentLifecycleEffects(), trackedDriver(driver), canRun);
}
/** Budget evaluation owns the transaction and supplies the aggregate agent decision. */
export function applyAgentBudgetHoldInTransaction(tx: Db, companyId: string, agentId: string, blocked: boolean) {
return applyAgentPolicyHoldsInTransaction(tx, createAgentLifecycleEffects(), companyId, agentId, { budget: blocked });
}
export function fenceActiveCompanyAgentStopsInTransaction(tx: Db, companyId: string, status: "paused" | "archived") {
return fenceActiveCompanyStops(tx, createAgentLifecycleEffects(), companyId, status);
}
-2
View File
@@ -1,5 +1,4 @@
import { approvalResolutionRecords, resolvableApprovalStatuses as resolvableStatuses } from "../lib/approval-records.js";
import { budgetService } from "./budgets.js";
import { notifyHireApproved } from "./hire-hook.js";
import { and, asc, eq, inArray, sql, type SQL } from "drizzle-orm";
@@ -177,7 +176,6 @@ export function approvalService(db: Db) {
const result = await (status === "approved" ? lifecycle.approveHire : lifecycle.rejectHire)(target, userId, note);
if (result?.hireApprovedAgentId && result.approval) {
await reconcileApprovedBuiltInAgent(result.approval.companyId, result.approval.payload, db);
await budgetService(db).deliverPendingEnforcement(result.approval.companyId);
void notifyHireApproved(db, { companyId: result.approval.companyId, agentId: result.hireApprovedAgentId,
source: "approval", sourceId: result.approval.id, approvedAt: result.approval.decidedAt ?? new Date() }).catch(() => {});
}
@@ -7,7 +7,7 @@ import { withAccountingReadSnapshot, withAccountingTransaction } from "./account
import { receiptFingerprint } from "./receipt-fingerprint.js";
import { logActivity } from "./activity-log.js";
import { updateMonthlySpendProjections } from "./costs.js";
import { budgetServiceInTransaction, deliverBudgetEnforcement, type BudgetServiceHooks } from "./budgets.js";
import { budgetServiceInTransaction, type BudgetServiceHooks } from "./budgets.js";
async function invoiceCostCandidates(tx: Db, companyId: string, biller: string, line: {
costEventId?: string | null; runId?: string | null; providerRequestId?: string | null; model?: string | null;
@@ -180,7 +180,6 @@ export function billingReconciliationService(db: Db, hooks: BudgetServiceHooks =
details: { adjustmentId: adjustment.id, previousCents: input.expectedCents, correctedCents: input.correctedCents, reason: input.reason, pricing: input.pricing } }, publications);
return adjustment;
});
await deliverBudgetEnforcement(db, hooks, companyId);
return result;
},
};
+6 -2
View File
@@ -117,6 +117,7 @@ export function browserUseClient(
path: string,
method = "GET",
body?: unknown,
beforeDispatch?: () => Promise<void>,
): Promise<unknown> {
const until = backoff.get(scope) ?? 0;
if (until > Date.now())
@@ -126,6 +127,9 @@ export function browserUseClient(
until - Date.now(),
true,
);
// The owning service commits its final admission fence immediately before
// this one paid POST. Gate failures never become uncertain provider work.
if (beforeDispatch) await beforeDispatch();
let response: Response;
try {
const init: RequestInit = {
@@ -207,8 +211,8 @@ export function browserUseClient(
allowResizing?: true;
};
maxCostUsd?: number;
}) {
return created.parse(await call("/runs", "POST", body));
}, beforeDispatch?: () => Promise<void>) {
return created.parse(await call("/runs", "POST", body, beforeDispatch));
},
async status(runId: string) {
return z
+208 -103
View File
@@ -16,7 +16,8 @@ import {
issues,
heartbeatRuns,
agents,
costEvents,
companies,
toolCatalogEntries,
} from "@paperclipai/db";
import {
BROWSER_USE_IDLE_MS,
@@ -24,10 +25,13 @@ import {
type BrowserUseSettings,
type BrowserUseViewportRequest,
type TaskBrowser,
centsToUnits,
unitsToCents,
centsToUsd,
usdToUnits,
} from "@paperclipai/shared";
import {
browserUseClient,
browserUseCostCap,
BrowserUseError,
browserUseTerminal,
browserUseViewerUrl,
@@ -36,15 +40,17 @@ import {
type BrowserUseRequest,
} from "./browser-use-client.js";
import { secretService } from "./secrets.js";
import { budgetService, type BudgetServiceHooks } from "./budgets.js";
import { budgetService, budgetServiceInTransaction, type BudgetServiceHooks } from "./budgets.js";
import { issueTreeControlService } from "./issue-tree-control.js";
import { accessService } from "./access.js";
import { toolAccessPolicyService } from "./tool-access-policy.js";
import { createCostEventInTransaction } from "./costs.js";
import { createFinanceEventInTransaction } from "./finance.js";
import { withAccountingTransaction } from "./accounting-transaction.js";
import { logActivity } from "./activity-log.js";
import { logActivity, type ActivityPublication } from "./activity-log.js";
import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notifications.js";
import { browserUseViewports } from "./browser-use-viewport.js";
import { withBudgetSpan } from "./budget-tracing.js";
import { forbidden, notFound, conflict } from "../errors.js";
type Session = typeof sessions.$inferSelect;
@@ -54,6 +60,7 @@ type Binding = {
runId: string | null;
issueId: string | null;
approvedInvocationId?: string;
gatewayId?: string | null;
};
type Grant = typeof connectionGrants.$inferSelect;
const argsSchema = z
@@ -67,6 +74,17 @@ const argsSchema = z
const terminalSessions = ["closed", "failed"];
const liveSessions = ["starting", "running", "idle", "stopping"];
const requestMarker = (invocationId: string) => `\n\n[Paperclip request: ${invocationId}]`;
type BrowserBudgetHold = { invocationId: string; amountCents: string; projectId: string | null;
state: "held" | "released" | "settled"; uncapped?: boolean;
dispatchCommittedAt?: string; accountedCentsExact?: string; billingUnknown?: boolean };
const resultObject = (value: unknown): Record<string, unknown> =>
value && typeof value === "object" && !Array.isArray(value) ? value as Record<string, unknown> : {};
const budgetHold = (value: unknown) => resultObject(value).paperclipBudgetHold as BrowserBudgetHold | undefined;
function publicResult(value: unknown) {
if (!value || typeof value !== "object" || Array.isArray(value)) return value;
const { paperclipBudgetHold: _internal, ...visible } = resultObject(value);
return Object.keys(visible).length ? visible : null;
}
export function browserUseService(
db: Db,
@@ -114,6 +132,10 @@ export function browserUseService(
r.configPath.toLowerCase() === "headers.x-browser-use-api-key",
);
if (!ref) throw forbidden("Reconnect Browser Use to restore its API key.");
const [credentialSnapshot] = await db.select().from(companySecrets).where(and(
eq(companySecrets.companyId, s.companyId), eq(companySecrets.id, ref.secretId)));
if (!credentialSnapshot || (!cleanup && credentialSnapshot.status !== "active"))
throw forbidden("Browser credential is unavailable.");
if (row.grant.kind === "user") {
const [secret] = await db
.select()
@@ -154,6 +176,7 @@ export function browserUseService(
);
return {
...row,
credentialSnapshot,
client: browserUseClient(
{ "X-Browser-Use-API-Key": key },
request,
@@ -271,7 +294,7 @@ export function browserUseService(
status: s.status,
runStatus: run?.status ?? null,
progress: run?.progress ?? null,
result: run?.result ?? null,
result: publicResult(run?.result ?? null),
costCents: run?.accountedCents ?? 0,
browsers: ownedBrowsers,
error: s.error,
@@ -307,7 +330,8 @@ export function browserUseService(
connectionId: grant.connectionId,
grantId: grant.id,
};
const { client } = await credentials(scope);
const selected = await credentials(scope);
const { client } = selected;
if (name === "browser_profiles") {
// A safe Test action that creates no paid work and exposes no unapproved profile.
if (!config.allowedProfileIds.length) return { profiles: [] };
@@ -403,27 +427,89 @@ export function browserUseService(
.limit(1);
if (existing[0])
return snapshot(await owned(binding, grant, existing[0].sessionId));
const issue = await assertBinding(binding, approved);
const { policies } = await budgets.overview(binding.companyId);
const caps = policies
.filter(
(p) =>
p.isActive &&
p.hardStopEnabled &&
p.metric === "billed_cents" &&
p.amount > 0 &&
(p.scopeType === "company" ||
(p.scopeType === "agent" && p.scopeId === binding.agentId) ||
(p.scopeType === "project" && p.scopeId === issue.projectId)),
)
.map((p) => p.remainingAmount / 100);
const maxCostUsd = browserUseCostCap(
config.maxCostUsd,
args.maxCostUsd,
...caps,
);
const admissionCheck = async (tx: Db, publications: ActivityPublication[], hold?: BrowserBudgetHold) => {
const [parent] = await tx.select().from(heartbeatRuns).where(and(
eq(heartbeatRuns.id, binding.runId!), eq(heartbeatRuns.companyId, binding.companyId),
eq(heartbeatRuns.agentId, binding.agentId!))).for("update");
const [issue] = await tx.select().from(issues).where(and(eq(issues.id, binding.issueId!), eq(issues.companyId, binding.companyId)));
const [currentAgent] = await tx.select().from(agents).where(and(eq(agents.id, binding.agentId!), eq(agents.companyId, binding.companyId)));
const [company] = await tx.select({ status: companies.status }).from(companies).where(eq(companies.id, binding.companyId));
if (!parent || (!approved && parent.status !== "running") || parent.resultJson?.cancellation || parent.resultJson?.startupCancellation
|| !issue || issue.assigneeAgentId !== binding.agentId || issue.status === "cancelled"
|| !currentAgent || currentAgent.lifecycleState !== "ready" || ["paused", "terminated", "pending_approval"].includes(currentAgent.status)
|| company?.status !== "active") return { reason: "The browser's owning task, company or agent run is no longer active." };
const [invocation] = await tx.select().from(toolInvocations).where(and(
eq(toolInvocations.id, invocationId), eq(toolInvocations.companyId, binding.companyId)));
if (invocation && (invocation.agentId !== binding.agentId || invocation.runId !== binding.runId
|| invocation.issueId !== binding.issueId || invocation.connectionId !== grant.connectionId
|| invocation.status !== "executing" || !["agent", "user", "system", "plugin"].includes(invocation.actorType)))
return { reason: "The browser invocation is no longer authorized." };
if (approved && invocation?.approvalState !== "approved") return { reason: "The browser approval is no longer authorized." };
if (await issueTreeControlService(tx).getActivePauseHoldGate(binding.companyId, issue.id)) return { reason: "This task is paused." };
const [current] = await tx.select({ grant: connectionGrants, connection: toolConnections }).from(connectionGrants)
.innerJoin(toolConnections, eq(toolConnections.id, connectionGrants.connectionId)).where(and(
eq(connectionGrants.id, grant.id), eq(connectionGrants.companyId, binding.companyId),
eq(toolConnections.id, grant.connectionId), eq(toolConnections.companyId, binding.companyId))).for("share");
const [secret] = await tx.select().from(companySecrets).where(and(eq(companySecrets.id, selected.credentialSnapshot.id),
eq(companySecrets.companyId, binding.companyId))).for("share");
// Reuse the gateway's policy evaluator, including profile precedence,
// exclusions and explicit denies. REST gallery tools do not need MCP
// installation rows, and a matching binding alone is not authorization.
const [entry] = await tx.select().from(toolCatalogEntries).where(and(
eq(toolCatalogEntries.companyId, binding.companyId), eq(toolCatalogEntries.connectionId, grant.connectionId),
eq(toolCatalogEntries.toolName, name))).limit(1);
if (!entry || (invocation?.catalogEntryId && invocation.catalogEntryId !== entry.id)
|| (invocation?.upstreamToolName && invocation.upstreamToolName !== name))
return { reason: "Browser tool is no longer available." };
const access = await toolAccessPolicyService(tx).decide({
companyId: binding.companyId, actor: { actorType: (invocation?.actorType ?? "agent") as "agent" | "user" | "system" | "plugin",
actorId: invocation?.actorId ?? binding.agentId!, agentId: binding.agentId },
runContext: { heartbeatRunId: binding.runId, issueId: issue.id, projectId: issue.projectId, gatewayId: binding.gatewayId ?? invocation?.gatewayId },
request: { applicationId: entry.applicationId, connectionId: grant.connectionId, catalogEntryId: entry.id,
toolName: invocation?.toolName ?? entry.name, upstreamToolName: name, riskLevel: entry.riskLevel,
applicationKey: invocation?.applicationKey, providerType: invocation?.providerType ?? "rest_api", arguments: value, sideEffecting: true },
consumeRateLimit: false, writeAuditEvent: false,
}, { admittedInvocationId: invocation?.id });
if (!access.allowed && !(approved && access.decision === "require_approval")) return { reason: access.explanation };
if (!current || current.grant.status !== "active" || current.connection.status !== "active" || !current.connection.enabled
|| !isBrowserUseConnection(current.connection)
|| current.grant.updatedAt.getTime() !== selected.grant.updatedAt.getTime()
|| current.connection.updatedAt.getTime() !== selected.connection.updatedAt.getTime()
|| JSON.stringify(current.grant.credentialSecretRefs) !== JSON.stringify(selected.grant.credentialSecretRefs)
|| !secret || secret.status !== "active" || secret.latestVersion !== selected.credentialSnapshot.latestVersion)
return { reason: "Browser access or its credential changed before dispatch." };
const admission = await budgetServiceInTransaction(tx, publications).evaluateAdmission(binding.companyId, binding.agentId, { projectId: issue.projectId });
if (admission.block) return { reason: admission.block.reason };
if (hold && centsToUnits(hold.amountCents) === 0n && admission.policies.length)
return { reason: "Browser budget changed before dispatch; try again with a bounded cost cap." };
const [currentConfig] = await tx.select().from(settings).where(and(eq(settings.companyId, binding.companyId), eq(settings.grantId, grant.id)));
if (args.profileId && !currentConfig?.allowedProfileIds.includes(args.profileId)) return { reason: "This profile is not allowed for this credential." };
let amount: bigint | undefined;
for (const cap of [currentConfig?.maxCostUsd, args.maxCostUsd]) {
if (cap === null || cap === undefined) continue;
const value = usdToUnits(cap);
amount = amount === undefined || value < amount ? value : amount;
}
for (const [index, policy] of admission.policies.entries()) {
const observed = admission.observations[index];
const committed = centsToUnits(observed.ledgerCents) + centsToUnits(observed.effectiveHeldCents);
const available = centsToUnits(policy.amount) - committed;
if (hold ? available < 0n || (centsToUnits(hold.amountCents) === 0n && available === 0n) : available <= 0n)
return { reason: "No browser budget remains; available capacity is reserved by unfinished work." };
if (!hold) amount = amount === undefined || available < amount ? available : amount;
}
if (amount !== undefined && amount <= 0n) return { reason: "No browser budget remains." };
if (hold && amount !== undefined && (centsToUnits(hold.amountCents) === 0n || amount < centsToUnits(hold.amountCents)))
return { reason: "Browser cost settings changed before dispatch; try again." };
if (hold && issue.projectId !== hold.projectId) return { reason: "The browser task's project changed before dispatch." };
return { amountCents: unitsToCents(amount ?? 0n), maxCostUsd: amount === undefined ? undefined : Number(centsToUsd(unitsToCents(amount))), projectId: issue.projectId };
};
const sid = args.sessionId;
const s = await db.transaction(async (tx) => {
const admitted = await withAccountingTransaction(db, binding.companyId, async (tx, publications) => {
const [duplicate] = await tx.select().from(runs).where(and(eq(runs.companyId, binding.companyId), eq(runs.invocationId, invocationId)));
if (duplicate) return { existing: duplicate.sessionId };
const check = await admissionCheck(tx, publications);
if (check.reason) return { reason: check.reason };
await notifyDeliveryWork(tx, DELIVERY_QUEUES.browser);
let row: Session;
if (sid) {
@@ -463,7 +549,10 @@ export function browserUseService(
})
.returning();
}
const hold: BrowserBudgetHold = { invocationId, amountCents: check.amountCents!, projectId: check.projectId!,
state: "held", uncapped: check.maxCostUsd === undefined };
await tx.insert(runs).values({
result: { paperclipBudgetHold: hold },
companyId: binding.companyId,
sessionId: row.id,
heartbeatRunId: binding.runId!,
@@ -472,8 +561,13 @@ export function browserUseService(
? new Date(Date.now() + BROWSER_USE_IDLE_MS)
: null,
});
return row;
return { session: row, hold, maxCostUsd: check.maxCostUsd };
});
if (admitted.reason) throw forbidden(admitted.reason);
if (admitted.existing) return snapshot(await owned(binding, grant, admitted.existing));
const s = admitted.session!;
const maxCostUsd = admitted.maxCostUsd;
let providerAttempted = false;
let result: Awaited<ReturnType<typeof client.start>>;
try {
result = await client.start({
@@ -488,6 +582,21 @@ export function browserUseService(
...(args.profileId ? { profileId: args.profileId } : {}),
},
...(maxCostUsd === undefined ? {} : { maxCostUsd }),
}, async () => {
const reason = await withAccountingTransaction(db, binding.companyId, async (tx, publications) => {
const [current] = await tx.select().from(runs).where(and(eq(runs.invocationId, invocationId), eq(runs.companyId, binding.companyId))).for("update");
const [currentSession] = await tx.select().from(sessions).where(and(eq(sessions.id, s.id), eq(sessions.companyId, binding.companyId))).for("update");
const hold = budgetHold(current?.result);
if (!currentSession || currentSession.status !== "starting" || currentSession.stopRequested) return "Browser session was stopped before dispatch.";
if (!current || current.status !== "creating" || current.providerRunId || hold?.state !== "held" || hold.dispatchCommittedAt)
return "Browser provider dispatch no longer owns this invocation.";
const check = await admissionCheck(tx, publications, hold);
if (check.reason) return check.reason;
await tx.update(runs).set({ result: { ...resultObject(current.result), paperclipBudgetHold: { ...hold, dispatchCommittedAt: new Date().toISOString() } } }).where(eq(runs.id, current.id));
return null;
});
if (reason) throw new BrowserUseError(403, reason, 0, true);
providerAttempted = true;
});
await db.transaction(async (tx) => {
await notifyDeliveryWork(tx, DELIVERY_QUEUES.browser);
@@ -508,12 +617,19 @@ export function browserUseService(
});
} catch (error) {
// A lost create response can mean paid work exists. Never submit it again.
const rejected = error instanceof BrowserUseError && error.requestRejected;
const message = rejected ? error.message :
const rejected = !providerAttempted || (error instanceof BrowserUseError && error.requestRejected);
const message = rejected ? (error instanceof Error ? error.message : "Browser dispatch was rejected.") :
"Browser run creation could not be confirmed. Paperclip is locating and stopping possible provider work. Do not start another run yet.";
await db.transaction(async tx => {
await withAccountingTransaction(db, binding.companyId, async (tx, publications) => {
await notifyDeliveryWork(tx, DELIVERY_QUEUES.browser);
const [current] = await tx.select().from(sessions).where(eq(sessions.id, s.id)).for("update");
if (rejected) {
await tx.update(runs).set({ result: sql`jsonb_set(${runs.result}, '{paperclipBudgetHold,state}', '"released"'::jsonb)` }).where(eq(runs.invocationId, invocationId));
await budgetServiceInTransaction(tx, publications).evaluateAffectedScopes(binding.companyId, binding.agentId, admitted.hold!.projectId);
} else {
await tx.update(runs).set({ result: sql`jsonb_set(${runs.result}, '{paperclipBudgetHold,billingUnknown}', 'true'::jsonb)` }).where(eq(runs.invocationId, invocationId));
await budgetServiceInTransaction(tx, publications).evaluateAffectedScopes(binding.companyId, binding.agentId, admitted.hold!.projectId);
}
// A restart must not leave a drained rejection attached to a starting
// session: that combination has neither recovery work nor an idle timer.
await tx.update(runs)
@@ -917,75 +1033,47 @@ export function browserUseService(
ReturnType<ReturnType<typeof browserUseClient>["summary"]>
>,
) {
const cents = Math.max(0, Math.round(Number(summary.totalCostUsd) * 100));
if (!Number.isSafeInteger(cents))
throw new BrowserUseError(502, "Browser Use returned invalid usage.");
const event = await withAccountingTransaction(db, s.companyId, async (tx, publications) => {
const [issue] = await tx
.select({ projectId: issues.projectId })
.from(issues)
let amount: bigint;
try { amount = usdToUnits(summary.totalCostUsd); }
catch { throw new BrowserUseError(502, "Browser Use returned invalid usage."); }
if (amount < 0n || !browserUseTerminal(summary.status))
throw new BrowserUseError(502, "Browser Use returned invalid terminal usage.");
const cents = unitsToCents(amount);
await withBudgetSpan("budget.accounting.settlement", () => withAccountingTransaction(db, s.companyId, async (tx, publications) => {
const [current] = await tx.select().from(runs).where(and(eq(runs.id, run.id), eq(runs.companyId, s.companyId))).for("update");
const hold = budgetHold(current.result);
const previous = centsToUnits(hold?.accountedCentsExact ?? current.accountedCents);
const [issue] = await tx.select({ projectId: issues.projectId }).from(issues)
.where(and(eq(issues.id, s.issueId), eq(issues.companyId, s.companyId)));
const [current] = await tx
.select()
.from(runs)
.where(eq(runs.id, run.id))
.for("update");
if (cents <= current.accountedCents) return null;
const projectId = hold ? hold.projectId : issue?.projectId ?? null;
// Replace the separate browser hold with its ledger charge atomically.
// Provider output never gets to overwrite the private accounting state.
const settled: BrowserBudgetHold = { ...(hold ?? { invocationId: current.invocationId, amountCents: "0.0000000", projectId }), state: "settled",
accountedCentsExact: unitsToCents(amount > previous ? amount : previous) };
await tx.update(runs).set({
result: { ...resultObject(sanitizeBrowserUse({ result: summary.result, output: summary.output })),
paperclipBudgetHold: settled },
eventsDrained: 1,
accountedCents: Math.min(2_147_483_647, Math.round(Number(unitsToCents(amount > previous ? amount : previous)))),
}).where(eq(runs.id, current.id));
if (amount <= previous) {
if (hold?.state === "held") await budgetServiceInTransaction(tx, publications).evaluateAffectedScopes(s.companyId, s.agentId, projectId);
return;
}
const e = await createCostEventInTransaction(tx, s.companyId, {
agentId: s.agentId,
issueId: s.issueId,
projectId: issue?.projectId,
agentId: s.agentId, issueId: s.issueId, projectId,
heartbeatRunId: run.heartbeatRunId,
provider: "browser-use-cloud",
biller: "browser-use-cloud",
billingType: "metered_api",
model: summary.model,
costCents: cents - current.accountedCents,
occurredAt: new Date(),
billingCode: `browser-use-cloud:${run.id}`,
idempotencyKey: `browser-use-cloud:${run.id}:${cents}`,
provider: "browser-use-cloud", biller: "browser-use-cloud", billingType: "metered_api", model: summary.model,
costCents: unitsToCents(amount - previous), occurredAt: new Date(),
billingCode: `browser-use-cloud:${run.id}`, idempotencyKey: `browser-use-cloud:${run.id}:${cents}`,
}, publications);
await createFinanceEventInTransaction(tx, publications, s.companyId, {
agentId: e.agentId,
issueId: e.issueId,
projectId: e.projectId,
heartbeatRunId: e.heartbeatRunId,
costEventId: e.id,
billingCode: e.billingCode,
eventKind: "inference_charge",
direction: "debit",
biller: "browser-use-cloud",
provider: "browser-use-cloud",
model: e.model,
amountCents: e.costCents,
currency: "USD",
description: "Browser Use agent run",
occurredAt: e.occurredAt,
agentId: e.agentId, issueId: e.issueId, projectId: e.projectId, heartbeatRunId: e.heartbeatRunId,
costEventId: e.id, billingCode: e.billingCode, eventKind: "inference_charge", direction: "debit",
biller: "browser-use-cloud", provider: "browser-use-cloud", model: e.model,
amountCents: e.costCentsExact, currency: "USD", description: "Browser Use agent run", occurredAt: e.occurredAt,
});
await tx
.update(runs)
.set({ accountedCents: cents })
.where(eq(runs.id, run.id));
return e;
});
// Re-evaluate on retries too: accounting may have committed immediately
// before a crash interrupted budget enforcement.
const lastEvent =
event ??
(
await db
.select()
.from(costEvents)
.where(
and(
eq(costEvents.companyId, s.companyId),
eq(costEvents.billingCode, `browser-use-cloud:${run.id}`),
),
)
.orderBy(desc(costEvents.createdAt))
.limit(1)
)[0];
if (lastEvent) await budgets.evaluateCostEvent(lastEvent);
}));
}
async function reconcile(s: Session) {
const { grant, connection, client } = await credentials(s, true);
@@ -1046,6 +1134,33 @@ export function browserUseService(
active = true;
continue;
}
if (budgetHold(run.result)?.state === "held" && !budgetHold(run.result)?.dispatchCommittedAt) {
const repaired = await withAccountingTransaction(db, s.companyId, async (tx, publications) => {
const [current] = await tx.select().from(runs).where(and(eq(runs.id, run.id), eq(runs.companyId, s.companyId))).for("update");
const hold = budgetHold(current?.result);
if (!current || current.providerRunId || current.status !== "creating" || hold?.state !== "held" || hold.dispatchCommittedAt
|| Date.now() - current.createdAt.getTime() <= 60_000) return null;
await notifyDeliveryWork(tx, DELIVERY_QUEUES.browser);
await tx.update(runs).set({ status: "failed", eventsDrained: 1,
result: { ...resultObject(current.result), paperclipBudgetHold: { ...hold, state: "released" } } }).where(eq(runs.id, run.id));
const [session] = await tx.select().from(sessions).where(eq(sessions.id, s.id)).for("update");
const stopping = session.stopRequested === "end";
const [updated] = await tx.update(sessions).set({
status: stopping ? "stopping" : session.providerSessionId ? "idle" : "failed",
idleDeadline: !stopping && session.providerSessionId ? new Date(Date.now() + BROWSER_USE_IDLE_MS) : null,
updatedAt: new Date(),
}).where(eq(sessions.id, s.id)).returning();
await budgetServiceInTransaction(tx, publications).evaluateAffectedScopes(s.companyId, s.agentId, hold.projectId);
return updated;
});
if (repaired) {
Object.assign(s, repaired);
run.eventsDrained = 1;
run.status = "failed";
if (terminalSessions.includes(s.status)) return;
continue;
}
}
end = true;
await db.update(sessions).set({ status: "stopping", stopRequested: "end" })
.where(eq(sessions.id, s.id));
@@ -1154,16 +1269,6 @@ export function browserUseService(
if (browserUseTerminal(state.status) && drained) {
const summary = await client.summary(run.providerRunId);
await account(s, run, summary);
await db
.update(runs)
.set({
result: sanitizeBrowserUse({
result: summary.result,
output: summary.output,
}),
eventsDrained: 1,
})
.where(eq(runs.id, run.id));
await signal(s, "browser.run_finished");
}
}
@@ -0,0 +1,86 @@
import type { Db } from "@paperclipai/db";
import { withBudgetSpan } from "./budget-tracing.js";
import type { createDeliveryWorkCoordinator } from "./delivery-work-coordinator.js";
import { DELIVERY_QUEUES } from "./delivery-work-notifications.js";
import {
deliverBudgetEnforcement,
hasPendingBudgetEnforcement,
nextBudgetReconciliationAt,
reconcileDueBudgetPolicies,
type BudgetServiceHooks,
} from "./budgets.js";
import { readCostAccountingWork } from "./run-cost-accounting.js";
import { subscribeAllCompanyLiveEvents } from "./live-events.js";
type Coordinator = Pick<ReturnType<typeof createDeliveryWorkCoordinator>, "register">;
const accountingWakes = new WeakMap<Db, () => void>();
/** A durable filesystem receipt can need recovery even when its database
* attempt failed before a transaction identity existed. Startup also replays it. */
export function wakeCostAccountingWork(db: Db): void {
accountingWakes.get(db)?.();
}
/** The version pair is durable cancellation intent. Signals only request an
* inspection; startup and concrete deadlines recover work after interruption. */
export function registerBudgetEnforcementWork(db: Db, coordinator: Coordinator, hooks: BudgetServiceHooks) {
let startup = true;
let reconciliationCursor: string | null = null;
let reconciling = false;
let nextDeadline: number | null = null;
return coordinator.register(DELIVERY_QUEUES.budgetEnforcement, {
retryMs: 1_000,
runDuringDrain: true,
async run(signal) {
if (signal.aborted) return;
if (startup || reconciling || (nextDeadline !== null && nextDeadline <= Date.now())) {
const page = await reconcileDueBudgetPolicies(db, { startup, ...(reconciliationCursor ? { cursor: reconciliationCursor } : {}) });
reconciliationCursor = page.nextCursor;
reconciling = reconciliationCursor !== null;
if (!reconciling) startup = false;
}
if (signal.aborted) return;
await withBudgetSpan("budget.enforcement.deliver", () => deliverBudgetEnforcement(db, hooks));
},
async hasPending() {
const [pending, next] = await Promise.all([
hasPendingBudgetEnforcement(db), nextBudgetReconciliationAt(db),
]);
nextDeadline = next?.getTime() ?? null;
return pending || reconciling;
},
nextRunAt: () => nextDeadline,
});
}
/** Accounting repairs admitted work and never dispatches another paid turn.
* A waiting partial receipt stays held without polling unchanged debt. */
export function registerCostAccountingWork(db: Db, coordinator: Coordinator, reconcile: () => Promise<unknown>) {
let nextDeadline: number | null = null;
const worker = coordinator.register(DELIVERY_QUEUES.accounting, {
retryMs: 5_000,
runDuringDrain: true,
async run(signal) {
if (signal.aborted) return;
await withBudgetSpan("budget.accounting.recover", reconcile);
},
async hasPending() {
const work = await readCostAccountingWork(db);
nextDeadline = work.nextRunAt;
return work.pending;
},
nextRunAt: () => nextDeadline,
});
accountingWakes.set(db, worker.wake);
// The existing heartbeat terminal event is published after its persisted
// transition. Receipt writes have transaction-aware accounting signals too.
const unsubscribe = subscribeAllCompanyLiveEvents(event => {
if (event.type === "heartbeat.run.status" &&
["succeeded", "failed", "timed_out", "cancelled", "interrupted"].includes(String(event.payload.status))) worker.wake();
});
const stop = () => {
unsubscribe();
if (accountingWakes.get(db) === worker.wake) accountingWakes.delete(db);
};
return { ...worker, stop };
}
+167
View File
@@ -0,0 +1,167 @@
import { addCents, normalizeCents, usdToUnits, unitsToCents } from "@paperclipai/shared";
import { and, eq, gte, inArray, lt, or, sql, type SQL } from "drizzle-orm";
import { budgetPolicies, budgetReservations, browserUseRuns, browserUseSessions, costEvents, heartbeatRuns, issues, nativeRunFinalizations, type Db } from "@paperclipai/db";
import { withBudgetSpan } from "./budget-tracing.js";
export interface BudgetExposure {
ledgerCents: string;
liveObservedCents: string;
effectiveHeldCents: string;
pendingTerminalRunCount: number;
unpricedEventCount: number;
/** Unconfirmed POSTs or unfinished uncapped browser work, separate from ledger data quality. */
unknownBrowserRunCount: number;
}
export type SpendPolicy = Pick<typeof budgetPolicies.$inferSelect, "companyId" | "scopeType" | "scopeId" | "windowKind" | "metric">;
export type ObservedSpend = BudgetExposure & {
windowStart: Date; windowEnd: Date; total: number; totalExact: string;
pendingRunCount: number; recoveringRunCount: number;
};
export function budgetWindow(kind: string, now = new Date()) {
return kind === "lifetime" ? { start: new Date(0), end: new Date(Date.UTC(9999, 0, 1)) }
: { start: new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth(), 1)), end: new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth() + 1, 1)) };
}
/** Only normalized per-run prices are exposure. An incomplete receipt can
* contain a known partial price while still being unpriced for settlement. */
export function knownRunCostCents(usage: Record<string, unknown>): string | null {
if (usage.billingType === "subscription_included" || usage.billingType === "subscription") return "0";
if (usage.usageSource === "session_cumulative" || usage.usageBasis === "session_cumulative") return null;
for (const value of [usage.cacheAdjustedCostUsd, usage.costUsdExact, usage.costUsd]) {
if ((typeof value !== "number" || !Number.isFinite(value) || value < 0) && typeof value !== "string") continue;
try {
const amount = usdToUnits(value);
if (amount >= 0n) return unitsToCents(amount);
} catch { /* An older malformed field cannot erase a reliable later field. */ }
}
return null;
}
/** SQL compatibility for active runs written before the exposure projection.
* Numeric text is bounded and checked before casting legacy JSON fields. */
function normalizedNumericCents(value: SQL, usd = false) {
const scale = usd ? 1_000_000_000 : 10_000_000;
return sql`case when length(${value}) <= 160 and ${value} ~ '^[+]?[0-9]+([.][0-9]+)?([eE][+-]?[0-9]{1,3})?$'
then case when abs(coalesce(substring(${value} from '[eE]([+-]?[0-9]+)$'), '0')::int) <= 100
then case when round((${value})::numeric * ${scale}) <= 999999999999999999999999::numeric
then round((${value})::numeric * ${scale}) / 10000000::numeric else null end
else null end else null end`;
}
const usage = heartbeatRuns.usageJson;
const projected = normalizedNumericCents(sql`${usage}->'budgetExposure'->>'costCents'`);
const hasProjection = sql`jsonb_typeof(${usage}->'budgetExposure') = 'object'
and jsonb_typeof(${usage}->'budgetExposure'->'sourceId') = 'string'
and jsonb_typeof(${usage}->'budgetExposure'->'sequence') = 'number'`;
const legacy = sql`case when coalesce(${usage}->>'usageSource', 'per_run') <> 'session_cumulative'
and coalesce(${usage}->>'usageBasis', 'per_run') <> 'session_cumulative' then
coalesce(${normalizedNumericCents(sql`${usage}->>'cacheAdjustedCostUsd'`, true)},
${normalizedNumericCents(sql`${usage}->>'costUsdExact'`, true)}, ${normalizedNumericCents(sql`${usage}->>'costUsd'`, true)}) else null end`;
// The projection already separates attempts and cumulative/session counters.
// Raw compatibility fields are a fallback, never stronger monetary evidence.
// A canonical unknown projection must remain unknown, not inherit an unsafe
// cumulative aggregate retained only for old accounting compatibility.
const liveCost = sql`case when ${heartbeatRuns.costAccountedAt} is null and ${heartbeatRuns.id} is not null
and coalesce(${usage}->>'billingType', '') not in ('subscription_included', 'subscription')
then case when ${hasProjection} then coalesce(${projected}, 0)
else coalesce(${projected}, ${legacy}, 0) end else 0 end`;
/** Batched ledger, held-operation and unfinished-run aggregates. Never scans
* agents or receipt history; each policy retains its own scope and window. */
async function readObservedSpends(db: Db, policies: SpendPolicy[], now: Date): Promise<ObservedSpend[]> {
const windows = policies.map(policy => budgetWindow(policy.windowKind, now));
const billable = policies.map((policy, index) => ({ policy, index })).filter(({ policy }) => policy.metric === "billed_cents");
const ledgerColumns: Record<string, SQL> = {}, heldColumns: Record<string, SQL> = {}, runColumns: Record<string, SQL> = {}, browserColumns: Record<string, SQL> = {};
const ledgerPredicates: SQL[] = [], heldPredicates: SQL[] = [], runPredicates: SQL[] = [], browserPredicates: SQL[] = [];
for (const { policy, index } of billable) {
const { start, end } = windows[index];
const ledger = and(eq(costEvents.companyId, policy.companyId),
policy.scopeType === "agent" ? eq(costEvents.agentId, policy.scopeId) : undefined,
policy.scopeType === "project" ? eq(costEvents.projectId, policy.scopeId) : undefined,
policy.windowKind === "calendar_month_utc" ? and(gte(costEvents.occurredAt, start), lt(costEvents.occurredAt, end)) : undefined)!;
const held = and(eq(budgetReservations.companyId, policy.companyId), eq(budgetReservations.state, "held"),
policy.scopeType === "agent" ? eq(budgetReservations.agentId, policy.scopeId) : undefined,
policy.scopeType === "project" ? eq(budgetReservations.projectId, policy.scopeId) : undefined)!;
// Hosted browser work is a separate paid operation from its parent run.
// Its durable hold survives unknown POST outcomes and monthly rollover.
const browserHold = sql`${browserUseRuns.result}->'paperclipBudgetHold'`;
const hasBrowserHold = sql`coalesce(jsonb_typeof(${browserHold}) = 'object'
and ${browserHold}->>'invocationId' = ${browserUseRuns.invocationId}::text, false)`;
// Earlier browser rows did not persist a cap/dispatch projection. Their
// unfinished work remains unknown until a terminal receipt is accounted.
const legacyBrowserPending = sql`not ${hasBrowserHold} and ${browserUseRuns.eventsDrained} = 0
and (${browserUseRuns.providerRunId} is not null or ${browserUseRuns.status} in ('creating', 'unknown', 'queued', 'dispatching', 'running'))`;
const browser = and(eq(browserUseRuns.companyId, policy.companyId), eq(browserUseSessions.companyId, policy.companyId),
sql`((${hasBrowserHold} and ${browserHold}->>'state' = 'held') or (${legacyBrowserPending}))`,
policy.scopeType === "agent" ? eq(browserUseSessions.agentId, policy.scopeId) : undefined,
policy.scopeType === "project" ? sql`case when ${hasBrowserHold} then ${browserHold}->>'projectId'
else ${issues.projectId}::text end = ${policy.scopeId}` : undefined)!;
browserPredicates.push(browser);
browserColumns[`held${index}`] = sql`coalesce(sum(${normalizedNumericCents(sql`${browserHold}->>'amountCents'`)}) filter (where ${browser}), 0)::text`;
// A successful uncapped POST still has no bounded monetary exposure until
// terminal accounting. Introducing a budget cannot treat that work as free.
// Older holds had no explicit flag: zero meant an uncapped dispatch. New
// capped zero holds are distinct, and unstarted reservations are not spend.
const uncapped = sql`case when ${browserHold}->>'uncapped' = 'false' then false
when ${browserHold}->>'uncapped' = 'true' then true
else coalesce(${normalizedNumericCents(sql`${browserHold}->>'amountCents'`)}, 0) = 0 end`;
const dispatched = sql`(${browserUseRuns.providerRunId} is not null or ${browserHold}->>'dispatchCommittedAt' is not null)`;
browserColumns[`unknown${index}`] = sql`count(*) filter (where ${browser} and
((${legacyBrowserPending}) or ${browserHold}->>'billingUnknown' = 'true' or (${uncapped} and ${dispatched})))::int`;
const inWindow = policy.windowKind === "calendar_month_utc"
? sql`coalesce(${heartbeatRuns.finishedAt}, ${now.toISOString()}::timestamptz) >= ${start.toISOString()}::timestamptz
and coalesce(${heartbeatRuns.finishedAt}, ${now.toISOString()}::timestamptz) < ${end.toISOString()}::timestamptz` : sql`true`;
const run = and(eq(heartbeatRuns.companyId, policy.companyId), eq(heartbeatRuns.costAccountingPending, true),
policy.scopeType === "agent" ? eq(heartbeatRuns.agentId, policy.scopeId) : undefined,
policy.scopeType === "project" ? sql`${usage}->'ledgerScope'->>'projectId' = ${policy.scopeId}` : undefined)!;
const terminal = and(run, inArray(heartbeatRuns.status, ["succeeded", "failed", "timed_out", "cancelled", "interrupted"]), inWindow)!;
// Keep correlated subqueries in nested SQL chunks. Drizzle strips column
// qualifiers from a single-table selection's direct column chunks.
const hasReservation = sql`exists (select 1 from ${budgetReservations}
where ${budgetReservations.companyId} = ${heartbeatRuns.companyId}
and ${budgetReservations.runId} = ${heartbeatRuns.id} and ${budgetReservations.state} = 'held')`;
const recovering = sql`exists (select 1 from ${nativeRunFinalizations}
where ${nativeRunFinalizations.runId} = ${heartbeatRuns.id}
and ${nativeRunFinalizations.companyId} = ${heartbeatRuns.companyId}
and ${nativeRunFinalizations.resultId} is null and ${nativeRunFinalizations.phase} <> 'terminal_failure')`;
ledgerPredicates.push(ledger); heldPredicates.push(held); runPredicates.push(run);
ledgerColumns[`total${index}`] = sql`coalesce(sum(${costEvents.costCents}) filter (where ${ledger}), 0)::text`;
ledgerColumns[`unpriced${index}`] = sql`count(*) filter (where ${ledger} and ${costEvents.costStatus} = 'unpriced' and ${costEvents.billingType} <> 'subscription_included')::int`;
runColumns[`live${index}`] = sql`coalesce(sum(${liveCost}) filter (where ${run} and ${inWindow}), 0)::text`;
runColumns[`unreserved${index}`] = sql`coalesce(sum(${liveCost}) filter (where ${run} and not ${hasReservation}), 0)::text`;
// Reservations carry over UTC windows. A reservation and its running
// observation describe one operation, so count their maximum, not both.
heldColumns[`held${index}`] = sql`coalesce(sum(greatest(${budgetReservations.amountCents}, ${liveCost})) filter (where ${held}), 0)::text`;
runColumns[`pending${index}`] = sql`count(*) filter (where ${terminal})::int`;
runColumns[`recovering${index}`] = sql`count(*) filter (where ${terminal} and ${heartbeatRuns.runtimeMode} = 'native' and ${recovering})::int`;
}
const [ledger = {}] = billable.length ? await db.select(ledgerColumns).from(costEvents).where(or(...ledgerPredicates)) : [];
const [held = {}] = billable.length ? await db.select(heldColumns).from(budgetReservations)
.leftJoin(heartbeatRuns, and(eq(heartbeatRuns.id, budgetReservations.runId), eq(heartbeatRuns.companyId, budgetReservations.companyId)))
.where(or(...heldPredicates)) : [];
const [runs = {}] = billable.length ? await db.select(runColumns).from(heartbeatRuns).where(or(...runPredicates)) : [];
const [browser = {}] = billable.length ? await db.select(browserColumns).from(browserUseRuns)
.innerJoin(browserUseSessions, and(eq(browserUseSessions.id, browserUseRuns.sessionId), eq(browserUseSessions.companyId, browserUseRuns.companyId)))
.leftJoin(issues, and(eq(issues.id, browserUseSessions.issueId), eq(issues.companyId, browserUseSessions.companyId)))
.where(or(...browserPredicates)) : [];
return policies.map((_, index) => {
const ledgerCents = normalizeCents(String(ledger[`total${index}`] ?? 0));
const liveObservedCents = normalizeCents(String(runs[`live${index}`] ?? 0));
const totalExact = addCents(ledgerCents, liveObservedCents);
const pendingRunCount = Number(runs[`pending${index}`] ?? 0);
return { windowStart: windows[index].start, windowEnd: windows[index].end,
ledgerCents, liveObservedCents, effectiveHeldCents: addCents(String(held[`held${index}`] ?? 0), String(runs[`unreserved${index}`] ?? 0), String(browser[`held${index}`] ?? 0)),
total: Number(totalExact), totalExact, pendingRunCount, pendingTerminalRunCount: pendingRunCount,
unpricedEventCount: Number(ledger[`unpriced${index}`] ?? 0), unknownBrowserRunCount: Number(browser[`unknown${index}`] ?? 0), recoveringRunCount: Number(runs[`recovering${index}`] ?? 0) };
});
}
export async function computeObservedSpend(db: Db, policy: SpendPolicy) {
return (await computeObservedSpends(db, [policy]))[0];
}
export function computeObservedSpends(db: Db, policies: SpendPolicy[], now = new Date()): Promise<ObservedSpend[]> {
return withBudgetSpan("budget.exposure.evaluate", () => readObservedSpends(db, policies, now), { "budget.policy_count": policies.length });
}
@@ -0,0 +1,115 @@
import type { AdapterUsageCheckpoint } from "@paperclipai/adapter-utils";
import { describe, expect, it } from "vitest";
import { createLiveRunUsageProjection, liveRunUsageCostCents } from "./budget-live-usage.js";
const openai: AdapterUsageCheckpoint = {
complete: false, usageBasis: "per_run", provider: "openai", biller: "openai",
billingType: "metered_api", model: "gpt-6.1-sol", costStatus: "unpriced",
usage: { inputTokens: 1000, cachedInputTokens: 1000, outputTokens: 1000 },
};
const anthropic: AdapterUsageCheckpoint = {
complete: false, usageBasis: "per_run", provider: "anthropic", biller: "anthropic",
billingType: "metered_api", model: "claude-sonnet-5", costStatus: "unpriced",
usage: { inputTokens: 22462, cachedInputTokens: 208218, cacheWriteTokens: 22450, outputTokens: 2450 },
};
describe("live run budget pricing", () => {
it("projects partial OpenAI token usage while preserving final-accounting uncertainty", () => {
const receipt = Object.freeze({ ...openai, usage: Object.freeze({ ...openai.usage! }) });
expect(liveRunUsageCostCents(receipt)).toBe("1.2100000");
expect(receipt).toMatchObject({ complete: false, costStatus: "unpriced" });
expect(receipt.costUsdExact).toBeUndefined();
});
it("projects partial Anthropic token usage using the reviewed cache-write bound", () => {
expect(liveRunUsageCostCents(anthropic)).toBe("15.5967600");
expect(anthropic).toMatchObject({ complete: false, costStatus: "unpriced" });
});
it("uses explicit provider cost before estimates even on incomplete unpriced receipts", () => {
expect(liveRunUsageCostCents({ ...openai, costUsdExact: "0.000000001" })).toBe("0.0000001");
expect(liveRunUsageCostCents({ ...openai, costUsd: 0 })).toBe("0.0000000");
expect(liveRunUsageCostCents({ ...openai, costUsd: 5, cacheAdjustedCostUsd: 0.123456789 })).toBe("12.3456789");
expect(liveRunUsageCostCents({ ...openai, model: "unsupported", costUsd: 0.05 })).toBe("5.0000000");
});
it("applies supported OpenAI context and processing-tier guards", () => {
expect(liveRunUsageCostCents({ ...openai, pricingContext: { serviceTier: "priority", contextTier: "long" } }))
.toBe("3.8400000");
expect(liveRunUsageCostCents({ ...openai, pricingContext: { serviceTier: "unrecognized" } })).toBeNull();
});
it.each([
{ usageBasis: "session_cumulative" }, { usageBasis: undefined }, { provider: "unknown" }, { biller: "openrouter" },
{ billingType: "unknown" }, { model: "unsupported" }, { model: "constructor" }, { usage: undefined },
{ usage: { inputTokens: 1, cacheWriteTokens: 2, outputTokens: 1 } },
{ usageByModel: [{ model: "gpt-6.1-sol", usage: { inputTokens: 1, outputTokens: 1 }, costUsd: 1 }] },
] as Partial<AdapterUsageCheckpoint>[])("does not invent a price for unsupported OpenAI usage: %j", override => {
expect(liveRunUsageCostCents({ ...openai, ...override })).toBeNull();
});
it.each([
{ biller: "aws_bedrock" }, { model: "claude-sonnet-latest" }, { pricingContext: { serviceTier: "priority" } },
{ usage: { inputTokens: 1, outputTokens: 1 } },
{ usage: { inputTokens: 1, cacheWriteTokens: 2, cachedInputTokens: 0, outputTokens: 1 } },
] as Partial<AdapterUsageCheckpoint>[])("retains Anthropic rate-card restrictions: %j", override => {
expect(liveRunUsageCostCents({ ...anthropic, ...override })).toBeNull();
});
it.each(["subscription", "subscription_included"] as const)("keeps %s usage outside billed exposure", billingType => {
expect(liveRunUsageCostCents({ ...openai, billingType, costUsd: 10 })).toBe("0");
});
it("does not use cumulative explicit cost as a per-run price", () => {
expect(liveRunUsageCostCents({ ...openai, usageBasis: "session_cumulative", costUsd: 10 })).toBeNull();
});
});
describe("live exposure across provider attempts", () => {
it("adds a token-priced retry after an explicitly free rejected attempt", () => {
const observe = createLiveRunUsageProjection();
expect(observe("rejected", { ...openai, complete: true, costUsd: 0, usage: { inputTokens: 0, outputTokens: 0 } }))
.toBe("0.0000000");
expect(observe("retry", openai)).toBe("1.2100000");
});
it("prices different provider attempts separately instead of applying one model to aggregate counters", () => {
const observe = createLiveRunUsageProjection();
expect(observe("openai", openai)).toBe("1.2100000");
expect(observe("anthropic", anthropic)).toBe("16.8067600");
});
it("preserves a known attempt price when its next partial checkpoint omits it", () => {
const observe = createLiveRunUsageProjection();
observe("first", { ...openai, model: "unsupported", costUsd: 0.8 });
observe("first", { ...openai, model: "unsupported" });
expect(observe("second", { ...openai, model: "unsupported", costUsd: 0.6 })).toBe("140.0000000");
expect(observe("second", { ...openai, model: "unsupported", costUsd: 0.4 })).toBe("140.0000000");
});
it("accepts a complete attempt correction without changing other attempts", () => {
const observe = createLiveRunUsageProjection();
observe("first", { ...openai, costUsd: 0.8 });
observe("second", { ...openai, costUsd: 0.6 });
expect(observe("first", { ...openai, complete: true, costUsd: 0.5 })).toBe("110.0000000");
});
it("restores the durable aggregate projection after same-run recorder replacement", () => {
const observe = createLiveRunUsageProjection({ receipt: { ...openai, costUsd: 0 }, costCents: "140" });
expect(observe("default", { ...openai, model: "unsupported" })).toBe("140.0000000");
expect(observe("next", { ...openai, costUsd: 0.3 })).toBe("170.0000000");
});
it("restores explicit unknown exposure without pricing unsafe compatibility aggregates", () => {
const observe = createLiveRunUsageProjection({ receipt: { ...openai, costUsd: 100 }, costCents: null });
expect(observe("default", { ...openai, model: "unsupported" })).toBeNull();
});
it("keeps unpriced attempts unknown while retaining known portions", () => {
const observe = createLiveRunUsageProjection();
expect(observe("unknown", { ...openai, usageBasis: "session_cumulative", costUsd: 100 })).toBeNull();
expect(observe("known", { ...openai, costUsd: 0.1 })).toBe("10.0000000");
expect(observe("subscription", { ...anthropic, billingType: "subscription_included", costUsd: 100 }))
.toBe("10.0000000");
});
});
+36
View File
@@ -0,0 +1,36 @@
import type { AdapterUsageCheckpoint } from "@paperclipai/adapter-utils";
import { priceAnthropicReceipt } from "./anthropic-pricing.js";
import { knownRunCostCents } from "./budget-exposure.js";
import { priceCodexReceipt } from "./codex-pricing.js";
import { addCents, compareCents, normalizeCents } from "@paperclipai/shared";
/** Live admission can use a reviewed rate-card estimate before a receipt is
* final. The original receipt stays incomplete/unpriced for ledger settlement. */
export function liveRunUsageCostCents(receipt: AdapterUsageCheckpoint): string | null {
const reported = knownRunCostCents({ ...receipt, usageSource: receipt.usageBasis });
if (reported !== null) return reported;
if (receipt.usageBasis !== "per_run") return null;
const projection = priceAnthropicReceipt(priceCodexReceipt({ ...receipt, complete: true }));
return knownRunCostCents({ ...projection, usageSource: projection.usageBasis });
}
/** Attempt receipts can omit a later price without retracting earlier spend.
* Keep that knowledge per attempt before combining different models/billers. */
export function createLiveRunUsageProjection(initial?: {
receipt: AdapterUsageCheckpoint;
costCents?: string | null;
}) {
const costs = new Map<string, string>();
if (initial) {
const cost = initial.costCents !== undefined ? initial.costCents : liveRunUsageCostCents(initial.receipt);
if (cost !== null) costs.set("default", normalizeCents(cost));
}
return (attemptId: string, receipt: AdapterUsageCheckpoint): string | null => {
const reported = liveRunUsageCostCents(receipt);
const previous = costs.get(attemptId);
if (reported !== null && (receipt.complete || previous === undefined || compareCents(reported, previous) > 0)) {
costs.set(attemptId, normalizeCents(reported));
}
return costs.size ? addCents(...costs.values()) : null;
};
}
+88 -24
View File
@@ -1,15 +1,78 @@
import { and, eq, sql } from "drizzle-orm";
import { agents, budgetPolicies, budgetReservations, heartbeatRuns, nativeRunFinalizations, type Db } from "@paperclipai/db";
import { agents, budgetReservations, heartbeatRuns, nativeRunFinalizations, type Db } from "@paperclipai/db";
import { centsToUnits, unitsToCents } from "@paperclipai/shared";
import { conflict, notFound } from "../errors.js";
import { withAccountingTransaction } from "./accounting-transaction.js";
import { budgetServiceInTransaction, budgetPoliciesForRun, computeObservedSpend } from "./budgets.js";
import { budgetServiceInTransaction } from "./budgets.js";
import { withBudgetSpan } from "./budget-tracing.js";
import type { ActivityPublication } from "./activity-log.js";
/** Native finalization takes its coordinator before issue/run locks. Recovery
* admission uses the same order after the company accounting lock, so its
* later ownership checks never wait for a coordinator while holding a run. */
export async function lockRunBudgetRecoveryOwnerInTransaction(
tx: Db, companyId: string, runId: string, recoveryLeaseOwner?: string,
) {
if (!recoveryLeaseOwner) return;
const [coordinator] = await tx.select({ leaseOwner: nativeRunFinalizations.leaseOwner, leaseExpiresAt: nativeRunFinalizations.leaseExpiresAt }).from(nativeRunFinalizations)
.where(and(eq(nativeRunFinalizations.companyId, companyId), eq(nativeRunFinalizations.runId, runId))).for("update");
if (!coordinator || coordinator.leaseOwner !== recoveryLeaseOwner || !coordinator.leaseExpiresAt || coordinator.leaseExpiresAt.getTime() <= Date.now())
throw conflict("Native recovery no longer owns provider dispatch");
}
/** Called by the final dispatch gate with the company and run already locked.
* Admission may have waited through preparation; that earlier reservation is
* evidence of ownership, not permission to bypass a newer charge or pause.
* Recovery callers prelock their coordinator before taking issue/run locks. */
export async function recheckRunBudgetReservationInTransaction(
tx: Db,
publications: ActivityPublication[],
run: typeof heartbeatRuns.$inferSelect,
reservationId: string,
recoveryLeaseOwner?: string,
): Promise<string | null> {
if (run.resultJson?.cancellation || run.resultJson?.startupCancellation)
return "Run cancellation already requested";
if (run.costAccountedAt || !["queued", "running"].includes(run.status))
return "Run can no longer start provider work";
const [reservation] = await tx.select().from(budgetReservations).where(and(
eq(budgetReservations.id, reservationId), eq(budgetReservations.companyId, run.companyId),
eq(budgetReservations.runId, run.id), eq(budgetReservations.agentId, run.agentId),
eq(budgetReservations.state, "held"),
)).for("update");
if (!reservation) return "Provider dispatch no longer owns a held budget reservation";
if (recoveryLeaseOwner) {
const [owner] = run.runtimeMode === "native" ? await tx.select({ runId: nativeRunFinalizations.runId })
.from(nativeRunFinalizations).where(and(eq(nativeRunFinalizations.companyId, run.companyId),
eq(nativeRunFinalizations.runId, run.id), eq(nativeRunFinalizations.leaseOwner, recoveryLeaseOwner),
sql`${nativeRunFinalizations.leaseExpiresAt} > clock_timestamp()`)) : [];
if (!owner) return "Native recovery no longer owns provider dispatch";
}
const { block, policies, observations } = await budgetServiceInTransaction(tx, publications)
.evaluateAdmission(run.companyId, run.agentId, { projectId: reservation.projectId });
if (block) return block.reason;
const [agent] = await tx.select({ status: agents.status, lifecycleState: agents.lifecycleState }).from(agents)
.where(and(eq(agents.companyId, run.companyId), eq(agents.id, run.agentId)));
if (!agent || agent.lifecycleState !== "ready" || ["paused", "terminated", "pending_approval"].includes(agent.status))
return "Agent is paused or its lifecycle is not ready for new work";
for (const [index, policy] of policies.entries()) {
const observed = observations[index];
// This operation's hold is already included in exposure. Do not propose it
// twice, or reject a nonzero reservation that exactly fills the budget.
const committed = centsToUnits(observed.ledgerCents) + centsToUnits(observed.effectiveHeldCents);
const limit = centsToUnits(policy.amount);
if (committed > limit || (centsToUnits(reservation.amountCents) === 0n && committed >= limit))
return "Available budget is reserved by unfinished runs";
}
return null;
}
/** Reserve before dispatch, under the same company lock as charges and policy
* changes. Estimates constrain admission; they cannot cap a provider's bill.
* A reservation survives timeouts and restarts until accounting proves closure. */
export async function reserveRunBudget(db: Db, companyId: string, runId: string, projectId: string | null, ledgerScope: Record<string, unknown> = {}, recoveryLeaseOwner?: string) {
return withAccountingTransaction(db, companyId, async (tx, publications) => {
const result = await withBudgetSpan("budget.admission.reserve", () => withAccountingTransaction(db, companyId, async (tx, publications) => {
await lockRunBudgetRecoveryOwnerInTransaction(tx, companyId, runId, recoveryLeaseOwner);
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");
@@ -22,40 +85,37 @@ export async function reserveRunBudget(db: Db, companyId: string, runId: string,
&& existing.projectId === projectId ? await tx.select({ runId: nativeRunFinalizations.runId })
.from(nativeRunFinalizations).where(and(eq(nativeRunFinalizations.companyId, companyId),
eq(nativeRunFinalizations.runId, runId), eq(nativeRunFinalizations.leaseOwner, recoveryLeaseOwner),
sql`${nativeRunFinalizations.leaseExpiresAt} > now()`)).for("update") : [];
sql`${nativeRunFinalizations.leaseExpiresAt} > clock_timestamp()`)) : [];
if (owner) {
// Recovery has already made this run queued/running, so its own pending
// accounting is excluded. Current overruns, other unfinished runs and
// scope pauses still apply before any more paid work starts.
const block = await budgetServiceInTransaction(tx, publications).getInvocationBlock(companyId, run.agentId, { projectId });
if (block) throw conflict(block.reason);
const [agent] = await tx.select({ status: agents.status }).from(agents)
if (block) return { admissionBlockReason: block.reason };
const [agent] = await tx.select({ status: agents.status, lifecycleState: agents.lifecycleState }).from(agents)
.where(and(eq(agents.companyId, companyId), eq(agents.id, run.agentId)));
if (agent?.status === "paused") throw conflict("Agent is paused and cannot start new work");
if (!agent || agent.lifecycleState !== "ready" || ["paused", "terminated", "pending_approval"].includes(agent.status))
return { admissionBlockReason: "Agent is paused and cannot start new work" };
return { ...existing, reused: true };
}
throw conflict("Provider dispatch has already reserved this run");
}
const block = await budgetServiceInTransaction(tx, publications).getInvocationBlock(companyId, run.agentId, { projectId });
if (block) throw conflict(block.reason);
const policies = await tx.select().from(budgetPolicies).where(and(
eq(budgetPolicies.companyId, companyId), eq(budgetPolicies.isActive, true), eq(budgetPolicies.hardStopEnabled, true),
eq(budgetPolicies.metric, "billed_cents"), sql`${budgetPolicies.amount} > 0`,
budgetPoliciesForRun(companyId, run.agentId, projectId),
));
const { block, policies, observations } = await budgetServiceInTransaction(tx, publications).evaluateAdmission(companyId, run.agentId, { projectId });
if (block) return { admissionBlockReason: block.reason };
// A budget release can start asynchronous resume hooks. Capacity alone
// does not authorize dispatch before that lifecycle operation is ready.
const [agent] = await tx.select({ status: agents.status, lifecycleState: agents.lifecycleState }).from(agents)
.where(and(eq(agents.companyId, companyId), eq(agents.id, run.agentId)));
if (!agent || agent.lifecycleState !== "ready" || ["paused", "terminated", "pending_approval"].includes(agent.status))
return { admissionBlockReason: "Agent is paused or its lifecycle is not ready for new work" };
const amount = policies.reduce((max, policy) => { const next = centsToUnits(policy.reservationCents); return next > max ? next : max; }, 0n);
for (const policy of policies) {
const [held] = await tx.select({ amount: sql<string>`coalesce(sum(${budgetReservations.amountCents}), 0)::text` }).from(budgetReservations).where(and(
eq(budgetReservations.companyId, companyId), eq(budgetReservations.state, "held"),
policy.scopeType === "agent" ? eq(budgetReservations.agentId, run.agentId) : undefined,
policy.scopeType === "project" ? eq(budgetReservations.projectId, projectId!) : undefined,
));
const observed = await computeObservedSpend(tx, policy);
for (const [index, policy] of policies.entries()) {
const observed = observations[index];
// Even zero-estimate runs cannot enter a scope whose available capacity
// is completely reserved. Outstanding reservations carry across months.
const committed = centsToUnits(observed.totalExact) + centsToUnits(held.amount);
const committed = centsToUnits(observed.ledgerCents) + centsToUnits(observed.effectiveHeldCents);
const limit = centsToUnits(policy.amount);
if (committed >= limit || committed + amount > limit) throw conflict("Available budget is reserved by unfinished runs");
if (committed >= limit || committed + amount > limit) return { admissionBlockReason: "Available budget is reserved by unfinished runs" };
}
// A zero-valued row also fences duplicate dispatch when estimates are off.
const [reservation] = await tx.insert(budgetReservations).values({ companyId, runId, agentId: run.agentId, projectId,
@@ -64,5 +124,9 @@ export async function reserveRunBudget(db: Db, companyId: string, runId: string,
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, reused: false };
});
}));
// Commit required block/resume transitions even when dispatch is denied.
// Throwing inside the transaction would discard those durable obligations.
if ("admissionBlockReason" in result) throw conflict(result.admissionBlockReason ?? "Budget admission blocked");
return result;
}
@@ -0,0 +1,31 @@
import { afterEach, expect, it, vi } from "vitest";
import { withBudgetSpan } from "./budget-tracing.js";
const start = vi.hoisted(() => vi.fn());
vi.mock("@opentelemetry/api", () => ({ trace: { getTracer: () => ({ startActiveSpan: start }) } }));
afterEach(() => { vi.unstubAllEnvs(); vi.resetAllMocks(); });
it("leaves tracing unused without an operator endpoint", async () => {
vi.stubEnv("OTEL_EXPORTER_OTLP_ENDPOINT", "");
await expect(withBudgetSpan("budget.accounting.transaction", async () => 42)).resolves.toBe(42);
expect(start).not.toHaveBeenCalled();
});
it.each(["before", "after", "end", "rejected", "missing"])("keeps accounting authoritative when tracing fails %s", async failure => {
vi.stubEnv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://collector.test");
start.mockImplementation((_name, _options, callback) => {
if (failure === "before") throw new Error("before callback");
if (failure === "rejected") return Promise.reject(new Error("tracer rejected"));
if (failure === "missing") return undefined;
const result = callback({ end: () => { if (failure === "end") throw new Error("end failed"); } });
if (failure === "after") throw new Error("after callback");
return result;
});
const work = vi.fn(async () => 42);
await expect(withBudgetSpan("budget.accounting.transaction", work)).resolves.toBe(42);
expect(work).toHaveBeenCalledOnce();
const error = new Error("original accounting failure");
const failed = vi.fn(async () => { throw error; });
await expect(withBudgetSpan("budget.accounting.transaction", failed)).rejects.toBe(error);
expect(failed).toHaveBeenCalledOnce();
});
+18
View File
@@ -0,0 +1,18 @@
import { trace, type Span } from "@opentelemetry/api";
/** Operator-gated OpenTelemetry only; these spans contain no prompt, receipt,
* credential, or provider error text. Diagnostics never change accounting. */
export async function withBudgetSpan<T>(name: string, work: () => Promise<T>, attributes?: Record<string, number | boolean>) {
if (!process.env.OTEL_EXPORTER_OTLP_ENDPOINT?.trim()) return work();
let operation: Promise<T> | undefined;
const run = (span?: Span) => operation ??= Promise.resolve().then(work).finally(() => {
try { span?.end(); } catch { /* Diagnostics cannot change accounting. */ }
});
try {
const traced = trace.getTracer("paperclip.budgets").startActiveSpan(name, { attributes }, span => run(span));
// Some external SDK implementations can fail before or after the callback.
// Own the accounting promise so either failure still executes work once.
void Promise.resolve(traced).catch(() => undefined);
} catch { /* Preserve the accounting operation if span creation fails. */ }
return operation ?? run();
}
+303 -202
View File
@@ -1,6 +1,5 @@
import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notifications.js";
import { compareCents, normalizeCents } from "@paperclipai/shared";
import { and, desc, eq, gte, inArray, lt, ne, sql, or, isNull, type SQL } from "drizzle-orm";
import { addCents, compareCents, normalizeCents } from "@paperclipai/shared";
import { and, desc, eq, inArray, sql, or, isNull } from "drizzle-orm";
import type { Db } from "@paperclipai/db";
import {
@@ -8,10 +7,10 @@ import {
approvals,
budgetIncidents,
budgetPolicies,
budgetReservations,
companies,
costEvents,
heartbeatRuns,
nativeRunFinalizations,
projects,
} from "@paperclipai/db";
import type {
@@ -31,12 +30,19 @@ 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";
import { computeObservedSpend, computeObservedSpends, type ObservedSpend } from "./budget-exposure.js";
import { notifyDeliveryWork, DELIVERY_QUEUES } from "./delivery-work-notifications.js";
import { listPendingBudgetCancellationScopes } from "./heartbeat/budget-cancellation.js";
import { withBudgetSpan } from "./budget-tracing.js";
export { computeObservedSpend, computeObservedSpends } from "./budget-exposure.js";
export type { BudgetExposure } from "./budget-exposure.js";
type ScopeRecord = {
companyId: string;
name: string;
paused: boolean;
pauseReason: PauseReason | null;
budgetHeld: boolean;
};
type PolicyRow = typeof budgetPolicies.$inferSelect;
@@ -109,6 +115,7 @@ async function resolveScopeRecord(db: Db, scopeType: BudgetScopeType, scopeId: s
name: row.name,
paused: row.status === "paused" || Boolean(row.pausedAt),
pauseReason: (row.pauseReason as ScopeRecord["pauseReason"]) ?? null,
budgetHeld: row.pauseReason === "budget",
};
}
@@ -119,6 +126,7 @@ async function resolveScopeRecord(db: Db, scopeType: BudgetScopeType, scopeId: s
name: agents.name,
status: agents.status,
pauseReason: agents.pauseReason,
lifecycleHolds: agents.lifecycleHolds,
})
.from(agents)
.where(eq(agents.id, scopeId))
@@ -129,6 +137,7 @@ async function resolveScopeRecord(db: Db, scopeType: BudgetScopeType, scopeId: s
name: row.name,
paused: row.status === "paused",
pauseReason: (row.pauseReason as ScopeRecord["pauseReason"]) ?? null,
budgetHeld: row.lifecycleHolds.includes("budget"),
};
}
@@ -148,6 +157,7 @@ async function resolveScopeRecord(db: Db, scopeType: BudgetScopeType, scopeId: s
name: row.name,
paused: Boolean(row.pausedAt),
pauseReason: (row.pauseReason as ScopeRecord["pauseReason"]) ?? null,
budgetHeld: row.pauseReason === "budget",
};
}
@@ -169,68 +179,50 @@ export function budgetPoliciesForRun(companyId: string, agentId: string | null,
);
}
type SpendPolicy = Pick<PolicyRow, "companyId" | "scopeType" | "scopeId" | "windowKind" | "metric">;
type ObservedSpend = { windowStart: Date; windowEnd: Date; total: number; totalExact: string; unpricedEventCount: number; pendingRunCount: number; recoveringRunCount: number };
/** One ledger scan and one pending-run scan for the affected policies. The
* company lock still protects decisions; nothing is cached across mutations.
* Each aggregate keeps its own scope/window, including mixed lifetime and
* monthly policies. An outer predicate excludes rows no policy needs. */
async function computeObservedSpends(db: Db, policies: SpendPolicy[]): Promise<ObservedSpend[]> {
const now = new Date();
const windows = policies.map(policy => resolveWindow(policy.windowKind as BudgetWindowKind, now));
const billable = policies.map((policy, index) => ({ policy, index })).filter(({ policy }) => policy.metric === "billed_cents");
const ledgerColumns: Record<string, SQL> = {};
const pendingColumns: Record<string, SQL> = {};
const ledgerPredicates: SQL[] = [], pendingPredicates: SQL[] = [];
for (const { policy, index } of billable) {
const { start, end } = windows[index];
const ledger = and(
eq(costEvents.companyId, policy.companyId),
policy.scopeType === "agent" ? eq(costEvents.agentId, policy.scopeId) : undefined,
policy.scopeType === "project" ? eq(costEvents.projectId, policy.scopeId) : undefined,
policy.windowKind === "calendar_month_utc" ? and(gte(costEvents.occurredAt, start), lt(costEvents.occurredAt, end)) : undefined,
)!;
const pending = and(
eq(heartbeatRuns.companyId, policy.companyId), eq(heartbeatRuns.costAccountingPending, true),
inArray(heartbeatRuns.status, ["succeeded", "failed", "timed_out", "cancelled", "interrupted"]),
policy.scopeType === "agent" ? eq(heartbeatRuns.agentId, policy.scopeId) : undefined,
policy.scopeType === "project" ? sql`${heartbeatRuns.usageJson}->'ledgerScope'->>'projectId' = ${policy.scopeId}` : undefined,
policy.windowKind === "calendar_month_utc" ? sql`coalesce(${heartbeatRuns.finishedAt}, ${heartbeatRuns.createdAt}) >= ${start.toISOString()}::timestamptz
and coalesce(${heartbeatRuns.finishedAt}, ${heartbeatRuns.createdAt}) < ${end.toISOString()}::timestamptz` : undefined,
)!;
ledgerPredicates.push(ledger); pendingPredicates.push(pending);
ledgerColumns[`total${index}`] = sql`coalesce(sum(${costEvents.costCents}) filter (where ${ledger}), 0)::text`;
ledgerColumns[`unpriced${index}`] = sql`count(*) filter (where ${ledger} and ${costEvents.costStatus} = 'unpriced' and ${costEvents.billingType} <> 'subscription_included')::int`;
pendingColumns[`pending${index}`] = sql`count(*) filter (where ${pending})::int`;
pendingColumns[`recovering${index}`] = sql`count(*) filter (where ${pending} and ${heartbeatRuns.runtimeMode} = 'native'
and exists (select 1 from ${nativeRunFinalizations} where ${nativeRunFinalizations.runId} = ${heartbeatRuns.id}
and ${nativeRunFinalizations.companyId} = ${heartbeatRuns.companyId}
and ${nativeRunFinalizations.resultId} is null and ${nativeRunFinalizations.phase} <> 'terminal_failure'))::int`;
}
const [ledger = {}] = billable.length ? await db.select(ledgerColumns).from(costEvents).where(or(...ledgerPredicates)) : [];
const [pending = {}] = billable.length ? await db.select(pendingColumns).from(heartbeatRuns).where(or(...pendingPredicates)) : [];
return policies.map((_, index) => ({
windowStart: windows[index].start, windowEnd: windows[index].end,
total: Number(ledger[`total${index}`] ?? 0), totalExact: normalizeCents(String(ledger[`total${index}`] ?? 0)),
unpricedEventCount: Number(ledger[`unpriced${index}`] ?? 0), pendingRunCount: Number(pending[`pending${index}`] ?? 0),
recoveringRunCount: Number(pending[`recovering${index}`] ?? 0),
}));
}
export async function computeObservedSpend(db: Db, policy: SpendPolicy) {
return (await computeObservedSpends(db, [policy]))[0];
function unpricedUsageBlocks(policy: PolicyRow, observed: ObservedSpend) {
return policy.unpricedUsagePolicy !== "allow"
&& (observed.unpricedEventCount > 0 || observed.unknownBrowserRunCount > 0);
}
function observedBlocks(policy: PolicyRow, observed: Awaited<ReturnType<typeof computeObservedSpend>>) {
if (!policy.isActive || !policy.hardStopEnabled || policy.amount <= 0) return false;
return compareCents(observed.totalExact, policy.amount) >= 0 || (observed.pendingRunCount - observed.recoveringRunCount) > 0 || (policy.unpricedUsagePolicy !== "allow" && observed.unpricedEventCount > 0);
return compareCents(observed.totalExact, policy.amount) >= 0 || (observed.pendingRunCount - observed.recoveringRunCount) > 0 || unpricedUsageBlocks(policy, observed);
}
export async function policyBlocks(db: Db, policy: PolicyRow) {
if (!policy.isActive || !policy.hardStopEnabled || policy.amount <= 0) return false;
return observedBlocks(policy, await computeObservedSpend(db, policy));
}
/** Supplemental handoff check after authoritative accounting admission commits.
* Read only: the caller can hold issue/run ownership, so this must never take
* the company lock, reconcile holds, or deliver effects. A later concurrent
* charge targets already-admitted work through the normal durable Stop path. */
export async function recheckAdmittedRunBudgetReadOnly(db: Db, input: {
companyId: string; agentId: string; runId: string; projectId: string | null; reservationId: string;
}): Promise<string | null> {
return withBudgetSpan("budget.admission.evaluate", async () => {
const [reservation] = await db.select({ amountCents: budgetReservations.amountCents }).from(budgetReservations).where(and(
eq(budgetReservations.id, input.reservationId), eq(budgetReservations.companyId, input.companyId),
eq(budgetReservations.runId, input.runId), eq(budgetReservations.agentId, input.agentId),
sql`${budgetReservations.projectId} is not distinct from ${input.projectId}::uuid`,
eq(budgetReservations.state, "held"),
));
if (!reservation) return "Provider dispatch no longer owns a held budget reservation";
const policies = await db.select().from(budgetPolicies).where(and(eq(budgetPolicies.companyId, input.companyId),
eq(budgetPolicies.isActive, true), eq(budgetPolicies.hardStopEnabled, true), eq(budgetPolicies.metric, "billed_cents"),
sql`${budgetPolicies.amount} > 0`, budgetPoliciesForRun(input.companyId, input.agentId, input.projectId)));
const observations = await computeObservedSpends(db, policies);
for (const [index, policy] of policies.entries()) {
const observed = observations[index];
if (observedBlocks(policy, observed)) return "Applicable budget no longer permits provider work";
const capacity = compareCents(addCents(observed.ledgerCents, observed.effectiveHeldCents), policy.amount);
if (capacity > 0 || (compareCents(reservation.amountCents, 0) === 0 && capacity >= 0))
return "Available budget is reserved by unfinished runs";
}
return null;
});
}
function buildApprovalPayload(input: {
policy: PolicyRow;
scopeName: string;
@@ -283,7 +275,11 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
const recordActivity = (input: LogActivityInput) => logActivity(db, input, publications);
async function pauseScopeForBudget(policy: PolicyRow) {
const now = new Date();
if (policy.scopeType === "agent") return;
if (policy.scopeType === "agent") {
const { applyAgentBudgetHoldInTransaction } = await import("./agent-lifecycle.js");
await applyAgentBudgetHoldInTransaction(db, policy.companyId, policy.scopeId, true);
return;
}
if (policy.scopeType === "project") {
await db
@@ -293,10 +289,13 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
pausedAt: now,
updatedAt: now,
})
.where(and(eq(projects.id, policy.scopeId), or(isNull(projects.pausedAt), eq(projects.pauseReason, "budget"))));
.where(and(eq(projects.id, policy.scopeId), isNull(projects.pausedAt)));
return;
}
const scope = await resolveScopeRecord(db, "company", policy.scopeId);
if (scope.paused) return;
await notifyDeliveryWork(db, DELIVERY_QUEUES.agentLifecycle, policy.companyId);
await notifyDeliveryWork(db, DELIVERY_QUEUES.skillSourceSync);
await db
.update(companies)
@@ -306,13 +305,20 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
pausedAt: now,
updatedAt: now,
})
.where(and(eq(companies.id, policy.scopeId), or(eq(companies.status, "active"), and(eq(companies.status, "paused"), eq(companies.pauseReason, "budget")))));
.where(and(eq(companies.id, policy.scopeId), eq(companies.status, "active")));
}
async function pauseAndCancelScopeForBudget(policy: PolicyRow) {
async function pauseAndCancelScopeForBudget(policy: PolicyRow, requestCancellation = true) {
await pauseScopeForBudget(policy);
if (!requestCancellation) return;
// A pending generation already covers all admitted work in this scope.
// Rechecks must not invalidate a worker that is delivering that intent.
await notifyDeliveryWork(db, DELIVERY_QUEUES.budgetEnforcement, policy.companyId);
await db.update(budgetPolicies).set({
enforcementVersion: sql`${budgetPolicies.enforcementVersion} + 1`,
enforcementVersion: sql`case when ${budgetPolicies.enforcementVersion} > ${budgetPolicies.enforcementDeliveredVersion}
then ${budgetPolicies.enforcementVersion} else ${budgetPolicies.enforcementVersion} + 1 end`,
updatedAt: sql`case when ${budgetPolicies.enforcementVersion} > ${budgetPolicies.enforcementDeliveredVersion}
then ${budgetPolicies.updatedAt} else now() end`,
}).where(eq(budgetPolicies.id, policy.id));
}
@@ -321,11 +327,14 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
eq(budgetPolicies.companyId, policy.companyId), eq(budgetPolicies.scopeType, policy.scopeType),
eq(budgetPolicies.scopeId, policy.scopeId), eq(budgetPolicies.isActive, true), eq(budgetPolicies.hardStopEnabled, true),
));
for (const candidate of policies) {
if (await policyBlocks(db, candidate)) return;
}
const observations = await computeObservedSpends(db, policies);
if (policies.some((candidate, index) => observedBlocks(candidate, observations[index]))) return;
const now = new Date();
if (policy.scopeType === "agent") return;
if (policy.scopeType === "agent") {
const { applyAgentBudgetHoldInTransaction } = await import("./agent-lifecycle.js");
await applyAgentBudgetHoldInTransaction(db, policy.companyId, policy.scopeId, false);
return;
}
if (policy.scopeType === "project") {
await db
@@ -340,6 +349,9 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
}
await notifyDeliveryWork(db, DELIVERY_QUEUES.skillSourceSync);
const scope = await resolveScopeRecord(db, "company", policy.scopeId);
if (scope.pauseReason !== "budget") return;
await notifyDeliveryWork(db, DELIVERY_QUEUES.agentLifecycle, policy.companyId);
await db
.update(companies)
.set({
@@ -559,23 +571,25 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
});
}
async function reconcileScope(companyId: string, scopeType: BudgetScopeType, scopeId: string) {
const scope = await resolveScopeRecord(db, scopeType, scopeId);
async function reconcileScope(companyId: string, scopeType: BudgetScopeType, scopeId: string,
prepared?: { scope: ScopeRecord; policies: PolicyRow[]; observations: ObservedSpend[] }) {
const scope = prepared?.scope ?? await resolveScopeRecord(db, scopeType, scopeId);
if (scope.companyId !== companyId) throw notFound("Budget scope not found");
const policies = await db.select().from(budgetPolicies).where(and(
const policies = prepared?.policies ?? await db.select().from(budgetPolicies).where(and(
eq(budgetPolicies.companyId, companyId), eq(budgetPolicies.scopeType, scopeType), eq(budgetPolicies.scopeId, scopeId),
));
let recoveringPolicy: PolicyRow | undefined;
for (const policy of policies) {
const observed = await computeObservedSpend(db, policy);
const observations = prepared?.observations ?? await computeObservedSpends(db, policies);
for (const [index, policy] of policies.entries()) {
const observed = observations[index];
if (observedBlocks(policy, observed)) {
await createIncidentIfNeeded(policy, "hard", observed);
await pauseAndCancelScopeForBudget(policy);
await pauseAndCancelScopeForBudget(policy, !scope.budgetHeld);
const title = scopeType === "company" ? "Company" : scopeType === "agent" ? "Agent" : "Project";
if (observed.pendingRunCount > 0) {
return { scopeType, scopeId, scopeName: scope.name, reason: `${title} cannot start work while completed runs await accounting.` };
}
if (policy.unpricedUsagePolicy !== "allow" && observed.unpricedEventCount > 0) {
if (unpricedUsageBlocks(policy, observed)) {
return { scopeType, scopeId, scopeName: scope.name, reason: `${title} cannot start work because recorded usage has no reliable price.` };
}
return { scopeType, scopeId, scopeName: scope.name, reason: scopeType === "project"
@@ -600,12 +614,112 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
if (scope.paused && scope.pauseReason === "budget" && !policies.length) {
return { scopeType, scopeId, scopeName: scope.name, reason: "Budget pause requires a policy or an explicit operator resume." };
}
if (scope.pauseReason === "budget" && policies[0]) await resumeScopeFromBudget(policies[0]);
if (scope.budgetHeld && policies[0]) await resumeScopeFromBudget(policies[0]);
if (recoveringPolicy) return { scopeType, scopeId, scopeName: scope.name,
reason: "New work must wait while a native run recovers its unfinished accounting." };
return null;
}
async function evaluateAffectedScopes(companyId: string, agentId: string | null, projectId: string | null) {
const candidatePolicies = await db
.select()
.from(budgetPolicies)
.where(
and(
eq(budgetPolicies.companyId, companyId),
eq(budgetPolicies.isActive, true),
budgetPoliciesForRun(companyId, agentId, projectId),
),
);
const policies = candidatePolicies.filter(policy => policy.metric === "billed_cents" && policy.amount > 0);
const observations = await computeObservedSpends(db, policies);
for (const [index, policy] of policies.entries()) {
const observed = observations[index];
const observedAmount = observed.total;
const softThreshold = Math.ceil((policy.amount * policy.warnPercent) / 100);
if (policy.notifyEnabled && compareCents(observed.totalExact, softThreshold) >= 0 && (!policy.hardStopEnabled || compareCents(observed.totalExact, policy.amount) < 0)) {
const softIncident = await createIncidentIfNeeded(policy, "soft", observed);
if (softIncident?.created) {
await recordActivity({
companyId: policy.companyId,
actorType: "system",
actorId: "budget_service",
action: "budget.soft_threshold_crossed",
entityType: "budget_incident",
entityId: softIncident.incident.id,
details: {
scopeType: policy.scopeType,
scopeId: policy.scopeId,
amountObserved: observedAmount,
amountLimit: policy.amount,
},
});
}
}
if (observedBlocks(policy, observed)) {
await resolveOpenSoftIncidents(policy.id);
const hardIncident = await createIncidentIfNeeded(policy, "hard", observed);
await pauseAndCancelScopeForBudget(policy);
if (hardIncident?.created) {
await recordActivity({
companyId: policy.companyId,
actorType: "system",
actorId: "budget_service",
action: "budget.hard_threshold_crossed",
entityType: "budget_incident",
entityId: hardIncident.incident.id,
details: {
scopeType: policy.scopeType,
scopeId: policy.scopeId,
amountObserved: observedAmount,
amountLimit: policy.amount,
approvalId: hardIncident.incident.approvalId ?? null,
},
});
}
}
}
// A final receipt or price correction can reduce a partial observation.
// Release only this scope's budget-owned hold after all its policies agree.
const healthyScopes = new Map<string, PolicyRow>();
for (const [index, policy] of policies.entries()) {
if (!observedBlocks(policy, observations[index])) healthyScopes.set(`${policy.scopeType}:${policy.scopeId}`, policy);
}
for (const policy of healthyScopes.values()) {
const scope = await resolveScopeRecord(db, policy.scopeType as BudgetScopeType, policy.scopeId);
if (scope.budgetHeld) await resumeScopeFromBudget(policy);
}
}
async function evaluateAdmissionInTransaction(companyId: string, agentId: string | null,
context?: { issueId?: string | null; projectId?: string | null }) {
const scopes: Array<{ scopeType: BudgetScopeType; scopeId: string }> = [
{ scopeType: "company", scopeId: companyId }, ...(agentId ? [{ scopeType: "agent" as const, scopeId: agentId }] : []),
...(context?.projectId ? [{ scopeType: "project" as const, scopeId: context.projectId }] : []),
];
const records = await Promise.all(scopes.map(scope => resolveScopeRecord(db, scope.scopeType, scope.scopeId)));
if (records.some(scope => scope.companyId !== companyId)) throw notFound("Budget scope not found");
const policies = await db.select().from(budgetPolicies).where(and(eq(budgetPolicies.companyId, companyId),
budgetPoliciesForRun(companyId, agentId, context?.projectId ?? null)));
const observations = await computeObservedSpends(db, policies);
const reservationIndexes = policies.map((policy, i) => policy.isActive && policy.hardStopEnabled && policy.metric === "billed_cents" && policy.amount > 0 ? i : -1).filter(i => i >= 0);
const admission = (block: Awaited<ReturnType<typeof reconcileScope>>) => ({ block,
policies: reservationIndexes.map(i => policies[i]), observations: reservationIndexes.map(i => observations[i]) });
for (const [index, { scopeType, scopeId }] of scopes.entries()) {
const indexes = policies.map((policy, i) => policy.scopeType === scopeType && policy.scopeId === scopeId ? i : -1).filter(i => i >= 0);
const block = await reconcileScope(companyId, scopeType, scopeId, { scope: records[index],
policies: indexes.map(i => policies[i]), observations: indexes.map(i => observations[i]) });
if (block) return admission(block);
}
return admission(null);
}
const evaluateAdmission = (...args: Parameters<typeof evaluateAdmissionInTransaction>) =>
withBudgetSpan("budget.admission.evaluate", () => evaluateAdmissionInTransaction(...args));
return {
reconcileScope,
listPolicies: async (companyId: string): Promise<BudgetPolicy[]> => {
@@ -652,6 +766,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
// setting; editing a disabled positive policy must preserve that choice.
const nextIsActive = amount > 0 && (input.isActive ?? (existing?.amount === 0 ? true : existing?.isActive) ?? true);
const now = new Date();
await notifyDeliveryWork(db, DELIVERY_QUEUES.budgetEnforcement, companyId);
const row = existing
? await db
.update(budgetPolicies)
@@ -665,6 +780,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
isActive: nextIsActive,
updatedByUserId: actorUserId,
updatedAt: now,
enforcementVersion: sql`${budgetPolicies.enforcementVersion} + 1`,
})
.where(eq(budgetPolicies.id, existing.id))
.returning()
@@ -773,88 +889,12 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
};
},
evaluateCostEvent: async (event: typeof costEvents.$inferSelect) => {
const candidatePolicies = await db
.select()
.from(budgetPolicies)
.where(
and(
eq(budgetPolicies.companyId, event.companyId),
eq(budgetPolicies.isActive, true),
budgetPoliciesForRun(event.companyId, event.agentId, event.projectId),
),
);
evaluateAffectedScopes,
evaluateCostEvent: (event: typeof costEvents.$inferSelect) => evaluateAffectedScopes(event.companyId, event.agentId, event.projectId),
const policies = candidatePolicies.filter(policy => policy.metric === "billed_cents" && policy.amount > 0);
const observations = await computeObservedSpends(db, policies);
for (const [index, policy] of policies.entries()) {
const observed = observations[index];
const observedAmount = observed.total;
const softThreshold = Math.ceil((policy.amount * policy.warnPercent) / 100);
if (policy.notifyEnabled && compareCents(observed.totalExact, softThreshold) >= 0 && (!policy.hardStopEnabled || compareCents(observed.totalExact, policy.amount) < 0)) {
const softIncident = await createIncidentIfNeeded(policy, "soft", observed);
if (softIncident?.created) {
await recordActivity({
companyId: policy.companyId,
actorType: "system",
actorId: "budget_service",
action: "budget.soft_threshold_crossed",
entityType: "budget_incident",
entityId: softIncident.incident.id,
details: {
scopeType: policy.scopeType,
scopeId: policy.scopeId,
amountObserved: observedAmount,
amountLimit: policy.amount,
},
});
}
}
if (observedBlocks(policy, observed)) {
await resolveOpenSoftIncidents(policy.id);
const hardIncident = await createIncidentIfNeeded(policy, "hard", observed);
await pauseAndCancelScopeForBudget(policy);
if (hardIncident?.created) {
await recordActivity({
companyId: policy.companyId,
actorType: "system",
actorId: "budget_service",
action: "budget.hard_threshold_crossed",
entityType: "budget_incident",
entityId: hardIncident.incident.id,
details: {
scopeType: policy.scopeType,
scopeId: policy.scopeId,
amountObserved: observedAmount,
amountLimit: policy.amount,
approvalId: hardIncident.incident.approvalId ?? null,
},
});
}
}
}
},
getInvocationBlock: async (
companyId: string,
agentId: string | null,
context?: { issueId?: string | null; projectId?: string | null },
) => {
const scopes: Array<{ scopeType: BudgetScopeType; scopeId: string }> = [
{ scopeType: "company", scopeId: companyId }, ...(agentId ? [{ scopeType: "agent" as const, scopeId: agentId }] : []),
...(context?.projectId ? [{ scopeType: "project" as const, scopeId: context.projectId }] : []),
];
for (const { scopeType, scopeId } of scopes) {
if ((await resolveScopeRecord(db, scopeType, scopeId)).companyId !== companyId) throw notFound("Budget scope not found");
}
for (const { scopeType, scopeId } of scopes) {
const block = await reconcileScope(companyId, scopeType, scopeId);
if (block) return block;
}
return null;
},
evaluateAdmission,
getInvocationBlock: async (companyId: string, agentId: string | null, context?: { issueId?: string | null; projectId?: string | null }) =>
(await evaluateAdmission(companyId, agentId, context)).block,
resolveIncident: async (
companyId: string,
@@ -874,7 +914,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
if (input.action === "raise_budget_and_resume") {
const nextAmount = Math.max(0, Math.floor(input.amount ?? 0));
const observed = await computeObservedSpend(db, policy);
if (policy.unpricedUsagePolicy !== "allow" && observed.unpricedEventCount > 0) {
if (unpricedUsageBlocks(policy, observed)) {
throw unprocessable("Resolve unpriced usage or explicitly allow it in the budget policy before resuming");
}
if (observed.pendingRunCount > 0) throw unprocessable("Completed runs must finish accounting before resuming");
@@ -883,6 +923,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
}
const now = new Date();
await notifyDeliveryWork(db, DELIVERY_QUEUES.budgetEnforcement, companyId);
await db
.update(budgetPolicies)
.set({
@@ -890,6 +931,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
isActive: true,
updatedByUserId: actorUserId,
updatedAt: now,
enforcementVersion: sql`${budgetPolicies.enforcementVersion} + 1`,
})
.where(eq(budgetPolicies.id, policy.id));
@@ -972,64 +1014,123 @@ export async function withCurrentBudgetEnforcement<T>(db: Db, scope: BudgetEnfor
});
}
/** 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, agentId?: string | null) {
const { createAgentLifecycle } = await import("./agent-lifecycle.js");
// Global recovery uses the bounded lifecycle sweep for policy holds.
if (companyId) await createAgentLifecycle(db).reconcilePolicyHolds(companyId, agentId);
async function deliverScopes<T>(items: T[], deliver: (item: T) => Promise<void>) {
let next = 0;
await Promise.all(Array.from({ length: Math.min(4, items.length) }, async () => {
while (next < items.length) await deliver(items[next++]);
}));
}
/** Cancellation is an at-least-once external effect. Only the delivery worker
* invokes this: request paths commit intent and return without awaiting it. */
export async function deliverBudgetEnforcement(db: Db, hooks: BudgetServiceHooks, companyId?: string, _agentId?: string | null) {
if (!hooks.cancelWorkForScope) return;
const attempted = new Set<string>();
const scopeKey = (scope: BudgetEnforcementScope) => `${scope.companyId}:${scope.scopeType}:${scope.scopeId}`;
const pending = await db.select().from(budgetPolicies).where(and(
sql`${budgetPolicies.enforcementVersion} > ${budgetPolicies.enforcementDeliveredVersion}`,
companyId ? eq(budgetPolicies.companyId, companyId) : undefined,
));
for (const policy of pending) {
try {
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(), enforcement: { policyId: policy.id, version: policy.enforcementVersion } };
});
if (cancellation) await hooks.cancelWorkForScope(cancellation);
await db.update(budgetPolicies).set({ enforcementDeliveredVersion: policy.enforcementVersion })
.where(and(eq(budgetPolicies.id, policy.id), eq(budgetPolicies.enforcementVersion, policy.enforcementVersion)));
} catch (error) {
logger.warn({ err: error, policyId: policy.id }, "Budget cancellation delivery pending; accounting is committed");
}
}
)).orderBy(budgetPolicies.updatedAt, budgetPolicies.id).limit(100);
await deliverScopes(pending, policy => withBudgetSpan("budget.enforcement.scope", async () => {
try {
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(), enforcement: { policyId: policy.id, version: policy.enforcementVersion } };
});
if (cancellation) {
attempted.add(scopeKey(cancellation));
await hooks.cancelWorkForScope!(cancellation);
}
await db.update(budgetPolicies).set({ enforcementDeliveredVersion: policy.enforcementVersion })
.where(and(eq(budgetPolicies.id, policy.id), eq(budgetPolicies.enforcementVersion, policy.enforcementVersion)));
} catch (error) {
logger.warn({ err: error, policyId: policy.id }, "Budget cancellation delivery pending; accounting is committed");
// Rotate a failed scope behind other pending work. A poisoned delivery
// must not monopolize the bounded batch on every retry.
await db.update(budgetPolicies).set({ updatedAt: new Date() }).where(and(eq(budgetPolicies.id, policy.id),
eq(budgetPolicies.enforcementVersion, policy.enforcementVersion))).catch(() => undefined);
}
}, { "budget.enforcement.queue_delay_ms": Math.max(0, Date.now() - policy.updatedAt.getTime()) }));
// A committed run Stop outlives a policy edit or deletion. It remains owed
// even if no policy generation now authorizes a fresh cancellation.
const committed = (await listPendingBudgetCancellationScopes(db))
.filter(scope => (!companyId || scope.companyId === companyId) && !attempted.has(scopeKey(scope)));
await deliverScopes(committed, async scope => {
try { await hooks.cancelWorkForScope!(scope); }
catch (error) { logger.warn({ err: error, companyId: scope.companyId, scopeType: scope.scopeType, scopeId: scope.scopeId }, "Committed budget Stop remains pending"); }
});
}
export function budgetService(db: Db, hooks: BudgetServiceHooks = {}) {
const reads = budgetServiceInTransaction(db);
async function mutate<T>(companyId: string, work: (service: ReturnType<typeof budgetServiceInTransaction>) => Promise<T>, agentId?: string | null) {
const result = await withAccountingTransaction(db, companyId, (tx, publications) => work(budgetServiceInTransaction(tx, publications)));
await deliverBudgetEnforcement(db, hooks, companyId, agentId);
return result;
export async function hasPendingBudgetEnforcement(db: Db): Promise<boolean> {
const [pending] = await db.select({ id: budgetPolicies.id }).from(budgetPolicies)
.where(sql`${budgetPolicies.enforcementVersion} > ${budgetPolicies.enforcementDeliveredVersion}`).limit(1);
return Boolean(pending) || (await listPendingBudgetCancellationScopes(db, 1)).length > 0;
}
const budgetOwnedScope = or(
and(eq(budgetPolicies.scopeType, "company"), sql`exists (select 1 from ${companies} where ${companies.id} = ${budgetPolicies.scopeId} and ${companies.pauseReason} = 'budget')`),
and(eq(budgetPolicies.scopeType, "agent"), sql`exists (select 1 from ${agents} where ${agents.id} = ${budgetPolicies.scopeId} and ${agents.lifecycleHolds} @> '["budget"]'::jsonb)`),
and(eq(budgetPolicies.scopeType, "project"), sql`exists (select 1 from ${projects} where ${projects.id} = ${budgetPolicies.scopeId} and ${projects.pauseReason} = 'budget')`),
)!;
const openIncidentScope = sql`exists (select 1 from ${budgetIncidents}
where ${budgetIncidents.policyId} = ${budgetPolicies.id} and ${budgetIncidents.status} = 'open')`;
const unfinishedRunScope = sql`exists (select 1 from ${heartbeatRuns}
where ${heartbeatRuns.companyId} = ${budgetPolicies.companyId} and ${heartbeatRuns.costAccountingPending} = true
and ${heartbeatRuns.costAccountedAt} is null and (${budgetPolicies.scopeType} = 'company'
or (${budgetPolicies.scopeType} = 'agent' and ${heartbeatRuns.agentId} = ${budgetPolicies.scopeId})
or (${budgetPolicies.scopeType} = 'project' and ${heartbeatRuns.usageJson}->'ledgerScope'->>'projectId' = ${budgetPolicies.scopeId}::text)))`;
export async function nextBudgetReconciliationAt(db: Db): Promise<Date | null> {
const [held] = await db.select({ id: budgetPolicies.id }).from(budgetPolicies)
.where(and(eq(budgetPolicies.windowKind, "calendar_month_utc"), or(budgetOwnedScope, openIncidentScope))).limit(1);
return held ? currentUtcMonthWindow().end : null;
}
/** Startup repairs configured policy scopes, not every agent. Afterwards only
* held monthly scopes owe a timed reevaluation; writes reconcile their own scopes. */
export async function reconcileDueBudgetPolicies(db: Db, options: { startup?: boolean; cursor?: string | null } = {}): Promise<{ nextCursor: string | null }> {
const actionable = options.startup ? or(budgetOwnedScope, openIncidentScope,
sql`${budgetPolicies.enforcementVersion} > ${budgetPolicies.enforcementDeliveredVersion}`,
and(eq(budgetPolicies.isActive, true), eq(budgetPolicies.hardStopEnabled, true), sql`${budgetPolicies.amount} > 0`, unfinishedRunScope))
: and(eq(budgetPolicies.windowKind, "calendar_month_utc"), or(budgetOwnedScope, openIncidentScope));
const rows = await db.select({ id: budgetPolicies.id, companyId: budgetPolicies.companyId, scopeType: budgetPolicies.scopeType, scopeId: budgetPolicies.scopeId })
.from(budgetPolicies).where(and(actionable, options.cursor ? sql`${budgetPolicies.id} > ${options.cursor}::uuid` : undefined))
.orderBy(budgetPolicies.id).limit(100);
const scopes = new Map(rows.map(scope => [`${scope.companyId}:${scope.scopeType}:${scope.scopeId}`, scope]));
for (const scope of scopes.values()) {
try {
await withAccountingTransaction(db, scope.companyId, (tx, publications) => budgetServiceInTransaction(tx, publications)
.reconcileScope(scope.companyId, scope.scopeType as BudgetScopeType, scope.scopeId));
} catch (error) {
if (error instanceof HttpError && error.status === 404) continue;
throw error;
}
}
return { nextCursor: rows.length === 100 ? rows[rows.length - 1].id : null };
}
export function budgetService(db: Db, _hooks: BudgetServiceHooks = {}) {
const reads = budgetServiceInTransaction(db);
const mutate = <T>(companyId: string, work: (service: ReturnType<typeof budgetServiceInTransaction>) => Promise<T>) =>
withAccountingTransaction(db, companyId, (tx, publications) => work(budgetServiceInTransaction(tx, publications)));
return {
...reads,
deliverPendingEnforcement: (companyId?: string) => deliverBudgetEnforcement(db, hooks, companyId),
reconcilePolicies: async () => {
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");
}
}
},
evaluateAdmission: (companyId: string, agentId: string | null, context?: { issueId?: string | null; projectId?: string | null }) =>
mutate(companyId, service => service.evaluateAdmission(companyId, agentId, context)),
evaluateAffectedScopes: (companyId: string, agentId: string | null, projectId: string | null) =>
mutate(companyId, service => service.evaluateAffectedScopes(companyId, agentId, projectId)),
reconcileScope: (companyId: string, scopeType: BudgetScopeType, scopeId: string) =>
mutate(companyId, service => service.reconcileScope(companyId, scopeType, scopeId)),
reconcilePolicies: () => reconcileDueBudgetPolicies(db, { startup: true }),
upsertPolicy: (companyId: string, input: BudgetPolicyUpsertInput, actorUserId: string | null) =>
mutate(companyId, (service) => service.upsertPolicy(companyId, input, actorUserId)),
evaluateCostEvent: (event: typeof costEvents.$inferSelect) =>
mutate(event.companyId, (service) => service.evaluateCostEvent(event)),
mutate(companyId, service => service.upsertPolicy(companyId, input, actorUserId)),
evaluateCostEvent: (event: typeof costEvents.$inferSelect) => mutate(event.companyId, service => service.evaluateCostEvent(event)),
getInvocationBlock: (companyId: string, agentId: string | null, context?: { issueId?: string | null; projectId?: string | null }) =>
mutate(companyId, (service) => service.getInvocationBlock(companyId, agentId, context), agentId),
mutate(companyId, service => service.getInvocationBlock(companyId, agentId, context)),
resolveIncident: (companyId: string, incidentId: string, input: BudgetIncidentResolutionInput, actorUserId: string) =>
mutate(companyId, (service) => service.resolveIncident(companyId, incidentId, input, actorUserId)),
mutate(companyId, service => service.resolveIncident(companyId, incidentId, input, actorUserId)),
};
}
+24 -25
View File
@@ -1,9 +1,9 @@
import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notifications.js";
import { isRunnerApiResponseAsset } from "./asset-content-access.js";
import { deleteCompany } from "./company-deletion.js";
import { createAgentLifecycle } from "./agent-lifecycle.js";
import { fenceActiveCompanyAgentStopsInTransaction } from "./agent-lifecycle.js";
import { publishAccountingActivities } from "./accounting-transaction.js";
import { budgetServiceInTransaction, deliverBudgetEnforcement, type BudgetServiceHooks } from "./budgets.js";
import { budgetServiceInTransaction, type BudgetServiceHooks } from "./budgets.js";
import { and, count, eq, gte, inArray, isNull, lt, notInArray, sql } from "drizzle-orm";
import type { Db } from "@paperclipai/db";
import {
@@ -28,7 +28,6 @@ import {
rekeyCompanyIssueIdentifiers,
} from "./issue-prefix.js";
import { environmentService } from "./environments.js";
import { heartbeatService } from "./heartbeat.js";
import { logActivity, type ActivityPublication } from "./activity-log.js";
import { builtInAgentService } from "./built-in-agents.js";
@@ -48,24 +47,22 @@ const SYSTEM_COMPANY_ACTOR: CompanyActivityActor = {
runId: null,
};
export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
export function companyService(db: Db, _budgetHooks: BudgetServiceHooks = {}) {
const environmentsSvc = environmentService(db);
const heartbeat = heartbeatService(db);
const builtInAgents = builtInAgentService(db);
type CompanyTx = Parameters<Parameters<typeof db.transaction>[0]>[0];
async function applyArchiveCascadeInTx(tx: CompanyTx, id: string) {
const pausedAgentRows = await tx.select({ id: agents.id }).from(agents).where(and(eq(agents.companyId, id), notInArray(agents.status, ["paused", "terminated", "pending_approval"])));
const [agentCount] = await tx.select({ value: count() }).from(agents).where(and(eq(agents.companyId, id), notInArray(agents.status, ["paused", "terminated", "pending_approval"])));
const activeRunIds = await tx
.select({ id: heartbeatRuns.id })
const [runCount] = await tx
.select({ value: count() })
.from(heartbeatRuns)
.where(and(
eq(heartbeatRuns.companyId, id),
inArray(heartbeatRuns.status, ["queued", "running"]),
))
.then((rows) => rows.map((row) => row.id));
inArray(heartbeatRuns.status, ["queued", "running", "scheduled_retry"]),
));
await tx
.update(agentWakeupRequests)
@@ -81,18 +78,16 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
isNull(agentWakeupRequests.runId),
));
return { agentsPaused: pausedAgentRows.length, activeRunIds };
return { agentsPauseRequested: agentCount.value, runsCancellationRequested: runCount.value };
}
async function finalizeArchive(
id: string,
actor: CompanyActivityActor,
cascade: { agentsPaused: number; activeRunIds: string[] },
cascade: { agentsPauseRequested: number; runsCancellationRequested: number },
) {
for (const runId of cascade.activeRunIds) {
await heartbeat.cancelRun(runId, "Cancelled because the company was archived");
}
// The committed company status blocks dispatch immediately. Lifecycle
// projection owns provider termination, including paused agents with work.
await logActivity(db, {
companyId: id,
actorType: actor.actorType,
@@ -103,8 +98,8 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
entityType: "company",
entityId: id,
details: {
agentsPaused: cascade.agentsPaused,
runsCancelled: cascade.activeRunIds.length,
agentsPauseRequested: cascade.agentsPauseRequested,
runsCancellationRequested: cascade.runsCancellationRequested,
},
});
}
@@ -304,14 +299,14 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
) => {
const budgetPublications: ActivityPublication[] = [];
const result = await db.transaction(async (tx) => {
if (data.status !== undefined || data.budgetMonthlyCents !== undefined) {
await tx.select({ id: companies.id }).from(companies).where(eq(companies.id, id)).for("no key update");
}
const existing = await getCompanyQuery(tx)
.where(eq(companies.id, id))
.then((rows) => rows[0] ?? null);
if (!existing) return null;
if (data.budgetMonthlyCents !== undefined) {
await tx.select({ id: companies.id }).from(companies).where(eq(companies.id, id)).for("no key update");
}
const { logoAssetId, ...companyPatch } = data;
const willReactivate = existing.status !== "active" && companyPatch.status === "active";
const willArchive = existing.status !== "archived" && companyPatch.status === "archived";
@@ -335,6 +330,7 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
if (companyPatch.status !== undefined && companyPatch.status !== existing.status) {
await notifyDeliveryWork(tx, DELIVERY_QUEUES.skillSourceSync);
await notifyDeliveryWork(tx, DELIVERY_QUEUES.agentLifecycle, id);
}
const updated = await tx
.update(companies)
@@ -375,6 +371,9 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
}
const archiveCascade = willArchive ? await applyArchiveCascadeInTx(tx, id) : null;
if (companyPatch.status !== existing.status && (companyPatch.status === "paused" || companyPatch.status === "archived")) {
await fenceActiveCompanyAgentStopsInTransaction(tx as unknown as Db, id, companyPatch.status);
}
if (logoAssetId === null) {
await tx.delete(companyLogos).where(eq(companyLogos.companyId, id));
@@ -422,9 +421,7 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
};
});
if (!result) return null;
await createAgentLifecycle(db).reconcilePolicyHolds(id);
publishAccountingActivities(id, budgetPublications);
if (data.budgetMonthlyCents !== undefined) await deliverBudgetEnforcement(db, budgetHooks, id);
// Post-commit, fire-and-forget, and BEFORE any finalization that
// could throw: a Cloud-pinned primary company that crossed the
// archived boundary (either direction) rings the harness so the
@@ -476,6 +473,7 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
.select({ status: companies.status })
.from(companies)
.where(eq(companies.id, id))
.for("no key update")
.then((rows) => rows[0] ?? null);
if (!existing) return null;
@@ -483,6 +481,7 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
if (!wasAlreadyArchived) {
await notifyDeliveryWork(tx, DELIVERY_QUEUES.skillSourceSync);
await notifyDeliveryWork(tx, DELIVERY_QUEUES.agentLifecycle, id);
await tx
.update(companies)
.set({ status: "archived", updatedAt: new Date() })
@@ -490,6 +489,7 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
}
const cascade = wasAlreadyArchived ? null : await applyArchiveCascadeInTx(tx, id);
if (!wasAlreadyArchived) await fenceActiveCompanyAgentStopsInTransaction(tx as unknown as Db, id, "archived");
const row = await getCompanyQuery(tx)
.where(eq(companies.id, id))
@@ -502,7 +502,6 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
};
});
if (!result) return null;
await createAgentLifecycle(db).reconcilePolicyHolds(id);
// Same doorbell rule as update(): the archive is committed, so ring
// before finalization, which can throw without undoing it.
+2 -2
View File
@@ -3,7 +3,7 @@ import { agentObservabilityService } from "./agent-observability.js";
import { withAccountHomeSecretMutationLock } from "@paperclipai/adapter-codex-local/server";
import type { Db } from "@paperclipai/db";
import { eq, inArray } from "drizzle-orm";
import { createAgentLifecycle, scheduleAgentLifecycle } from "./agent-lifecycle.js";
import { createAgentLifecycle, processAgentLifecycle } from "./agent-lifecycle.js";
import { agentLifecycleCompanyDeletion } from "../modules/agent-lifecycle/company-deletion.js";
import {
companies,
@@ -50,7 +50,7 @@ export async function deleteCompany(db: Db, id: string) {
const companyAgents = await db.select({ id: agents.id }).from(agents).where(eq(agents.companyId, id));
for (const agent of companyAgents) {
await lifecycle.terminateAgent(agent.id);
await scheduleAgentLifecycle(db, agent.id);
await processAgentLifecycle(db, agent.id);
}
return withAccountHomeSecretMutationLock(undefined, id, () => db.transaction(async (tx) => {
// Exclude accounting writers before taking child locks. KEY SHARE must
+7 -6
View File
@@ -5,7 +5,7 @@ import { alias } from "drizzle-orm/pg-core";
import type { Db } from "@paperclipai/db";
import { activityLog, agents, agentRuntimeState, authUsers, companies, companyMemberships, costEvents, heartbeatRuns, issues, projects, goals, aiSubscriptions } from "@paperclipai/db";
import { notFound, unprocessable, conflict } from "../errors.js";
import { budgetService, budgetServiceInTransaction, type BudgetServiceHooks } from "./budgets.js";
import { budgetServiceInTransaction, type BudgetServiceHooks } from "./budgets.js";
import { logActivity, type LogActivityInput, type ActivityPublication } from "./activity-log.js";
import { withAccountingReadSnapshot, withAccountingTransaction } from "./accounting-transaction.js";
import { visibleIssueCondition } from "./issue-visibility.js";
@@ -82,7 +82,8 @@ export async function getMonthlySpendTotal(
return Number(row?.total ?? 0);
}
export async function createCostEventInTransaction(db: Db, companyId: string, data: Omit<typeof costEvents.$inferInsert, "companyId" | "receiptHash" | "costCents"> & { costCents: MoneyInput }, publications: ActivityPublication[] = [], actor?: Pick<LogActivityInput, "actorType" | "actorId" | "agentId">) {
export async function createCostEventInTransaction(db: Db, companyId: string, data: Omit<typeof costEvents.$inferInsert, "companyId" | "receiptHash" | "costCents"> & { costCents: MoneyInput }, publications: ActivityPublication[] = [], actor?: Pick<LogActivityInput, "actorType" | "actorId" | "agentId">,
options: { deferBudgetEvaluation?: boolean } = {}) {
const parsed = ((data.usageKind === "decision" || data.usageKind === "fast_response" || data.usageKind === "prompt_injection") ? createServiceCostEventSchema : createCostEventSchema).safeParse({ ...data, occurredAt: data.occurredAt.toISOString() });
if (!parsed.success) throw unprocessable("Invalid cost receipt", parsed.error.flatten());
const values = {
@@ -152,7 +153,9 @@ export async function createCostEventInTransaction(db: Db, companyId: string, da
updatedAt: new Date(),
}).where(and(eq(agentRuntimeState.agentId, event.agentId), eq(agentRuntimeState.companyId, companyId)));
}
await budgetServiceInTransaction(db, publications).evaluateCostEvent(event);
// A multi-model run receipt inserts all parts before one scope evaluation.
// The owning accounting transaction must perform that deferred evaluation.
if (!options.deferBudgetEvaluation) await budgetServiceInTransaction(db, publications).evaluateCostEvent(event);
await logActivity(db, {
companyId, actorType: actor?.actorType ?? "system", actorId: actor?.actorId ?? "cost_accounting", agentId: actor?.agentId ?? event.agentId,
...((event.usageKind === "decision" || event.usageKind === "fast_response" || event.usageKind === "prompt_injection") ? { responsibleUserIdOverride: event.responsibleUserId } : {}),
@@ -162,13 +165,11 @@ export async function createCostEventInTransaction(db: Db, companyId: string, da
return { ...event, costCentsExact: values.costCents };
}
export function costService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
const budgets = budgetService(db, budgetHooks);
export function costService(db: Db, _budgetHooks: BudgetServiceHooks = {}) {
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;
},
+11 -15
View File
@@ -1,7 +1,8 @@
import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notifications.js";
import { randomUUID } from "node:crypto";
import { and, desc, eq, gte, lt, lte, or, sql } from "drizzle-orm";
import {
agents, authUsers, budgetPolicies, budgetReservations, companies, companyMemberships, companyDecisionModels, costEvents,
agents, authUsers, budgetReservations, companies, companyMemberships, companyDecisionModels, costEvents,
decisionInvocations, heartbeatRuns, issues, projects, companySecrets, type Db,
} from "@paperclipai/db";
import {
@@ -14,7 +15,8 @@ import { aiConnectionService } from "./ai-connections.js";
import { accessService } from "./access.js";
import { type AuthorizationActor } from "./authorization.js";
import { withAccountingTransaction } from "./accounting-transaction.js";
import { budgetPoliciesForRun, budgetService, budgetServiceInTransaction, computeObservedSpend, type BudgetServiceHooks } from "./budgets.js";
import { withBudgetSpan } from "./budget-tracing.js";
import { budgetServiceInTransaction, type BudgetServiceHooks } from "./budgets.js";
import { createCostEventInTransaction } from "./costs.js";
import { logActivity } from "./activity-log.js";
import { decisionReceipt, runDecisionProvider, type DecisionProviderOutcome } from "./decision-model-provider.js";
@@ -138,7 +140,7 @@ export function decisionModelService(db: Db, options: { provider?: typeof runDec
return (context: DecisionContext) => { let result = cache.get(context); if (!result) { result = availability(context); cache.set(context, result); } return result; };
}
async function settle(invocation: typeof decisionInvocations.$inferSelect, outcome: DecisionProviderOutcome) {
await withAccountingTransaction(db, invocation.companyId, async (tx, publications) => {
await withBudgetSpan("budget.accounting.settlement", () => withAccountingTransaction(db, invocation.companyId, async (tx, publications) => {
const [current] = await tx.select().from(decisionInvocations).where(and(eq(decisionInvocations.companyId, invocation.companyId), eq(decisionInvocations.id, invocation.id))).for("update");
if (!current || current.status !== "running") return;
const receipt = outcome.receipt;
@@ -161,6 +163,7 @@ export function decisionModelService(db: Db, options: { provider?: typeof runDec
providerRequestId: receipt.providerRequestId, pricingProvenance: receipt.pricingProvenance, occurredAt: invocation.startedAt,
}, publications, { actorType: invocation.actorType as "system" | "user" | "agent", actorId: invocation.actorId, agentId: links.agentId });
const finishedAt = new Date();
await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, invocation.companyId);
await tx.update(decisionInvocations).set({ status: outcome.answers ? "succeeded" : receipt.costStatus === "unpriced" ? "unknown" : "failed",
errorCode: outcome.errorCode ?? null, costEventId: event.id, providerRequestId: receipt.providerRequestId,
inputTokens: receipt.inputTokens, outputTokens: receipt.outputTokens, finishedAt, durationMs: Math.max(0, finishedAt.getTime() - invocation.startedAt.getTime()),
@@ -168,8 +171,7 @@ export function decisionModelService(db: Db, options: { provider?: typeof runDec
if (receipt.costStatus !== "unpriced") await tx.update(budgetReservations).set({ state: outcome.noProviderWork ? "released" : "settled", settledAt: finishedAt }).where(and(
eq(budgetReservations.companyId, invocation.companyId), eq(budgetReservations.decisionInvocationId, invocation.id), eq(budgetReservations.state, "held"),
));
});
await budgetService(db, options.budgetHooks).deliverPendingEnforcement(invocation.companyId);
}));
}
async function decide(context: DecisionContext, request: DecisionRequest, signal?: AbortSignal): Promise<DecisionResult> {
const parsed = decisionRequestSchema.safeParse(request);
@@ -202,23 +204,17 @@ export function decisionModelService(db: Db, options: { provider?: typeof runDec
if (!currentSecret || currentSecret.status !== "active" || currentSecret.version !== credentialSnapshot.version)
return { reason: "connection_unavailable" as const };
const budgets = budgetServiceInTransaction(tx, publications);
const block = await budgets.getInvocationBlock(context.companyId, current.agentId, { projectId: current.projectId });
const { block, policies, observations } = await budgets.evaluateAdmission(context.companyId, current.agentId, { projectId: current.projectId });
const [company] = await tx.select({ status: companies.status }).from(companies).where(eq(companies.id, context.companyId));
const [agent] = current.agentId ? await tx.select({ status: agents.status }).from(agents).where(and(eq(agents.id, current.agentId), eq(agents.companyId, context.companyId))) : [];
if (block || company.status !== "active" || (current.agentId && (!agent || agent.status === "paused" || agent.status === "terminated"))) return { reason: "budget_blocked" as const };
const policies = await tx.select().from(budgetPolicies).where(and(eq(budgetPolicies.companyId, context.companyId), eq(budgetPolicies.isActive, true),
eq(budgetPolicies.hardStopEnabled, true), eq(budgetPolicies.metric, "billed_cents"), sql`${budgetPolicies.amount} > 0`, budgetPoliciesForRun(context.companyId, current.agentId, current.projectId)));
const amount = policies.reduce((max, p) => { const n = centsToUnits(p.reservationCents); return n > max ? n : max; }, 0n);
for (const policy of policies) {
const [held] = await tx.select({ total: sql<string>`coalesce(sum(${budgetReservations.amountCents}),0)::text` }).from(budgetReservations).where(and(
eq(budgetReservations.companyId, context.companyId), eq(budgetReservations.state, "held"),
policy.scopeType === "agent" ? eq(budgetReservations.agentId, current.agentId!) : undefined,
policy.scopeType === "project" ? eq(budgetReservations.projectId, current.projectId!) : undefined,
));
const committed = centsToUnits((await computeObservedSpend(tx, policy)).totalExact) + centsToUnits(held.total);
for (const [index, policy] of policies.entries()) {
const committed = centsToUnits(observations[index].ledgerCents) + centsToUnits(observations[index].effectiveHeldCents);
if (committed >= centsToUnits(policy.amount) || committed + amount > centsToUnits(policy.amount)) return { reason: "budget_blocked" as const };
}
signal?.throwIfAborted();
await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, context.companyId);
const [invocation] = await tx.insert(decisionInvocations).values({ id: randomUUID(), companyId: context.companyId, feature: context.feature.id,
actorType: current.actorType, actorId: current.actorId, responsibleUserId: current.userId, agentId: current.agentId,
issueId: current.issueId, projectId: current.projectId, runId: current.runId, identityContextId: current.identityContextId,
@@ -29,6 +29,9 @@ export function createDeliveryWorkCoordinator(input: {
// An empty queue's maintenance may wait until dispatch admission reopens.
// Notifications and unresolved producer writes still reconcile in drain.
pausePassiveMaintenance?: boolean;
// Accounting and stop delivery finish already-admitted obligations;
// they must continue during admission drain, never warm standby.
runDuringDrain?: boolean;
}) {
if (stopped) throw new Error("Delivery coordinator is stopped");
if (workers.has(queue)) throw new Error(`Delivery worker already registered: ${queue}`);
@@ -64,7 +67,7 @@ export function createDeliveryWorkCoordinator(input: {
const attempt = new AbortController();
controller = attempt;
running = Promise.resolve().then(() => reconcileDatabaseWork(input.owner, queue))
.then(() => workerStopped || !input.canRun() ? undefined : task.run(attempt.signal)).then(() => workerStopped ? false : task.hasPending()).then(pending => {
.then(() => workerStopped || !(input.canRun() || (task.runDuringDrain && input.canReconcile?.())) ? undefined : task.run(attempt.signal)).then(() => workerStopped ? false : task.hasPending()).then(pending => {
pendingWork = pending || databaseWorkPending(input.owner, queue);
if (pendingWork) schedule(task.retryMs);
else { finishIntent?.(); finishIntent = undefined; }
@@ -2,6 +2,9 @@ import { signalDatabaseWork, subscribeDatabaseWork } from "@paperclipai/db";
/** Topics for existing durable queues; register intent before their writes. */
export const DELIVERY_QUEUES = {
budgetEnforcement: "budget-enforcement",
accounting: "cost-accounting",
agentLifecycle: "agent-lifecycle",
chatBoardMessages: "chat-board-messages",
chatTaskStarts: "chat-task-starts",
chatReceiptReactions: "chat-receipt-reactions",
+9 -59
View File
@@ -20,7 +20,6 @@ import {
issueAttachments,
projects,
authUsers,
budgetPolicies,
budgetReservations,
companies,
companyFastResponses,
@@ -60,11 +59,9 @@ import {
type ActivityPublication,
} from "./activity-log.js";
import { withAccountingTransaction } from "./accounting-transaction.js";
import { withBudgetSpan } from "./budget-tracing.js";
import {
budgetPoliciesForRun,
budgetService,
budgetServiceInTransaction,
computeObservedSpend,
type BudgetServiceHooks,
} from "./budgets.js";
import { createCostEventInTransaction } from "./costs.js";
@@ -567,7 +564,7 @@ export function fastResponseService(
);
}
async function settle(request: Request, outcome: FastResponseOutcome) {
await withAccountingTransaction(
await withBudgetSpan("budget.accounting.settlement", () => withAccountingTransaction(
db,
request.companyId,
async (tx, publications) => {
@@ -660,10 +657,8 @@ export function fastResponseService(
),
);
},
);
await budgetService(db, options.budgetHooks).deliverPendingEnforcement(
request.companyId,
);
));
}
async function process(
request: Request,
@@ -853,60 +848,15 @@ export function fastResponseService(
(!agent || ["paused", "terminated"].includes(agent.status)))
)
return null;
if (
await budgetServiceInTransaction(tx, []).getInvocationBlock(
request.companyId,
request.agentId,
{ projectId: current.issue?.projectId },
)
)
return null;
const policies = await tx
.select()
.from(budgetPolicies)
.where(
and(
eq(budgetPolicies.companyId, request.companyId),
eq(budgetPolicies.isActive, true),
eq(budgetPolicies.hardStopEnabled, true),
eq(budgetPolicies.metric, "billed_cents"),
sql`${budgetPolicies.amount} > 0`,
budgetPoliciesForRun(
request.companyId,
request.agentId,
current.issue?.projectId ?? null,
),
),
);
const { block, policies, observations } = await budgetServiceInTransaction(tx, []).evaluateAdmission(
request.companyId, request.agentId, { projectId: current.issue?.projectId });
if (block) return null;
const amount = policies.reduce((max, p) => {
const n = centsToUnits(p.reservationCents);
return n > max ? n : max;
}, 0n);
for (const policy of policies) {
const [held] = await tx
.select({
total: sql<string>`coalesce(sum(${budgetReservations.amountCents}),0)::text`,
})
.from(budgetReservations)
.where(
and(
eq(budgetReservations.companyId, request.companyId),
eq(budgetReservations.state, "held"),
policy.scopeType === "agent"
? eq(budgetReservations.agentId, request.agentId!)
: undefined,
policy.scopeType === "project"
? eq(
budgetReservations.projectId,
current.issue!.projectId!,
)
: undefined,
),
);
const committed =
centsToUnits(
(await computeObservedSpend(tx, policy)).totalExact,
) + centsToUnits(held.total);
for (const [index, policy] of policies.entries()) {
const committed = centsToUnits(observations[index].ledgerCents) + centsToUnits(observations[index].effectiveHeldCents);
if (
committed >= centsToUnits(policy.amount) ||
committed + amount > centsToUnits(policy.amount)
+26 -27
View File
@@ -330,7 +330,7 @@ import {
projects,
projectWorkspaces,
} from "@paperclipai/db";
import { HttpError, notFound } from "../errors.js";
import { conflict, HttpError, notFound } from "../errors.js";
import { logger } from "../middleware/logger.js";
import { publishLiveEvent } from "./live-events.js";
@@ -929,6 +929,14 @@ export function heartbeatService(
const budgetHooks = {
cancelWorkForScope: cancelBudgetScopeWork,
};
const reconcileCostAccounting = createCostAccountingReconciler(db, budgetHooks, async () => {
const outcomes = await Promise.allSettled([
decisionModelService(db, { budgetHooks }).recoverInterrupted(),
promptInjectionPreventionService(db, { budgetHooks }).recoverInterrupted(),
]);
const failed = outcomes.find(outcome => outcome.status === "rejected");
if (failed?.status === "rejected") throw failed.reason;
});
const budgets = budgetService(db, budgetHooks);
const recovery = recoveryService(db, {
enqueueWakeup,
@@ -2740,12 +2748,13 @@ export function heartbeatService(
}
const dispatchResolvedInteractionContinuationAfterAdmission = async <T>(
dispatch: (markDispatchStarted: () => void) => Promise<T>,
reservation: { id: string; reused: boolean },
): Promise<
{ dispatched: true; resultPromise: Promise<T> } | { dispatched: false }
> => {
await controllerLease.assertOwned("dispatching");
// Recheck after workspace/credential preparation, immediately before the
// provider handoff. Never hold validation locks while adapter code runs.
// Recheck after preparation. The final locked gate commits a dispatch
// claim before entering adapter bootstrap or contacting a provider.
await authorizeFailedChatRetryExecution();
if (
!(await withChatControlRecoveryGate(
@@ -2771,33 +2780,27 @@ export function heartbeatService(
}
return { dispatched: false };
}
if (
!issueId ||
(!isResolvedInteractionContinuationWakeContext(context) &&
typeof context.workspaceReconstructionActionId !== "string" &&
!["native_safe_replacement", "native_provider_overloaded"].includes(run.scheduledRetryReason ?? ""))
) {
return { dispatched: true, resultPromise: dispatch(() => {}) };
const enforceContinuationOwnership = Boolean(issueId && (isResolvedInteractionContinuationWakeContext(context) ||
typeof context.workspaceReconstructionActionId === "string" ||
["native_safe_replacement", "native_provider_overloaded"].includes(run.scheduledRetryReason ?? "")));
if (issueId && enforceContinuationOwnership) {
await options.beforeResolvedInteractionContinuationDispatchCheck?.({ runId: run.id, issueId });
await options.afterResolvedInteractionContinuationDispatchCheck?.({ runId: run.id, issueId });
}
await options.beforeResolvedInteractionContinuationDispatchCheck?.({
runId: run.id,
issueId,
});
await options.afterResolvedInteractionContinuationDispatchCheck?.({
runId: run.id,
issueId,
});
const gate = await runDispatch.dispatchResolvedInteractionIfCurrent({
runId: run.id,
companyId: run.companyId,
expectedStatus: "running",
// Synchronous handoff under the ownership lock; the gate commits
// without awaiting the adapter's asynchronous bootstrap or finalizer.
reservationId: reservation.id,
enforceContinuationOwnership,
...(runOptions.nativeLeaseOwner ? { recoveryLeaseOwner: runOptions.nativeLeaseOwner } : {}),
// The gate persists admission under the ownership lock, then invokes
// this callback after commit so provider I/O holds no accounting lock.
dispatch,
});
if (gate.dispatched) return gate;
if ("admissionBlockReason" in gate) throw conflict(gate.admissionBlockReason);
if (gate.cancellation.outcome === "cancelled") {
applyRunDispatchPostCommitEffects(
gate.cancellation.postCommitEffects,
@@ -2818,7 +2821,7 @@ export function heartbeatService(
return await dispatchResolvedInteractionContinuationAfterAdmission((markDispatchStarted) => {
entered = true;
return dispatch(markDispatchStarted);
});
}, reservation);
} finally {
if (!entered && !reservation.reused) {
await db.update(heartbeatRuns).set({ costAccountingPending: true,
@@ -4092,11 +4095,7 @@ export function heartbeatService(
scanSilentActiveRuns,
reconcileTaskWatchdogs,
reconcileCostAccounting: async () => {
await decisionModelService(db, { budgetHooks }).recoverInterrupted();
await promptInjectionPreventionService(db, { budgetHooks }).recoverInterrupted();
return createCostAccountingReconciler(db, budgetHooks)();
},
reconcileCostAccounting,
buildRunOutputSilence,
@@ -0,0 +1,98 @@
import type { Db } from "@paperclipai/db";
import { agents, companies, createDb, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS, heartbeatRuns } from "@paperclipai/db";
import { sql } from "drizzle-orm";
import { afterAll, afterEach, beforeAll, beforeEach, vi, type Mock } from "vitest";
import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../../../__tests__/helpers/embedded-postgres.js";
import { runningProcesses } from "../../../adapters/index.js";
import { terminateLocalService } from "../../local-service-supervisor.js";
import { createHeartbeatRunState } from "../run-state.js";
import { createHeartbeatRunPreparation } from "../run-preparation.js";
import { createHeartbeatLifecycle, type HeartbeatLifecycleDependencies } from "../run-lifecycle.js";
import type { HeartbeatRunControlDependencies } from "../run-control.js";
type TestCallbacks = HeartbeatRunControlDependencies & ReturnType<typeof createHeartbeatLifecycle> & {
getRun: Mock<HeartbeatRunControlDependencies["getRun"]>;
getAgent: Mock<HeartbeatRunControlDependencies["getAgent"]>;
setRunStatusFromLive: Mock<HeartbeatRunControlDependencies["setRunStatusFromLive"]>;
getSchedulingSuppression: Mock<HeartbeatRunControlDependencies["getSchedulingSuppression"]>;
options: { closeWarmNativeSessionsForRun: Mock<NonNullable<HeartbeatRunControlDependencies["options"]["closeWarmNativeSessionsForRun"]>> };
};
/** Shared real lifecycle callbacks; each focused suite controls its provider mocks. */
export function callbacks(db: Db): TestCallbacks {
const state = createHeartbeatRunState(db);
const lifecycleDependencies = {
...state,
...createHeartbeatRunPreparation(db),
issuesSvc: {
addComment: vi.fn<HeartbeatLifecycleDependencies["issuesSvc"]["addComment"]>(),
listReviewAttention: vi.fn<HeartbeatLifecycleDependencies["issuesSvc"]["listReviewAttention"]>(),
},
recovery: { escalateStrandedAssignedIssue: vi.fn<HeartbeatLifecycleDependencies["recovery"]["escalateStrandedAssignedIssue"]>() },
companySkills: { completeTestRunForIssue: vi.fn<HeartbeatLifecycleDependencies["companySkills"]["completeTestRunForIssue"]>() },
budgets: { getInvocationBlock: vi.fn(async () => null) },
treeControlSvc: { getActivePauseHoldGate: vi.fn(async () => null) },
enqueueWakeup: vi.fn<HeartbeatLifecycleDependencies["enqueueWakeup"]>(),
getCurrentUserRedactionOptions: vi.fn(async () => ({})),
getAgentInvokability: vi.fn<HeartbeatLifecycleDependencies["getAgentInvokability"]>(async () => ({ invokable: true })),
emitTerminalAgentTaskRun: vi.fn<HeartbeatLifecycleDependencies["emitTerminalAgentTaskRun"]>(),
budgetHooks: {},
} satisfies HeartbeatLifecycleDependencies;
const lifecycle = createHeartbeatLifecycle(db, lifecycleDependencies);
return {
...state,
...lifecycle,
getRun: vi.fn(state.getRun),
getAgent: vi.fn(state.getAgent),
setRunStatusFromLive: vi.fn(lifecycle.setRunStatusFromLive),
setRunStatus: vi.fn(lifecycle.setRunStatus),
setWakeupStatus: vi.fn(lifecycle.setWakeupStatus),
appendRunEvent: vi.fn(lifecycle.appendRunEvent),
finalizeAgentStatus: vi.fn<HeartbeatRunControlDependencies["finalizeAgentStatus"]>(),
options: { closeWarmNativeSessionsForRun: vi.fn(async () => ({ closed: 1, busy: 0, failed: 0 })) },
envOrchestrator: { releaseForRun: vi.fn<HeartbeatRunControlDependencies["envOrchestrator"]["releaseForRun"]>(async () => ({ released: [], errors: [] })) },
enqueueWakeup: lifecycleDependencies.enqueueWakeup,
startNextQueuedRunForAgent: vi.fn<HeartbeatRunControlDependencies["startNextQueuedRunForAgent"]>(),
sweepPendingCleanupLeases: vi.fn<HeartbeatRunControlDependencies["sweepPendingCleanupLeases"]>(),
timerClaimWasFirstHeartbeat: vi.fn<HeartbeatRunControlDependencies["timerClaimWasFirstHeartbeat"]>(() => undefined),
wakeQueue: { releaseIssueExecution: vi.fn<HeartbeatRunControlDependencies["wakeQueue"]["releaseIssueExecution"]>(async () => ({ outcome: { kind: "released" }, postCommitEffects: [] })) },
applyWakeQueuePostCommitEffects: vi.fn(async () => undefined),
getSchedulingSuppression: vi.fn(async () => ({ suppressed: false })),
activeRunExecutions: new Set<string>(),
processRunCancellationSettlements: new Map<string, { settled: Promise<void>; failed: boolean; error?: unknown }>(),
failedProcessRunCancellations: new Map<string, { settled: Promise<void>; failed: boolean; error?: unknown }>(),
} satisfies HeartbeatRunControlDependencies;
}
export const runControlPostgresSupport = await getEmbeddedPostgresTestSupport();
if (!runControlPostgresSupport.supported) console.warn(`Skipping run-control database tests: ${runControlPostgresSupport.reason}`);
export function setupRunControlDatabase(name: string, ready: (db: Db) => void) {
let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>>;
let db: Db;
beforeAll(async () => {
database = await startEmbeddedPostgresTestDatabase(name);
db = createDb(database.connectionString);
ready(db);
}, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS);
afterAll(async () => {
await db?.$client.end();
await database?.cleanup();
});
beforeEach(async () => {
vi.mocked(terminateLocalService).mockReset().mockResolvedValue(undefined);
await db.execute(sql`truncate table companies restart identity cascade`);
});
afterEach(() => runningProcesses.clear());
async function fixture() {
const [company] = await db.insert(companies).values({ name: "Run control", issuePrefix: "RC" }).returning();
const [agent] = await db.insert(agents).values({ companyId: company.id, name: "Controlled agent", adapterType: "codex_local" }).returning();
return { company, agent };
}
async function liveRun(companyId: string, agentId: string) {
const [run] = await db.insert(heartbeatRuns).values({ companyId, agentId, status: "running", runtimeMode: "legacy", runtimeModeResolvedAt: new Date() }).returning();
runningProcesses.set(run.id, { child: { pid: 424242 } as never, graceSec: 4, processGroupId: 424242 });
return run;
}
return { fixture, liveRun };
}
@@ -0,0 +1,469 @@
import { randomUUID } from "node:crypto";
import { spawn } from "node:child_process";
import { once } from "node:events";
import type { Db } from "@paperclipai/db";
import { agents, agentWakeupRequests, browserUseRuns, browserUseSessions, budgetPolicies, companies, connectionGrants, costEvents, environmentLeases, heartbeatRuns, issues, projects, toolApplications, toolConnections } from "@paperclipai/db";
import { eq, inArray } from "drizzle-orm";
import { describe, expect, it, vi } from "vitest";
import { runningProcesses } from "../../adapters/index.js";
import { terminateLocalService } from "../local-service-supervisor.js";
import { persistUsageReceipt } from "../usage-receipts.js";
import { registerBudgetEnforcementWork } from "../budget-enforcement-worker.js";
import { createDeliveryWorkCoordinator } from "../delivery-work-coordinator.js";
import { DELIVERY_QUEUES, subscribeDeliveryWork } from "../delivery-work-notifications.js";
import { PROCESS_START_REQUESTED } from "../native-local-process-stop.js";
import { createPostgresRunDispatchAdapter } from "../../modules/run-dispatch/adapters/postgres.js";
import { createHeartbeatRunControl } from "./run-control.js";
import { listPendingBudgetCancellationScopes } from "./budget-cancellation.js";
import { callbacks, runControlPostgresSupport, setupRunControlDatabase } from "./__tests__/run-control-fixtures.js";
vi.mock("../live-events.js", () => ({ publishLiveEvent: vi.fn() }));
vi.mock("../local-service-supervisor.js", async (importOriginal) => ({
...await importOriginal<typeof import("../local-service-supervisor.js")>(),
terminateLocalService: vi.fn(async () => undefined),
}));
describe.skipIf(!runControlPostgresSupport.supported)("scoped budget cancellation", () => {
let db: Db;
const { fixture, liveRun } = setupRunControlDatabase("paperclip-heartbeat-budget-cancellation", value => { db = value; });
it("rechecks stale budget enforcement before cancelling runs or queued wakes", async () => {
const { company, agent } = await fixture();
const [run] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id, status: "queued", runtimeMode: "legacy" }).returning();
const [wake] = await db.insert(agentWakeupRequests).values({ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued" }).returning();
const deps = callbacks(db);
await createHeartbeatRunControl(db, deps).cancelBudgetScopeWork({ companyId: company.id, scopeType: "agent", scopeId: agent.id, enforcement: { policyId: randomUUID(), version: 1 } });
expect((await deps.getRun(run.id))?.status).toBe("queued");
expect((await db.select().from(agentWakeupRequests).where(eq(agentWakeupRequests.id, wake.id)))[0].status).toBe("queued");
expect(deps.setRunStatusFromLive).not.toHaveBeenCalled();
expect(deps.finalizeAgentStatus).not.toHaveBeenCalled();
expect(terminateLocalService).not.toHaveBeenCalled();
});
it("does not apply an earlier enforcement generation after a budget increase", async () => {
const { company, agent } = await fixture();
const [policy] = await db.insert(budgetPolicies).values({ companyId: company.id,
scopeType: "agent", scopeId: agent.id, windowKind: "lifetime", amount: 10_000,
enforcementVersion: 2, enforcementDeliveredVersion: 1 }).returning();
const run = await liveRun(company.id, agent.id);
const [wake] = await db.insert(agentWakeupRequests).values({ companyId: company.id, agentId: agent.id,
source: "assignment", status: "queued" }).returning();
await createHeartbeatRunControl(db, callbacks(db)).cancelBudgetScopeWork({ companyId: company.id,
scopeType: "agent", scopeId: agent.id, enforcement: { policyId: policy.id, version: 1 } });
expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, run.id)))[0]).toMatchObject({ status: "running", resultJson: null });
expect((await db.select().from(agentWakeupRequests).where(eq(agentWakeupRequests.id, wake.id)))[0].status).toBe("queued");
expect(terminateLocalService).not.toHaveBeenCalled();
});
it("fences every affected run and queued wake before bounded provider delivery", async () => {
const { company, agent } = await fixture();
const runs = await db.insert(heartbeatRuns).values(Array.from({ length: 5 }, () => ({
companyId: company.id, agentId: agent.id, status: "running", runtimeMode: "legacy", runtimeModeResolvedAt: new Date(),
}))).returning();
const [unrelatedAgent] = await db.insert(agents).values({ companyId: company.id, name: "Unrelated" }).returning();
const unrelated = await db.insert(heartbeatRuns).values({ companyId: company.id,
agentId: unrelatedAgent.id, status: "running", runtimeMode: "legacy" }).returning().then(rows => rows[0]);
const wakes = await db.insert(agentWakeupRequests).values([
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued" },
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "deferred_issue_execution" },
{ companyId: company.id, agentId: unrelatedAgent.id, source: "assignment", status: "queued" },
]).returning();
runs.forEach((run, i) => runningProcesses.set(run.id, {
child: { pid: 425000 + i } as never, graceSec: 4, processGroupId: 425000 + i,
}));
const releases = new Map<number, () => void>();
let active = 0, maximumActive = 0, releaseAll = false;
vi.mocked(terminateLocalService).mockImplementation(async input => {
active += 1;
maximumActive = Math.max(maximumActive, active);
if (!releaseAll) await new Promise<void>(resolve => releases.set(input.pid, resolve));
active -= 1;
});
const deps = callbacks(db);
const delivery = createHeartbeatRunControl(db, deps).cancelBudgetScopeWork({ companyId: company.id,
scopeType: "agent", scopeId: agent.id });
try {
await vi.waitFor(() => expect(releases.size).toBe(4));
const persisted = await db.select().from(heartbeatRuns).where(inArray(heartbeatRuns.id, runs.map(run => run.id)));
expect(persisted.every(run => run.status === "running" && Boolean(run.resultJson?.startupCancellation))).toBe(true);
const pendingWakes = await db.select().from(agentWakeupRequests).where(inArray(agentWakeupRequests.id, wakes.map(wake => wake.id)));
expect(wakes.map(wake => pendingWakes.find(row => row.id === wake.id)?.status)).toEqual(["cancelled", "cancelled", "queued"]);
expect((await deps.getRun(unrelated.id))?.resultJson).toBeNull();
// A slow first Stop must not stop another slot from making progress.
const firstPid = releases.keys().next().value!;
const otherPid = [...releases.keys()].find(pid => pid !== firstPid)!;
releases.get(otherPid)!();
await vi.waitFor(() => expect(releases.size).toBe(5));
expect(maximumActive).toBe(4);
expect(active).toBe(4);
} finally {
releaseAll = true;
releases.forEach(release => release());
await delivery;
}
expect((await db.select().from(heartbeatRuns).where(inArray(heartbeatRuns.id, runs.map(run => run.id))))
.every(run => run.status === "cancelled" && (run.resultJson?.executionCancellation as { state?: string })?.state === "acknowledged")).toBe(true);
expect((await deps.getRun(unrelated.id))?.status).toBe("running");
});
it("wakes scoped Stop delivery from a partial usage overrun without waiting for adapter termination", async () => {
const { company, agent } = await fixture();
const [project, otherProject] = await db.insert(projects).values([
{ companyId: company.id, name: "Over budget" }, { companyId: company.id, name: "Still admitted" },
]).returning();
const sourceId = randomUUID();
const run = await liveRun(company.id, agent.id);
await db.update(heartbeatRuns).set({ contextSnapshot: { projectId: project.id }, costAccountingPending: true,
usageJson: { accountingReceiptSourceId: sourceId, accountingReceiptSequence: 0,
accountingReceiptReady: false, ledgerScope: { projectId: project.id } },
}).where(eq(heartbeatRuns.id, run.id));
const [queued, unrelated] = await db.insert(heartbeatRuns).values([
{ companyId: company.id, agentId: agent.id, status: "queued", runtimeMode: "legacy", contextSnapshot: { projectId: project.id } },
{ companyId: company.id, agentId: agent.id, status: "running", runtimeMode: "legacy", contextSnapshot: { projectId: otherProject.id } },
]).returning();
const wakes = await db.insert(agentWakeupRequests).values([
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued", payload: { projectId: project.id } },
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "deferred_issue_execution", payload: { projectId: project.id } },
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued", payload: { projectId: otherProject.id } },
]).returning();
const [policy] = await db.insert(budgetPolicies).values({ companyId: company.id, scopeType: "project",
scopeId: project.id, windowKind: "lifetime", amount: 100, notifyEnabled: false }).returning();
let releaseTermination!: () => void;
const termination = new Promise<void>(resolve => { releaseTermination = resolve; });
vi.mocked(terminateLocalService).mockImplementation(async () => termination);
const deps = callbacks(db);
const control = createHeartbeatRunControl(db, deps);
const errors = vi.fn();
const coordinator = createDeliveryWorkCoordinator({ owner: db, canRun: () => true, onError: errors });
const worker = registerBudgetEnforcementWork(db, coordinator, { cancelWorkForScope: control.cancelBudgetScopeWork });
try {
// Drain startup recovery first: only the receipt's committed notification
// may wake the idle worker below. Do not call wake or deliver explicitly.
await worker.ready;
expect(terminateLocalService).not.toHaveBeenCalled();
expect(coordinator.nextWakeAt()).toBeNull();
const startedAt = performance.now();
await persistUsageReceipt(db, { schema: "paperclip/accounting-receipt/v1", id: randomUUID(),
companyId: company.id, runId: run.id, sourceId, sequence: 1, receivedAt: new Date().toISOString(),
adapterType: agent.adapterType, receipt: { complete: false, usageBasis: "per_run", costUsd: 1.1,
provider: "test", biller: "test", model: "budget-stop-fixture", billingType: "metered_api", costStatus: "unpriced" },
});
const receiptReturnedAt = performance.now();
await vi.waitFor(async () => {
const persisted = await deps.getRun(run.id);
expect(persisted).toMatchObject({ status: "running", usageJson: { budgetExposure: { sequence: 1, costCents: "110.0000000" } },
resultJson: { startupCancellation: { budgetEnforcement: { policyId: policy.id, scopeType: "project", scopeId: project.id } } } });
expect(terminateLocalService).toHaveBeenCalledTimes(1);
}, { interval: 10, timeout: 3_000 });
const fenceObservedAt = performance.now();
const fenced = await deps.getRun(run.id);
expect(fenced?.resultJson?.executionCancellation).toBeUndefined();
expect((await deps.getRun(queued.id))?.resultJson?.startupCancellation).toBeTruthy();
const persistedWakes = await db.select().from(agentWakeupRequests).where(inArray(agentWakeupRequests.id, wakes.map(wake => wake.id)));
expect(wakes.map(wake => persistedWakes.find(row => row.id === wake.id)?.status)).toEqual(["cancelled", "cancelled", "queued"]);
const [pending] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.id));
expect(pending.enforcementVersion).toBeGreaterThan(pending.enforcementDeliveredVersion);
const requestStartedAt = performance.now();
let responseTimeout: ReturnType<typeof setTimeout> | undefined;
try {
const rows = await Promise.race([
db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, unrelated.id)),
new Promise<never>((_, reject) => { responseTimeout = setTimeout(() => reject(new Error("Unrelated database request blocked by Stop delivery")), 1_000); }),
]);
expect(rows[0]).toMatchObject({ status: "running", resultJson: null });
} finally { clearTimeout(responseTimeout); }
const unrelatedRequestMs = performance.now() - requestStartedAt;
const terminationReleasedAt = performance.now();
releaseTermination();
await vi.waitFor(async () => {
expect(await deps.getRun(run.id)).toMatchObject({ status: "cancelled", resultJson: { executionCancellation: { state: "acknowledged" } } });
const [delivered] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.id));
expect(delivered.enforcementDeliveredVersion).toBe(delivered.enforcementVersion);
}, { interval: 10, timeout: 3_000 });
expect((await deps.getRun(unrelated.id))?.status).toBe("running");
expect(errors).not.toHaveBeenCalled();
console.info("Partial usage budget Stop timing", { receiptMs: Math.round(receiptReturnedAt - startedAt),
receiptReturnToFenceMs: Math.round(fenceObservedAt - receiptReturnedAt), unrelatedRequestMs: Math.round(unrelatedRequestMs),
controlledTerminationWaitMs: Math.round(terminationReleasedAt - fenceObservedAt),
terminationReleaseToAcknowledgementMs: Math.round(performance.now() - terminationReleasedAt) });
} finally {
releaseTermination();
await coordinator.stop();
}
});
it("drains other provider stops after a failed delivery and retries the unfinished run", async () => {
const { company, agent } = await fixture();
const failed = await liveRun(company.id, agent.id);
const [other] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id,
status: "running", runtimeMode: "legacy", runtimeModeResolvedAt: new Date() }).returning();
runningProcesses.set(other.id, { child: { pid: 434343 } as never, graceSec: 4, processGroupId: 434343 });
vi.mocked(terminateLocalService).mockImplementation(async input => {
if (input.pid === 424242) throw new Error("provider unreachable");
});
const deps = callbacks(db);
const control = createHeartbeatRunControl(db, deps);
const scope = { companyId: company.id, scopeType: "agent" as const, scopeId: agent.id };
await expect(control.cancelBudgetScopeWork(scope)).rejects.toThrow("Budget cancellation delivery remains pending");
expect((await deps.getRun(other.id))?.status).toBe("cancelled");
expect(await deps.getRun(failed.id)).toMatchObject({ status: "running", resultJson: { startupCancellation: expect.any(Object) } });
expect(runningProcesses.has(failed.id)).toBe(true);
vi.mocked(terminateLocalService).mockResolvedValue(undefined);
await control.cancelBudgetScopeWork(scope);
expect(await deps.getRun(failed.id)).toMatchObject({ status: "cancelled", resultJson: { executionCancellation: { state: "acknowledged" } } });
});
it("shares four Stop slots across independent scopes and service instances", async () => {
const { company, agent } = await fixture();
const [otherAgent] = await db.insert(agents).values({ companyId: company.id, name: "Another scope" }).returning();
const runs = await db.insert(heartbeatRuns).values([agent.id, otherAgent.id].flatMap(agentId =>
Array.from({ length: 3 }, () => ({ companyId: company.id, agentId,
status: "running", runtimeMode: "legacy", runtimeModeResolvedAt: new Date() })))).returning();
runs.forEach((run, i) => runningProcesses.set(run.id, {
child: { pid: 435000 + i } as never, graceSec: 4, processGroupId: 435000 + i,
}));
let active = 0, maximumActive = 0, releaseAll = false;
const releases = new Map<number, () => void>();
vi.mocked(terminateLocalService).mockImplementation(async input => {
active += 1;
maximumActive = Math.max(maximumActive, active);
if (!releaseAll) await new Promise<void>(resolve => releases.set(input.pid, resolve));
active -= 1;
});
const deliveries = [agent, otherAgent].map(selected => createHeartbeatRunControl(db, callbacks(db))
.cancelBudgetScopeWork({ companyId: company.id, scopeType: "agent", scopeId: selected.id }));
try {
await vi.waitFor(() => expect(releases.size).toBe(4));
expect(maximumActive).toBe(4);
expect((await db.select().from(heartbeatRuns).where(inArray(heartbeatRuns.id, runs.map(run => run.id))))
.every(run => Boolean(run.resultJson?.startupCancellation))).toBe(true);
releases.values().next().value!();
await vi.waitFor(() => expect(releases.size).toBe(5));
expect(maximumActive).toBe(4);
} finally {
releaseAll = true;
releases.forEach(release => release());
await Promise.all(deliveries);
}
});
it("blocks retry promotion after the budget batch commits but before Stop delivery", async () => {
const { company, agent } = await fixture();
const [run] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id,
status: "scheduled_retry", runtimeMode: "legacy", runtimeModeResolvedAt: new Date(),
scheduledRetryAt: new Date(), scheduledRetryReason: "transient_failure" }).returning();
const deps = callbacks(db);
const getAgent = deps.getAgent.getMockImplementation()!;
let entered!: () => void, release!: () => void;
const deliveryEntered = new Promise<void>(resolve => { entered = resolve; });
const continueDelivery = new Promise<void>(resolve => { release = resolve; });
deps.getAgent.mockImplementationOnce(async id => { entered(); await continueDelivery; return getAgent(id); });
const cancellation = createHeartbeatRunControl(db, deps).cancelBudgetScopeWork({ companyId: company.id,
scopeType: "agent", scopeId: agent.id });
try {
await deliveryEntered;
expect(await createPostgresRunDispatchAdapter(db).promoteOrCancelDueRetry({
companyId: company.id, runId: run.id, now: new Date(),
})).toMatchObject({ outcome: "not_promoted" });
expect(await deps.getRun(run.id)).toMatchObject({ status: "scheduled_retry",
resultJson: { startupCancellation: expect.any(Object) } });
} finally {
release();
await cancellation;
}
expect(await deps.getRun(run.id)).toMatchObject({ status: "cancelled" });
});
it("recovers a committed Stop after its policy was deleted without fencing newly admitted work", async () => {
const { company, agent } = await fixture();
const oldRun = await liveRun(company.id, agent.id);
const [newRun] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id,
status: "queued", runtimeMode: "legacy" }).returning();
const intent = { companyId: company.id, scopeType: "agent", scopeId: agent.id,
policyId: randomUUID(), version: 1 };
await db.update(heartbeatRuns).set({ resultJson: { startupCancellation: {
requestedAt: new Date().toISOString(), budgetEnforcement: intent,
} } }).where(eq(heartbeatRuns.id, oldRun.id));
const scopes = await listPendingBudgetCancellationScopes(db);
expect(scopes).toHaveLength(1);
const deps = callbacks(db);
await createHeartbeatRunControl(db, deps).cancelBudgetScopeWork(scopes[0]);
expect(await deps.getRun(oldRun.id)).toMatchObject({ status: "cancelled",
resultJson: { startupCancellation: { budgetEnforcement: { deliveredAt: expect.any(String) } } } });
expect(await deps.getRun(newRun.id)).toMatchObject({ status: "queued", resultJson: null });
expect(await listPendingBudgetCancellationScopes(db)).toEqual([]);
});
it("retains a terminal budget Stop until actual stop acknowledgement exists", async () => {
const { company, agent } = await fixture();
const [run] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id,
status: "cancelled", runtimeMode: "legacy", startedAt: new Date(), resultJson: {
startupCancellation: { budgetEnforcement: { companyId: company.id, scopeType: "agent", scopeId: agent.id,
policyId: randomUUID(), version: 1 } },
} }).returning();
const control = createHeartbeatRunControl(db, callbacks(db));
const [scope] = await listPendingBudgetCancellationScopes(db);
await expect(control.cancelBudgetScopeWork(scope)).rejects.toThrow("Budget cancellation delivery remains pending");
expect(await listPendingBudgetCancellationScopes(db)).toHaveLength(1);
expect(terminateLocalService).not.toHaveBeenCalled();
await db.update(heartbeatRuns).set({ resultJson: { ...run.resultJson,
executionCancellation: { state: "acknowledged", acknowledgedAt: new Date().toISOString() },
} }).where(eq(heartbeatRuns.id, run.id));
await control.cancelBudgetScopeWork(scope);
expect(await listPendingBudgetCancellationScopes(db)).toEqual([]);
});
it.each(["local exit", "remote receipt", "unproven", "later launch"] as const)("reconciles normal completion before queued budget Stop using %s evidence", async evidence => {
const { company, agent } = await fixture();
const [run] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id,
status: "running", runtimeMode: "native", startedAt: new Date() }).returning();
const deps = callbacks(db);
if (evidence === "local exit" || evidence === "later launch") {
const child = spawn(process.execPath, ["-e", "setTimeout(() => process.exit(0), 50)"], { detached: true, stdio: "ignore" });
const exited = once(child, "exit");
try {
expect(child.pid).toBeGreaterThan(1);
await deps.persistRunProcessMetadata(run.id, { pid: child.pid!, processGroupId: child.pid!,
startedAt: new Date().toISOString(), targetKind: "local" });
await exited;
} finally { if (child.exitCode === null) child.kill("SIGKILL"); }
if (evidence === "later launch") await deps.appendRunEvent(run, {
eventType: PROCESS_START_REQUESTED, stream: "system", level: "info", message: "A later launch invalidated the earlier exit proof",
});
}
if (evidence === "remote receipt") {
const leaseId = randomUUID(), providerLeaseId = randomUUID();
await db.insert(environmentLeases).values({ id: leaseId, companyId: company.id, heartbeatRunId: run.id,
provider: "daytona", providerLeaseId, status: "released", releasedAt: new Date(), cleanupStatus: "success",
metadata: { remoteExecutionTermination: { schema: "paperclip.remote-termination.v1", companyId: company.id,
runId: run.id, leaseId, provider: "daytona", providerLeaseId, state: "stopped", confirmedAt: new Date().toISOString() } },
});
}
const getRun = deps.getRun.getMockImplementation()!;
deps.getRun.mockImplementationOnce(async id => {
// Budget delivery persisted its fence, then ordinary native completion
// won the race before this run acquired a provider Stop delivery slot.
expect((await getRun(id))?.resultJson?.startupCancellation).toBeTruthy();
await db.update(heartbeatRuns).set({ status: "succeeded", finishedAt: new Date(), nativePhase: "committed",
processPid: null, processGroupId: null, processStartedAt: null }).where(eq(heartbeatRuns.id, id));
return getRun(id);
});
const policyId = randomUUID();
await db.insert(budgetPolicies).values({ id: policyId, companyId: company.id, scopeType: "agent",
scopeId: agent.id, windowKind: "lifetime", amount: 1, enforcementVersion: 1 });
await db.insert(costEvents).values({ companyId: company.id, agentId: agent.id, provider: "test",
model: "overrun", inputTokens: 1, outputTokens: 1, costCents: 2, occurredAt: new Date() });
const scope = { companyId: company.id, scopeType: "agent" as const, scopeId: agent.id,
enforcement: { policyId, version: 1 } };
const control = createHeartbeatRunControl(db, deps);
const succeeded = evidence === "local exit" || evidence === "remote receipt";
const delivery = control.cancelBudgetScopeWork(scope);
if (succeeded) await delivery;
else await expect(delivery).rejects.toThrow("Budget cancellation delivery remains pending");
const persisted = await getRun(run.id);
expect(persisted).toMatchObject({ status: "succeeded", nativePhase: "committed" });
expect(persisted?.resultJson?.executionCancellation).toBeUndefined();
const intent = (persisted?.resultJson?.startupCancellation as { budgetEnforcement?: Record<string, unknown> })?.budgetEnforcement;
if (succeeded) expect(intent).toMatchObject({ deliveredAt: expect.any(String),
terminationProof: evidence === "local exit" ? "local_process_exit" : "provider_termination_receipt" });
else expect(intent?.deliveredAt).toBeUndefined();
expect(await listPendingBudgetCancellationScopes(db)).toHaveLength(succeeded ? 0 : 1);
expect(terminateLocalService).not.toHaveBeenCalled();
expect(deps.setRunStatusFromLive).not.toHaveBeenCalled();
});
it("moves an attempted failed Stop behind another pending scope", async () => {
const { company, agent } = await fixture();
const [otherAgent] = await db.insert(agents).values({ companyId: company.id, name: "Later recovery" }).returning();
await db.insert(heartbeatRuns).values([agent, otherAgent].map((selected, index) => ({
companyId: company.id, agentId: selected.id, status: "cancelled", runtimeMode: "legacy",
resultJson: { startupCancellation: { requestedAt: `2020-01-0${index + 1}T00:00:00.000Z`,
budgetEnforcement: { companyId: company.id, scopeType: "agent", scopeId: selected.id,
policyId: randomUUID(), version: 1 } } },
})));
const [first] = await listPendingBudgetCancellationScopes(db, 1);
expect(first.scopeId).toBe(agent.id);
await expect(createHeartbeatRunControl(db, callbacks(db)).cancelBudgetScopeWork(first))
.rejects.toThrow("Budget cancellation delivery remains pending");
const [next] = await listPendingBudgetCancellationScopes(db, 1);
expect(next.scopeId).toBe(otherAgent.id);
const [attempted] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.agentId, agent.id));
expect((attempted.resultJson?.startupCancellation as { budgetEnforcement?: { lastAttemptAt?: string } })
?.budgetEnforcement?.lastAttemptAt).toEqual(expect.any(String));
});
it("limits project budget cancellation by company, effective project, and admission cutoff", async () => {
const { company, agent } = await fixture();
const [otherCompany] = await db.insert(companies).values({ name: "Other company", issuePrefix: "OC" }).returning();
const [otherAgent] = await db.insert(agents).values({ companyId: otherCompany.id, name: "Other company agent" }).returning();
const projectRows = await db.insert(projects).values([
{ companyId: company.id, name: "Target" }, { companyId: company.id, name: "Override" },
]).returning();
const [issue] = await db.insert(issues).values({ companyId: company.id, projectId: projectRows[0].id, title: "Bound task", issueNumber: 1, identifier: "RC-1" }).returning();
const old = new Date("2026-01-01T00:00:00Z");
const cutoff = new Date("2026-01-02T00:00:00Z");
const wakes = await db.insert(agentWakeupRequests).values([
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued", createdAt: old, payload: { issueId: issue.id } },
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued", createdAt: old, payload: { issueId: issue.id, projectId: projectRows[1].id } },
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued", createdAt: cutoff, payload: { projectId: projectRows[0].id } },
{ companyId: otherCompany.id, agentId: otherAgent.id, source: "assignment", status: "queued", createdAt: old, payload: { projectId: projectRows[0].id } },
]).returning();
await createHeartbeatRunControl(db, callbacks(db)).cancelBudgetScopeWork({ companyId: company.id, scopeType: "project", scopeId: projectRows[0].id, createdBefore: cutoff });
const persisted = await db.select().from(agentWakeupRequests);
expect(wakes.map(wake => persisted.find(row => row.id === wake.id)?.status)).toEqual(["cancelled", "queued", "queued", "queued"]);
});
it("fences detached browsers after their parent run finishes using admitted project scope and cutoff", async () => {
const { company, agent } = await fixture();
const [project, otherProject] = await db.insert(projects).values([
{ companyId: company.id, name: "Stopped project" }, { companyId: company.id, name: "Other project" },
]).returning();
const [movedIssue, targetIssue] = await db.insert(issues).values([
{ companyId: company.id, projectId: otherProject.id, title: "Moved after browser admission" },
{ companyId: company.id, projectId: project.id, title: "Current project" },
]).returning();
const [parent] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id,
status: "succeeded", finishedAt: new Date(), runtimeMode: "legacy" }).returning();
const [application] = await db.insert(toolApplications).values({ companyId: company.id, name: "Browser fixture", type: "rest_api" }).returning();
const [connection] = await db.insert(toolConnections).values({ companyId: company.id, applicationId: application.id,
name: "Browser fixture", uid: "budget-browser-fixture", transport: "rest_api" }).returning();
const [grant] = await db.insert(connectionGrants).values({ companyId: company.id, connectionId: connection.id, kind: "organization" }).returning();
const cutoff = new Date();
const old = new Date(cutoff.getTime() - 1_000);
const later = new Date(cutoff.getTime() + 60_000);
const [admitted, other, legacy, newer, noProject] = await db.insert(browserUseSessions).values([
{ issueId: movedIssue.id, stopRequested: "cancel" }, { issueId: targetIssue.id },
{ issueId: targetIssue.id }, { issueId: targetIssue.id, createdAt: later }, { issueId: targetIssue.id },
].map(values => ({ companyId: company.id, agentId: agent.id, connectionId: connection.id, grantId: grant.id,
providerSessionId: randomUUID(), status: "running", createdAt: old, nextPollAt: later,
leaseToken: randomUUID(), leaseUntil: later, error: "retained cleanup context", ...values }))).returning();
await db.insert(browserUseRuns).values([admitted, other, legacy, noProject].map((session, index) => {
const invocationId = randomUUID();
return { companyId: company.id, sessionId: session.id, heartbeatRunId: parent.id, invocationId,
status: "running", detachedUntil: later, createdAt: old,
result: index === 2 ? null : { paperclipBudgetHold: { invocationId, amountCents: "10", state: "held",
projectId: index === 0 ? project.id : index === 1 ? otherProject.id : null } },
};
}));
const wake = vi.fn();
const unsubscribe = subscribeDeliveryWork(db, DELIVERY_QUEUES.browser, wake);
const deps = callbacks(db);
try {
await createHeartbeatRunControl(db, deps).cancelBudgetScopeWork({ companyId: company.id,
scopeType: "project", scopeId: project.id, createdBefore: cutoff });
const persisted = await db.select().from(browserUseSessions).where(eq(browserUseSessions.companyId, company.id));
for (const original of [admitted, legacy]) {
const stopped = persisted.find(session => session.id === original.id)!;
expect(stopped).toMatchObject({ status: "stopping", stopRequested: "end", providerSessionId: original.providerSessionId,
leaseToken: original.leaseToken, leaseUntil: original.leaseUntil, error: original.error });
expect(stopped.nextPollAt.getTime()).toBeLessThan(later.getTime());
}
for (const original of [other, newer, noProject]) expect(persisted.find(session => session.id === original.id)).toEqual(original);
expect(await deps.getRun(parent.id)).toMatchObject({ status: "succeeded", resultJson: null });
expect(wake).toHaveBeenCalledTimes(1);
expect(terminateLocalService).not.toHaveBeenCalled();
expect(deps.setRunStatusFromLive).not.toHaveBeenCalled();
} finally { unsubscribe(); }
});
});
@@ -0,0 +1,43 @@
import { heartbeatRuns, type Db } from "@paperclipai/db";
import { isUuidLike } from "@paperclipai/shared";
import { and, eq, sql } from "drizzle-orm";
import type { BudgetEnforcementScope } from "../budgets.js";
/** Committed Stop intent survives policy edits, deletion, and terminal status.
* Run control writes deliveredAt only after its ordinary stop proof succeeds. */
export function pendingBudgetRunCancellationPredicate() {
return and(
sql`jsonb_typeof(${heartbeatRuns.resultJson}->'startupCancellation'->'budgetEnforcement') = 'object'`,
sql`${heartbeatRuns.resultJson}->'startupCancellation'->'budgetEnforcement'->>'deliveredAt' is null`,
);
}
export function budgetCancellationScopePredicate(scope: Pick<BudgetEnforcementScope, "companyId" | "scopeType" | "scopeId">) {
return and(eq(heartbeatRuns.companyId, scope.companyId),
sql`${heartbeatRuns.resultJson}->'startupCancellation'->'budgetEnforcement'->>'companyId' = ${scope.companyId}`,
sql`${heartbeatRuns.resultJson}->'startupCancellation'->'budgetEnforcement'->>'scopeType' = ${scope.scopeType}`,
sql`${heartbeatRuns.resultJson}->'startupCancellation'->'budgetEnforcement'->>'scopeId' = ${scope.scopeId}`);
}
export async function listPendingBudgetCancellationScopes(db: Db, limit = 100): Promise<BudgetEnforcementScope[]> {
const intent = sql`${heartbeatRuns.resultJson}->'startupCancellation'->'budgetEnforcement'`;
const scopeType = sql<string>`${intent}->>'scopeType'`, scopeId = sql<string>`${intent}->>'scopeId'`;
const attemptAt = sql`coalesce((${intent}->>'lastAttemptAt')::timestamptz,
(${heartbeatRuns.resultJson}->'startupCancellation'->>'requestedAt')::timestamptz, ${heartbeatRuns.updatedAt})`;
const rows = await db.select({ companyId: heartbeatRuns.companyId,
intent: sql<Record<string, unknown>>`(array_agg(${intent} order by ${attemptAt}, ${heartbeatRuns.id}))[1]`,
}).from(heartbeatRuns).where(pendingBudgetRunCancellationPredicate())
.groupBy(heartbeatRuns.companyId, scopeType, scopeId)
.orderBy(sql`min(${attemptAt})`, heartbeatRuns.companyId, scopeType, scopeId)
.limit(limit);
return rows.flatMap(({ companyId, intent }) => {
if (intent.companyId !== companyId || !["company", "agent", "project"].includes(String(intent.scopeType))
|| typeof intent.scopeId !== "string" || !isUuidLike(intent.scopeId)
|| typeof intent.policyId !== "string" || !isUuidLike(intent.policyId)
|| !Number.isSafeInteger(intent.version) || Number(intent.version) < 0) return [];
const cutoff = typeof intent.createdBefore === "string" ? new Date(intent.createdBefore) : null;
return [{ companyId, scopeType: intent.scopeType as BudgetEnforcementScope["scopeType"], scopeId: intent.scopeId,
enforcement: { policyId: intent.policyId, version: Number(intent.version) },
...(cutoff && Number.isFinite(cutoff.getTime()) ? { createdBefore: cutoff } : {}) }];
});
}
@@ -1,9 +1,8 @@
import { randomUUID } from "node:crypto";
import type { Db } from "@paperclipai/db";
import { activityLog, agents, agentWakeupRequests, caseEvents, caseHandoffs, cases, companies, createDb, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS, environmentLeases, heartbeatRuns, issues, nativeRunFinalizations, projects } from "@paperclipai/db";
import { eq, inArray, sql } from "drizzle-orm";
import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../../__tests__/helpers/embedded-postgres.js";
import { activityLog, agents, agentWakeupRequests, caseEvents, caseHandoffs, cases, companies, environmentLeases, heartbeatRuns, issues, nativeRunFinalizations } from "@paperclipai/db";
import { eq, inArray } from "drizzle-orm";
import { afterEach, describe, expect, it, vi } from "vitest";
import { runningProcesses } from "../../adapters/index.js";
import {
adapterExecutionControls,
@@ -16,14 +15,13 @@ import { HttpError } from "../../errors.js";
import { captureException } from "../../sentry.js";
import { terminateLocalService } from "../local-service-supervisor.js";
import { createHeartbeatRunState } from "./run-state.js";
import { createHeartbeatRunPreparation } from "./run-preparation.js";
import { createHeartbeatLifecycle, type HeartbeatLifecycleDependencies } from "./run-lifecycle.js";
import { createHeartbeatRunControl, type HeartbeatRunControlDependencies } from "./run-control.js";
import { createHeartbeatRunControl } from "./run-control.js";
import { preserveCancelledNativeWorkspaces } from "../native-runtime/native-failed-workspace-recovery.js";
import { createPostgresRunDispatchAdapter } from "../../modules/run-dispatch/adapters/postgres.js";
import { caseHandoffService } from "../case-handoffs.js";
import { instanceSettingsService } from "../instance-settings.js";
import { issueService } from "../issues.js";
import { callbacks, runControlPostgresSupport, setupRunControlDatabase } from "./__tests__/run-control-fixtures.js";
vi.mock("../live-events.js", () => ({ publishLiveEvent: vi.fn() }));
vi.mock("../adapter-execution-control.js", async (importOriginal) => {
@@ -40,50 +38,6 @@ vi.mock("../local-service-supervisor.js", async (importOriginal) => ({
}));
type Run = typeof heartbeatRuns.$inferSelect;
function callbacks(db: Db) {
const state = createHeartbeatRunState(db);
const lifecycleDependencies = {
...state,
...createHeartbeatRunPreparation(db),
issuesSvc: {
addComment: vi.fn<HeartbeatLifecycleDependencies["issuesSvc"]["addComment"]>(),
listReviewAttention: vi.fn<HeartbeatLifecycleDependencies["issuesSvc"]["listReviewAttention"]>(),
},
recovery: { escalateStrandedAssignedIssue: vi.fn<HeartbeatLifecycleDependencies["recovery"]["escalateStrandedAssignedIssue"]>() },
companySkills: { completeTestRunForIssue: vi.fn<HeartbeatLifecycleDependencies["companySkills"]["completeTestRunForIssue"]>() },
budgets: { getInvocationBlock: vi.fn(async () => null) },
treeControlSvc: { getActivePauseHoldGate: vi.fn(async () => null) },
enqueueWakeup: vi.fn<HeartbeatLifecycleDependencies["enqueueWakeup"]>(),
getCurrentUserRedactionOptions: vi.fn(async () => ({})),
getAgentInvokability: vi.fn<HeartbeatLifecycleDependencies["getAgentInvokability"]>(async () => ({ invokable: true })),
emitTerminalAgentTaskRun: vi.fn<HeartbeatLifecycleDependencies["emitTerminalAgentTaskRun"]>(),
budgetHooks: {},
} satisfies HeartbeatLifecycleDependencies;
const lifecycle = createHeartbeatLifecycle(db, lifecycleDependencies);
return {
...state,
...lifecycle,
getRun: vi.fn(state.getRun),
getAgent: vi.fn(state.getAgent),
setRunStatusFromLive: vi.fn(lifecycle.setRunStatusFromLive),
setRunStatus: vi.fn(lifecycle.setRunStatus),
setWakeupStatus: vi.fn(lifecycle.setWakeupStatus),
appendRunEvent: vi.fn(lifecycle.appendRunEvent),
finalizeAgentStatus: vi.fn<HeartbeatRunControlDependencies["finalizeAgentStatus"]>(),
options: { closeWarmNativeSessionsForRun: vi.fn(async () => ({ closed: 1, busy: 0, failed: 0 })) },
envOrchestrator: { releaseForRun: vi.fn<HeartbeatRunControlDependencies["envOrchestrator"]["releaseForRun"]>(async () => ({ released: [], errors: [] })) },
enqueueWakeup: lifecycleDependencies.enqueueWakeup,
startNextQueuedRunForAgent: vi.fn<HeartbeatRunControlDependencies["startNextQueuedRunForAgent"]>(),
sweepPendingCleanupLeases: vi.fn<HeartbeatRunControlDependencies["sweepPendingCleanupLeases"]>(),
timerClaimWasFirstHeartbeat: vi.fn<HeartbeatRunControlDependencies["timerClaimWasFirstHeartbeat"]>(() => undefined),
wakeQueue: { releaseIssueExecution: vi.fn<HeartbeatRunControlDependencies["wakeQueue"]["releaseIssueExecution"]>(async () => ({ outcome: { kind: "released" }, postCommitEffects: [] })) },
applyWakeQueuePostCommitEffects: vi.fn(async () => undefined),
getSchedulingSuppression: vi.fn(async () => ({ suppressed: false })),
activeRunExecutions: new Set<string>(),
processRunCancellationSettlements: new Map<string, { settled: Promise<void>; failed: boolean; error?: unknown }>(),
failedProcessRunCancellations: new Map<string, { settled: Promise<void>; failed: boolean; error?: unknown }>(),
} satisfies HeartbeatRunControlDependencies;
}
function guardedDatabase() {
const access = vi.fn(() => { throw new Error("Unexpected run-control database access"); });
return { db: new Proxy({}, { get: access }) as Db, access };
@@ -171,35 +125,9 @@ describe("heartbeat run-control module boundary", () => {
});
});
const support = await getEmbeddedPostgresTestSupport();
if (!support.supported) console.warn(`Skipping run-control database tests: ${support.reason}`);
describe.skipIf(!support.supported)("heartbeat run-control database wiring", () => {
let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>>;
describe.skipIf(!runControlPostgresSupport.supported)("heartbeat run-control database wiring", () => {
let db: Db;
beforeAll(async () => {
database = await startEmbeddedPostgresTestDatabase("paperclip-heartbeat-run-control-module");
db = createDb(database.connectionString);
}, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS);
afterAll(async () => {
await db?.$client.end();
await database?.cleanup();
});
beforeEach(async () => {
vi.mocked(terminateLocalService).mockReset().mockResolvedValue(undefined);
await db.execute(sql`truncate table companies restart identity cascade`);
});
afterEach(() => runningProcesses.clear());
async function fixture() {
const [company] = await db.insert(companies).values({ name: "Run control", issuePrefix: "RC" }).returning();
const [agent] = await db.insert(agents).values({ companyId: company.id, name: "Controlled agent", adapterType: "codex_local" }).returning();
return { company, agent };
}
async function liveRun(companyId: string, agentId: string) {
const [run] = await db.insert(heartbeatRuns).values({ companyId, agentId, status: "running", runtimeMode: "legacy", runtimeModeResolvedAt: new Date() }).returning();
runningProcesses.set(run.id, { child: { pid: 424242 } as never, graceSec: 4, processGroupId: 424242 });
return run;
}
const { fixture, liveRun } = setupRunControlDatabase("paperclip-heartbeat-run-control-module", value => { db = value; });
it.each(["queued", "running", "cancelled", "succeeded"] as const)("does not supersede a retry that becomes %s between lookup and the cancellation fence", async status => {
const { company, agent } = await fixture();
@@ -704,37 +632,4 @@ describe.skipIf(!support.supported)("heartbeat run-control database wiring", ()
expect(wakes.map(wake => persisted.find(row => row.id === wake.id)?.status)).toEqual(["cancelled", "cancelled", "claimed", "queued", "queued"]);
});
it("rechecks stale budget enforcement before cancelling runs or queued wakes", async () => {
const { company, agent } = await fixture();
const [run] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id, status: "queued", runtimeMode: "legacy" }).returning();
const [wake] = await db.insert(agentWakeupRequests).values({ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued" }).returning();
const deps = callbacks(db);
await createHeartbeatRunControl(db, deps).cancelBudgetScopeWork({ companyId: company.id, scopeType: "agent", scopeId: agent.id, enforcement: { policyId: randomUUID(), version: 1 } });
expect((await deps.getRun(run.id))?.status).toBe("queued");
expect((await db.select().from(agentWakeupRequests).where(eq(agentWakeupRequests.id, wake.id)))[0].status).toBe("queued");
expect(deps.setRunStatusFromLive).not.toHaveBeenCalled();
expect(deps.finalizeAgentStatus).not.toHaveBeenCalled();
expect(terminateLocalService).not.toHaveBeenCalled();
});
it("limits project budget cancellation by company, effective project, and admission cutoff", async () => {
const { company, agent } = await fixture();
const [otherCompany] = await db.insert(companies).values({ name: "Other company", issuePrefix: "OC" }).returning();
const [otherAgent] = await db.insert(agents).values({ companyId: otherCompany.id, name: "Other company agent" }).returning();
const projectRows = await db.insert(projects).values([
{ companyId: company.id, name: "Target" }, { companyId: company.id, name: "Override" },
]).returning();
const [issue] = await db.insert(issues).values({ companyId: company.id, projectId: projectRows[0].id, title: "Bound task", issueNumber: 1, identifier: "RC-1" }).returning();
const old = new Date("2026-01-01T00:00:00Z");
const cutoff = new Date("2026-01-02T00:00:00Z");
const wakes = await db.insert(agentWakeupRequests).values([
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued", createdAt: old, payload: { issueId: issue.id } },
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued", createdAt: old, payload: { issueId: issue.id, projectId: projectRows[1].id } },
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued", createdAt: cutoff, payload: { projectId: projectRows[0].id } },
{ companyId: otherCompany.id, agentId: otherAgent.id, source: "assignment", status: "queued", createdAt: old, payload: { projectId: projectRows[0].id } },
]).returning();
await createHeartbeatRunControl(db, callbacks(db)).cancelBudgetScopeWork({ companyId: company.id, scopeType: "project", scopeId: projectRows[0].id, createdBefore: cutoff });
const persisted = await db.select().from(agentWakeupRequests);
expect(wakes.map(wake => persisted.find(row => row.id === wake.id)?.status)).toEqual(["cancelled", "queued", "queued", "queued"]);
});
});
+295 -85
View File
@@ -1,5 +1,5 @@
import { readTaskRetryIntent, reconcileTaskChatRetryReceipts } from "../task-retry.js";
import { recordedLocalExecutionStopped } from "../native-local-process-stop.js";
import { PROCESS_IDENTITY_RECORDED, PROCESS_START_REQUESTED, readNativeLocalProcessStop, recordedLocalExecutionStopped } from "../native-local-process-stop.js";
import { isHeartbeatRunTerminalStatus } from "./run-lifecycle.js";
import { pendingCaseHandoffForRun, cancelCaseHandoffsForIssue, isCaseHandoffSourceReleased } from "./case-handoff-lifecycle.js";
import { publishActivity, type ActivityPublication } from "../activity-log.js";
@@ -57,6 +57,8 @@ import {
waitForAdapterStop,
} from "../adapter-execution-control.js";
import { AdapterStopTimeoutError } from "../adapter-stop-timeout.js";
import { budgetCancellationScopePredicate, pendingBudgetRunCancellationPredicate } from "./budget-cancellation.js";
import { withBudgetSpan } from "../budget-tracing.js";
import { captureException } from "../../sentry.js";
import { randomUUID } from "node:crypto";
import {
@@ -82,9 +84,12 @@ import {
} from "@paperclipai/shared";
import {
agentWakeupRequests,
browserUseRuns,
browserUseSessions,
companies,
environmentLeases,
heartbeatRuns,
heartbeatRunEvents,
issueComments,
issueRecoveryActions,
issues,
@@ -96,6 +101,7 @@ import {
notFound,
} from "../../errors.js";
import { logger } from "../../middleware/logger.js";
import { DELIVERY_QUEUES, notifyDeliveryWork } from "../delivery-work-notifications.js";
import {
queuedCommentIdsFromWakePayload,
withQueuedCommentIdsInWakePayload,
@@ -134,6 +140,27 @@ type ProcessRunCancellationSettlement = {
error?: unknown;
};
// This bounds external Stop concurrency across scopes and heartbeat service
// instances sharing a database. Durable cancellation rows own recovery.
const budgetStopDeliveryLimits = new WeakMap<Db, { active: number; waiting: Array<() => void> }>();
async function withBudgetStopDelivery<T>(db: Db, work: () => Promise<T>): Promise<T> {
let limit = budgetStopDeliveryLimits.get(db);
if (!limit) {
limit = { active: 0, waiting: [] };
budgetStopDeliveryLimits.set(db, limit);
}
if (limit.active >= 4) await new Promise<void>(resolve => limit.waiting.push(resolve));
else limit.active += 1;
try {
return await withBudgetSpan("budget.cancellation.deliver", work);
} finally {
const next = limit.waiting.shift();
if (next) next();
else limit.active -= 1;
}
}
/** Cancellation shares executor ownership and delegates lifecycle and admission effects. */
export interface HeartbeatRunControlDependencies extends Pick<ReturnType<typeof createHeartbeatRunState>,
"getRun" | "getAgent">,
@@ -255,6 +282,43 @@ export type CancelRunOptions = {
suppressImmediateRecovery?: boolean;
};
/** The same durable fence protects both individual Stop and budget fan-out.
* It records intent only; provider termination still needs its own receipt. */
function runCancellationFenceSql(input: {
cancellation: ReturnType<typeof requestedRunCancellation>;
errorCode: string;
reassignmentStopRequested?: boolean;
preserveCancellation?: boolean;
budgetEnforcement?: BudgetEnforcementScope;
}) {
const cancellation = JSON.stringify(input.cancellation);
const startupCancellation = sql`jsonb_build_object(
'requestedAt', ${new Date().toISOString()}::text,
'errorCode', ${input.errorCode}::text,
'beforeNativeSelection', ${heartbeatRuns.runtimeMode} = 'legacy'
and ${heartbeatRuns.runtimeModeResolvedAt} is null
and ${heartbeatRuns.executionStage} = 'preparing'
and coalesce(${heartbeatRuns.runnerProfileJson}->'adapterDispatch'->>'adapterType' = 'paperclip_runner', false)
)`;
let startupFence = input.preserveCancellation
? sql`coalesce(${heartbeatRuns.resultJson}->'startupCancellation', ${startupCancellation})`
: startupCancellation;
if (input.budgetEnforcement) {
const intent = JSON.stringify({ companyId: input.budgetEnforcement.companyId,
scopeType: input.budgetEnforcement.scopeType, scopeId: input.budgetEnforcement.scopeId,
policyId: input.budgetEnforcement.enforcement?.policyId, version: input.budgetEnforcement.enforcement?.version,
createdBefore: input.budgetEnforcement.createdBefore?.toISOString() });
startupFence = sql`${startupFence} || jsonb_build_object('budgetEnforcement',
coalesce(nullif(${heartbeatRuns.resultJson}->'startupCancellation'->'budgetEnforcement', 'null'::jsonb), ${intent}::jsonb))`;
}
return sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) ||
${JSON.stringify(input.reassignmentStopRequested ? { reassignmentStopRequested: true } : {})}::jsonb ||
jsonb_build_object('cancellation', ${input.preserveCancellation
? sql`coalesce(${heartbeatRuns.resultJson}->'cancellation', ${cancellation}::jsonb)`
: sql`${cancellation}::jsonb`}) ||
jsonb_build_object('startupCancellation', ${startupFence})`;
}
/** Run control binds cancellation, resource release, and saved-input resumption. */
export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunControlDependencies) {
const {
@@ -793,7 +857,7 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont
}
}
async function listProjectScopedRunIds(companyId: string, projectId: string, createdBefore?: Date) {
async function listProjectScopedRunIds(companyId: string, projectId: string, createdBefore?: Date, database: Db = db) {
const runIssueId = sql<
string | null
>`${heartbeatRuns.contextSnapshot} ->> 'issueId'`;
@@ -801,7 +865,7 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont
string | null
>`coalesce(${heartbeatRuns.contextSnapshot} ->> 'projectId', ${issues.projectId}::text)`;
const rows = await db
const rows = await database
.selectDistinctOn([heartbeatRuns.id], { id: heartbeatRuns.id })
.from(heartbeatRuns)
.leftJoin(
@@ -866,69 +930,68 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont
async function cancelPendingWakeupsForBudgetScope(
scope: BudgetEnforcementScope,
tx: Db,
) {
return withCurrentBudgetEnforcement(db, scope, async (tx) => {
const now = new Date();
let wakeupIds: string[] = [];
const now = new Date();
let wakeupIds: string[] = [];
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 (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 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,
));
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;
}
function cancellationTerminationGraceMs(
@@ -946,6 +1009,81 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont
runId: string,
reason = "Cancelled by control plane",
options: CancelRunOptions = {},
) {
return cancelRun(runId, reason, options);
}
async function acknowledgeBudgetCancellation(run: typeof heartbeatRuns.$inferSelect, terminationProof?: string) {
if (!parseObject(parseObject(run.resultJson?.startupCancellation).budgetEnforcement).scopeType) return;
const acknowledged = await db.update(heartbeatRuns).set({ resultJson: sql`jsonb_set(${heartbeatRuns.resultJson},
'{startupCancellation,budgetEnforcement}', ${heartbeatRuns.resultJson}->'startupCancellation'->'budgetEnforcement'
|| ${JSON.stringify({ deliveredAt: new Date().toISOString(), ...(terminationProof ? { terminationProof } : {}) })}::jsonb)` })
.where(and(eq(heartbeatRuns.companyId, run.companyId), eq(heartbeatRuns.id, run.id),
pendingBudgetRunCancellationPredicate(),
terminationProof ? and(eq(heartbeatRuns.status, run.status), eq(heartbeatRuns.nextEventSeq, run.nextEventSeq),
sql`${heartbeatRuns.processPid} is not distinct from ${run.processPid}`,
sql`${heartbeatRuns.processGroupId} is not distinct from ${run.processGroupId}`,
sql`${heartbeatRuns.processStartedAt} is not distinct from ${run.processStartedAt?.toISOString() ?? null}::timestamptz`) : undefined,
)).returning({ id: heartbeatRuns.id });
if (terminationProof && !acknowledged.length)
throw conflict("Execution evidence changed before acknowledging budget Stop", { runId: run.id });
}
function hasRecordedStopAcknowledgement(run: typeof heartbeatRuns.$inferSelect) {
if (parseObject(run.resultJson?.executionCancellation).state === "acknowledged") return true;
const intent = parseObject(run.resultJson?.nativeCancellation);
return intent.schema === "paperclip.native-cancellation.v1" && intent.runId === run.id
&& intent.companyId === run.companyId && intent.issueId === run.nativeIssueId
&& intent.scope === "run" && intent.reasonCode === "cancellation_run_only"
&& intent.dispatchState === "acknowledged" && intent.dispatched === true
&& typeof intent.intentAuditId === "string" && typeof intent.acknowledgementAuditId === "string";
}
async function completedRunTerminationProof(run: typeof heartbeatRuns.$inferSelect): Promise<string | null> {
// Logical completion can win the race with queued Stop delivery. It does
// not prove a provider exited, and an in-process owner must still settle.
if (!isHeartbeatRunTerminalStatus(run.status) || activeRunExecutions.has(run.id) ||
runningProcesses.has(run.id) || adapterExecutionControls.has(run.id) || processRunCancellationSettlements.has(run.id)) return null;
const leases = await db.select().from(environmentLeases).where(and(
eq(environmentLeases.companyId, run.companyId), eq(environmentLeases.heartbeatRunId, run.id),
));
if (leases.some(lease => lease.provider !== "local"))
return leases.every(hasRemoteTerminationReceipt) ? "provider_termination_receipt" : null;
if (run.processPid || run.processGroupId || run.processStartedAt)
return await recordedLocalExecutionStopped(db, run) ? "local_process_exit" : null;
if (await readNativeLocalProcessStop(db, run.companyId, run.id)) return "local_process_stop_receipt";
// Normal native finalization clears mutable process fields. Rebind the
// latest server-owned launch identity before checking that exact local
// process and group. Any newer start invalidates the older observation.
const [identity] = await db.select({ eventType: heartbeatRunEvents.eventType, payload: heartbeatRunEvents.payload })
.from(heartbeatRunEvents).where(and(eq(heartbeatRunEvents.companyId, run.companyId),
eq(heartbeatRunEvents.runId, run.id), isNull(heartbeatRunEvents.sourceEventId),
inArray(heartbeatRunEvents.eventType, [PROCESS_START_REQUESTED, PROCESS_IDENTITY_RECORDED, "native.local_process_stopped"])))
.orderBy(desc(heartbeatRunEvents.seq)).limit(1);
const recorded = identity?.payload;
if (identity?.eventType !== PROCESS_IDENTITY_RECORDED || typeof recorded?.processPid !== "number" ||
!Number.isSafeInteger(recorded.processPid) || recorded.processPid <= 1 ||
(recorded.processGroupId !== null && (typeof recorded.processGroupId !== "number" ||
!Number.isSafeInteger(recorded.processGroupId) || recorded.processGroupId <= 1)) ||
typeof recorded.processStartedAt !== "string") return null;
const processStartedAt = new Date(recorded.processStartedAt);
if (Number.isNaN(processStartedAt.getTime())) return null;
return await recordedLocalExecutionStopped(db, { ...run, processPid: recorded.processPid,
processGroupId: recorded.processGroupId, processStartedAt }) ? "local_process_exit" : null;
}
function deliveredBudgetCancellationResult(run: typeof heartbeatRuns.$inferSelect, deliveredAt: Date) {
const startup = parseObject(run.resultJson?.startupCancellation);
const intent = parseObject(startup.budgetEnforcement);
return intent.scopeType ? { startupCancellation: { ...startup,
budgetEnforcement: { ...intent, deliveredAt: deliveredAt.toISOString() } } } : {};
}
async function cancelRun(
runId: string,
reason: string,
options: CancelRunOptions,
committedBudgetCancellation = false,
) {
let run = await getRun(runId);
if (!run) throw notFound("Heartbeat run not found");
@@ -1041,11 +1179,19 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont
!CANCELLABLE_HEARTBEAT_RUN_STATUSES.includes(
run.status as (typeof CANCELLABLE_HEARTBEAT_RUN_STATUSES)[number],
)
)
) {
if (committedBudgetCancellation) {
const terminationProof = hasRecordedStopAcknowledgement(run) ? undefined : await completedRunTerminationProof(run);
if (terminationProof === null) {
throw conflict("Budget Stop is committed, but provider termination has not been acknowledged", { runId: run.id });
}
await acknowledgeBudgetCancellation(run, terminationProof);
}
return run;
}
const agent = await getAgent(run.agentId);
const errorCode = options.errorCode ?? "cancelled";
const cancellation = (options.expectedStatus && run.resultJson?.startupCancellation
const cancellation = ((options.expectedStatus || committedBudgetCancellation) && run.resultJson?.startupCancellation
? readRunCancellation(run.resultJson) : null) ?? requestedRunCancellation(options.resultJson ?? {}, reason);
const pendingProcessCancellation = processRunCancellationSettlements.get(
@@ -1069,23 +1215,13 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont
// 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 (options.expectedStatus || options.budgetEnforcement || (!options.cancellationRequestId && (run.runtimeMode === "native" || (!run.runtimeModeResolvedAt && !running && !control)))) {
if (!committedBudgetCancellation && (options.expectedStatus || options.budgetEnforcement || (!options.cancellationRequestId && (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,
'errorCode', ${errorCode}::text,
'beforeNativeSelection', ${heartbeatRuns.runtimeMode} = 'legacy'
and ${heartbeatRuns.runtimeModeResolvedAt} is null
and ${heartbeatRuns.executionStage} = 'preparing'
and coalesce(${heartbeatRuns.runnerProfileJson}->'adapterDispatch'->>'adapterType' = 'paperclip_runner', false)
))`,
resultJson: runCancellationFenceSql({ cancellation, errorCode,
reassignmentStopRequested: options.errorCode === "issue_reassigned" && options.resultJson?.reassignmentStopConfirmed === true }),
}).where(and(eq(heartbeatRuns.id, runId), inArray(heartbeatRuns.status,
options.expectedStatus ? [options.expectedStatus] : pendingNativeRetry ? [...CANCELLABLE_HEARTBEAT_RUN_STATUSES, "failed"] : [...CANCELLABLE_HEARTBEAT_RUN_STATUSES],
))).returning();
@@ -1229,6 +1365,7 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont
}
// The owned adapter already finalized this run and its lifecycle.
// Do not replay the process cancellation side effects below.
await acknowledgeBudgetCancellation(stopped);
return { run: stopped, updated: false };
}
}
@@ -1271,6 +1408,10 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont
errorCode, errorMessage: reason,
})
: {}),
// Commit this receipt with the terminal transition. A
// crash after provider Stop cannot strand an undispatched
// queued run whose normal path needs no process receipt.
...deliveredBudgetCancellationResult(run, finishedAt),
// The native cancellation helper may have advanced a durable
// pending intent to its acknowledged state after `run` was
// first read. Never let that stale snapshot overwrite the
@@ -1428,17 +1569,86 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont
};
}
async function stopBrowserSessionsForBudgetScope(tx: Db, scope: BudgetEnforcementScope) {
// Browser work can outlive its terminal parent run. Its admitted project
// snapshot remains authoritative even if the task moves to another project.
const activeHold = sql`(select ${browserUseRuns.result}->'paperclipBudgetHold'
from ${browserUseRuns} where ${browserUseRuns.companyId} = ${browserUseSessions.companyId}
and ${browserUseRuns.sessionId} = ${browserUseSessions.id}
and ${browserUseRuns.result}->'paperclipBudgetHold'->>'state' = 'held'
and ${browserUseRuns.result}->'paperclipBudgetHold'->>'invocationId' = ${browserUseRuns.invocationId}::text
order by ${browserUseRuns.createdAt} desc, ${browserUseRuns.id} desc limit 1)`;
const project = sql`case when ${activeHold} ? 'projectId' then ${activeHold}->>'projectId'
else (select ${issues.projectId}::text from ${issues}
where ${issues.id} = ${browserUseSessions.issueId} and ${issues.companyId} = ${scope.companyId}) end`;
const sessions = await tx.select({ id: browserUseSessions.id }).from(browserUseSessions).where(and(
eq(browserUseSessions.companyId, scope.companyId),
inArray(browserUseSessions.status, ["starting", "running", "idle", "stopping"]),
scope.scopeType === "agent" ? eq(browserUseSessions.agentId, scope.scopeId) : undefined,
scope.scopeType === "project" ? sql`${project} = ${scope.scopeId}` : undefined,
scope.createdBefore ? lt(browserUseSessions.createdAt, scope.createdBefore) : undefined,
)).for("update");
if (!sessions.length) return;
await notifyDeliveryWork(tx, DELIVERY_QUEUES.browser, scope.companyId);
await tx.update(browserUseSessions).set({ status: "stopping", stopRequested: "end",
nextPollAt: new Date(), updatedAt: new Date(),
}).where(and(eq(browserUseSessions.companyId, scope.companyId),
inArray(browserUseSessions.id, sessions.map(session => session.id))));
}
async function cancelBudgetScopeWork(scope: BudgetEnforcementScope) {
const runIds = scope.scopeType === "project"
? await listProjectScopedRunIds(scope.companyId, scope.scopeId, scope.createdBefore)
: await db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(and(
eq(heartbeatRuns.companyId, scope.companyId),
scope.scopeType === "agent" ? eq(heartbeatRuns.agentId, scope.scopeId) : undefined,
scope.createdBefore ? lt(heartbeatRuns.createdAt, scope.createdBefore) : undefined,
const reason = "Cancelled due to budget pause";
const newlyFenced = await withCurrentBudgetEnforcement(db, scope, async (tx) => {
// Queue admission and every selected run lose authority in this commit,
// before any provider can delay cancellation delivery.
await cancelPendingWakeupsForBudgetScope(scope, tx);
await stopBrowserSessionsForBudgetScope(tx, scope);
const candidates = scope.scopeType === "project"
? await listProjectScopedRunIds(scope.companyId, scope.scopeId, scope.createdBefore, tx)
: await tx.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(and(
eq(heartbeatRuns.companyId, scope.companyId),
scope.scopeType === "agent" ? eq(heartbeatRuns.agentId, scope.scopeId) : undefined,
scope.createdBefore ? lt(heartbeatRuns.createdAt, scope.createdBefore) : undefined,
inArray(heartbeatRuns.status, [...CANCELLABLE_HEARTBEAT_RUN_STATUSES]),
)).then((rows) => rows.map((row) => row.id));
if (!candidates.length) return [];
return tx.update(heartbeatRuns).set({ resultJson: runCancellationFenceSql({
cancellation: requestedRunCancellation({}, reason), errorCode: "cancelled", preserveCancellation: true, budgetEnforcement: scope,
}) }).where(and(eq(heartbeatRuns.companyId, scope.companyId), inArray(heartbeatRuns.id, candidates),
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", { budgetEnforcement: scope });
await cancelPendingWakeupsForBudgetScope(scope);
scope.createdBefore ? lt(heartbeatRuns.createdAt, scope.createdBefore) : undefined,
)).returning({ id: heartbeatRuns.id }).then((rows) => rows.map((row) => row.id));
});
// A later policy change revokes permission to fence new work, not the Stop
// already committed for older work. Recovery must deliver those intents.
const pending = await db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(and(
pendingBudgetRunCancellationPredicate(), budgetCancellationScopePredicate(scope),
));
const runIds = [...new Set([...(newlyFenced ?? []), ...pending.map(run => run.id)])];
if (!runIds.length) return;
const failures: unknown[] = [];
let nextRun = 0;
await Promise.all(Array.from({ length: Math.min(4, runIds.length) }, async () => {
while (nextRun < runIds.length) {
const runId = runIds[nextRun++]!;
try {
// The committed intent owns these old runs even if a later budget
// edit admits new work. Do not revoke or repeat that earlier fence.
await withBudgetStopDelivery(db, async () => {
const [attempted] = await db.update(heartbeatRuns).set({ resultJson: sql`jsonb_set(${heartbeatRuns.resultJson},
'{startupCancellation,budgetEnforcement,lastAttemptAt}', ${JSON.stringify(new Date().toISOString())}::jsonb)` })
.where(and(eq(heartbeatRuns.companyId, scope.companyId), eq(heartbeatRuns.id, runId),
pendingBudgetRunCancellationPredicate())).returning({ id: heartbeatRuns.id });
if (attempted) await cancelRun(runId, reason, {}, true);
});
} catch (error) {
failures.push(error);
}
}
}));
// Drain the batch before rejecting; one failed provider must not prevent
// the remaining runs from receiving their already-committed Stop.
if (failures.length) throw new AggregateError(failures, "Budget cancellation delivery remains pending");
}
return {
@@ -9,6 +9,7 @@ import { getHeartbeatRunRuntimeStatus } from "../heartbeat-run-runtime-status.js
import { createHeartbeatRunState } from "./run-state.js";
import { createHeartbeatRunPreparation } from "./run-preparation.js";
import { createHeartbeatLifecycle, persistHeartbeatRunProcessMetadata, type HeartbeatLifecycleDependencies } from "./run-lifecycle.js";
import { DELIVERY_QUEUES, subscribeDeliveryWork } from "../delivery-work-notifications.js";
vi.mock("../live-events.js", () => ({ publishLiveEvent: vi.fn() }));
@@ -161,6 +162,43 @@ describe.skipIf(!support.supported)("heartbeat lifecycle database wiring", () =>
expect(persisted.usageJson).toEqual({ accountingReceiptReady: false, usageSource: "per_run", inputTokens: 8, outputTokens: 3 });
});
it.each(["setRunStatus", "setRunStatusIfRunning"] as const)("signals committed terminal accounting before %s publishes a failed live event", async operation => {
const { run } = await fixture({ costAccountingPending: true });
const settled = vi.fn();
const unsubscribe = subscribeDeliveryWork(db, DELIVERY_QUEUES.accounting, settled);
vi.mocked(publishLiveEvent).mockImplementationOnce(() => { throw new Error("Live publication unavailable"); });
try {
await expect(createHeartbeatLifecycle(db, callbacks(db))[operation](run.id, "succeeded", { finishedAt: new Date() }))
.rejects.toThrow("Live publication unavailable");
const [persisted] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, run.id));
expect(persisted.status).toBe("succeeded");
expect(persisted.costAccountingPending).toBe(true);
expect(settled).toHaveBeenCalledOnce();
} finally { unsubscribe(); }
});
it("does not request accounting for an unchanged running status", async () => {
const { run } = await fixture();
const settled = vi.fn();
const unsubscribe = subscribeDeliveryWork(db, DELIVERY_QUEUES.accounting, settled);
try {
await createHeartbeatLifecycle(db, callbacks(db)).setRunStatus(run.id, "running");
expect(settled).not.toHaveBeenCalled();
} finally { unsubscribe(); }
});
it("signals legacy terminal accounting after the nested terminalizer transaction commits", async () => {
const { run } = await fixture({ runtimeMode: "legacy", costAccountingPending: true });
const settled = vi.fn();
const unsubscribe = subscribeDeliveryWork(db, DELIVERY_QUEUES.accounting, settled);
try {
const result = await createHeartbeatLifecycle(db, callbacks(db)).setRunStatusIfRunning(run.id, "failed", { finishedAt: new Date() });
expect(result.updated).toBe(true);
expect(result.run?.status).toBe("failed");
expect(settled).toHaveBeenCalledOnce();
} finally { unsubscribe(); }
});
it.each(["claimed", "cancelled"] as const)("settles a %s wake without reviving a revoked request", async (status) => {
const { wake } = await fixture();
await db.update(agentWakeupRequests).set({ status }).where(eq(agentWakeupRequests.id, wake.id));
+23 -13
View File
@@ -26,7 +26,6 @@ import {
} from "../native-runtime/native-cancellation-request.js";
import { cancellationResultJson } from "../run-cancellation.js";
import { CONTROLLER_LEASE_LOST_FAILURE_CODE } from "../legacy-controller-lease.js";
import { accountRunCost } from "../run-cost-accounting.js";
import { isWaitingConversation } from "../agent-conversations.js";
import { PROCESS_IDENTITY_RECORDED } from "../native-local-process-stop.js";
import {
@@ -39,6 +38,7 @@ import {
isConversationAdapter,
} from "../conversation-continuation.js";
import { terminalizeLegacyExecution } from "../legacy-execution-recovery.js";
import { DELIVERY_QUEUES, notifyDeliveryWork } from "../delivery-work-notifications.js";
import { buildHeartbeatRunStatusLiveEventPayload } from "../heartbeat-run-status-payload.js";
import { randomUUID } from "node:crypto";
import {
@@ -1010,6 +1010,17 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl
return recordHeartbeatRunRuntimeProgress(currentRun, update, issueId);
}
async function writeRunStatus<T>(status: string, companyId: string | undefined, write: (tx: Db) => Promise<T>) {
if (!isHeartbeatRunTerminalStatus(status)) return write(db);
// Register on the outer transaction before the terminal write. An unknown
// commit cannot strand newly-actionable accounting after the live event.
return db.transaction(async transaction => {
const tx = transaction as unknown as Db;
await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, companyId);
return write(tx);
});
}
async function setRunStatus(
runId: string,
status: string,
@@ -1050,16 +1061,16 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl
},
};
}
const updated =
const updated = await writeRunStatus(status, previousStatus?.companyId, writer =>
previousStatus && previousStatus.runtimeMode === "legacy" && isHeartbeatRunTerminalStatus(status)
? await terminalizeLegacyExecution({
db,
? terminalizeLegacyExecution({
db: writer,
run: previousStatus,
status,
patch,
reconcileIfNeeded: true,
})
: await db
: writer
.update(heartbeatRuns)
.set({
status,
@@ -1069,7 +1080,7 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl
})
.where(eq(heartbeatRuns.id, runId))
.returning()
.then((rows) => rows[0] ?? null);
.then((rows) => rows[0] ?? null));
if (updated) {
publishLiveEvent({
@@ -1146,10 +1157,10 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl
},
};
}
const updated =
const updated = await writeRunStatus(status, previousStatus?.companyId, writer =>
previousStatus && previousStatus.runtimeMode === "legacy" && isHeartbeatRunTerminalStatus(status)
? await terminalizeLegacyExecution({
db,
? terminalizeLegacyExecution({
db: writer,
run: previousStatus,
status,
patch,
@@ -1157,7 +1168,7 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl
reconcileIfNeeded: true,
writeConditions: [nativeRunnerOwnershipNotHeldCondition(), ...(cancellationCondition ? [cancellationCondition] : [])],
})
: await db
: writer
.update(heartbeatRuns)
.set({
status,
@@ -1174,7 +1185,7 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl
// issue promotion, or release of the unfinished source lease.
...(status === "failed" && patch?.errorCode === "process_lost"
? [
not(pendingNativeWorkspaceFinalizationCondition(db)),
not(pendingNativeWorkspaceFinalizationCondition(writer)),
// Result publication locks and updates this heartbeat in
// its transaction. If this UPDATE waited behind it, the
// joined tables can still use the older statement snapshot;
@@ -1193,7 +1204,7 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl
.catch((error: unknown) => {
if (cancellationCondition) rethrowNativeCancellationLockConflict(error);
throw error;
});
}));
if (updated) {
publishLiveEvent({
@@ -2975,7 +2986,6 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl
normalizedUsage?: UsageTotals | null,
) {
await ensureRuntimeState(agent);
await accountRunCost(db, run.id, budgetHooks);
await db
.update(agentRuntimeState)
.set({
+10
View File
@@ -6,6 +6,8 @@ import { connectionDeliveryWaitingForResponse, finishedIdleRun as finishedRun, f
import type { Db } from "@paperclipai/db";
import { sql } from "drizzle-orm";
import { verifiedResourceCleanup } from "./cleanup-verification-proof.js";
import { companyHoldMismatch } from "../modules/agent-lifecycle/index.js";
import { pendingBudgetRunCancellationPredicate } from "./heartbeat/budget-cancellation.js";
export interface IdleSleepDrainStatus {
ownerId?: string;
@@ -52,6 +54,11 @@ const WORK_CHECKS = [
OR process_stopped_at IS NULL OR next_attempt_at IS NOT NULL
OR receipt ? 'baseline' OR receipt->>'cleanupPending' = 'true'`,
`SELECT 1 FROM budget_reservations WHERE state NOT IN ('settled', 'released')`,
`SELECT 1 FROM budget_policies WHERE enforcement_version > enforcement_delivered_version`,
`SELECT 1 FROM agents WHERE lifecycle_state IN
('preparing', 'verifying', 'pausing', 'resuming', 'terminating', 'cleaning_up')`,
`SELECT 1 FROM decision_invocations WHERE status = 'running'`,
`SELECT 1 FROM prompt_injection_invocations WHERE status = 'running'`,
// Persistent machine credentials can receive work without a human opening
// the instance. They need an ingress wake contract before we can sleep.
`SELECT 1 FROM agent_api_keys WHERE revoked_at IS NULL`,
@@ -275,6 +282,9 @@ export async function readIdleSleepSafety(
(!scheduledRoutines || query !== `SELECT 1 FROM routines WHERE status NOT IN ('paused', 'archived')`)
&& (!scheduledIssueMonitors || query !== `SELECT 1 FROM issues WHERE status NOT IN ('done', 'cancelled')`))
.map((query) => sql`EXISTS (${sql.raw(query)})`);
checks.push(sql`EXISTS (SELECT 1 FROM agents JOIN companies ON companies.id = agents.company_id
WHERE ${companyHoldMismatch()})`);
checks.push(sql`EXISTS (SELECT 1 FROM heartbeat_runs WHERE ${pendingBudgetRunCancellationPredicate()})`);
// Only the current worker's owned runtime drain can exempt a plugin.
// Missing workers, old SDKs, and newly installed plugins still block.
checks.push(quietPlugins.length === 0
@@ -2,6 +2,7 @@ import { recoverLegacyUnsafeWorkspaceExports } from "./native-workspace-export-r
import { hasPendingNativeChildCompletion } from "./native-child-completion-delivery.js";
import { dismissAutomaticCompletionReviews, decisionHasRetiredAutomaticReview } from "./automatic-completion-reviews.js";
import { logger } from "../../middleware/logger.js";
import { DELIVERY_QUEUES, notifyDeliveryWork } from "../delivery-work-notifications.js";
import { createHash, randomUUID } from "node:crypto";
import { and, asc, desc, eq, gt, inArray, isNotNull, isNull, lt, lte, ne, not, notInArray, or, sql, type SQL } from "drizzle-orm";
import type { Db } from "@paperclipai/db";
@@ -489,6 +490,7 @@ export async function claimNativeSessionResumptions(input: {
|| (!!checkpoint && !checkpointHasProviderIdentity)
) {
const failureCode = originalFailureCode;
await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, row.run.companyId);
await tx.update(nativeRunFinalizations).set({
phase: "terminal_failure",
leaseOwner: null,
@@ -62,6 +62,7 @@ import { waitForPendingRunFailureReports } from "../run-failure-report.js";
import { deliverExecutionStatuses } from "../execution-status-delivery.js";
import { commitNativeStatusDecision } from "./status-decision-committer.js";
import { NATIVE_STATUS_ARBITER_POLICY_VERSION, type NativeStatusDecision } from "./status-arbiter.js";
import { DELIVERY_QUEUES, subscribeDeliveryWork } from "../delivery-work-notifications.js";
const embeddedPostgresSupport = await getEmbeddedPostgresTestSupport();
const describeEmbeddedPostgres = embeddedPostgresSupport.supported ? describe : describe.skip;
@@ -946,6 +947,26 @@ describeEmbeddedPostgres("native run finalizer / status decision committer — a
expect(alreadyEmittedByCommittedDecision).toBe(true);
expect(agentTaskRunCalls(callsBefore)).toHaveLength(1);
});
it("signals durable accounting when native commitment or permanent failure makes a receipt actionable", async () => {
const settled = vi.fn();
const unsubscribe = subscribeDeliveryWork(db, DELIVERY_QUEUES.accounting, settled);
try {
const committed = await seedNativeRun();
await driveToCompleteResult(committed);
const beforeCommit = settled.mock.calls.length;
await finalizeNativeRun({ db, runId: committed.runId, workspaceFinalizeStatus: "succeeded", projectRunStatus: true });
expect((await db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, committed.runId)))[0].phase).toBe("committed");
expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, committed.runId)))[0].status).toBe("succeeded");
expect(settled.mock.calls.length).toBeGreaterThan(beforeCommit);
const failed = await seedNativeRun();
await driveToCompleteResult(failed);
const beforeFailure = settled.mock.calls.length;
await recordNativeFinalizationFailure({ db, runId: failed.runId, error: new Error("native_fixture_permanent_failure"), permanent: true });
expect((await db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, failed.runId)))[0].phase).toBe("terminal_failure");
expect(settled.mock.calls.length).toBeGreaterThan(beforeFailure);
} finally { unsubscribe(); }
});
});
@@ -64,6 +64,7 @@ import {
restoreNativeChatReviewPresentationInTransaction,
} from "./native-chat-review-presentation.js";
import { logger } from "../../middleware/logger.js";
import { DELIVERY_QUEUES, notifyDeliveryWork } from "../delivery-work-notifications.js";
import {
CHAT_RUN_PRESENTATION_AUTHORIZATION_REASON,
findHeartbeatRunCompletionComment,
@@ -418,6 +419,7 @@ async function recordRetryableFailure(input: {
? "native_workspace_sync_out_retry_exhausted"
: "native_finalization_retry_exhausted"
: input.failureCode;
if (phase === "terminal_failure") await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, input.run.companyId);
await tx
.update(nativeRunFinalizations)
.set({
@@ -660,69 +662,74 @@ async function projectCommittedRun(input: {
const projectedStatus = projectNativeTerminalRunStatus(
terminalState as "succeeded" | "failed" | "cancelled",
);
const now = new Date();
const [updatedRun] = await input.db
.update(heartbeatRuns)
.set({
executionStatusDeliveryId: randomUUID(),
status: projectedStatus,
finishedAt: input.run.finishedAt ?? now,
nativePhase: "committed",
nativePhaseUpdatedAt: now,
...(terminalState === "succeeded"
? {
error: null,
errorCode: null,
// Capture the row being updated, not the earlier admission read:
// cleanup may have recorded a new diagnostic in the meantime.
// Other result metadata and all physical-owner evidence stay put.
resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb)
|| ${recoveredExecutionFailureMetadata()}
|| jsonb_build_object('finalizationPhase', 'committed', 'failureCode', null, 'originalFailureCode', null, 'nextAttemptAt', null)`,
}
: { resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) || jsonb_build_object('finalizationPhase', 'committed', 'failureCode', null, 'originalFailureCode', null, 'nextAttemptAt', null)` }),
updatedAt: now,
})
.where(
const needsRepair = and(
eq(heartbeatRuns.id, input.run.id),
eq(heartbeatRuns.runtimeMode, "native"),
or(
inArray(heartbeatRuns.status, ["queued", "running", "failed"]),
and(
eq(heartbeatRuns.id, input.run.id),
eq(heartbeatRuns.runtimeMode, "native"),
eq(heartbeatRuns.status, "succeeded"),
or(
inArray(heartbeatRuns.status, ["queued", "running", "failed"]),
and(
eq(heartbeatRuns.status, "succeeded"),
or(
isNotNull(heartbeatRuns.error),
isNotNull(heartbeatRuns.errorCode),
isNull(heartbeatRuns.finishedAt),
sql`${heartbeatRuns.nativePhase} is distinct from 'committed'`,
sql`${heartbeatRuns.resultJson}->>'finalizationPhase' is distinct from 'committed'`,
sql`${heartbeatRuns.resultJson}->>'nextAttemptAt' is not null`,
sql`${heartbeatRuns.resultJson}->>'failureCode' is not null`,
sql`${heartbeatRuns.resultJson}->>'originalFailureCode' is not null`,
),
),
),
// Reconciliation revisits committed results periodically. Only repair
// a changed projection; rewriting an unchanged failed run would mint a
// fresh status delivery (and failure toast) on every sweep. Check the
// current row so concurrent replays cannot both queue the same repair.
or(
sql`${heartbeatRuns.status} is distinct from ${projectedStatus}`,
isNotNull(heartbeatRuns.error),
isNotNull(heartbeatRuns.errorCode),
isNull(heartbeatRuns.finishedAt),
sql`${heartbeatRuns.nativePhase} is distinct from 'committed'`,
sql`${heartbeatRuns.resultJson}->>'finalizationPhase' is distinct from 'committed'`,
sql`${heartbeatRuns.resultJson}->>'nextAttemptAt' is not null`,
sql`${heartbeatRuns.resultJson}->>'failureCode' is not null`,
sql`${heartbeatRuns.resultJson}->>'originalFailureCode' is not null`,
...(terminalState === "succeeded"
? [isNotNull(heartbeatRuns.error), isNotNull(heartbeatRuns.errorCode)]
: []),
),
nativeRunnerOwnershipNotHeldCondition(),
),
)
.returning();
),
// Reconciliation revisits committed results periodically. Only repair
// a changed projection; rewriting an unchanged failed run would mint a
// fresh status delivery (and failure toast) on every sweep. Check the
// current row so concurrent replays cannot both queue the same repair.
or(
sql`${heartbeatRuns.status} is distinct from ${projectedStatus}`,
isNull(heartbeatRuns.finishedAt),
sql`${heartbeatRuns.nativePhase} is distinct from 'committed'`,
sql`${heartbeatRuns.resultJson}->>'finalizationPhase' is distinct from 'committed'`,
sql`${heartbeatRuns.resultJson}->>'nextAttemptAt' is not null`,
sql`${heartbeatRuns.resultJson}->>'failureCode' is not null`,
sql`${heartbeatRuns.resultJson}->>'originalFailureCode' is not null`,
...(terminalState === "succeeded"
? [isNotNull(heartbeatRuns.error), isNotNull(heartbeatRuns.errorCode)]
: []),
),
nativeRunnerOwnershipNotHeldCondition(),
);
const [candidate] = await input.db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(needsRepair).limit(1);
if (!candidate) return;
const now = new Date();
const [updatedRun] = await input.db.transaction(async transaction => {
const tx = transaction as unknown as Db;
await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, input.run.companyId);
return tx
.update(heartbeatRuns)
.set({
executionStatusDeliveryId: randomUUID(),
status: projectedStatus,
finishedAt: input.run.finishedAt ?? now,
nativePhase: "committed",
nativePhaseUpdatedAt: now,
...(terminalState === "succeeded"
? {
error: null,
errorCode: null,
// Capture the row being updated, not the earlier admission read:
// cleanup may have recorded a new diagnostic in the meantime.
// Other result metadata and all physical-owner evidence stay put.
resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb)
|| ${recoveredExecutionFailureMetadata()}
|| jsonb_build_object('finalizationPhase', 'committed', 'failureCode', null, 'originalFailureCode', null, 'nextAttemptAt', null)`,
}
: { resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) || jsonb_build_object('finalizationPhase', 'committed', 'failureCode', null, 'originalFailureCode', null, 'nextAttemptAt', null)` }),
updatedAt: now,
})
.where(needsRepair)
.returning();
});
// Metadata repairs can preserve the terminal status. Only a genuine status
// transition should emit another terminal event. A failed run that leaves a
// pending retry still needs its first failure report.
@@ -1579,7 +1586,7 @@ export async function finalizeNativeRun(input: {
statusVersionAfter: Number(committed.issue.statusVersion),
workspaceFinalizeStatus: input.workspaceFinalizeStatus,
};
const [updatedRun] = await input.db
const projectRun = (tx: Db) => tx
.update(heartbeatRuns)
.set({
executionStatusDeliveryId: randomUUID(),
@@ -1610,6 +1617,13 @@ export async function finalizeNativeRun(input: {
})
.where(eq(heartbeatRuns.id, run.id))
.returning();
const [updatedRun] = input.projectRunStatus
? await input.db.transaction(async transaction => {
const tx = transaction as unknown as Db;
await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, run.companyId);
return projectRun(tx);
})
: await projectRun(input.db);
if (
input.projectRunStatus &&
!alreadyEmittedByCommittedDecision &&
@@ -79,12 +79,22 @@ import * as nativeSessionResume from "./native-session-resume.js";
import { buildNativeHeartbeatPreparationSpans } from "./native-run-trace.js";
import { NativeRunnerOwnershipUnverifiedError } from "./native-runner-ownership.js";
import type { AdapterRuntimeEvent } from "../../adapters/index.js";
import type { AdapterUsageCheckpoint } from "@paperclipai/adapter-utils";
vi.mock("node:fs", async (importOriginal) => {
const actual = await importOriginal<typeof import("node:fs")>();
return { ...actual, readFileSync: vi.fn(actual.readFileSync) };
});
const usageRecorder = vi.hoisted(() => ({
create: vi.fn<typeof import("../usage-receipts.js")["createRunUsageRecorder"]>(),
capture: vi.fn<(receipt: AdapterUsageCheckpoint) => Promise<AdapterUsageCheckpoint>>(),
}));
vi.mock("../usage-receipts.js", async (importOriginal) => ({
...await importOriginal<typeof import("../usage-receipts.js")>(),
createRunUsageRecorder: usageRecorder.create,
}));
const githubAccess = vi.hoisted(() => ({
activate: vi.fn((_binding: { runId: string }) => vi.fn()),
stop: vi.fn(async () => undefined),
@@ -376,6 +386,9 @@ import {
} from "./native-session-executor.js";
beforeEach(() => {
usageRecorder.capture.mockReset().mockImplementation(async receipt => receipt);
usageRecorder.create.mockReset().mockResolvedValue({ capture: usageRecorder.capture,
complete: vi.fn(async () => ({ complete: true })), persistFailure: vi.fn(async () => undefined) });
state.bundledPackManifestPath = null;
state.bundledRunnerBinary = null;
state.createAssignedMcpTools.mockReset();
@@ -6183,21 +6196,42 @@ describe("native terminal-turn accounting", () => {
});
await expect(executePaperclipNativeSession({ db: leaseDb(), execution: { ...execution, provider: { kind: "codex", model: "gpt-6-astra" } } as NativeExecutionInput, runnerInstanceId: "runner", onUsage })).rejects.toBe(failure);
expect(onUsage).toHaveBeenLastCalledWith(expect.objectContaining({ complete: true }));
expect(usageRecorder.create).not.toHaveBeenCalled();
});
it("clears provisional unpriced status when direct database persistence receives complete usage", async () => {
it("uses the canonical recorder for partial and completed usage when no onUsage callback is supplied", async () => {
const updates: Array<{ table: unknown; values: Record<string, unknown> }> = [];
state.execute.mockReset().mockImplementationOnce(async () => {
await accountingEvents.committed!({ eventType: "item.completed", turnId: "turn", emittedAt: new Date().toISOString(), payload: { kind: "usage", usage: { runDelta: { inputTokens: 10, outputTokens: 2, providerCostUsd: 0.25 } } } } as unknown as PrpEvent);
return { result: { summary: "cancelled" }, terminal: { runTerminalState: "cancelled" },
turnId: "turn", normalizedSessionId: "session", providerSessionId: null, driverKind: "test", driverVersion: "1", nativeEventCount: 1, highestContiguousSourceSeq: 1, usage: null };
});
await executePaperclipNativeSession({ db: leaseDb(execution, {}, {}, updates), execution, runnerInstanceId: "runner" });
const dialect = new PgDialect();
const writes = updates.filter(update => update.values.usageJson).map(update =>
JSON.parse(dialect.sqlToQuery(update.values.usageJson as SQL).params[0] as string) as Record<string, unknown>);
expect(writes[0]).toMatchObject({ accountingReceiptReady: false, costStatus: "unpriced" });
// Apply the same JSON merge as the database: omitted keys retain the old value.
expect(Object.assign({}, ...writes)).toMatchObject({ accountingReceiptReady: true, costStatus: null, costUsd: 0.25 });
const database = leaseDb(execution, {}, {}, updates);
await executePaperclipNativeSession({ db: database, execution, runnerInstanceId: "runner" });
expect(usageRecorder.create).toHaveBeenCalledExactlyOnceWith(database, {
companyId: execution.binding.companyId, runId: execution.binding.runId, adapterType: "paperclip_runner",
});
expect(usageRecorder.capture).toHaveBeenNthCalledWith(1, expect.objectContaining({
complete: false, costStatus: "unpriced", costUsd: 0.25, usageBasis: "per_run",
}));
expect(usageRecorder.capture).toHaveBeenLastCalledWith(expect.objectContaining({ complete: true, costUsd: 0.25 }));
expect(usageRecorder.capture.mock.calls.at(-1)![0].costStatus).toBeUndefined();
expect(updates.filter(update => update.values.usageJson)).toEqual([]);
});
it.each(["recorder_creation", "receipt_capture"] as const)("does not hide a failed canonical accounting %s", async stage => {
const failure = new Error(`accounting_${stage}_failed`);
if (stage === "recorder_creation") usageRecorder.create.mockRejectedValueOnce(failure);
else usageRecorder.capture.mockRejectedValueOnce(failure);
state.execute.mockReset().mockImplementationOnce(async () => {
await accountingEvents.committed!({ eventType: "item.completed", turnId: "turn", emittedAt: new Date().toISOString(),
payload: { kind: "usage", usage: { runDelta: { inputTokens: 10, outputTokens: 2, providerCostUsd: 0.25 } } },
} as unknown as PrpEvent);
throw new Error("Execution must not advance past failed accounting capture");
});
await expect(executePaperclipNativeSession({ db: leaseDb(), execution, runnerInstanceId: "runner" })).rejects.toBe(failure);
expect(usageRecorder.create).toHaveBeenCalledOnce();
if (stage === "recorder_creation") expect(usageRecorder.capture).not.toHaveBeenCalled();
else expect(usageRecorder.capture).toHaveBeenCalledOnce();
});
it.each(["stale_warm_zero", "partial", "fresh", "fresh_zero", "new_turn", "later_partial", "merged_partial"])("requires fresh complete usage at native completion (%s)", async scenario => {
@@ -1,3 +1,4 @@
import { createRunUsageRecorder } from "../usage-receipts.js";
import { applyObservabilityEnvironment, type ProviderObservabilityLaunch } from "@paperclipai/adapter-utils/agent-observability";
import { museRunnerBroker } from "../muse-runner-broker.js";
import { nativeCaseHandoffResult } from "./native-case-handoff.js";
@@ -8217,6 +8218,7 @@ async function executePaperclipNativeSessionWithinScope(
)).orderBy(heartbeatRunEvents.seq);
for (const row of history) turnAccounting.observe(record(row.payload).prpEvent as PrpEvent);
}
let fallbackUsageRecorder: ReturnType<typeof createRunUsageRecorder> | undefined;
const persistAccountingUsage = async (usage: Record<string, unknown> | null, complete: boolean) => {
const accountingUsage = normalizeNativeUsage(usage, { inputIncludesCacheReads: input.execution.provider.kind === "codex" });
const accountingCost = accountingUsage ? nativeUsageCostUsd(usage, input.execution.provider) : undefined;
@@ -8226,15 +8228,14 @@ async function executePaperclipNativeSessionWithinScope(
: "unpriced" as const;
const billing = resolveNativeBilling(input.execution.provider, input.runnerEnvironment, input.billingIdentity);
const costUsdExact = typeof usage?.accountingCostUsdExact === "string" ? usage.accountingCostUsdExact : undefined;
await input.onUsage?.({ usage: accountingUsage, ...billing, complete, usageBasis: "per_run", costUsdExact,
model: input.execution.provider.model ?? "unknown", costUsd: accountingCost ?? null, costStatus });
if (!input.onUsage) await input.db.update(heartbeatRuns).set({
costAccountingPending: true,
usageJson: sql`coalesce(${heartbeatRuns.usageJson}, '{}'::jsonb) || ${JSON.stringify({
...accountingUsage, ...billing, accountingReceiptReady: complete, model: input.execution.provider.model ?? "unknown",
costUsd: accountingCost ?? null, costUsdExact: costUsdExact ?? null, costStatus: costStatus ?? null, usageSource: "per_run",
})}::jsonb`,
}).where(and(eq(heartbeatRuns.id, input.execution.binding.runId), eq(heartbeatRuns.companyId, input.execution.binding.companyId), isNull(heartbeatRuns.costAccountedAt)));
const checkpoint: AdapterUsageCheckpoint = { usage: accountingUsage, ...billing, complete, usageBasis: "per_run", costUsdExact,
model: input.execution.provider.model ?? "unknown", costUsd: accountingCost ?? null, costStatus };
if (input.onUsage) await input.onUsage(checkpoint);
else {
fallbackUsageRecorder ??= createRunUsageRecorder(input.db, { companyId: input.execution.binding.companyId,
runId: input.execution.binding.runId, adapterType: "paperclip_runner" });
await (await fallbackUsageRecorder).capture(checkpoint);
}
};
// appendEvent awaits this receipt before the runtime may commit completeRun.
// Replayed events must rebuild it too after a stop between event and receipt.
@@ -17,6 +17,8 @@ import { preserveNativeFailedWorkspace } from "./native-failed-workspace-recover
import { settleStopOnlyCleanup } from "../sandbox-stop-and-retain.js";
import { createRunDispatch } from "../../modules/run-dispatch/index.js";
import { drainHeartbeatRunsToQuiescence } from "../../__tests__/helpers/drain-heartbeat-runs.js";
import { reserveRunBudget } from "../budget-reservations.js";
import { accountRunCost } from "../run-cost-accounting.js";
const providerExecute = vi.hoisted(() => vi.fn(async () => ({ exitCode: 0, signal: null, timedOut: false,
errorMessage: null, summary: "Bounded reconstruction completed.", provider: "test", model: "test" })));
vi.mock("../../adapters/index.js", async () => ({
@@ -173,6 +175,7 @@ describe("reconstruction after confirmed native workspace loss", () => {
expect(successorId).not.toBe("undefined");
await db.update(heartbeatRuns).set({ status: "running" }).where(eq(heartbeatRuns.id, successorId));
await db.update(issues).set({ status: "in_progress", executionRunId: successorId }).where(eq(issues.id, f.issueId));
const reservation = await reserveRunBudget(db, f.companyId, successorId, null);
if (kind === "changed_evidence") {
const [action] = await db.select().from(issueRecoveryActions).where(eq(issueRecoveryActions.id, later.actionId!));
await db.update(issueRecoveryActions).set({ evidence: { ...action.evidence, executionReconciliation: {
@@ -182,7 +185,7 @@ describe("reconstruction after confirmed native workspace loss", () => {
if (kind === "extra_reconciled") await interveningCancelledRun(f);
const dispatch = vi.fn(async () => "provider entered");
const result = await createRunDispatch(db).dispatchResolvedInteractionIfCurrent({ companyId: f.companyId,
runId: successorId, expectedStatus: "running", dispatch });
runId: successorId, reservationId: reservation.id, expectedStatus: "running", dispatch });
expect(result.dispatched).toBe(kind === "unchanged");
expect(dispatch).toHaveBeenCalledTimes(kind === "unchanged" ? 1 : 0);
});
@@ -286,16 +289,17 @@ describe("reconstruction after confirmed native workspace loss", () => {
expect(successorId).not.toBe("undefined");
await db.update(heartbeatRuns).set({ status: "running" }).where(eq(heartbeatRuns.id, successorId));
await db.update(issues).set({ status: "in_progress", executionRunId: successorId }).where(eq(issues.id, f.issueId));
const reservation = await reserveRunBudget(db, f.companyId, successorId, null);
if (kind === "blocked") await db.update(issues).set({ status: "blocked" }).where(eq(issues.id, f.issueId));
if (kind === "attachment") await db.delete(issueAttachments).where(eq(issueAttachments.id, f.attachmentId));
if (kind === "paused") await db.update(agents).set({ status: "paused" }).where(eq(agents.id, f.agentId));
if (kind === "forged_successor") await db.update(issueRecoveryActions).set({ evidence: { ...action.evidence, continuationRunId: randomUUID() } }).where(eq(issueRecoveryActions.id, f.actionId));
const dispatch = vi.fn(async () => "provider entered");
const outcome = await createRunDispatch(db).dispatchResolvedInteractionIfCurrent({ companyId: f.companyId,
runId: successorId, expectedStatus: "running", dispatch });
runId: successorId, reservationId: reservation.id, expectedStatus: "running", dispatch });
expect(outcome.dispatched).toBe(kind === "unchanged");
expect(dispatch).toHaveBeenCalledTimes(kind === "unchanged" ? 1 : 0);
if (!outcome.dispatched) expect(outcome.cancellation).toMatchObject({ outcome: "cancelled" });
if (!outcome.dispatched) expect(outcome).toMatchObject({ cancellation: { outcome: "cancelled" } });
});
it.each(["unchanged", "attachment", "blocked", "paused", "forged_wake", "foreign_parent", "unrelated_run", "uncertain_parent", "new_recovery"] as const)("keeps bounded reconstruction retries under current authority: %s", async kind => {
const f = await seed();
@@ -314,6 +318,7 @@ describe("reconstruction after confirmed native workspace loss", () => {
if (retry.outcome !== "scheduled" || !retry.run) throw new Error("Expected a bounded bootstrap retry");
await db.update(heartbeatRuns).set({ status: "running" }).where(eq(heartbeatRuns.id, retry.run.id));
await db.update(issues).set({ executionRunId: retry.run.id }).where(eq(issues.id, f.issueId));
const reservation = await reserveRunBudget(db, f.companyId, retry.run.id, null);
if (kind === "attachment") await db.delete(issueAttachments).where(eq(issueAttachments.id, f.attachmentId));
if (kind === "blocked") await db.update(issues).set({ status: "blocked" }).where(eq(issues.id, f.issueId));
if (kind === "paused") await db.update(agents).set({ status: "paused" }).where(eq(agents.id, f.agentId));
@@ -330,21 +335,23 @@ describe("reconstruction after confirmed native workspace loss", () => {
fingerprint: randomUUID(), evidence: { runId: firstId }, nextAction: "Reconcile the new execution", status: "active" });
const dispatch = vi.fn(async () => "provider entered");
const outcome = await createRunDispatch(db).dispatchResolvedInteractionIfCurrent({ companyId: f.companyId,
runId: retry.run.id, expectedStatus: "running", dispatch });
runId: retry.run.id, reservationId: reservation.id, expectedStatus: "running", dispatch });
expect(outcome.dispatched).toBe(kind === "unchanged");
expect(dispatch).toHaveBeenCalledTimes(kind === "unchanged" ? 1 : 0);
if (kind === "unchanged") {
expect(retry.run.scheduledRetryAttempt).toBe(1);
await db.update(heartbeatRuns).set({ status: "failed", finishedAt: new Date(), errorCode: "adapter_failed",
resultJson: { executionRecovery: { kind: "bootstrap", providerWorkStarted: false } } }).where(eq(heartbeatRuns.id, retry.run.id));
await accountRunCost(db, retry.run.id);
const second = await heartbeat.scheduleBoundedRetry(retry.run.id, { now: new Date(), random: () => 0, delayMs: 0 });
expect(second.outcome).toBe("scheduled");
if (second.outcome !== "scheduled" || !second.run) throw new Error("Expected the second bounded retry");
expect(second.run.scheduledRetryAttempt).toBe(2);
await db.update(heartbeatRuns).set({ status: "running" }).where(eq(heartbeatRuns.id, second.run.id));
await db.update(issues).set({ executionRunId: second.run.id }).where(eq(issues.id, f.issueId));
const secondReservation = await reserveRunBudget(db, f.companyId, second.run.id, null);
expect(await createRunDispatch(db).dispatchResolvedInteractionIfCurrent({ companyId: f.companyId,
runId: second.run.id, expectedStatus: "running", dispatch })).toMatchObject({ dispatched: true });
runId: second.run.id, reservationId: secondReservation.id, expectedStatus: "running", dispatch })).toMatchObject({ dispatched: true });
await db.update(heartbeatRuns).set({ status: "failed", finishedAt: new Date(),
resultJson: { executionRecovery: { kind: "bootstrap", providerWorkStarted: false } } }).where(eq(heartbeatRuns.id, second.run.id));
expect(await heartbeat.scheduleBoundedRetry(second.run.id)).toMatchObject({ outcome: "retry_exhausted" });
@@ -57,6 +57,7 @@ import {
import { emitAgentTaskRun } from "../agent-task-run-telemetry.js";
import { materializeNativeProviderCapacityRetry } from "./native-provider-capacity-retry.js";
import { NATIVE_PROVIDER_OVERLOADED_CODE } from "./native-provider-failure.js";
import { DELIVERY_QUEUES, notifyDeliveryWork } from "../delivery-work-notifications.js";
export class NativeStatusRaceError extends Error {
readonly code = "native_status_race" as const;
@@ -1056,6 +1057,7 @@ async function materializeDecisionEffect(input: {
),
);
}
await notifyDeliveryWork(input.tx, DELIVERY_QUEUES.accounting, input.companyId);
const [run] = await input.tx
.update(heartbeatRuns)
.set({
@@ -1845,6 +1847,7 @@ export async function commitNativeStatusDecision(input: {
const finalizationError = input.decision.effects.find(
effect => effect.kind === "record_finalization_error",
);
await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, input.companyId);
await tx
.update(nativeRunFinalizations)
.set({
@@ -1,6 +1,7 @@
import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notifications.js";
import { randomUUID } from "node:crypto";
import { and, eq, lt, sql } from "drizzle-orm";
import { agents, budgetPolicies, budgetReservations, companies, companyPromptInjectionPrevention,
import { agents, budgetReservations, companies, companyPromptInjectionPrevention,
promptInjectionInvocations, companySecrets, heartbeatRuns, runIdentityContexts, issues, projects, type Db } from "@paperclipai/db";
import { PROMPT_INJECTION_DEFAULT_MODEL, PROMPT_INJECTION_POLICY_VERSION, decisionProviderForConnection,
updatePromptInjectionPreventionSchema, centsToUnits, unitsToCents,
@@ -9,7 +10,8 @@ import { PROMPT_INJECTION_DEFAULT_MODEL, PROMPT_INJECTION_POLICY_VERSION, decisi
import { HttpError, forbidden } from "../errors.js";
import { aiConnectionService } from "./ai-connections.js";
import { withAccountingTransaction } from "./accounting-transaction.js";
import { budgetPoliciesForRun, budgetService, budgetServiceInTransaction, computeObservedSpend, type BudgetServiceHooks } from "./budgets.js";
import { withBudgetSpan } from "./budget-tracing.js";
import { budgetServiceInTransaction, type BudgetServiceHooks } from "./budgets.js";
import { createCostEventInTransaction } from "./costs.js";
import { logActivity } from "./activity-log.js";
import { runPromptInjectionProvider, promptInjectionInput, promptInjectionRequestFits, type PromptInjectionProviderOutcome } from "./prompt-injection-provider.js";
@@ -102,7 +104,7 @@ export function promptInjectionPreventionService(db: Db, options: {
return { config, connection, userId, agentId, runId, identityContextId, issueId, projectId, actorType: context.actorType, actorId: context.actorId };
}
async function settle(invocation: typeof promptInjectionInvocations.$inferSelect, outcome: PromptInjectionProviderOutcome) {
await withAccountingTransaction(db, invocation.companyId, async (tx, publications) => {
await withBudgetSpan("budget.accounting.settlement", () => withAccountingTransaction(db, invocation.companyId, async (tx, publications) => {
const [current] = await tx.select().from(promptInjectionInvocations).where(and(eq(promptInjectionInvocations.companyId, invocation.companyId), eq(promptInjectionInvocations.id, invocation.id))).for("update");
if (!current || current.status !== "running") return;
const receipt = outcome.receipt;
@@ -125,6 +127,7 @@ export function promptInjectionPreventionService(db: Db, options: {
providerRequestId: receipt.providerRequestId, pricingProvenance: receipt.pricingProvenance, occurredAt: invocation.startedAt,
}, publications, { actorType: invocation.actorType as "system" | "user" | "agent", actorId: invocation.actorId, agentId: links.agentId });
const finishedAt = new Date();
await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, invocation.companyId);
await tx.update(promptInjectionInvocations).set({ status: outcome.verdict ? "succeeded" : receipt.costStatus === "unpriced" ? "unknown" : "failed",
violation: outcome.verdict?.violation ?? null, findings: outcome.verdict?.findings ?? null, errorCode: outcome.errorCode ?? null, costEventId: event.id, providerRequestId: receipt.providerRequestId,
inputTokens: receipt.inputTokens, outputTokens: receipt.outputTokens, finishedAt, durationMs: Math.max(0, finishedAt.getTime() - invocation.startedAt.getTime()),
@@ -132,8 +135,7 @@ export function promptInjectionPreventionService(db: Db, options: {
if (receipt.costStatus !== "unpriced") await tx.update(budgetReservations).set({ state: outcome.noProviderWork ? "released" : "settled", settledAt: finishedAt }).where(and(
eq(budgetReservations.companyId, invocation.companyId), eq(budgetReservations.promptInjectionInvocationId, invocation.id), eq(budgetReservations.state, "held"),
));
});
await budgetService(db, options.budgetHooks).deliverPendingEnforcement(invocation.companyId);
}));
}
async function classify(context: PromptInjectionContext, content: string, signal?: AbortSignal): Promise<PromptInjectionModelOutcome & { usage?: DecisionUsage }> {
signal?.throwIfAborted();
@@ -171,23 +173,17 @@ export function promptInjectionPreventionService(db: Db, options: {
if (!currentSecret || currentSecret.status !== "active" || currentSecret.version !== credentialSnapshot.version)
return { reason: "connection_unavailable" as const };
const budgets = budgetServiceInTransaction(tx, publications);
const block = await budgets.getInvocationBlock(context.companyId, current.agentId, { projectId: current.projectId });
const { block, policies, observations } = await budgets.evaluateAdmission(context.companyId, current.agentId, { projectId: current.projectId });
const [company] = await tx.select({ status: companies.status }).from(companies).where(eq(companies.id, context.companyId));
const [agent] = current.agentId ? await tx.select({ status: agents.status }).from(agents).where(and(eq(agents.id, current.agentId), eq(agents.companyId, context.companyId))) : [];
if (block || company.status !== "active" || (current.agentId && (!agent || agent.status === "paused" || agent.status === "terminated"))) return { reason: "budget_blocked" as const };
const policies = await tx.select().from(budgetPolicies).where(and(eq(budgetPolicies.companyId, context.companyId), eq(budgetPolicies.isActive, true),
eq(budgetPolicies.hardStopEnabled, true), eq(budgetPolicies.metric, "billed_cents"), sql`${budgetPolicies.amount} > 0`, budgetPoliciesForRun(context.companyId, current.agentId, current.projectId)));
const amount = policies.reduce((max, p) => { const n = centsToUnits(p.reservationCents); return n > max ? n : max; }, 0n);
for (const policy of policies) {
const [held] = await tx.select({ total: sql<string>`coalesce(sum(${budgetReservations.amountCents}),0)::text` }).from(budgetReservations).where(and(
eq(budgetReservations.companyId, context.companyId), eq(budgetReservations.state, "held"),
policy.scopeType === "agent" ? eq(budgetReservations.agentId, current.agentId!) : undefined,
policy.scopeType === "project" ? eq(budgetReservations.projectId, current.projectId!) : undefined,
));
const committed = centsToUnits((await computeObservedSpend(tx, policy)).totalExact) + centsToUnits(held.total);
for (const [index, policy] of policies.entries()) {
const committed = centsToUnits(observations[index].ledgerCents) + centsToUnits(observations[index].effectiveHeldCents);
if (committed >= centsToUnits(policy.amount) || committed + amount > centsToUnits(policy.amount)) return { reason: "budget_blocked" as const };
}
signal?.throwIfAborted();
await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, context.companyId);
const [invocation] = await tx.insert(promptInjectionInvocations).values({ id: randomUUID(), companyId: context.companyId, source: context.source, toolInvocationId: context.toolInvocationId ?? null, toolName: context.toolName,
actorType: current.actorType, actorId: current.actorId, responsibleUserId: current.userId, agentId: current.agentId,
issueId: current.issueId, projectId: current.projectId, runId: current.runId, identityContextId: current.identityContextId,
@@ -0,0 +1,114 @@
import { randomUUID } from "node:crypto";
import { promises as fs } from "node:fs";
import os from "node:os";
import path from "node:path";
import { and, eq, sql } from "drizzle-orm";
import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest";
import { agents, companies, costEvents, createDb, heartbeatRuns, promptInjectionInvocations, type Db } from "@paperclipai/db";
import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../__tests__/helpers/embedded-postgres.js";
import { accountRunCost, createCostAccountingReconciler, readCostAccountingWork } from "./run-cost-accounting.js";
const support = await getEmbeddedPostgresTestSupport();
(support.supported ? describe : describe.skip)("run cost accounting work", () => {
let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>>;
let db: Db;
let directory: string;
beforeAll(async () => {
database = await startEmbeddedPostgresTestDatabase("run-cost-work-");
db = createDb(database.connectionString);
directory = await fs.mkdtemp(path.join(os.tmpdir(), "run-cost-work-"));
}, 60_000);
afterEach(async () => { await db.transaction(async tx => {
await tx.execute(sql`SET LOCAL client_min_messages = warning`);
await tx.execute(sql`TRUNCATE companies CASCADE`);
}); });
afterAll(async () => { if (directory) await fs.rm(directory, { force: true, recursive: true }); await database?.cleanup(); });
async function seed(change: Partial<typeof heartbeatRuns.$inferInsert> = {}) {
const companyId = randomUUID(), agentId = randomUUID(), runId = randomUUID();
await db.insert(companies).values({ id: companyId, name: "Accounting", issuePrefix: "ACCT" });
await db.insert(agents).values({ id: agentId, companyId, name: "Accounting", role: "engineer", status: "idle", adapterType: "process" });
await db.insert(heartbeatRuns).values({ id: runId, companyId, agentId, invocationSource: "on_demand", status: "succeeded",
costAccountingPending: true, finishedAt: new Date(),
usageJson: { accountingReceiptReady: true, provider: "test", model: "test", billingType: "metered_api", costUsd: 0.01 }, ...change });
return { companyId, agentId, runId };
}
it("returns from settled or running calls without opening an accounting transaction", async () => {
const { runId } = await seed({ costAccountingPending: false });
const transaction = vi.fn(db.transaction.bind(db));
const observed = new Proxy(db, { get(target, property) {
if (property === "transaction") return transaction;
const value = Reflect.get(target, property);
return typeof value === "function" ? value.bind(target) : value;
} }) as Db;
expect(await accountRunCost(observed, runId)).toBe(false);
await db.update(heartbeatRuns).set({ status: "running", costAccountingPending: true }).where(eq(heartbeatRuns.id, runId));
expect(await accountRunCost(observed, runId)).toBe(false);
expect(transaction).not.toHaveBeenCalled();
});
it("settles once without delivering cancellation in the caller", async () => {
const { companyId, runId } = await seed();
const cancelWorkForScope = vi.fn(async () => { throw new Error("Cancellation must be asynchronous"); });
expect(await accountRunCost(db, runId, { cancelWorkForScope })).toBe(true);
expect(await accountRunCost(db, runId, { cancelWorkForScope })).toBe(false);
const events = await db.select().from(costEvents).where(and(eq(costEvents.companyId, companyId), eq(costEvents.heartbeatRunId, runId)));
expect(events).toHaveLength(1);
expect(events[0].costCents).toBe(1);
expect(cancelWorkForScope).not.toHaveBeenCalled();
expect(await readCostAccountingWork(db, Date.now(), directory)).toEqual({ pending: false, nextRunAt: null });
});
it("repairs a stale pending marker without creating another charge", async () => {
const { runId } = await seed({ costAccountedAt: new Date(), usageJson: { accountingReceiptReady: false } });
expect((await readCostAccountingWork(db, Date.now(), directory)).pending).toBe(true);
expect(await accountRunCost(db, runId)).toBe(false);
const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId));
expect(run.costAccountingPending).toBe(false);
expect(await db.select().from(costEvents)).toHaveLength(0);
});
it("waits for a missing final receipt without scheduling unchanged debt", async () => {
const { runId } = await seed({ usageJson: { accountingReceiptReady: false, inputTokens: 5 } });
expect(await readCostAccountingWork(db, Date.now(), directory)).toEqual({ pending: false, nextRunAt: null });
const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId));
expect(run.costAccountingPending).toBe(true);
await db.update(heartbeatRuns).set({ usageJson: { accountingReceiptReady: true, costUsd: 0.01 } }).where(eq(heartbeatRuns.id, runId));
expect((await readCostAccountingWork(db, Date.now(), directory)).pending).toBe(true);
});
it("reconstructs the actual interrupted-service deadline and disarms after settlement", async () => {
const { companyId } = await seed({ costAccountingPending: false });
const startedAt = new Date();
const [invocation] = await db.insert(promptInjectionInvocations).values({ companyId, source: "settings_test", actorType: "user", actorId: "test",
connectionId: randomUUID(), grantId: randomUUID(), policyVersion: "tool-result-v1", model: "openai/gpt-oss-safeguard-20b", startedAt }).returning();
const work = await readCostAccountingWork(db, startedAt.getTime(), directory);
expect(work).toEqual({ pending: false, nextRunAt: startedAt.getTime() + 120_001 });
expect((await readCostAccountingWork(db, startedAt.getTime() + 120_001, directory)).pending).toBe(true);
await db.update(promptInjectionInvocations).set({ status: "unknown" }).where(eq(promptInjectionInvocations.id, invocation.id));
expect(await readCostAccountingWork(db, startedAt.getTime() + 120_001, directory)).toEqual({ pending: false, nextRunAt: null });
});
it("shares service and run recovery across overlapping signals", async () => {
let release!: () => void;
const recover = vi.fn(() => new Promise<void>(resolve => { release = resolve; }));
const reconcile = createCostAccountingReconciler(db, {}, recover);
const first = reconcile(), second = reconcile();
expect(first).toBe(second);
expect(recover).toHaveBeenCalledOnce();
release();
await first;
expect(recover).toHaveBeenCalledOnce();
});
it("recovers run receipts even when an unrelated service invocation cannot settle", async () => {
const { runId } = await seed();
const reconcile = createCostAccountingReconciler(db, {}, async () => { throw new Error("Unrelated service receipt failed"); });
await expect(reconcile()).rejects.toThrow("Unrelated service receipt failed");
const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId));
expect(run.costAccountedAt).not.toBeNull();
expect(run.costAccountingPending).toBe(false);
expect(await db.select().from(costEvents)).toHaveLength(1);
});
});
+65 -27
View File
@@ -1,10 +1,11 @@
import { usdToUnits, unitsToCents } from "@paperclipai/shared";
import { and, asc, eq, inArray, sql } from "drizzle-orm";
import { accountingRuntimeBaselines, budgetReservations, agentRuntimeState, agents, costEvents, heartbeatRuns, nativeRunFinalizations, issues, projects, type Db } from "@paperclipai/db";
import { and, asc, eq, inArray, or, sql } from "drizzle-orm";
import { accountingRuntimeBaselines, budgetReservations, agentRuntimeState, agents, costEvents, heartbeatRuns, nativeRunFinalizations, issues, projects, decisionInvocations, promptInjectionInvocations, type Db } from "@paperclipai/db";
import { withAccountingTransaction } from "./accounting-transaction.js";
import { withBudgetSpan } from "./budget-tracing.js";
import { conflict } from "../errors.js";
import { createCostEventInTransaction } from "./costs.js";
import { budgetService, deliverBudgetEnforcement, type BudgetServiceHooks } from "./budgets.js";
import { budgetServiceInTransaction, type BudgetServiceHooks } from "./budgets.js";
import { logger } from "../middleware/logger.js";
import { promises as fs } from "node:fs";
import { replayUsageReceipts, indexPendingUsageReceipts, recoverPendingRunUsageReceipts, type UsageReceiptIndex } from "./usage-receipts.js";
@@ -21,13 +22,13 @@ export async function accountRunCost(db: Db, runId: string, hooks: BudgetService
pending: heartbeatRuns.costAccountingPending, accountedAt: heartbeatRuns.costAccountedAt,
issueId: sql<string | null>`${heartbeatRuns.usageJson}->'ledgerScope'->>'issueId'`,
}).from(heartbeatRuns).where(eq(heartbeatRuns.id, runId));
if (!identity) return false;
// Finalization can call twice. Avoid disk work for settled or active runs,
// but retain the locked recheck, stale-marker cleanup, and stop delivery.
if (!identity || !identity.pending || !terminalStatuses.includes(identity.status)) return false;
// Finalization can call twice. Settled or active runs owe no accounting work.
// Pending acknowledgements still receive the locked stale-marker cleanup.
const index = identity.pending && !identity.accountedAt && terminalStatuses.includes(identity.status)
? receiptIndex ?? await indexPendingUsageReceipts() : undefined;
let recoveredFiles: string[] = [];
const accounted = await withAccountingTransaction(db, identity.companyId, async (tx, publications) => {
const accounted = await withBudgetSpan("budget.accounting.settlement", () => withAccountingTransaction(db, identity.companyId, async (tx, publications) => {
// Native recovery owns the issue before updating its run. Acquire the
// charge's issue FK lock first so inserting it cannot reverse that order.
const issueId = text(identity.issueId);
@@ -38,6 +39,8 @@ export async function accountRunCost(db: Db, runId: string, hooks: BudgetService
if (!run?.costAccountingPending || !terminalStatuses.includes(run.status)) return false;
if (run.costAccountedAt) {
await tx.update(heartbeatRuns).set({ costAccountingPending: false }).where(eq(heartbeatRuns.id, run.id));
await budgetServiceInTransaction(tx, publications).evaluateAffectedScopes(run.companyId, run.agentId,
text(object(object(run.usageJson).ledgerScope).projectId));
return false;
}
// The bounded global replay can leave a newer receipt behind an older
@@ -69,7 +72,11 @@ export async function accountRunCost(db: Db, runId: string, hooks: BudgetService
await tx.update(heartbeatRuns).set({ costAccountingPending: false, costAccountedAt: new Date() }).where(eq(heartbeatRuns.id, run.id));
await tx.update(budgetReservations).set({ state: preProviderFailure ? "released" : "settled", settledAt: new Date() })
.where(and(eq(budgetReservations.companyId, run.companyId), eq(budgetReservations.runId, run.id), eq(budgetReservations.state, "held")));
if (preProviderFailure) return true;
if (preProviderFailure) {
await budgetServiceInTransaction(tx, publications).evaluateAffectedScopes(run.companyId, run.agentId,
text(object(object(run.usageJson).ledgerScope).projectId));
return true;
}
const usage = object(run.usageJson);
const scope = object(usage.ledgerScope);
// Deleted attribution targets become unallocated; a deletion must never
@@ -121,11 +128,12 @@ export async function accountRunCost(db: Db, runId: string, hooks: BudgetService
...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, undefined, { deferBudgetEvaluation: true });
}
} else {
await createCostEventInTransaction(tx, run.companyId, receipt, publications);
await createCostEventInTransaction(tx, run.companyId, receipt, publications, undefined, { deferBudgetEvaluation: true });
}
await budgetServiceInTransaction(tx, publications).evaluateAffectedScopes(run.companyId, run.agentId, text(scope.projectId));
const [event] = await tx.select({
inputTokens: sql<string>`coalesce(sum(${costEvents.inputTokens}),0)::text`,
cachedInputTokens: sql<string>`coalesce(sum(${costEvents.cachedInputTokens}),0)::text`,
@@ -145,26 +153,52 @@ export async function accountRunCost(db: Db, runId: string, hooks: BudgetService
}).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));
return true;
});
}));
for (const file of recoveredFiles) {
try { await fs.rm(file, { force: true }); }
catch (error) { logger.warn({ err: error, runId }, "Accounting committed; recovered receipt cleanup will retry"); }
}
await deliverBudgetEnforcement(db, hooks, identity.companyId);
return accounted;
}
export async function reconcileRunCosts(db: Db, hooks: BudgetServiceHooks = {}) {
const pending = await db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(and(
function actionableRunAccounting() {
return and(
eq(heartbeatRuns.costAccountingPending, true), inArray(heartbeatRuns.status, terminalStatuses),
sql`${heartbeatRuns.usageJson}->>'accountingCaptureFailed' is distinct from 'true'`,
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);
or(sql`${heartbeatRuns.costAccountedAt} is not null`, and(
sql`${heartbeatRuns.usageJson}->>'accountingCaptureFailed' is distinct from 'true'`,
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')`,
)),
);
}
/** Only unfinished, actionable debt warrants retries. Incomplete provider
* receipts retain their accounting/idle hold and wake when evidence arrives. */
export async function readCostAccountingWork(db: Db, now = Date.now(), receiptDirectory?: string) {
const [pending] = await db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(actionableRunAccounting()).limit(1);
const [deadline] = await db.execute<{ at: Date | string | null }>(sql`
select min(started_at) + interval '120 seconds 1 millisecond' as at from (
select ${decisionInvocations.startedAt} as started_at from ${decisionInvocations}
where ${decisionInvocations.status} = 'running'
union all
select ${promptInjectionInvocations.startedAt} from ${promptInjectionInvocations}
where ${promptInjectionInvocations.status} = 'running'
) running_service_invocations`);
const nextRunAt = deadline?.at == null ? null : new Date(deadline.at).getTime();
const receipts = await indexPendingUsageReceipts(receiptDirectory);
return {
pending: !!pending || receipts.size > 0 || (nextRunAt !== null && nextRunAt <= now),
nextRunAt,
};
}
export async function reconcileRunCosts(db: Db, hooks: BudgetServiceHooks = {}) {
const pending = await db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(actionableRunAccounting())
.orderBy(asc(heartbeatRuns.updatedAt), asc(heartbeatRuns.id)).limit(100);
let accounted = 0;
const receiptIndex = pending.length > 0 ? await indexPendingUsageReceipts() : undefined;
for (const run of pending) {
@@ -176,20 +210,24 @@ export async function reconcileRunCosts(db: Db, hooks: BudgetServiceHooks = {})
logger.error({ err: error, runId: run.id }, "Run accounting remains pending for retry");
}
}
await budgetService(db, hooks).reconcilePolicies();
return { scanned: pending.length, accounted };
}
/** Share one full sweep across overlapping scheduler ticks. Do not enqueue
* another sweep: the next tick retries after either success or failure. */
export function createCostAccountingReconciler(db: Db, hooks: BudgetServiceHooks = {}) {
/** Share recovery across overlapping signals. The owning delivery worker
* schedules retries only while a durable receipt or actionable debt remains. */
export function createCostAccountingReconciler(db: Db, hooks: BudgetServiceHooks = {}, recoverServices?: () => Promise<void>) {
let inFlight: Promise<Awaited<ReturnType<typeof reconcileRunCosts>>> | undefined;
return () => {
if (!inFlight) {
inFlight = (async () => {
try {
let serviceRecoveryFailure: { error: unknown } | undefined;
try { await recoverServices?.(); }
catch (error) { serviceRecoveryFailure = { error }; }
await replayUsageReceipts(db);
return await reconcileRunCosts(db, hooks);
const result = await reconcileRunCosts(db, hooks);
if (serviceRecoveryFailure) throw serviceRecoveryFailure.error;
return result;
} finally {
inFlight = undefined;
}
+56 -2
View File
@@ -188,6 +188,11 @@ function sha256(value: unknown): string {
return createHash("sha256").update(stableStringify(value)).digest("hex");
}
type ConsumedRateAdmission = { policyId: string; versionHash: string };
function ratePolicyVersionHash(policy: typeof toolPolicies.$inferSelect): string {
return sha256({ ...policy, createdAt: policy.createdAt.toISOString(), updatedAt: policy.updatedAt.toISOString() });
}
function isoDateOrNull(value: unknown): string | null {
if (!value) return null;
if (value instanceof Date) return Number.isNaN(value.getTime()) ? null : value.toISOString();
@@ -1336,6 +1341,32 @@ export function toolAccessPolicyService(db: Db) {
});
}
async function admittedRateLimitAdmissions(ctx: ToolAccessContext, invocationId: string): Promise<ConsumedRateAdmission[]> {
const [invocation] = await db.select({
actorType: toolInvocations.actorType, actorId: toolInvocations.actorId, agentId: toolInvocations.agentId,
runId: toolInvocations.runId, issueId: toolInvocations.issueId, gatewayId: toolInvocations.gatewayId,
applicationId: toolInvocations.applicationId, connectionId: toolInvocations.connectionId,
catalogEntryId: toolInvocations.catalogEntryId, toolName: toolInvocations.toolName,
providerType: toolInvocations.providerType, applicationKey: toolInvocations.applicationKey,
upstreamToolName: toolInvocations.upstreamToolName, argumentsHash: toolInvocations.argumentsHash,
status: toolInvocations.status, policyDecision: toolInvocations.policyDecision, approvalState: toolInvocations.approvalState,
policyExplanation: toolInvocations.policyExplanation,
}).from(toolInvocations).where(and(eq(toolInvocations.companyId, ctx.companyId), eq(toolInvocations.id, invocationId)));
const authorized = invocation && invocation.status === "executing"
&& (invocation.policyDecision === "allow" || (invocation.policyDecision === "require_approval" && invocation.approvalState === "approved"))
&& invocation.actorType === ctx.actorType && invocation.actorId === ctx.actorId && invocation.agentId === ctx.agentId
&& invocation.runId === ctx.heartbeatRunId && invocation.issueId === ctx.issueId && invocation.gatewayId === ctx.gatewayId
&& invocation.applicationId === ctx.applicationId && invocation.connectionId === ctx.connectionId
&& invocation.catalogEntryId === ctx.catalogEntryId && invocation.toolName === ctx.toolName
&& invocation.providerType === ctx.providerType && invocation.applicationKey === ctx.applicationKey
&& invocation.upstreamToolName === ctx.upstreamToolName
&& invocation.argumentsHash === ctx.argumentsHash;
const stored = authorized ? invocation.policyExplanation?.consumedRateAdmissions : null;
return Array.isArray(stored) ? stored.filter((item): item is ConsumedRateAdmission => isRecord(item)
&& typeof item.policyId === "string" && CANONICAL_UUID_RE.test(item.policyId)
&& typeof item.versionHash === "string" && /^[a-f0-9]{64}$/.test(item.versionHash)) : [];
}
/**
* Decides one tool access request. Pass the same `cache` to every call of a
* batch that shares one actor and run context (for example, one tools/list
@@ -1343,14 +1374,32 @@ export function toolAccessPolicyService(db: Db) {
*/
async function decide(
input: ToolAccessDecisionInput,
options: { cache?: ToolAccessDecisionCache } = {},
options: { cache?: ToolAccessDecisionCache; admittedInvocationId?: string } = {},
): Promise<ToolAccessDecision> {
const consumedRateAdmissions: ConsumedRateAdmission[] = [];
const result = await decideInternal(input, options, consumedRateAdmissions);
return consumedRateAdmissions.length ? { ...result,
policyExplanation: { ...result.policyExplanation, consumedRateAdmissions },
} : result;
}
async function decideInternal(
input: ToolAccessDecisionInput,
options: { cache?: ToolAccessDecisionCache; admittedInvocationId?: string },
consumedRateAdmissions: ConsumedRateAdmission[],
): Promise<ToolAccessDecision> {
// A decision that consumes a rate limit can also record a trust-rule hit
// from the policy row, so it always reads fresh rows.
const cache = input.consumeRateLimit === true ? undefined : options.cache;
const cache = input.consumeRateLimit === true || options.admittedInvocationId ? undefined : options.cache;
const loaded = await loadContext(input, cache);
if (!loaded.ok) return loaded.decision;
const { ctx, redaction } = loaded;
// A final handoff rechecks current permissions for work already admitted by
// the gateway. Its own last consumed slot cannot reject that same work.
// Only durable consumption of an unchanged rule retains that slot; new or
// edited rules still apply. This option never approves or consumes a call.
const admittedRates = input.consumeRateLimit === false && options.admittedInvocationId
? await admittedRateLimitAdmissions(ctx, options.admittedInvocationId) : [];
const profileState = await effectiveProfiles(ctx, cache);
const effectiveProfileIds = profileState.profiles.map((profile) => profile.id);
const permittedByProfile = profileState.profiles.some((profile) => {
@@ -1422,10 +1471,13 @@ export function toolAccessPolicyService(db: Db) {
{ redactionPlan: redaction.redactionPlan, policyExplanation },
);
}
const versionHash = ratePolicyVersionHash(policy);
if (admittedRates.some(admission => admission.policyId === policy.id && admission.versionHash === versionHash)) continue;
const state = await enforceRateLimit(policy, ctx, input.consumeRateLimit === true);
if (state?.limited) {
return decision("rate_limited", "rate_limited", "Tool access rate limit exceeded.", effectiveProfileIds, [policy.id], { rateLimitState: state, redactionPlan: redaction.redactionPlan, policyExplanation });
}
if (state && input.consumeRateLimit === true) consumedRateAdmissions.push({ policyId: policy.id, versionHash });
continue;
}
if (policy.policyType === "trust_rule") {
@@ -1608,6 +1660,7 @@ export function toolAccessPolicyService(db: Db) {
agentId: ctx.agentId,
issueId: ctx.issueId,
runId: ctx.heartbeatRunId,
gatewayId: ctx.gatewayId,
applicationId: ctx.applicationId,
connectionId: ctx.connectionId,
catalogEntryId: ctx.catalogEntryId,
@@ -1621,6 +1674,7 @@ export function toolAccessPolicyService(db: Db) {
argumentsHash,
argumentsSummary: redaction.summary,
policyDecision: accessDecision.decision,
policyExplanation: accessDecision.policyExplanation ?? null,
matchedPolicyIds: accessDecision.matchedPolicyIds,
approvalState: accessDecision.decision === "require_approval" ? "pending" : "not_required",
status,
+51 -5
View File
@@ -8,13 +8,17 @@ import { z } from "zod";
import { and, eq, isNull, sql } from "drizzle-orm";
import { heartbeatRuns, runUsageReceipts, type Db } from "@paperclipai/db";
import type { AdapterExecutionResult, AdapterUsageCheckpoint } from "@paperclipai/adapter-utils";
import { pricingProvenanceSchema, usdToUnits } from "@paperclipai/shared";
import { compareCents, normalizeCents, pricingProvenanceSchema, usdToUnits } from "@paperclipai/shared";
import { idleAccountingSpoolPath, registerIdleSpoolDirectory } from "./idle-local-work.js";
import { conflict, notFound } from "../errors.js";
import { logger } from "../middleware/logger.js";
import { receiptFingerprint } from "./receipt-fingerprint.js";
import { withAccountingTransaction } from "./accounting-transaction.js";
import { ensureDurableDirectory, syncDirectory } from "../lib/durable-directory.js";
import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notifications.js";
import { knownRunCostCents } from "./budget-exposure.js";
import { createLiveRunUsageProjection, liveRunUsageCostCents } from "./budget-live-usage.js";
import { budgetServiceInTransaction } from "./budgets.js";
const token = z.number().int().nonnegative().max(2_147_483_647);
const price = z.number().finite().nonnegative().nullable().optional();
@@ -36,6 +40,11 @@ const envelopeSchema = z.object({
schema: z.literal("paperclip/accounting-receipt/v1"), id: z.string().uuid(), companyId: z.string().uuid(),
runId: z.string().uuid(), sourceId: z.string().uuid(), sequence: z.number().int().positive(), receivedAt: z.iso.datetime(),
adapterType: z.string().max(100), receipt: checkpointSchema,
// Derived by the recorder, not a provider's final charge. Keeping it in the
// fsynced envelope preserves known partial attempts through database replay.
liveCostCents: z.string().max(160).refine(value => {
try { return compareCents(value, 0) >= 0; } catch { return false; }
}, "Invalid live cost").transform(normalizeCents).nullable().optional(),
});
export type UsageReceiptEnvelope = z.infer<typeof envelopeSchema>;
export const usageReceiptSpoolPath = idleAccountingSpoolPath;
@@ -58,7 +67,7 @@ export async function spoolUsageReceipt(envelope: UsageReceiptEnvelope, director
export async function persistUsageReceipt(db: Db, raw: UsageReceiptEnvelope) {
const envelope = envelopeSchema.parse(raw);
const hash = receiptFingerprint(envelope);
return withAccountingTransaction(db, envelope.companyId, async tx => {
return withAccountingTransaction(db, envelope.companyId, async (tx, publications) => {
const [run] = await tx.select().from(heartbeatRuns).where(and(eq(heartbeatRuns.id, envelope.runId), eq(heartbeatRuns.companyId, envelope.companyId))).for("update");
if (!run) throw notFound("Accounting run not found");
const [existing] = await tx.select().from(runUsageReceipts).where(and(eq(runUsageReceipts.companyId, envelope.companyId), eq(runUsageReceipts.runId, envelope.runId), eq(runUsageReceipts.sourceId, envelope.sourceId), eq(runUsageReceipts.sequence, envelope.sequence)));
@@ -73,6 +82,23 @@ export async function persistUsageReceipt(db: Db, raw: UsageReceiptEnvelope) {
if (run.costAccountedAt || previous.accountingReceiptSourceId !== envelope.sourceId || Number(previous.accountingReceiptSequence ?? 0) >= envelope.sequence) return;
const receipt = envelope.receipt;
const billingType = receipt.billingType === "api" ? "metered_api" : receipt.billingType === "subscription" ? "subscription_included" : receipt.billingType ?? "unknown";
const oldProjection = previous.budgetExposure && typeof previous.budgetExposure === "object"
? previous.budgetExposure as Record<string, unknown> : {};
const canonicalProjection = typeof oldProjection.sourceId === "string" && typeof oldProjection.sequence === "number";
let previousCost = canonicalProjection ? null : knownRunCostCents(previous);
if (typeof oldProjection.costCents === "string") {
try {
const projected = normalizeCents(oldProjection.costCents);
if (compareCents(projected, 0) >= 0 && (canonicalProjection || compareCents(projected, previousCost ?? "0") > 0)) previousCost = projected;
} catch { /* Older malformed projections are not billing evidence. */ }
}
const reportedCost = envelope.liveCostCents !== undefined ? envelope.liveCostCents
: liveRunUsageCostCents({ ...receipt, billingType } as AdapterUsageCheckpoint);
const observedCost = reportedCost === null ? previousCost
: !receipt.complete && previousCost !== null && compareCents(previousCost, reportedCost) > 0 ? previousCost : reportedCost;
// Running checkpoints cannot settle yet. Their terminal transition owns
// the accounting wake; only a late receipt needs to reactivate that worker.
if (!["queued", "running"].includes(run.status)) await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, envelope.companyId);
await tx.update(heartbeatRuns).set({
costAccountingPending: true,
accountingLastError: previous.accountingCaptureFailed === true ? "usage_capture_failed" : null,
@@ -83,9 +109,15 @@ export async function persistUsageReceipt(db: Db, raw: UsageReceiptEnvelope) {
accountingReceiptReady: previous.accountingCaptureFailed !== true && receipt.complete && receipt.usageBasis !== "session_cumulative", accountingReceiptId: envelope.id,
accountingReceiptSequence: envelope.sequence, accountingReceiptReceivedAt: envelope.receivedAt,
costStatus: receipt.costStatus ?? null,
budgetExposure: { sourceId: envelope.sourceId, sequence: envelope.sequence, costCents: observedCost },
pricingProvenance: receipt.pricingProvenance ?? { source: receipt.costUsd != null || receipt.costUsdExact != null ? "provider_reported" : "unknown", version: "accounting-receipt/v1" },
},
}).where(and(eq(heartbeatRuns.id, run.id), isNull(heartbeatRuns.costAccountedAt)));
if (observedCost !== null && compareCents(observedCost, previousCost ?? "0") !== 0) {
const scope = previous.ledgerScope && typeof previous.ledgerScope === "object" ? previous.ledgerScope as Record<string, unknown> : {};
await budgetServiceInTransaction(tx, publications).evaluateAffectedScopes(envelope.companyId, run.agentId,
typeof scope.projectId === "string" ? scope.projectId : null);
}
});
}
@@ -202,6 +234,7 @@ export async function createRunUsageRecorder(db: Db, input: { companyId: string;
let sequence = 0;
let currentAttempt = "default";
const attempts = new Map<string, AdapterUsageCheckpoint>();
let initialLiveUsage: Parameters<typeof createLiveRunUsageProjection>[0];
// 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.
@@ -210,8 +243,14 @@ export async function createRunUsageRecorder(db: Db, input: { companyId: 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);
if (previous) {
const restored = envelopeSchema.parse(previous.receiptJson);
const receipt = restored.receipt as AdapterUsageCheckpoint;
attempts.set("default", receipt);
initialLiveUsage = { receipt, costCents: restored.liveCostCents };
}
}
const projectLiveUsage = createLiveRunUsageProjection(initialLiveUsage);
let chain = Promise.resolve();
let captureFailed = bound[0].usageJson?.accountingCaptureFailed === true;
let finishFailedCapture: (() => void) | undefined;
@@ -253,6 +292,7 @@ export async function createRunUsageRecorder(db: Db, input: { companyId: string;
version: "accounting-receipt/v1",
};
attempts.set(attempt, priced);
const liveCostCents = projectLiveUsage(attempt, priced);
const parts = [...attempts.values()];
const latest = attempts.get(currentAttempt)!;
const receipt: AdapterUsageCheckpoint = { ...latest, complete: latest.complete && !captureFailed, usage: {
@@ -289,10 +329,16 @@ export async function createRunUsageRecorder(db: Db, input: { companyId: string;
// can be considered complete, even if a later attempt succeeded.
receipt.complete = !captureFailed && parts.every(part => part.complete);
}
const envelope: UsageReceiptEnvelope = { schema: "paperclip/accounting-receipt/v1", id: randomUUID(), ...input, sourceId, sequence: ++sequence, receivedAt: new Date().toISOString(), receipt };
const envelope: UsageReceiptEnvelope = { schema: "paperclip/accounting-receipt/v1", id: randomUUID(), ...input, sourceId, sequence: ++sequence, receivedAt: new Date().toISOString(), receipt, liveCostCents };
const file = await spoolUsageReceipt(envelope, directory);
try { await persistUsageReceipt(db, envelope); await fs.rm(file, { force: true }); }
catch (error) { logger.warn({ err: error, runId: input.runId, receiptId: envelope.id }, "Usage receipt is durable on disk; database replay pending"); }
catch (error) {
// The fsynced receipt is durable even when the DB transaction failed.
// Wake the owned spool recovery path; never certify that failed write as free.
const { wakeCostAccountingWork } = await import("./budget-enforcement-worker.js");
wakeCostAccountingWork(db);
logger.warn({ err: error, runId: input.runId, receiptId: envelope.id }, "Usage receipt is durable on disk; database replay pending");
}
return receipt;
}
return {