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

@@ -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({