mirror of
https://github.com/paperclipai/paperclip.git
synced 2026-11-10 12:37:55 +01:00
fix(budgets): raise_budget_and_resume continues the work its budget stopped, once (#20144)
## Thinking Path > - Paperclip is the open source app people use to manage AI agents for work > - A budget hard stop cancels the scope's running work; the board resolves it with Raise budget and resume > - In the restart-chaos H1 scenario, the cancelled tasks did not continue after resume and each needed a manual Retry > - Resume is the board's explicit decision to continue, so the work it stopped should continue once, as if the board clicked Retry > - This pull request files that Retry automatically, keyed by the cancelled run, once the run's stop has settled > - The benefit is that a budget raise restores the work without manual Retries, without repeating completed work, and without charging the run's attempt budget ## Linked Issues or Issue Description Refs #20143 (bug B; this PR builds on it), #20113 (bug A), #20101 (superseded: split into #20143 and this PR) **What happened?** Restart-chaos waves w-20261108T195944-s42 and w-20261108T210118-s9802 (H1). After `raise_budget_and_resume`, CHA-3's budget-cancelled run (enforcement stop) and CHA-1's run (agent lifecycle pause from the budget hold) stayed cancelled until someone clicked Retry. **Expected behavior** Resuming continues each task its budget stopped, exactly once, with no attempt charged, an activity entry, and no repeat of work that already continued. A run a person stopped stays stopped. **Steps to reproduce** Set a budget below spend while a task runs; let the stop cancel the run; resolve with Raise budget and resume. Before this change, no new run is queued. ## What Changed - New table `budget_resume_claims` (migration 0451). One row per cancelled run, outside the run's result JSON, so a Stop that settles during the resume cannot erase it. - `markBudgetCancelledWorkForResume` (called by `resolveIncident` before the incident closes): claims every task run its budget stopped since the incident opened. That covers runs budget enforcement fenced, including stops still in flight, and runs the agent lifecycle pause stopped for this budget (`pauseStop`, carried from the lifecycle operation's `pauseCause` so a resume before the Stop keeps the link). Runs a person stopped are excluded. Two indexed passes led by `(company_id, status)`. - `deliverBudgetResumeContinuations`: once a claimed run's stop settles, files the resolving board member's Retry (`requestTaskRetry`, idempotency key `budget-resume:<run>`, with the member's session source and the chat Retry callback). The Retry path checks, under the task lock, that the run is still the task's latest and the task is open, so continued or completed work is never repeated. Claim settlement and its activity entry commit together; retryable refusals and errors retry. - Delivery runs in the always-on budget-enforcement worker after each pass's Stop delivery. Its schedule includes deliverable claims and the earliest stale-claim expiry. The execution-continuation worker is a backstop. - `agent.paused` / `agent.resumed` logging is in #20113. `budget.cancelled_work_resume_requested` (incident) and `budget.cancelled_work_continued` / `_not_continued` (task) are recorded here. ## Verification - New `budget-resume-continuation.test.ts` (Postgres, real `resolveIncident` → claims → `requestTaskRetry` → real `resumeTaskRetry` into a queued run): - each task continues once, including a blocked task; repeat delivery or resolve saves nothing new; - a stop still in flight at resume, whose terminal write replaces the result JSON, is still continued; - a 30-task scope is fully continued; - a lifecycle-pause stop caused by the budget is continued, and one caused by a manual pause is not; - continued work, done tasks, a run that finished before its stop, and person-stopped runs (including ones enforcement also fenced) are not continued. - `lifecycle.test.ts`: the pause records its cause on the operation and on stopped runs, including when the budget is resumed before the pause's Stop. `budget-enforcement-worker.test.ts`: delivery after Stop delivery, and a wake at claim expiry. - 287 tests across the affected suites passed. Server typecheck (`typecheck-layers`, low-memory): no errors. - Not run: the H1 chaos rerun (Phase 3 runs it after merge). ## Risks - Resume now continues budget-stopped tasks automatically. A board that wants a task stopped can Stop it again; person-stopped runs are never continued. - A new table and one index (migration 0451). > For core feature work, check [`ROADMAP.md`](ROADMAP.md) first and discuss it in `#dev` before opening the PR. Feature PRs that overlap with planned core work may need to be redirected — check the roadmap first. See `CONTRIBUTING.md`. ## Model Used - Claude Opus 5.5 (`claude-opus-5-5`), Anthropic, via Claude Code with tool use (shell, file edits, test execution). Extended reasoning enabled. ## 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 - [ ] All Paperclip CI gates are green - [ ] Greptile is 5/5 with no open P2s, recommendations, or follow-ups - [x] I will address all Greptile and reviewer comments before requesting merge 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
1 parent
61fa13c6d8
commit
5a51a7685d
23 files changed
+1344
-18
No files matched your search
@@ -514,6 +514,8 @@ Another run's unfinished accounting or cleanup never refuses new work (see `doc/
|
||||
|
||||
A user's message refused by a pause (a paused agent, a paused company, or a budget pause that is not a reached hard stop) is saved as a deferred intent that names the refusal and its release trigger; it is never written as a terminal skip. The pause's own stop keeps it. Lifting the pause is the operator's explicit authorization to send what the pause refused, as Send is after a Stop: when the pause stopped the task's run, the lifted pause issues the task's next explicit wake, which adopts the saved message once. A task that never ran keeps the message in the task, and stranded-work recovery starts it.
|
||||
|
||||
A reached budget hard stop cancels the scope's running work. When a board member resolves the incident with Raise budget and resume, Paperclip claims every task run that the stop cancelled since the incident opened (runs the budget enforcement fenced, and runs the agent lifecycle pause stopped for this budget: each pause records why it started, and the runs it stops carry that cause), with that member and their session source, in one write before the incident closes, as rows in `budget_resume_claims`, outside the run's result JSON so a Stop settling at the same moment cannot erase them. This includes runs whose stop is still being delivered. Resolving again keeps existing claims. The execution-continuation worker delivers each claim once its run's stop has settled: it files that member's Retry for the task. The Retry clears the task's execution hold and is accepted only while the cancelled run is still the task's latest run and the task is open and assigned, checked under the task lock, so completed or continued work is never repeated. Its idempotency key names the cancelled run, so every pass, restart, or repeated resume returns the same single request. The claim and its activity entry settle in one commit. The Retry starts a new run through normal admission and charges no attempt to the cancelled run. Paperclip records `budget.cancelled_work_resume_requested` on the incident and `budget.cancelled_work_continued` (or `budget.cancelled_work_not_continued`, with the reason) on each task. A run a person stopped is never continued this way, even if budget enforcement also fenced it. A pause's stop is intentional, not a failure: a native run that the agent lifecycle pause or a budget pause cancelled keeps its task's status and never opens a reconciliation incident; resuming continues it. A retryable Retry refusal leaves the claim pending for the next pass.
|
||||
|
||||
The same holds for accepted work beyond a first assignment. A temporary refusal never terminally cancels or escalates accepted work: an unstarted queued run (a promoted comment, a Retry, or a recovery) stays queued, and stranded-work recovery leaves the task as it is instead of moving it to `blocked`. These waits are passive, so settling a run's accounting wakes queued-run dispatch, the saved-receipt worker, and stranded-work recovery for the company. The held work then starts through normal admission, once. When a native run reaches a terminal status through recovery, or releases its environment, each message saved behind it is hinted by its own receipt, so delivery does not wait out a company-wide rescan interval.
|
||||
|
||||
New comments received during an execution hold retain their individual deferred receipts and ordered comment ids. Release cannot drain those receipts while replay remains blocked; the next eligible wake can adopt them. Authorized external-chat requests also remain deferred with their exact durable receipt. They must use normal promotion and current authorization; a generic wake cannot adopt only their comment ids and discard their actor or session contract. A wait does not authorize replay, reset an incident retry budget, or bypass an interaction's delivery rules.
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
CREATE TABLE "budget_resume_claims" (
|
||||
"run_id" uuid PRIMARY KEY NOT NULL,
|
||||
"company_id" uuid NOT NULL,
|
||||
"incident_id" uuid NOT NULL,
|
||||
"actor_user_id" text NOT NULL,
|
||||
"actor_source" text,
|
||||
"state" text DEFAULT 'pending' NOT NULL,
|
||||
"claimed_at" timestamp with time zone,
|
||||
"settled_at" timestamp with time zone,
|
||||
"reason" text,
|
||||
"request_id" uuid,
|
||||
"created_at" timestamp with time zone DEFAULT now() NOT NULL
|
||||
);
|
||||
--> statement-breakpoint
|
||||
ALTER TABLE "budget_resume_claims" ADD CONSTRAINT "budget_resume_claims_run_id_heartbeat_runs_id_fk" FOREIGN KEY ("run_id") REFERENCES "public"."heartbeat_runs"("id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint
|
||||
ALTER TABLE "budget_resume_claims" ADD CONSTRAINT "budget_resume_claims_company_id_companies_id_fk" FOREIGN KEY ("company_id") REFERENCES "public"."companies"("id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint
|
||||
ALTER TABLE "budget_resume_claims" ADD CONSTRAINT "budget_resume_claims_incident_id_budget_incidents_id_fk" FOREIGN KEY ("incident_id") REFERENCES "public"."budget_incidents"("id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint
|
||||
CREATE INDEX "budget_resume_claims_pending_idx" ON "budget_resume_claims" USING btree ("company_id","run_id") WHERE "budget_resume_claims"."state" = 'pending';
|
||||
+308
-2
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"id": "782688f1-7843-4a93-8ea3-2cec362d269e",
|
||||
"prevId": "f0d8430b-7abf-4232-86a9-11a268a0d20e",
|
||||
"id": "b5699366-77c2-48c8-b7dc-658a75de7658",
|
||||
"prevId": "b8f43f1a-3230-4cda-991d-0e7b6323114b",
|
||||
"version": "7",
|
||||
"dialect": "postgresql",
|
||||
"tables": {
|
||||
@@ -4102,6 +4102,28 @@
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
},
|
||||
"agent_wakeup_requests_issue_blocker_stalled_uq": {
|
||||
"name": "agent_wakeup_requests_issue_blocker_stalled_uq",
|
||||
"columns": [
|
||||
{
|
||||
"expression": "company_id",
|
||||
"isExpression": false,
|
||||
"asc": true,
|
||||
"nulls": "last"
|
||||
},
|
||||
{
|
||||
"expression": "idempotency_key",
|
||||
"isExpression": false,
|
||||
"asc": true,
|
||||
"nulls": "last"
|
||||
}
|
||||
],
|
||||
"isUnique": true,
|
||||
"where": "\"agent_wakeup_requests\".\"idempotency_key\" LIKE 'issue_blocker_stalled:%' AND \"agent_wakeup_requests\".\"status\" IN ('queued', 'deferred_issue_execution', 'claimed', 'coalesced', 'completed')",
|
||||
"concurrently": true,
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
},
|
||||
"agent_wakeup_requests_company_payload_issue_idx": {
|
||||
"name": "agent_wakeup_requests_company_payload_issue_idx",
|
||||
"columns": [
|
||||
@@ -4252,6 +4274,18 @@
|
||||
"primaryKey": false,
|
||||
"notNull": false
|
||||
},
|
||||
"harness_verified_digest": {
|
||||
"name": "harness_verified_digest",
|
||||
"type": "text",
|
||||
"primaryKey": false,
|
||||
"notNull": false
|
||||
},
|
||||
"harness_verified_at": {
|
||||
"name": "harness_verified_at",
|
||||
"type": "timestamp with time zone",
|
||||
"primaryKey": false,
|
||||
"notNull": false
|
||||
},
|
||||
"reports_to": {
|
||||
"name": "reports_to",
|
||||
"type": "uuid",
|
||||
@@ -7848,6 +7882,150 @@
|
||||
},
|
||||
"isRLSEnabled": false
|
||||
},
|
||||
"public.budget_resume_claims": {
|
||||
"name": "budget_resume_claims",
|
||||
"schema": "",
|
||||
"columns": {
|
||||
"run_id": {
|
||||
"name": "run_id",
|
||||
"type": "uuid",
|
||||
"primaryKey": true,
|
||||
"notNull": true
|
||||
},
|
||||
"company_id": {
|
||||
"name": "company_id",
|
||||
"type": "uuid",
|
||||
"primaryKey": false,
|
||||
"notNull": true
|
||||
},
|
||||
"incident_id": {
|
||||
"name": "incident_id",
|
||||
"type": "uuid",
|
||||
"primaryKey": false,
|
||||
"notNull": true
|
||||
},
|
||||
"actor_user_id": {
|
||||
"name": "actor_user_id",
|
||||
"type": "text",
|
||||
"primaryKey": false,
|
||||
"notNull": true
|
||||
},
|
||||
"actor_source": {
|
||||
"name": "actor_source",
|
||||
"type": "text",
|
||||
"primaryKey": false,
|
||||
"notNull": false
|
||||
},
|
||||
"state": {
|
||||
"name": "state",
|
||||
"type": "text",
|
||||
"primaryKey": false,
|
||||
"notNull": true,
|
||||
"default": "'pending'"
|
||||
},
|
||||
"claimed_at": {
|
||||
"name": "claimed_at",
|
||||
"type": "timestamp with time zone",
|
||||
"primaryKey": false,
|
||||
"notNull": false
|
||||
},
|
||||
"settled_at": {
|
||||
"name": "settled_at",
|
||||
"type": "timestamp with time zone",
|
||||
"primaryKey": false,
|
||||
"notNull": false
|
||||
},
|
||||
"reason": {
|
||||
"name": "reason",
|
||||
"type": "text",
|
||||
"primaryKey": false,
|
||||
"notNull": false
|
||||
},
|
||||
"request_id": {
|
||||
"name": "request_id",
|
||||
"type": "uuid",
|
||||
"primaryKey": false,
|
||||
"notNull": false
|
||||
},
|
||||
"created_at": {
|
||||
"name": "created_at",
|
||||
"type": "timestamp with time zone",
|
||||
"primaryKey": false,
|
||||
"notNull": true,
|
||||
"default": "now()"
|
||||
}
|
||||
},
|
||||
"indexes": {
|
||||
"budget_resume_claims_pending_idx": {
|
||||
"name": "budget_resume_claims_pending_idx",
|
||||
"columns": [
|
||||
{
|
||||
"expression": "company_id",
|
||||
"isExpression": false,
|
||||
"asc": true,
|
||||
"nulls": "last"
|
||||
},
|
||||
{
|
||||
"expression": "run_id",
|
||||
"isExpression": false,
|
||||
"asc": true,
|
||||
"nulls": "last"
|
||||
}
|
||||
],
|
||||
"isUnique": false,
|
||||
"where": "\"budget_resume_claims\".\"state\" = 'pending'",
|
||||
"concurrently": false,
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
}
|
||||
},
|
||||
"foreignKeys": {
|
||||
"budget_resume_claims_run_id_heartbeat_runs_id_fk": {
|
||||
"name": "budget_resume_claims_run_id_heartbeat_runs_id_fk",
|
||||
"tableFrom": "budget_resume_claims",
|
||||
"tableTo": "heartbeat_runs",
|
||||
"columnsFrom": [
|
||||
"run_id"
|
||||
],
|
||||
"columnsTo": [
|
||||
"id"
|
||||
],
|
||||
"onDelete": "cascade",
|
||||
"onUpdate": "no action"
|
||||
},
|
||||
"budget_resume_claims_company_id_companies_id_fk": {
|
||||
"name": "budget_resume_claims_company_id_companies_id_fk",
|
||||
"tableFrom": "budget_resume_claims",
|
||||
"tableTo": "companies",
|
||||
"columnsFrom": [
|
||||
"company_id"
|
||||
],
|
||||
"columnsTo": [
|
||||
"id"
|
||||
],
|
||||
"onDelete": "cascade",
|
||||
"onUpdate": "no action"
|
||||
},
|
||||
"budget_resume_claims_incident_id_budget_incidents_id_fk": {
|
||||
"name": "budget_resume_claims_incident_id_budget_incidents_id_fk",
|
||||
"tableFrom": "budget_resume_claims",
|
||||
"tableTo": "budget_incidents",
|
||||
"columnsFrom": [
|
||||
"incident_id"
|
||||
],
|
||||
"columnsTo": [
|
||||
"id"
|
||||
],
|
||||
"onDelete": "cascade",
|
||||
"onUpdate": "no action"
|
||||
}
|
||||
},
|
||||
"compositePrimaryKeys": {},
|
||||
"uniqueConstraints": {},
|
||||
"policies": {},
|
||||
"checkConstraints": {},
|
||||
"isRLSEnabled": false
|
||||
},
|
||||
"public.built_in_managed_resources": {
|
||||
"name": "built_in_managed_resources",
|
||||
"schema": "",
|
||||
@@ -21965,6 +22143,112 @@
|
||||
"checkConstraints": {},
|
||||
"isRLSEnabled": false
|
||||
},
|
||||
"public.controller_instances": {
|
||||
"name": "controller_instances",
|
||||
"schema": "",
|
||||
"columns": {
|
||||
"boot_id": {
|
||||
"name": "boot_id",
|
||||
"type": "text",
|
||||
"primaryKey": true,
|
||||
"notNull": true
|
||||
},
|
||||
"epoch": {
|
||||
"name": "epoch",
|
||||
"type": "bigserial",
|
||||
"primaryKey": false,
|
||||
"notNull": true
|
||||
},
|
||||
"process_space": {
|
||||
"name": "process_space",
|
||||
"type": "text",
|
||||
"primaryKey": false,
|
||||
"notNull": true
|
||||
},
|
||||
"pid": {
|
||||
"name": "pid",
|
||||
"type": "integer",
|
||||
"primaryKey": false,
|
||||
"notNull": true
|
||||
},
|
||||
"started_at": {
|
||||
"name": "started_at",
|
||||
"type": "timestamp with time zone",
|
||||
"primaryKey": false,
|
||||
"notNull": true,
|
||||
"default": "now()"
|
||||
},
|
||||
"renewed_at": {
|
||||
"name": "renewed_at",
|
||||
"type": "timestamp with time zone",
|
||||
"primaryKey": false,
|
||||
"notNull": true,
|
||||
"default": "now()"
|
||||
},
|
||||
"lease_expires_at": {
|
||||
"name": "lease_expires_at",
|
||||
"type": "timestamp with time zone",
|
||||
"primaryKey": false,
|
||||
"notNull": true
|
||||
},
|
||||
"fenced_at": {
|
||||
"name": "fenced_at",
|
||||
"type": "timestamp with time zone",
|
||||
"primaryKey": false,
|
||||
"notNull": false
|
||||
},
|
||||
"fenced_by_epoch": {
|
||||
"name": "fenced_by_epoch",
|
||||
"type": "bigint",
|
||||
"primaryKey": false,
|
||||
"notNull": false
|
||||
},
|
||||
"stopped_at": {
|
||||
"name": "stopped_at",
|
||||
"type": "timestamp with time zone",
|
||||
"primaryKey": false,
|
||||
"notNull": false
|
||||
}
|
||||
},
|
||||
"indexes": {
|
||||
"controller_instances_epoch_idx": {
|
||||
"name": "controller_instances_epoch_idx",
|
||||
"columns": [
|
||||
{
|
||||
"expression": "epoch",
|
||||
"isExpression": false,
|
||||
"asc": true,
|
||||
"nulls": "last"
|
||||
}
|
||||
],
|
||||
"isUnique": true,
|
||||
"concurrently": false,
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
},
|
||||
"controller_instances_lease_idx": {
|
||||
"name": "controller_instances_lease_idx",
|
||||
"columns": [
|
||||
{
|
||||
"expression": "lease_expires_at",
|
||||
"isExpression": false,
|
||||
"asc": true,
|
||||
"nulls": "last"
|
||||
}
|
||||
],
|
||||
"isUnique": false,
|
||||
"concurrently": false,
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
}
|
||||
},
|
||||
"foreignKeys": {},
|
||||
"compositePrimaryKeys": {},
|
||||
"uniqueConstraints": {},
|
||||
"policies": {},
|
||||
"checkConstraints": {},
|
||||
"isRLSEnabled": false
|
||||
},
|
||||
"public.cost_events": {
|
||||
"name": "cost_events",
|
||||
"schema": "",
|
||||
@@ -34856,6 +35140,12 @@
|
||||
"primaryKey": false,
|
||||
"notNull": false
|
||||
},
|
||||
"controller_process": {
|
||||
"name": "controller_process",
|
||||
"type": "jsonb",
|
||||
"primaryKey": false,
|
||||
"notNull": false
|
||||
},
|
||||
"execution_stage": {
|
||||
"name": "execution_stage",
|
||||
"type": "text",
|
||||
@@ -35125,6 +35415,22 @@
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
},
|
||||
"heartbeat_runs_legacy_controller_lease_idx": {
|
||||
"name": "heartbeat_runs_legacy_controller_lease_idx",
|
||||
"columns": [
|
||||
{
|
||||
"expression": "controller_lease_expires_at",
|
||||
"isExpression": false,
|
||||
"asc": true,
|
||||
"nulls": "last"
|
||||
}
|
||||
],
|
||||
"isUnique": false,
|
||||
"where": "\"heartbeat_runs\".\"runtime_mode\" = 'legacy' AND \"heartbeat_runs\".\"status\" = 'running' AND \"heartbeat_runs\".\"controller_boot_id\" is not null",
|
||||
"concurrently": false,
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
},
|
||||
"heartbeat_runs_cost_accounting_pending_idx": {
|
||||
"name": "heartbeat_runs_cost_accounting_pending_idx",
|
||||
"columns": [
|
||||
+7
@@ -3144,6 +3144,13 @@
|
||||
"when": 1794170960477,
|
||||
"tag": "0450_petite_frightful_four",
|
||||
"breakpoints": true
|
||||
},
|
||||
{
|
||||
"idx": 451,
|
||||
"version": "7",
|
||||
"when": 1794185816322,
|
||||
"tag": "0451_budget_resume_claims",
|
||||
"breakpoints": true
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
import { sql } from "drizzle-orm";
|
||||
import { index, pgTable, text, timestamp, uuid } from "drizzle-orm/pg-core";
|
||||
import { budgetIncidents } from "./budget_incidents.js";
|
||||
import { companies } from "./companies.js";
|
||||
import { heartbeatRuns } from "./heartbeat_runs.js";
|
||||
|
||||
/**
|
||||
* A run a budget stop cancelled, owed one continuation because a board member
|
||||
* resolved the incident with Raise budget and resume. Kept outside the run's
|
||||
* result JSON, which terminal writers replace, so a Stop that settles at the
|
||||
* same moment cannot erase it.
|
||||
*/
|
||||
export const budgetResumeClaims = pgTable(
|
||||
"budget_resume_claims",
|
||||
{
|
||||
runId: uuid("run_id").primaryKey().references(() => heartbeatRuns.id, { onDelete: "cascade" }),
|
||||
companyId: uuid("company_id").notNull().references(() => companies.id, { onDelete: "cascade" }),
|
||||
incidentId: uuid("incident_id").notNull().references(() => budgetIncidents.id, { onDelete: "cascade" }),
|
||||
actorUserId: text("actor_user_id").notNull(),
|
||||
actorSource: text("actor_source"),
|
||||
state: text("state").notNull().default("pending"),
|
||||
claimedAt: timestamp("claimed_at", { withTimezone: true }),
|
||||
settledAt: timestamp("settled_at", { withTimezone: true }),
|
||||
reason: text("reason"),
|
||||
requestId: uuid("request_id"),
|
||||
createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(),
|
||||
},
|
||||
(table) => ({
|
||||
pendingIdx: index("budget_resume_claims_pending_idx").on(table.companyId, table.runId)
|
||||
.where(sql`${table.state} = 'pending'`),
|
||||
}),
|
||||
);
|
||||
@@ -22,6 +22,7 @@ export { invites } from "./invites.js";
|
||||
export { joinRequests } from "./join_requests.js";
|
||||
export { budgetPolicies } from "./budget_policies.js";
|
||||
export { budgetIncidents } from "./budget_incidents.js";
|
||||
export { budgetResumeClaims } from "./budget_resume_claims.js";
|
||||
export { agentConfigRevisions } from "./agent_config_revisions.js";
|
||||
export { agentApiKeys } from "./agent_api_keys.js";
|
||||
export { agentRuntimeState } from "./agent_runtime_state.js";
|
||||
|
||||
@@ -25,6 +25,9 @@ export interface AgentLifecycleOperation {
|
||||
verificationFailures?: number;
|
||||
attempts: number;
|
||||
responsibleUserId?: string | null;
|
||||
/** Why a pause started: its holds and the company's pause reason at that moment. Runs the
|
||||
* pause stops record it, so a budget resume continues only work its own budget stopped. */
|
||||
pauseCause?: { holds: string[]; companyPauseReason: string | null };
|
||||
}
|
||||
|
||||
export interface AgentLifecycleRequest {
|
||||
|
||||
@@ -17,6 +17,10 @@ vi.mock("../services/budgets.js", () => ({
|
||||
flushBudgetEvaluationWorkPage: work.evaluate,
|
||||
}));
|
||||
vi.mock("../services/budget-evaluation-work.js", () => ({ hasPendingBudgetEvaluation: work.evaluationPending }));
|
||||
const resumeWork = vi.hoisted(() => ({ expiry: vi.fn(async (): Promise<number | null> => null),
|
||||
pendingForRun: vi.fn(async (_db: unknown, runId: string) => runId !== "run-unclaimed") }));
|
||||
vi.mock("../services/budget-resume-continuation.js", () => ({ hasDeliverableBudgetResume: vi.fn(async () => false),
|
||||
hasPendingBudgetResumeForRun: resumeWork.pendingForRun, nextBudgetResumeClaimExpiryAt: resumeWork.expiry }));
|
||||
vi.mock("../services/run-cost-accounting.js", () => ({ readCostAccountingWork: work.accounting }));
|
||||
|
||||
const coordinators: ReturnType<typeof createDeliveryWorkCoordinator>[] = [];
|
||||
@@ -49,6 +53,50 @@ function setup() {
|
||||
}
|
||||
|
||||
describe("budget accounting delivery", () => {
|
||||
// H1 review: budget-resume claims are delivered by the always-on budget
|
||||
// worker, after the pass's Stop delivery, independent of scheduled agent work.
|
||||
it("delivers budget-resume continuations after Stop delivery in each pass", async () => {
|
||||
const s = setup();
|
||||
const order: string[] = [];
|
||||
work.deliver.mockImplementation(async () => { order.push("stops"); });
|
||||
const deliverResumeContinuations = vi.fn(async () => { order.push("continuations"); });
|
||||
await registerBudgetEnforcementWork(s.db, s.coordinator, { deliverResumeContinuations }).ready;
|
||||
expect(order).toEqual(["stops", "continuations"]);
|
||||
});
|
||||
|
||||
it("wakes resume delivery when a run's Stop settles", async () => {
|
||||
const s = setup();
|
||||
const deliverResumeContinuations = vi.fn(async () => undefined);
|
||||
const worker = registerBudgetEnforcementWork(s.db, s.coordinator, { deliverResumeContinuations });
|
||||
await worker.ready;
|
||||
expect(deliverResumeContinuations).toHaveBeenCalledTimes(1);
|
||||
// Another tenant's cancellation without a resume claim wakes nothing.
|
||||
publishLiveEvent({ companyId: "company-2", type: "heartbeat.run.status", payload: { runId: "run-unclaimed", status: "cancelled" } } as never);
|
||||
await vi.advanceTimersByTimeAsync(1_500);
|
||||
expect(deliverResumeContinuations).toHaveBeenCalledTimes(1);
|
||||
publishLiveEvent({ companyId: "company-1", type: "heartbeat.run.status", payload: { runId: "run-1", status: "cancelled" } } as never);
|
||||
await vi.advanceTimersByTimeAsync(1_500);
|
||||
expect(deliverResumeContinuations).toHaveBeenCalledTimes(2);
|
||||
// Stopping the worker removes the listener: later Stops wake nothing.
|
||||
await worker.stop();
|
||||
publishLiveEvent({ companyId: "company-1", type: "heartbeat.run.status", payload: { runId: "run-2", status: "cancelled" } } as never);
|
||||
await vi.advanceTimersByTimeAsync(1_500);
|
||||
expect(deliverResumeContinuations).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it("wakes when a failed dispatcher's resume claim goes stale", async () => {
|
||||
const s = setup();
|
||||
const expiry = Date.now() + 60_000;
|
||||
resumeWork.expiry.mockResolvedValue(expiry);
|
||||
const deliverResumeContinuations = vi.fn(async () => undefined);
|
||||
await registerBudgetEnforcementWork(s.db, s.coordinator, { deliverResumeContinuations }).ready;
|
||||
expect(deliverResumeContinuations).toHaveBeenCalledTimes(1);
|
||||
expect(s.coordinator.nextWakeAt()).toBe(expiry);
|
||||
resumeWork.expiry.mockResolvedValue(null);
|
||||
await vi.advanceTimersByTimeAsync(60_000);
|
||||
expect(deliverResumeContinuations).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it("continues cancellation delivery and retries a failed startup reconciliation", async () => {
|
||||
const s = setup();
|
||||
work.reconcile.mockResolvedValueOnce({ nextCursor: null, hadFailures: true });
|
||||
|
||||
@@ -214,6 +214,7 @@ import { createPluginToolDispatcher } from "./services/plugin-tool-dispatcher.js
|
||||
import { createToolGatewayService } from "./services/tool-gateway.js";
|
||||
import { toolAccessService } from "./services/tool-access.js";
|
||||
import { chatChannelService } from "./services/chat-channels.js";
|
||||
import type { BudgetResumeDelivery } from "./services/budget-resume-continuation.js";
|
||||
import { deliverNativeQuestionResponse } from "./services/native-runtime/native-question-bridge.js";
|
||||
import { nativeChatProgressWork } from "./services/chat-native-progress-work.js";
|
||||
import { chatInboundWakeNoticeWork } from "./services/chat-inbound-wake-notice-work.js";
|
||||
@@ -448,6 +449,8 @@ export async function createApp(
|
||||
databaseBackupService?: InstanceDatabaseBackupService;
|
||||
prepareIdleDatabaseBackup?: () => Promise<boolean>;
|
||||
idleDeadlineOwners?: IdleDeadlineOwners;
|
||||
/** Filled with the chat service's Retry callback for background budget-resume delivery. */
|
||||
budgetResumeDelivery?: BudgetResumeDelivery;
|
||||
databaseBackupHealth?: InspectDatabaseBackupHealthOptions;
|
||||
deploymentMode: DeploymentMode;
|
||||
deploymentExposure: DeploymentExposure;
|
||||
@@ -590,6 +593,7 @@ export async function createApp(
|
||||
deliverNativeQuestionResponse(db, interaction),
|
||||
storage: opts.storageService,
|
||||
});
|
||||
if (opts.budgetResumeDelivery) opts.budgetResumeDelivery.prepareChatRetry = chatChannels.prepareFailedChatRunRetry;
|
||||
// Voice capabilities are verified by the signed callback route, not by
|
||||
// board/agent bearer authentication. Mount before actorMiddleware so its
|
||||
// per-session Bearer cannot be mistaken for an agent API key.
|
||||
|
||||
@@ -92,6 +92,7 @@ import { embeddedPostgresOwnerPort } from "./embedded-postgres-owner.js";
|
||||
import { deliverExecutionStatuses } from "./services/execution-status-delivery.js";
|
||||
import { deliverReconciledExecutions, settleUnrecoverableExecutions } from "./services/execution-recovery-resolution.js";
|
||||
import { deliverPendingPauseResumes } from "./services/issue-tree-resume.js";
|
||||
import { deliverBudgetResumeContinuations, type BudgetResumeDelivery } from "./services/budget-resume-continuation.js";
|
||||
import { invalidateNativeReplacementDiscovery, reconcileSafeNativeReplacements } from "./services/native-runtime/native-safe-replacement.js";
|
||||
import { reconcileAbandonedExecutionControl } from "./services/execution-control-reconciliation.js";
|
||||
import { acknowledgeNativeStartupCleanup } from "./services/native-runtime/native-startup-cleanup.js";
|
||||
@@ -948,8 +949,11 @@ async function startServerWithDatabaseTeardown(
|
||||
logger.warn({ scope: cloudQualification.scope, blockerCodes: cloudQualification.blockers.map((blocker) => blocker.code) }, "Cloud configuration and package qualification has blockers");
|
||||
}
|
||||
const idleDeadlineOwners: IdleDeadlineOwners = {};
|
||||
// createApp fills the chat Retry callback; the continuation worker reads it.
|
||||
const budgetResumeDelivery: BudgetResumeDelivery = {};
|
||||
const app = await createApp(db as any, {
|
||||
idleDeadlineOwners,
|
||||
budgetResumeDelivery,
|
||||
prepareIdleDatabaseBackup: idleDatabaseBackupEnabled
|
||||
? async () => (await runServerDatabaseBackup("idle")) !== null
|
||||
: undefined,
|
||||
@@ -1396,13 +1400,19 @@ async function startServerWithDatabaseTeardown(
|
||||
{ queue: DELIVERY_QUEUES.executionContinuation, run: async () => {
|
||||
await deliverReconciledExecutions(db, heartbeat.wakeup);
|
||||
await deliverPendingPauseResumes(db, heartbeat.wakeup);
|
||||
await deliverBudgetResumeContinuations(db, { ...budgetResumeDelivery,
|
||||
resumeTaskRetry: (companyId, requestId) => heartbeat.resumeTaskRetry(companyId, requestId) });
|
||||
}, next: executionSchedule.continuation },
|
||||
] : []),
|
||||
], () => !heartbeatSchedulerStopped);
|
||||
|
||||
const budgetEnforcementWorker = registerBudgetEnforcementWork(db, deliveryWork, {
|
||||
cancelWorkForScope: accountingHeartbeat.cancelBudgetScopeWork,
|
||||
deliverResumeContinuations: () => deliverBudgetResumeContinuations(db, { ...budgetResumeDelivery,
|
||||
resumeTaskRetry: (companyId, requestId) => accountingHeartbeat.resumeTaskRetry(companyId, requestId) }),
|
||||
});
|
||||
// Its stop also removes the run-status listener that wakes resume delivery.
|
||||
server.on("close", () => { void budgetEnforcementWorker.stop(); });
|
||||
const accountingWorker = registerCostAccountingWork(db, deliveryWork, accountingHeartbeat.reconcileCostAccounting);
|
||||
server.on("close", accountingWorker.stop);
|
||||
skillSourceSync.start(deliveryWork);
|
||||
|
||||
@@ -371,8 +371,8 @@ const support = external ? { supported: true } : await getEmbeddedPostgresTestSu
|
||||
expect(await db.select().from(principalPermissionGrants).where(eq(principalPermissionGrants.companyId, companyId))).toHaveLength(0);
|
||||
});
|
||||
|
||||
// H1 bug B: a pause's stop is told apart from termination by the cause the
|
||||
// pause records on each run it stops.
|
||||
// H1 review: a budget resume continues only work its own budget stopped. The
|
||||
// pause records the holds and the company's pause reason on each run it stops.
|
||||
it("records what a pause stopped each active run for", async () => {
|
||||
const agent = await hire();
|
||||
await db.update(companies).set({ status: "paused", pauseReason: "budget" }).where(eq(companies.id, companyId));
|
||||
@@ -384,6 +384,46 @@ const support = external ? { supported: true } : await getEmbeddedPostgresTestSu
|
||||
expect(stopped.resultJson).toMatchObject({ pauseStop: { holds: ["company_paused"], companyPauseReason: "budget" } });
|
||||
});
|
||||
|
||||
// H1 review: a resume that removed the holds before the pause's Stop ran must
|
||||
// not erase why the pause stopped the run (the Stop itself stays owed).
|
||||
it("records the pause's original cause on runs it stops after a resume removed the holds", async () => {
|
||||
const agent = await hire();
|
||||
await db.update(agents).set({ lifecycleState: "pausing", lifecycleHolds: [],
|
||||
lifecycleOperation: { id: randomUUID(), hostComplete: false, completedPluginIds: [], attempts: 0,
|
||||
pauseCause: { holds: ["budget"], companyPauseReason: null } } }).where(eq(agents.id, agent.id));
|
||||
const [run] = await db.insert(heartbeatRuns).values({ companyId, agentId: agent.id, status: "queued" }).returning();
|
||||
await createLifecycleDriver(db, {} as never).runHost(await current(agent.id) as never);
|
||||
const [stopped] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, run.id));
|
||||
expect(stopped.status).toBe("cancelled");
|
||||
expect(stopped.resultJson).toMatchObject({ pauseStop: { holds: ["budget"], companyPauseReason: null } });
|
||||
});
|
||||
|
||||
it("keeps a budget pause's cause when the budget is resumed before the pause's Stop runs", async () => {
|
||||
const agent = await hire();
|
||||
await db.update(agents).set({ lifecycleState: "ready", status: "idle" }).where(eq(agents.id, agent.id));
|
||||
const [run] = await db.insert(heartbeatRuns).values({ companyId, agentId: agent.id, status: "queued" }).returning();
|
||||
const { applyAgentBudgetHoldInTransaction } = await import("../../../services/agent-lifecycle.js");
|
||||
// The budget hold pauses the agent, then the raise removes it before the
|
||||
// lifecycle worker has run the pause's Stop.
|
||||
await db.transaction(tx => applyAgentBudgetHoldInTransaction(tx as never, companyId, agent.id, true));
|
||||
await db.transaction(tx => applyAgentBudgetHoldInTransaction(tx as never, companyId, agent.id, false));
|
||||
const paused = await current(agent.id);
|
||||
expect(paused).toMatchObject({ lifecycleState: "pausing", lifecycleHolds: [] });
|
||||
await createLifecycleDriver(db, {} as never).runHost(paused as never);
|
||||
const [stopped] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, run.id));
|
||||
expect(stopped.resultJson).toMatchObject({ pauseStop: { holds: ["budget"] } });
|
||||
});
|
||||
|
||||
it("records why a budget pause started on its lifecycle operation", async () => {
|
||||
const agent = await hire();
|
||||
await db.update(agents).set({ lifecycleState: "ready", status: "idle" }).where(eq(agents.id, agent.id));
|
||||
await db.update(companies).set({ status: "paused", pauseReason: "budget" }).where(eq(companies.id, companyId));
|
||||
const { applyAgentBudgetHoldInTransaction } = await import("../../../services/agent-lifecycle.js");
|
||||
await db.transaction(tx => applyAgentBudgetHoldInTransaction(tx as never, companyId, agent.id, true));
|
||||
expect((await current(agent.id)).lifecycleOperation).toMatchObject({
|
||||
pauseCause: { holds: ["budget"], companyPauseReason: "budget" } });
|
||||
});
|
||||
|
||||
it("runs the saved harness test and keeps a failed configuration out of ready", async () => {
|
||||
const driver = createLifecycleDriver(db, {} as never);
|
||||
const work = worker(driver.runHost);
|
||||
|
||||
@@ -147,6 +147,11 @@ async function changeLockedAgent(tx: Db, effects: LifecycleEffects, agent: Lifec
|
||||
: agent.lifecycleOperation?.resumeState,
|
||||
};
|
||||
if (next === "ready") operation.resumeState = "ready";
|
||||
if (next === "pausing") {
|
||||
const [company] = await tx.select({ pauseReason: companies.pauseReason }).from(companies)
|
||||
.where(eq(companies.id, agent.companyId));
|
||||
operation.pauseCause = { holds, companyPauseReason: company?.pauseReason ?? null };
|
||||
}
|
||||
const status = compatibilityStatus(next, agent.status);
|
||||
if(!canConfigureAgentConnection(agent)&&canConfigureAgentConnection({lifecycleState:next,status}))await effects.beforeConnectionEligibilityRestore?.(tx,agent.companyId,agent.id);
|
||||
if (status !== agent.status) await effects.beforeStatusChange(tx, agent.companyId, agent.id, agent.status, status);
|
||||
|
||||
@@ -371,7 +371,8 @@ export function costRoutes(
|
||||
const companyId = req.params.companyId as string;
|
||||
const incidentId = req.params.incidentId as string;
|
||||
assertCompanyAccess(req, companyId);
|
||||
const incident = await budgets.resolveIncident(companyId, incidentId, req.body, req.actor.userId ?? "board");
|
||||
const incident = await budgets.resolveIncident(companyId, incidentId, req.body, req.actor.userId ?? "board",
|
||||
req.actor.source ?? null);
|
||||
res.json(incident);
|
||||
},
|
||||
);
|
||||
|
||||
@@ -63,12 +63,16 @@ export function createLifecycleDriver(db: Db, manager: PluginWorkerManager): Lif
|
||||
await db.select().from(agents).where(eq(agents.companyId, agent.companyId)))];
|
||||
const reason = agent.lifecycleHolds?.includes("company_archived")
|
||||
? "Cancelled because the company was archived" : "Agent lifecycle stop requested";
|
||||
// A pause records why it stopped each run (pauseStop). It marks the
|
||||
// stop as a pause, not termination or cleanup, which share this reason.
|
||||
// Record what this pause stopped the runs for, so a budget resume
|
||||
// continues only the work its own budget stopped, never a manual pause.
|
||||
if (agent.lifecycleState === "pausing") {
|
||||
// The cause recorded when the pause started survives a resume that
|
||||
// already removed the holds before this Stop ran.
|
||||
const cause = agent.lifecycleOperation?.pauseCause;
|
||||
const pauseStop = sql`jsonb_build_object('pauseStop', jsonb_build_object(
|
||||
'holds', ${JSON.stringify(agent.lifecycleHolds ?? [])}::jsonb,
|
||||
'companyPauseReason', (select ${companies.pauseReason} from ${companies} where ${companies.id} = ${agent.companyId})))`;
|
||||
'holds', ${JSON.stringify(cause?.holds ?? agent.lifecycleHolds ?? [])}::jsonb,
|
||||
'companyPauseReason', ${cause ? sql`${cause.companyPauseReason}::text`
|
||||
: sql`(select ${companies.pauseReason} from ${companies} where ${companies.id} = ${agent.companyId})`}))`;
|
||||
await db.update(heartbeatRuns).set({
|
||||
resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) || ${pauseStop}`,
|
||||
}).where(and(eq(heartbeatRuns.companyId, agent.companyId), eq(heartbeatRuns.agentId, agent.id),
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import { notifyMuseContactVerification } from "../modules/external-agents/index.js";
|
||||
import { notifyBudgetResumeAgentReady } from "./budget-resume-continuation.js";
|
||||
import { currentControllerProcessStamp } from "./controller-process-identity.js";
|
||||
import { notifyStrandedAgentStatus } from "./live-run-recovery-notifications.js";
|
||||
import { notifyDeferredExecutionEligibility } from "./deferred-execution-notifications.js";
|
||||
@@ -68,6 +69,7 @@ export function createAgentLifecycleEffects(_hooks: BudgetServiceHooks = {}): li
|
||||
await notifyNativeResponsePresentation(tx, { companyId, agentId });
|
||||
await notifyNativeReviewResponse(tx, { companyId, agentId });
|
||||
await notifyChatInboundWakeNotices(tx, { companyId });
|
||||
await notifyBudgetResumeAgentReady(tx, agentId);
|
||||
}
|
||||
},
|
||||
reportFailure(context, error) {
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import type { Db } from "@paperclipai/db";
|
||||
import { logger } from "../middleware/logger.js";
|
||||
import { withBudgetSpan } from "./budget-tracing.js";
|
||||
import type { createDeliveryWorkCoordinator } from "./delivery-work-coordinator.js";
|
||||
import { DELIVERY_QUEUES, subscribeDeliveryWork } from "./delivery-work-notifications.js";
|
||||
@@ -11,6 +12,7 @@ import {
|
||||
type BudgetServiceHooks,
|
||||
} from "./budgets.js";
|
||||
import { hasPendingBudgetEvaluation } from "./budget-evaluation-work.js";
|
||||
import { hasDeliverableBudgetResume, hasPendingBudgetResumeForRun, nextBudgetResumeClaimExpiryAt } from "./budget-resume-continuation.js";
|
||||
import { readCostAccountingWork } from "./run-cost-accounting.js";
|
||||
import { subscribeAllCompanyLiveEvents } from "./live-events.js";
|
||||
|
||||
@@ -32,7 +34,8 @@ export function registerBudgetEnforcementWork(db: Db, coordinator: Coordinator,
|
||||
let reconciling = false;
|
||||
let reconciliationFailed = false;
|
||||
let nextDeadline: number | null = null;
|
||||
return coordinator.register(DELIVERY_QUEUES.budgetEnforcement, {
|
||||
let resumeClaimDeadline: number | null = null;
|
||||
const worker = coordinator.register(DELIVERY_QUEUES.budgetEnforcement, {
|
||||
retryMs: 1_000,
|
||||
coalesceMs: 1_000,
|
||||
runDuringDrain: true,
|
||||
@@ -50,16 +53,44 @@ export function registerBudgetEnforcementWork(db: Db, coordinator: Coordinator,
|
||||
evaluationCursor = evaluation.nextCursor;
|
||||
if (signal.aborted) return;
|
||||
await withBudgetSpan("budget.enforcement.deliver", () => deliverBudgetEnforcement(db, hooks));
|
||||
if (signal.aborted) return;
|
||||
// Claims a resume saved, including runs whose Stop this pass just landed.
|
||||
// Registered with the always-on budget worker, not scheduled agent work.
|
||||
await hooks.deliverResumeContinuations?.().catch(error =>
|
||||
logger.warn({ err: error }, "Budget resume continuations remain pending"));
|
||||
},
|
||||
async hasPending() {
|
||||
const [pending, next, evaluation] = await Promise.all([
|
||||
const [pending, next, evaluation, resume, claimExpiry] = await Promise.all([
|
||||
hasPendingBudgetEnforcement(db), nextBudgetReconciliationAt(db), hasPendingBudgetEvaluation(db),
|
||||
// A deliverable resume claim (for example after a retryable refusal or
|
||||
// a database error) keeps this worker on its retry schedule.
|
||||
hooks.deliverResumeContinuations ? hasDeliverableBudgetResume(db) : Promise.resolve(false),
|
||||
// A claim held by a dispatcher that failed becomes deliverable when it
|
||||
// goes stale; wake for that moment without another budget signal.
|
||||
hooks.deliverResumeContinuations ? nextBudgetResumeClaimExpiryAt(db) : Promise.resolve(null),
|
||||
]);
|
||||
nextDeadline = next?.getTime() ?? null;
|
||||
return pending || evaluation || reconciling || startup || evaluationCursor !== null;
|
||||
resumeClaimDeadline = claimExpiry ?? null;
|
||||
return pending || evaluation || resume || reconciling || startup || evaluationCursor !== null;
|
||||
},
|
||||
nextRunAt: () => nextDeadline,
|
||||
nextRunAt: () => nextDeadline === null ? resumeClaimDeadline
|
||||
: resumeClaimDeadline === null ? nextDeadline : Math.min(nextDeadline, resumeClaimDeadline),
|
||||
});
|
||||
if (!hooks.deliverResumeContinuations) return worker;
|
||||
// A Stop that settles later (a pause's lifecycle Stop, a delivered budget
|
||||
// Stop) makes a pending resume claim deliverable. The terminal status event
|
||||
// is published after its persisted transition; wake for it, so delivery
|
||||
// never waits for unrelated budget work.
|
||||
// Only a run that carries a pending resume claim wakes the worker: one
|
||||
// primary-key lookup, so other cancellations never start an instance pass.
|
||||
const unsubscribe = subscribeAllCompanyLiveEvents(event => {
|
||||
if (event.type !== "heartbeat.run.status" || String(event.payload.status) !== "cancelled") return;
|
||||
const runId = event.payload.runId;
|
||||
if (typeof runId !== "string") return;
|
||||
void hasPendingBudgetResumeForRun(db, runId).then(pending => { if (pending) worker.wake(); })
|
||||
.catch(error => logger.warn({ err: error, runId }, "Budget resume wake check failed"));
|
||||
});
|
||||
return { ...worker, stop: async () => { unsubscribe(); await worker.stop(); } };
|
||||
}
|
||||
|
||||
/** Accounting repairs admitted work and never dispatches another paid turn.
|
||||
|
||||
@@ -0,0 +1,353 @@
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { and, eq, like } from "drizzle-orm";
|
||||
import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
|
||||
import {
|
||||
agents, agentWakeupRequests, budgetIncidents, budgetPolicies, budgetResumeClaims, companies, createDb, heartbeatRuns, issues,
|
||||
} from "@paperclipai/db";
|
||||
import { budgetService } from "./budgets.js";
|
||||
import { heartbeatService } from "./heartbeat.js";
|
||||
import { HttpError } from "../errors.js";
|
||||
import {
|
||||
deliverBudgetResumeContinuations, hasDeliverableBudgetResume, hasPendingBudgetResume, markBudgetCancelledWorkForResume,
|
||||
} from "./budget-resume-continuation.js";
|
||||
import {
|
||||
EMBEDDED_POSTGRES_TEST_TIMEOUT_MS, getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase,
|
||||
} from "../__tests__/helpers/embedded-postgres.js";
|
||||
|
||||
const mockLogActivity = vi.hoisted(() => vi.fn());
|
||||
vi.mock("./activity-log.js", async (importOriginal) => ({
|
||||
...await importOriginal<typeof import("./activity-log.js")>(),
|
||||
logActivity: mockLogActivity,
|
||||
}));
|
||||
|
||||
const support = await getEmbeddedPostgresTestSupport();
|
||||
|
||||
// H1 (restart chaos): a budget hard stop cancelled a run mid-task. After the
|
||||
// board raised the budget and resumed, the task stayed stopped until someone
|
||||
// clicked Retry. Resuming must continue that work, exactly once.
|
||||
describe.skipIf(!support.supported)("resuming after a budget hard stop", () => {
|
||||
let db!: ReturnType<typeof createDb>;
|
||||
let tempDb: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>> | null = null;
|
||||
|
||||
beforeAll(async () => {
|
||||
await import("./agent-lifecycle.js");
|
||||
tempDb = await startEmbeddedPostgresTestDatabase("paperclip-budget-resume-continuation-");
|
||||
db = createDb(tempDb.connectionString);
|
||||
}, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS);
|
||||
afterAll(async () => {
|
||||
await db?.$client.end();
|
||||
await tempDb?.cleanup();
|
||||
});
|
||||
|
||||
async function seedScope(label: string, scopeType: "company" | "agent" = "company") {
|
||||
const [company] = await db.insert(companies).values({ name: `Budget ${label}`,
|
||||
issuePrefix: `B${randomUUID().slice(0, 4).toUpperCase()}` }).returning();
|
||||
const [agent] = await db.insert(agents).values({ companyId: company.id, name: "Worker" }).returning();
|
||||
const scopeId = scopeType === "agent" ? agent.id : company.id;
|
||||
const [policy] = await db.insert(budgetPolicies).values({ companyId: company.id, scopeType,
|
||||
scopeId, windowKind: "calendar_month_utc", amount: 1, hardStopEnabled: true, isActive: true }).returning();
|
||||
const [incident] = await db.insert(budgetIncidents).values({ companyId: company.id, policyId: policy.id,
|
||||
scopeType, scopeId, metric: "billed_cents", windowKind: "calendar_month_utc",
|
||||
windowStart: new Date(Date.now() - 86_400_000), windowEnd: new Date(Date.now() + 86_400_000),
|
||||
thresholdType: "hard", amountLimit: 1, amountObserved: 0, createdAt: new Date(Date.now() - 60_000) }).returning();
|
||||
return { company, agent, policy, incident };
|
||||
}
|
||||
|
||||
async function seedTask(scope: Awaited<ReturnType<typeof seedScope>>, options: {
|
||||
budgetStop?: boolean; lifecycleStop?: "budget" | "agentBudget" | "manual"; issueStatus?: string; runStatus?: string; laterRun?: boolean;
|
||||
} = {}) {
|
||||
const { company, agent, policy } = scope;
|
||||
const [issue] = await db.insert(issues).values({ companyId: company.id, title: "Long task",
|
||||
status: options.issueStatus ?? "in_progress", assigneeAgentId: agent.id }).returning();
|
||||
const [run] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id,
|
||||
issueId: issue.id, status: options.runStatus ?? "cancelled", runtimeMode: "legacy",
|
||||
contextSnapshot: { issueId: issue.id }, createdAt: new Date(Date.now() - 50_000),
|
||||
resultJson: options.lifecycleStop ? {
|
||||
// H1: the agent's budget hold paused it, and the lifecycle stop cancelled the run.
|
||||
cancellation: { reason: "Agent lifecycle stop requested", source: "control_plane", expected: true,
|
||||
initiator: { type: "system" } },
|
||||
pauseStop: options.lifecycleStop === "budget" ? { holds: ["company_paused"], companyPauseReason: "budget" }
|
||||
: options.lifecycleStop === "agentBudget" ? { holds: ["budget"], companyPauseReason: null }
|
||||
: { holds: ["manual"], companyPauseReason: null },
|
||||
startupCancellation: { requestedAt: new Date(Date.now() - 40_000).toISOString(), errorCode: "cancelled" },
|
||||
} : options.budgetStop === false ? { stopReason: "cancelled" } : { startupCancellation: {
|
||||
requestedAt: new Date(Date.now() - 40_000).toISOString(), errorCode: "cancelled",
|
||||
budgetEnforcement: { companyId: company.id, scopeType: "company", scopeId: company.id,
|
||||
policyId: policy.id, version: 1 } } },
|
||||
}).returning();
|
||||
if (options.laterRun) await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id,
|
||||
issueId: issue.id, status: "succeeded", runtimeMode: "legacy", contextSnapshot: { issueId: issue.id },
|
||||
createdAt: new Date(Date.now() - 20_000) });
|
||||
return { issue, run };
|
||||
}
|
||||
|
||||
const resume = (scope: Awaited<ReturnType<typeof seedScope>>) => budgetService(db).resolveIncident(
|
||||
scope.company.id, scope.incident.id, { action: "raise_budget_and_resume", amount: 500 }, "local-board",
|
||||
"local_implicit");
|
||||
const continuations = (companyId: string) => db.select().from(agentWakeupRequests).where(and(
|
||||
eq(agentWakeupRequests.companyId, companyId), like(agentWakeupRequests.idempotencyKey, "task-retry:budget-resume:%")));
|
||||
const claimOf = async (runId: string) => (await db.select().from(budgetResumeClaims)
|
||||
.where(eq(budgetResumeClaims.runId, runId)))[0] as Record<string, unknown> | undefined;
|
||||
|
||||
it("continues each task the stop cancelled exactly once, through one continuation wake", async () => {
|
||||
const scope = await seedScope("continue");
|
||||
const working = await seedTask(scope);
|
||||
const blocked = await seedTask(scope, { issueStatus: "blocked" });
|
||||
await resume(scope);
|
||||
expect(await claimOf(working.run.id)).toMatchObject({ state: "pending", incidentId: scope.incident.id,
|
||||
actorUserId: "local-board", actorSource: "local_implicit" });
|
||||
expect(await hasPendingBudgetResume(db)).toBe(true);
|
||||
await deliverBudgetResumeContinuations(db);
|
||||
const saved = await continuations(scope.company.id);
|
||||
expect(saved.map(wake => wake.payload?.issueId).sort()).toEqual([working.issue.id, blocked.issue.id].sort());
|
||||
expect(saved.every(wake => wake.reason === "task_retry" && wake.requestedByActorId === "local-board")).toBe(true);
|
||||
expect(await claimOf(working.run.id)).toMatchObject({ state: "delivered", requestId: expect.any(String) });
|
||||
expect(await hasPendingBudgetResume(db)).toBe(false);
|
||||
expect(mockLogActivity).toHaveBeenCalledWith(expect.anything(), expect.objectContaining({
|
||||
action: "budget.cancelled_work_continued", entityId: working.issue.id, actorId: "local-board" }), expect.anything());
|
||||
// Delivering or resolving again files nothing new.
|
||||
await deliverBudgetResumeContinuations(db);
|
||||
await markBudgetCancelledWorkForResume(db, { companyId: scope.company.id, scopeType: "company",
|
||||
scopeId: scope.company.id, incidentId: scope.incident.id, incidentOpenedAt: scope.incident.createdAt,
|
||||
actorUserId: "local-board" });
|
||||
await deliverBudgetResumeContinuations(db);
|
||||
expect(await continuations(scope.company.id)).toHaveLength(2);
|
||||
});
|
||||
|
||||
it("continues a run whose stop was still being delivered when the board resumed", async () => {
|
||||
const scope = await seedScope("in-flight");
|
||||
const task = await seedTask(scope, { runStatus: "running" });
|
||||
await resume(scope);
|
||||
await deliverBudgetResumeContinuations(db);
|
||||
expect(await continuations(scope.company.id)).toEqual([]);
|
||||
expect(await claimOf(task.run.id)).toMatchObject({ state: "pending" });
|
||||
// The committed Stop lands after the resume. Its terminal write replaces
|
||||
// the run's result JSON with a copy built before the claim was saved; the
|
||||
// claim lives outside that JSON, so it survives.
|
||||
await db.update(heartbeatRuns).set({ status: "cancelled", finishedAt: new Date(),
|
||||
resultJson: { ...(task.run.resultJson ?? {}), stopReason: "cancelled" } }).where(eq(heartbeatRuns.id, task.run.id));
|
||||
await deliverBudgetResumeContinuations(db);
|
||||
expect(await continuations(scope.company.id)).toEqual([expect.objectContaining({
|
||||
payload: expect.objectContaining({ issueId: task.issue.id }) })]);
|
||||
});
|
||||
|
||||
it("dispatches the continuation through the real Retry path into a queued run", async () => {
|
||||
const scope = await seedScope("dispatch");
|
||||
const task = await seedTask(scope, { issueStatus: "blocked" });
|
||||
// Keep the new run queued: one run slot, already taken.
|
||||
await db.update(agents).set({ runtimeConfig: { heartbeat: { maxConcurrentRuns: 1 } } }).where(eq(agents.id, scope.agent.id));
|
||||
await db.insert(heartbeatRuns).values({ companyId: scope.company.id, agentId: scope.agent.id, status: "running" });
|
||||
await resume(scope);
|
||||
const heartbeat = heartbeatService(db);
|
||||
await deliverBudgetResumeContinuations(db, {
|
||||
resumeTaskRetry: (companyId, requestId) => heartbeat.resumeTaskRetry(companyId, requestId) });
|
||||
const queued = await db.select().from(heartbeatRuns).where(and(eq(heartbeatRuns.companyId, scope.company.id),
|
||||
eq(heartbeatRuns.status, "queued")));
|
||||
expect(queued, JSON.stringify(await db.select({ status: agentWakeupRequests.status, error: agentWakeupRequests.error,
|
||||
reason: agentWakeupRequests.reason, payload: agentWakeupRequests.payload }).from(agentWakeupRequests)
|
||||
.where(eq(agentWakeupRequests.companyId, scope.company.id)))).toHaveLength(1);
|
||||
expect(queued[0]).toMatchObject({ agentId: scope.agent.id });
|
||||
expect(queued[0].contextSnapshot).toMatchObject({ issueId: task.issue.id });
|
||||
});
|
||||
|
||||
it("continues every task of a large scope", async () => {
|
||||
const scope = await seedScope("large");
|
||||
for (let index = 0; index < 30; index += 1) await seedTask(scope);
|
||||
await resume(scope);
|
||||
await deliverBudgetResumeContinuations(db);
|
||||
expect(await continuations(scope.company.id)).toHaveLength(30);
|
||||
});
|
||||
|
||||
it.each([
|
||||
["work that already continued", { laterRun: true }, "work_continued_or_task_changed"],
|
||||
["a task that is now done", { issueStatus: "done" }, "work_continued_or_task_changed"],
|
||||
] as const)("does not continue %s", async (label, options, reason) => {
|
||||
const scope = await seedScope(label.replace(/\W+/g, "-"));
|
||||
const task = await seedTask(scope, options);
|
||||
await resume(scope);
|
||||
await deliverBudgetResumeContinuations(db);
|
||||
expect(await continuations(scope.company.id)).toEqual([]);
|
||||
expect(await claimOf(task.run.id)).toMatchObject({ state: "abandoned", reason });
|
||||
});
|
||||
|
||||
it("abandons a claim whose run finished before its stop landed", async () => {
|
||||
const scope = await seedScope("finished-first");
|
||||
const task = await seedTask(scope, { runStatus: "running" });
|
||||
await resume(scope);
|
||||
await db.update(heartbeatRuns).set({ status: "succeeded", finishedAt: new Date() }).where(eq(heartbeatRuns.id, task.run.id));
|
||||
await deliverBudgetResumeContinuations(db);
|
||||
expect(await continuations(scope.company.id)).toEqual([]);
|
||||
expect(await claimOf(task.run.id)).toMatchObject({ state: "abandoned", reason: "run_not_cancelled" });
|
||||
});
|
||||
|
||||
it("continues a run the agent's budget pause stopped through the lifecycle stop", async () => {
|
||||
const scope = await seedScope("lifecycle-stop");
|
||||
const task = await seedTask(scope, { lifecycleStop: "budget", issueStatus: "in_progress" });
|
||||
await resume(scope);
|
||||
expect(await claimOf(task.run.id)).toMatchObject({ state: "pending" });
|
||||
await deliverBudgetResumeContinuations(db);
|
||||
expect(await continuations(scope.company.id)).toEqual([expect.objectContaining({
|
||||
payload: expect.objectContaining({ issueId: task.issue.id }) })]);
|
||||
});
|
||||
|
||||
it("continues a run an agent-scope budget pause stopped through the lifecycle stop", async () => {
|
||||
const scope = await seedScope("agent-lifecycle-stop", "agent");
|
||||
const task = await seedTask(scope, { lifecycleStop: "agentBudget" });
|
||||
await resume(scope);
|
||||
await deliverBudgetResumeContinuations(db);
|
||||
expect(await continuations(scope.company.id)).toEqual([expect.objectContaining({
|
||||
payload: expect.objectContaining({ issueId: task.issue.id }) })]);
|
||||
});
|
||||
|
||||
it("continues a run whose budget pause Stop lands after the resume", async () => {
|
||||
const scope = await seedScope("late-pause-stop", "agent");
|
||||
// The budget pause started (its cause is on the lifecycle operation), but
|
||||
// the lifecycle worker has not stopped the run yet when the board resumes.
|
||||
await db.update(agents).set({ lifecycleState: "pausing", lifecycleHolds: [],
|
||||
lifecycleOperation: { id: randomUUID(), hostComplete: false, completedPluginIds: [], attempts: 0,
|
||||
pauseCause: { holds: ["budget"], companyPauseReason: null } } }).where(eq(agents.id, scope.agent.id));
|
||||
const task = await seedTask(scope, { budgetStop: false, runStatus: "running" });
|
||||
await resume(scope);
|
||||
expect(await claimOf(task.run.id)).toMatchObject({ state: "pending" });
|
||||
await db.update(heartbeatRuns).set({ status: "cancelled", finishedAt: new Date(), resultJson: {
|
||||
cancellation: { reason: "Agent lifecycle stop requested", source: "control_plane", expected: true, initiator: { type: "system" } },
|
||||
pauseStop: { holds: ["budget"], companyPauseReason: null } } }).where(eq(heartbeatRuns.id, task.run.id));
|
||||
await deliverBudgetResumeContinuations(db);
|
||||
expect(await continuations(scope.company.id)).toHaveLength(1);
|
||||
});
|
||||
|
||||
it("does not continue a run a person stopped after the resume claimed it", async () => {
|
||||
const scope = await seedScope("later-person-stop");
|
||||
const task = await seedTask(scope, { runStatus: "running" });
|
||||
await resume(scope);
|
||||
expect(await claimOf(task.run.id)).toMatchObject({ state: "pending" });
|
||||
await db.update(heartbeatRuns).set({ status: "cancelled", finishedAt: new Date(), resultJson: {
|
||||
...(task.run.resultJson ?? {}),
|
||||
cancellation: { reason: "Stopped by board", source: "operator", expected: true, initiator: { type: "user", id: "board-user" } } },
|
||||
}).where(eq(heartbeatRuns.id, task.run.id));
|
||||
await deliverBudgetResumeContinuations(db);
|
||||
expect(await continuations(scope.company.id)).toEqual([]);
|
||||
expect(await claimOf(task.run.id)).toMatchObject({ state: "abandoned", reason: "stopped_by_person" });
|
||||
});
|
||||
|
||||
it("releases the claim on a retryable Retry refusal, and settles it on a permanent one", async () => {
|
||||
const scope = await seedScope("retry-refusals");
|
||||
const task = await seedTask(scope);
|
||||
await resume(scope);
|
||||
const busy = vi.fn(async () => { throw new HttpError(409, "busy", { code: "task_retry_busy", retryable: true }); });
|
||||
await deliverBudgetResumeContinuations(db, { requestRetry: busy as never });
|
||||
expect(busy).toHaveBeenCalledOnce();
|
||||
expect(await claimOf(task.run.id)).toMatchObject({ state: "pending", claimedAt: null });
|
||||
const refused = vi.fn(async () => { throw new HttpError(409, "no", { code: "chat_failed_run_retry_not_authorized" }); });
|
||||
await deliverBudgetResumeContinuations(db, { requestRetry: refused as never });
|
||||
expect(await claimOf(task.run.id)).toMatchObject({ state: "abandoned", reason: "chat_failed_run_retry_not_authorized" });
|
||||
});
|
||||
|
||||
async function admissionStopped(label: string, cancellationReason: string) {
|
||||
const scope = await seedScope(label);
|
||||
const task = await seedTask(scope, { budgetStop: false });
|
||||
const reason = "Company cannot start work because its budget hard-stop is still exceeded.";
|
||||
await db.update(heartbeatRuns).set({ resultJson: { stopReason: "cancelled",
|
||||
cancellation: { reason: cancellationReason, source: "control_plane", expected: true, initiator: { type: "system" } },
|
||||
startupCancellation: { requestedAt: new Date(Date.now() - 30_000).toISOString(), errorCode: "cancelled" },
|
||||
budgetAdmissionStop: { scopeType: "company", scopeId: scope.company.id, reason } } }).where(eq(heartbeatRuns.id, task.run.id));
|
||||
return { scope, task, reason };
|
||||
}
|
||||
|
||||
it("continues a queued run that admission cancelled for this budget", async () => {
|
||||
const { scope, task } = await admissionStopped("admission-stop",
|
||||
"Company cannot start work because its budget hard-stop is still exceeded.");
|
||||
await resume(scope);
|
||||
await deliverBudgetResumeContinuations(db);
|
||||
expect(await continuations(scope.company.id)).toEqual([expect.objectContaining({
|
||||
payload: expect.objectContaining({ issueId: task.issue.id }) })]);
|
||||
});
|
||||
|
||||
it("never claims an admission-stamped run that another reason cancelled", async () => {
|
||||
// Admission saves its stamp before it cancels. A crash between them leaves
|
||||
// a queued, stamped run that may later be cancelled for something else.
|
||||
const { scope, task } = await admissionStopped("admission-other", "Cancelled because the agent no longer exists");
|
||||
await resume(scope);
|
||||
expect(await claimOf(task.run.id)).toBeUndefined();
|
||||
});
|
||||
|
||||
it("claims a queued run admission stamped, then follows how it was cancelled", async () => {
|
||||
// Admission saved its stamp, and the resume ran before its cancel landed.
|
||||
const reason = "Company cannot start work because its budget hard-stop is still exceeded.";
|
||||
for (const [label, cancelledFor, expected] of [
|
||||
["admission-pending", reason, "delivered"],
|
||||
["admission-pending-other", "Cancelled because the agent no longer exists", "abandoned"],
|
||||
] as const) {
|
||||
const scope = await seedScope(label);
|
||||
const task = await seedTask(scope, { budgetStop: false, runStatus: "queued" });
|
||||
await db.update(heartbeatRuns).set({ resultJson: {
|
||||
budgetAdmissionStop: { scopeType: "company", scopeId: scope.company.id, reason } } }).where(eq(heartbeatRuns.id, task.run.id));
|
||||
await resume(scope);
|
||||
expect(await claimOf(task.run.id)).toMatchObject({ state: "pending" });
|
||||
await db.update(heartbeatRuns).set({ status: "cancelled", resultJson: {
|
||||
budgetAdmissionStop: { scopeType: "company", scopeId: scope.company.id, reason },
|
||||
cancellation: { reason: cancelledFor, source: "control_plane", expected: true, initiator: { type: "system" } } } })
|
||||
.where(eq(heartbeatRuns.id, task.run.id));
|
||||
await deliverBudgetResumeContinuations(db);
|
||||
expect(await claimOf(task.run.id)).toMatchObject({ state: expected,
|
||||
...(expected === "abandoned" ? { reason: "cancelled_for_another_reason" } : {}) });
|
||||
}
|
||||
});
|
||||
|
||||
it("releases a stale claim while its agent is paused, so no pass is scheduled for it", async () => {
|
||||
const scope = await seedScope("stale-paused");
|
||||
const task = await seedTask(scope);
|
||||
await resume(scope);
|
||||
await db.update(budgetResumeClaims).set({ claimedAt: new Date(Date.now() - 120_000) })
|
||||
.where(eq(budgetResumeClaims.runId, task.run.id));
|
||||
await db.update(agents).set({ status: "paused" }).where(eq(agents.id, scope.agent.id));
|
||||
await deliverBudgetResumeContinuations(db);
|
||||
expect(await claimOf(task.run.id)).toMatchObject({ state: "pending", claimedAt: null });
|
||||
await db.update(agents).set({ status: "idle" }).where(eq(agents.id, scope.agent.id));
|
||||
await deliverBudgetResumeContinuations(db);
|
||||
expect(await claimOf(task.run.id)).toMatchObject({ state: "delivered" });
|
||||
});
|
||||
|
||||
it("waits for a paused agent to resume before filing its Retry", async () => {
|
||||
// A connected chat refuses a Retry while the resume is still lifting the
|
||||
// agent's hold; that refusal must not abandon the continuation.
|
||||
const scope = await seedScope("agent-resuming");
|
||||
const task = await seedTask(scope);
|
||||
await db.update(agents).set({ status: "paused" }).where(eq(agents.id, scope.agent.id));
|
||||
await resume(scope);
|
||||
await db.update(agents).set({ status: "paused" }).where(eq(agents.id, scope.agent.id));
|
||||
const requestRetry = vi.fn(async () => ({ accepted: true, requestId: null }));
|
||||
await deliverBudgetResumeContinuations(db, { requestRetry: requestRetry as never });
|
||||
expect(requestRetry).not.toHaveBeenCalled();
|
||||
expect(await hasDeliverableBudgetResume(db)).toBe(false);
|
||||
expect(await claimOf(task.run.id)).toMatchObject({ state: "pending", claimedAt: null });
|
||||
await db.update(agents).set({ status: "idle" }).where(eq(agents.id, scope.agent.id));
|
||||
expect(await hasDeliverableBudgetResume(db)).toBe(true);
|
||||
await deliverBudgetResumeContinuations(db, { requestRetry: requestRetry as never });
|
||||
expect(requestRetry).toHaveBeenCalledOnce();
|
||||
});
|
||||
|
||||
it("never claims a run a manual agent pause stopped", async () => {
|
||||
const scope = await seedScope("manual-pause");
|
||||
const task = await seedTask(scope, { lifecycleStop: "manual" });
|
||||
await resume(scope);
|
||||
expect(await claimOf(task.run.id)).toBeUndefined();
|
||||
});
|
||||
|
||||
it("never claims a run a person stopped while budget enforcement also fenced it", async () => {
|
||||
const scope = await seedScope("person-and-budget");
|
||||
const task = await seedTask(scope);
|
||||
await db.update(heartbeatRuns).set({ resultJson: { ...(task.run.resultJson ?? {}),
|
||||
cancellation: { reason: "Stopped by board", source: "operator", expected: true, initiator: { type: "user", id: "board-user" } },
|
||||
cancelledByActorType: "user" } }).where(eq(heartbeatRuns.id, task.run.id));
|
||||
await resume(scope);
|
||||
expect(await claimOf(task.run.id)).toBeUndefined();
|
||||
});
|
||||
|
||||
it("never claims a run a person stopped", async () => {
|
||||
const scope = await seedScope("person-stop");
|
||||
const task = await seedTask(scope, { budgetStop: false });
|
||||
await resume(scope);
|
||||
expect(await claimOf(task.run.id)).toBeUndefined();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,352 @@
|
||||
import { and, asc, eq, gt, inArray, isNull, lte, notInArray, or, sql } from "drizzle-orm";
|
||||
import { agents, budgetResumeClaims, heartbeatRuns, type Db } from "@paperclipai/db";
|
||||
import { HttpError } from "../errors.js";
|
||||
import { logger } from "../middleware/logger.js";
|
||||
import { logActivity, publishActivity, type ActivityPublication } from "./activity-log.js";
|
||||
import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notifications.js";
|
||||
import { budgetCancellationScopePredicate } from "./heartbeat/budget-cancellation.js";
|
||||
import { AGENT_LIFECYCLE_STOP_REASON } from "./governance-pause-stop.js";
|
||||
|
||||
/**
|
||||
* Resuming after a budget hard stop continues the work that the stop
|
||||
* cancelled, exactly once.
|
||||
*
|
||||
* When a board member resolves the incident with Raise budget and resume,
|
||||
* `markBudgetCancelledWorkForResume` saves a pending `budget_resume_claims` row
|
||||
* (the board member and their session source) for every task run this
|
||||
* scope's stop cancelled, before the incident closes. That includes runs whose
|
||||
* stop has not been delivered yet. Claims live outside the run's result JSON,
|
||||
* which terminal writers replace, so a Stop settling at the same moment cannot
|
||||
* erase one. The budget-enforcement worker (and the execution-continuation
|
||||
* worker as a backstop) calls `deliverBudgetResumeContinuations`: once a run's
|
||||
* stop has settled (`cancelled`), it files that board member's Retry.
|
||||
* - The Retry path clears the stopped task's execution hold, and accepts only
|
||||
* while the cancelled run is still the task's latest run and the task is
|
||||
* open and assigned, checked under the task lock. Work that already
|
||||
* continued or completed is never repeated.
|
||||
* - Its idempotency key names the cancelled run, so every pass, restart or
|
||||
* repeated resume returns the same single request.
|
||||
* - The claim and its activity entry settle in one commit.
|
||||
* - A Retry starts a new run through normal admission; nothing charges the
|
||||
* cancelled run's provider attempt or failure budget.
|
||||
*/
|
||||
|
||||
/** A claim older than this with no settled outcome belongs to an interrupted dispatcher. */
|
||||
export const BUDGET_RESUME_CLAIM_STALE_MS = 60_000;
|
||||
const PAGE_SIZE = 25;
|
||||
const ACTIVE_RUN_STATUSES = ["queued", "running", "scheduled_retry"] as const;
|
||||
/** Agent statuses that cannot take a Retry yet but may later: a resume still
|
||||
* lifting the agent's hold, or an approval. Their claims wait, unclaimed, and
|
||||
* the agent's status change wakes delivery (`notifyBudgetResumeAgentReady`). */
|
||||
const WAITING_AGENT_STATUSES = ["paused", "pending_approval"] as const;
|
||||
|
||||
type BudgetResumeClaim = {
|
||||
incidentId: string; actorUserId: string; actorSource: string | null; claimedAt: Date | null;
|
||||
};
|
||||
|
||||
type PrepareChatRetry = NonNullable<Parameters<typeof import("./task-retry.js").requestTaskRetry>[0]["prepareChatRetry"]>;
|
||||
|
||||
export type BudgetResumeDelivery = {
|
||||
/** The chat service's Retry callback, which connected conversations need. */
|
||||
prepareChatRetry?: PrepareChatRetry;
|
||||
/** Optional latency nudge; the deferred-receipt worker also dispatches the Retry. */
|
||||
resumeTaskRetry?: (companyId: string, requestId: string) => Promise<unknown>;
|
||||
/** The Retry path; defaults to requestTaskRetry. Tests substitute refusals. */
|
||||
requestRetry?: typeof import("./task-retry.js").requestTaskRetry;
|
||||
};
|
||||
|
||||
const runIssueId = sql<string | null>`coalesce(${heartbeatRuns.issueId}::text, ${heartbeatRuns.nativeIssueId}::text,
|
||||
${heartbeatRuns.contextSnapshot}->>'issueId')`;
|
||||
|
||||
/**
|
||||
* Stamp every task run this scope's budget stop cancelled since the incident
|
||||
* opened: runs the budget enforcement fenced, and runs the agent lifecycle
|
||||
* pause stopped for this budget (the run's `pauseStop` names the agent's
|
||||
* budget hold, or the company's budget pause). Runs that already carry a claim
|
||||
* keep it, so resolving again is a no-op. Arms the continuation worker in the
|
||||
* same commit. Two indexed passes: cancelled runs, and runs whose stop is
|
||||
* still being delivered.
|
||||
*/
|
||||
export async function markBudgetCancelledWorkForResume(db: Db, input: {
|
||||
companyId: string; scopeType: "company" | "agent" | "project"; scopeId: string;
|
||||
incidentId: string; incidentOpenedAt: Date; actorUserId: string; actorSource?: string | null;
|
||||
}): Promise<string[]> {
|
||||
return db.transaction(async tx => {
|
||||
// Register the work signals before the write they announce: the budget
|
||||
// worker delivers claims; the execution-continuation worker is a backstop.
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.budgetEnforcement, input.companyId);
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.executionContinuation);
|
||||
const openedAt = input.incidentOpenedAt.toISOString();
|
||||
// The pause records the holds it stopped the run for (pauseStop). Only a
|
||||
// stop this budget caused qualifies: the agent's own budget hold, or the
|
||||
// company's budget pause. A manual pause never does. The holds come from
|
||||
// the lifecycle operation's pauseCause, saved when the pause started, so a
|
||||
// resume that removed the holds before the Stop ran still leaves them here.
|
||||
const pauseStop = sql`${heartbeatRuns.resultJson}->'pauseStop'`;
|
||||
const lifecycleStop = input.scopeType === "project" ? sql`false` : and(
|
||||
sql`${heartbeatRuns.resultJson}->'cancellation'->>'reason' = ${AGENT_LIFECYCLE_STOP_REASON}`,
|
||||
sql`${heartbeatRuns.resultJson}->'cancellation'->'initiator'->>'type' = 'system'`,
|
||||
input.scopeType === "agent"
|
||||
? and(eq(heartbeatRuns.agentId, input.scopeId), sql`(${pauseStop}->'holds') ? 'budget'`)
|
||||
: sql`(${pauseStop}->'holds') ? 'company_paused' and ${pauseStop}->>'companyPauseReason' = 'budget'`)!;
|
||||
// A run still active on an agent whose pause this budget caused: the
|
||||
// pause's Stop may land after the incident closes. The lifecycle operation
|
||||
// recorded the cause when the pause started (pauseCause).
|
||||
const pauseCause = sql`a.lifecycle_operation->'pauseCause'`;
|
||||
const pendingPauseStop = input.scopeType === "project" ? sql`false` : sql`exists (select 1 from agents a
|
||||
where a.id = ${heartbeatRuns.agentId} and a.company_id = ${heartbeatRuns.companyId}
|
||||
and a.lifecycle_state = 'pausing'
|
||||
and ${input.scopeType === "agent"
|
||||
? sql`a.id = ${input.scopeId} and (${pauseCause}->'holds') ? 'budget'`
|
||||
: sql`(${pauseCause}->'holds') ? 'company_paused' and ${pauseCause}->>'companyPauseReason' = 'budget'`})`;
|
||||
// A queued run that admission cancelled for this budget's board-owned block.
|
||||
// Admission saves the stamp before it cancels; the recorded cancellation
|
||||
// must carry the stamp's reason, so a stamped run that later ran, or was
|
||||
// cancelled for another reason, never matches.
|
||||
const admissionStop = and(
|
||||
sql`${heartbeatRuns.resultJson}->'budgetAdmissionStop'->>'scopeType' = ${input.scopeType}`,
|
||||
sql`${heartbeatRuns.resultJson}->'budgetAdmissionStop'->>'scopeId' = ${input.scopeId}`,
|
||||
sql`${heartbeatRuns.resultJson}->'cancellation'->>'reason' = ${heartbeatRuns.resultJson}->'budgetAdmissionStop'->>'reason'`)!;
|
||||
const stoppedByThisScope = or(budgetCancellationScopePredicate(input), lifecycleStop, admissionStop)!;
|
||||
// A queued run admission stamped for this budget but has not cancelled
|
||||
// yet. Delivery waits for its Stop, and abandons the claim if the run was
|
||||
// cancelled for another reason.
|
||||
const pendingAdmissionStop = and(
|
||||
sql`${heartbeatRuns.resultJson}->'budgetAdmissionStop'->>'scopeType' = ${input.scopeType}`,
|
||||
sql`${heartbeatRuns.resultJson}->'budgetAdmissionStop'->>'scopeId' = ${input.scopeId}`)!;
|
||||
// A person's Stop stays stopped, even if budget enforcement also fenced
|
||||
// the run while that Stop was being delivered.
|
||||
const stoppedByPerson = sql`(coalesce(${heartbeatRuns.resultJson}->'cancellation'->'initiator'->>'type', '') = 'user'
|
||||
or coalesce(${heartbeatRuns.resultJson}->'cancellation'->>'source', '') = 'operator'
|
||||
or coalesce(${heartbeatRuns.resultJson}->>'cancelledByActorType', '') in ('user', 'board'))`;
|
||||
const claimable = and(
|
||||
eq(heartbeatRuns.companyId, input.companyId),
|
||||
sql`not ${stoppedByPerson}`,
|
||||
sql`${runIssueId} is not null`,
|
||||
);
|
||||
const stoppedSinceIncident = and(stoppedByThisScope,
|
||||
sql`coalesce((${heartbeatRuns.resultJson}->'startupCancellation'->>'requestedAt')::timestamptz,
|
||||
${heartbeatRuns.updatedAt}) >= ${openedAt}::timestamptz`)!;
|
||||
// Both passes lead with (company_id, status), which the company status
|
||||
// indexes serve. A run that finished any other way is owed nothing. An
|
||||
// existing claim is kept, so resolving again saves nothing new.
|
||||
const claim = (status: ReturnType<typeof eq>) => tx.insert(budgetResumeClaims).select(
|
||||
tx.select({ runId: heartbeatRuns.id, companyId: heartbeatRuns.companyId,
|
||||
incidentId: sql<string>`${input.incidentId}::uuid`.as("incident_id"),
|
||||
actorUserId: sql<string>`${input.actorUserId}`.as("actor_user_id"),
|
||||
actorSource: sql<string | null>`${input.actorSource ?? null}::text`.as("actor_source"),
|
||||
state: sql<string>`'pending'`.as("state"),
|
||||
claimedAt: sql<Date | null>`null::timestamptz`.as("claimed_at"),
|
||||
settledAt: sql<Date | null>`null::timestamptz`.as("settled_at"),
|
||||
reason: sql<string | null>`null::text`.as("reason"),
|
||||
requestId: sql<string | null>`null::uuid`.as("request_id"),
|
||||
createdAt: sql<Date>`now()`.as("created_at"),
|
||||
}).from(heartbeatRuns).where(and(claimable, status)),
|
||||
).onConflictDoNothing().returning({ id: budgetResumeClaims.runId });
|
||||
const finished = await claim(and(eq(heartbeatRuns.status, "cancelled"), stoppedSinceIncident) as ReturnType<typeof eq>);
|
||||
// Active runs are being stopped now: fenced by this budget's enforcement,
|
||||
// or on an agent whose budget pause has not delivered its Stop yet.
|
||||
const active = await claim(and(inArray(heartbeatRuns.status, [...ACTIVE_RUN_STATUSES]),
|
||||
or(stoppedSinceIncident, pendingPauseStop, pendingAdmissionStop)) as ReturnType<typeof eq>);
|
||||
return [...finished, ...active].map(row => row.id);
|
||||
});
|
||||
}
|
||||
|
||||
/** Compare-and-set the run's claim. Only the writer that wins dispatches. */
|
||||
async function setClaim(db: Pick<Db, "update">, input: {
|
||||
runId: string; expectedClaimedAt: Date | null;
|
||||
next: Partial<Pick<typeof budgetResumeClaims.$inferInsert, "state" | "claimedAt" | "settledAt" | "reason" | "requestId">>;
|
||||
}) {
|
||||
const [row] = await db.update(budgetResumeClaims).set(input.next).where(and(
|
||||
eq(budgetResumeClaims.runId, input.runId), eq(budgetResumeClaims.state, "pending"),
|
||||
input.expectedClaimedAt === null ? isNull(budgetResumeClaims.claimedAt)
|
||||
: sql`date_trunc('milliseconds', ${budgetResumeClaims.claimedAt}) = ${input.expectedClaimedAt.toISOString()}::timestamptz`,
|
||||
)).returning({ id: budgetResumeClaims.runId });
|
||||
return Boolean(row);
|
||||
}
|
||||
|
||||
/** True while a resumed budget stop still owes a continuation. */
|
||||
export async function hasPendingBudgetResume(db: Db) {
|
||||
const [row] = await db.select({ id: budgetResumeClaims.runId }).from(budgetResumeClaims)
|
||||
.where(eq(budgetResumeClaims.state, "pending")).limit(1);
|
||||
return Boolean(row);
|
||||
}
|
||||
|
||||
/** True when this run carries a pending resume claim (a primary-key lookup). */
|
||||
export async function hasPendingBudgetResumeForRun(db: Db, runId: string) {
|
||||
const [row] = await db.select({ id: budgetResumeClaims.runId }).from(budgetResumeClaims)
|
||||
.where(and(eq(budgetResumeClaims.runId, runId), eq(budgetResumeClaims.state, "pending"))).limit(1);
|
||||
return Boolean(row);
|
||||
}
|
||||
|
||||
/** True while a claim can be delivered now: its stop settled and no live
|
||||
* dispatcher holds it. A claim waiting for its Stop re-arms with that Stop. */
|
||||
export async function hasDeliverableBudgetResume(db: Db, now = new Date()) {
|
||||
const [row] = await db.select({ id: budgetResumeClaims.runId }).from(budgetResumeClaims)
|
||||
.innerJoin(heartbeatRuns, eq(heartbeatRuns.id, budgetResumeClaims.runId))
|
||||
.innerJoin(agents, eq(agents.id, heartbeatRuns.agentId))
|
||||
.where(and(eq(budgetResumeClaims.state, "pending"),
|
||||
notInArray(heartbeatRuns.status, [...ACTIVE_RUN_STATUSES]),
|
||||
notInArray(agents.status, [...WAITING_AGENT_STATUSES]),
|
||||
or(isNull(budgetResumeClaims.claimedAt),
|
||||
lte(budgetResumeClaims.claimedAt, new Date(now.getTime() - BUDGET_RESUME_CLAIM_STALE_MS))))).limit(1);
|
||||
return Boolean(row);
|
||||
}
|
||||
|
||||
/** When the earliest claimed, undelivered claim goes stale and becomes
|
||||
* deliverable again, so a worker can wake for it without another signal. */
|
||||
export async function nextBudgetResumeClaimExpiryAt(db: Db): Promise<number | null> {
|
||||
const [row] = await db.select({ at: sql<Date | null>`min(${budgetResumeClaims.claimedAt})` }).from(budgetResumeClaims)
|
||||
.where(and(eq(budgetResumeClaims.state, "pending"), sql`${budgetResumeClaims.claimedAt} is not null`));
|
||||
const claimedAt = row?.at ? new Date(row.at).getTime() : null;
|
||||
return claimedAt === null ? null : claimedAt + BUDGET_RESUME_CLAIM_STALE_MS;
|
||||
}
|
||||
|
||||
type PendingRow = {
|
||||
id: string; companyId: string; agentId: string; agentStatus: string; status: string; createdAt: Date;
|
||||
issueId: string | null; stoppedByPerson: boolean; otherCancellation: boolean; claim: BudgetResumeClaim | null;
|
||||
};
|
||||
|
||||
/**
|
||||
* Deliver every pending claim whose stop has settled. Pages through all
|
||||
* pending claims; a claim that waits (stop not settled, or a live
|
||||
* dispatcher's fresh claim) is skipped without a write.
|
||||
*/
|
||||
export async function deliverBudgetResumeContinuations(db: Db, delivery: BudgetResumeDelivery = {}, now = new Date()) {
|
||||
let after: string | null = null;
|
||||
for (;;) {
|
||||
const page = await db.select({
|
||||
id: heartbeatRuns.id, companyId: heartbeatRuns.companyId, agentId: heartbeatRuns.agentId,
|
||||
agentStatus: agents.status, status: heartbeatRuns.status, createdAt: heartbeatRuns.createdAt, issueId: runIssueId,
|
||||
incidentId: budgetResumeClaims.incidentId, actorUserId: budgetResumeClaims.actorUserId,
|
||||
actorSource: budgetResumeClaims.actorSource, claimedAt: budgetResumeClaims.claimedAt,
|
||||
stoppedByPerson: sql<boolean>`(coalesce(${heartbeatRuns.resultJson}->'cancellation'->'initiator'->>'type', '') = 'user'
|
||||
or coalesce(${heartbeatRuns.resultJson}->'cancellation'->>'source', '') = 'operator'
|
||||
or coalesce(${heartbeatRuns.resultJson}->>'cancelledByActorType', '') in ('user', 'board'))`,
|
||||
// Claimed only for its admission stamp, but cancelled for another reason.
|
||||
otherCancellation: sql<boolean>`(${heartbeatRuns.resultJson} ? 'budgetAdmissionStop'
|
||||
and coalesce(${heartbeatRuns.resultJson}->'cancellation'->>'reason', '')
|
||||
<> coalesce(${heartbeatRuns.resultJson}->'budgetAdmissionStop'->>'reason', '')
|
||||
and not coalesce(${heartbeatRuns.resultJson}->'startupCancellation' ? 'budgetEnforcement', false)
|
||||
and not ${heartbeatRuns.resultJson} ? 'pauseStop')`,
|
||||
}).from(budgetResumeClaims).innerJoin(heartbeatRuns, eq(heartbeatRuns.id, budgetResumeClaims.runId))
|
||||
.innerJoin(agents, eq(agents.id, heartbeatRuns.agentId))
|
||||
.where(and(eq(budgetResumeClaims.state, "pending"), after ? gt(budgetResumeClaims.runId, after) : undefined))
|
||||
.orderBy(asc(budgetResumeClaims.runId)).limit(PAGE_SIZE);
|
||||
if (!page.length) return;
|
||||
after = page[page.length - 1]!.id;
|
||||
const rows: PendingRow[] = page.map(({ incidentId, actorUserId, actorSource, claimedAt, ...run }) =>
|
||||
({ ...run, claim: { incidentId, actorUserId, actorSource, claimedAt } }));
|
||||
for (const run of rows) {
|
||||
// One task's failure must not stop the others. Its claim goes stale and
|
||||
// a later pass retries it.
|
||||
try { await deliverOne(db, delivery, now, run); }
|
||||
catch (error) { logger.warn({ err: error, runId: run.id }, "Budget resume continuation remains pending"); }
|
||||
}
|
||||
if (rows.length < PAGE_SIZE) return;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Settle a claim and record why, in one commit. A crash before it leaves the
|
||||
* claim pending; the Retry's idempotency key makes the next pass return the
|
||||
* same request, so the record is never lost or doubled.
|
||||
*/
|
||||
async function settle(db: Db, run: PendingRow, expectedClaimedAt: Date | null, outcome: {
|
||||
state: "delivered" | "abandoned"; reason?: string; requestId?: string | null;
|
||||
}) {
|
||||
const claim = run.claim!;
|
||||
const publications: ActivityPublication[] = [];
|
||||
const settled = await db.transaction(async tx => {
|
||||
if (!await setClaim(tx, { runId: run.id, expectedClaimedAt,
|
||||
next: { state: outcome.state, settledAt: new Date(),
|
||||
...(outcome.reason ? { reason: outcome.reason } : {}),
|
||||
...(outcome.requestId !== undefined ? { requestId: outcome.requestId } : {}) } })) return false;
|
||||
await logActivity(tx as unknown as Db, { companyId: run.companyId, actorType: "user", actorId: claim.actorUserId,
|
||||
action: outcome.state === "delivered" ? "budget.cancelled_work_continued" : "budget.cancelled_work_not_continued",
|
||||
entityType: "issue", entityId: run.issueId!,
|
||||
details: { budgetIncidentId: claim.incidentId, interruptedRunId: run.id,
|
||||
...(outcome.requestId ? { retryRequestId: outcome.requestId } : {}),
|
||||
...(outcome.reason ? { reason: outcome.reason } : {}) } }, publications);
|
||||
return true;
|
||||
});
|
||||
if (settled) for (const publication of publications) publishActivity(publication);
|
||||
return settled;
|
||||
}
|
||||
|
||||
async function deliverOne(db: Db, delivery: BudgetResumeDelivery, now: Date, run: PendingRow) {
|
||||
const claim = run.claim;
|
||||
if (!claim || !run.issueId) return;
|
||||
// The stop is still being delivered; revisit once it settles.
|
||||
if ((ACTIVE_RUN_STATUSES as readonly string[]).includes(run.status)) return;
|
||||
const claimedAt = claim.claimedAt ?? null;
|
||||
// A fresh claim belongs to a live dispatcher.
|
||||
if (claimedAt && now.getTime() - claimedAt.getTime() < BUDGET_RESUME_CLAIM_STALE_MS) return;
|
||||
// The run finished before its stop landed: nothing was cancelled.
|
||||
if (run.status !== "cancelled") { await settle(db, run, claimedAt, { state: "abandoned", reason: "run_not_cancelled" }); return; }
|
||||
// A person stopped the run after the resume claimed it: their Stop wins.
|
||||
if (run.stoppedByPerson) { await settle(db, run, claimedAt, { state: "abandoned", reason: "stopped_by_person" }); return; }
|
||||
if (run.otherCancellation) { await settle(db, run, claimedAt, { state: "abandoned", reason: "cancelled_for_another_reason" }); return; }
|
||||
// The resume is still lifting the agent's hold (a connected chat refuses a
|
||||
// Retry for a paused agent). Wait without claiming; the agent's status
|
||||
// change wakes delivery.
|
||||
// A stale claim is released first, so its expiry no longer schedules passes.
|
||||
if ((WAITING_AGENT_STATUSES as readonly string[]).includes(run.agentStatus)) {
|
||||
if (claimedAt) await setClaim(db, { runId: run.id, expectedClaimedAt: claimedAt, next: { claimedAt: null } });
|
||||
return;
|
||||
}
|
||||
const nextClaim = new Date(now.getTime());
|
||||
if (!await setClaim(db, { runId: run.id, expectedClaimedAt: claimedAt, next: { claimedAt: nextClaim } })) return;
|
||||
// The board member's own Retry: the path that clears a stopped task's
|
||||
// execution hold, and that accepts only while this run is still the task's
|
||||
// latest run and the task is open and assigned (checked under the task lock).
|
||||
const requestTaskRetry = delivery.requestRetry ?? (await import("./task-retry.js")).requestTaskRetry;
|
||||
let receipt: Awaited<ReturnType<typeof requestTaskRetry>>;
|
||||
try {
|
||||
receipt = await requestTaskRetry({ db, companyId: run.companyId, issueId: run.issueId,
|
||||
observedRunId: run.id, idempotencyKey: `budget-resume:${run.id}`, actorId: claim.actorUserId,
|
||||
...(claim.actorSource ? { actorSource: claim.actorSource } : {}),
|
||||
...(delivery.prepareChatRetry ? { prepareChatRetry: delivery.prepareChatRetry } : {}) });
|
||||
} catch (error) {
|
||||
// A refusal the Retry path names is permanent for this claim; record it.
|
||||
// Anything else (a database error) leaves the claim to go stale and retry.
|
||||
// A refusal the Retry path marks retryable (brief chat-authorization
|
||||
// contention) releases the claim for the next pass.
|
||||
if (error instanceof HttpError && (error.details as { retryable?: unknown } | undefined)?.retryable === true) {
|
||||
await setClaim(db, { runId: run.id, expectedClaimedAt: nextClaim, next: { claimedAt: null } });
|
||||
return;
|
||||
}
|
||||
if (error instanceof HttpError && error.status >= 400 && error.status < 500) {
|
||||
const code = (error.details as { code?: unknown } | undefined)?.code;
|
||||
await settle(db, run, nextClaim, { state: "abandoned",
|
||||
reason: typeof code === "string" ? code : `retry_refused_${error.status}` });
|
||||
return;
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
if (!receipt.accepted) {
|
||||
await settle(db, run, nextClaim, { state: "abandoned", reason: "work_continued_or_task_changed" });
|
||||
return;
|
||||
}
|
||||
if (!await settle(db, run, nextClaim, { state: "delivered", requestId: receipt.requestId ?? null })) return;
|
||||
if (receipt.requestId && !receipt.chatActionId && delivery.resumeTaskRetry)
|
||||
await delivery.resumeTaskRetry(run.companyId, receipt.requestId).catch(error =>
|
||||
logger.warn({ err: error, runId: run.id }, "Budget resume Retry remains queued"));
|
||||
}
|
||||
|
||||
/** Wake budget delivery when an agent becomes able to take work and a
|
||||
* resume claim waits on it. Runs inside the agent's status-change transaction. */
|
||||
export async function notifyBudgetResumeAgentReady(tx: Db, agentId: string) {
|
||||
const [pending] = await tx.select({ id: budgetResumeClaims.runId }).from(budgetResumeClaims)
|
||||
.innerJoin(heartbeatRuns, eq(heartbeatRuns.id, budgetResumeClaims.runId))
|
||||
.where(and(eq(heartbeatRuns.agentId, agentId), eq(budgetResumeClaims.state, "pending"))).limit(1);
|
||||
if (pending) await notifyDeliveryWork(tx, DELIVERY_QUEUES.budgetEnforcement);
|
||||
}
|
||||
|
||||
/** Wake the continuation worker after budget Stops land. */
|
||||
export async function notifyBudgetStopsSettled(db: Db, runIds: string[]) {
|
||||
if (!runIds.length) return;
|
||||
const [pending] = await db.select({ id: budgetResumeClaims.runId }).from(budgetResumeClaims).where(and(
|
||||
inArray(budgetResumeClaims.runId, runIds), eq(budgetResumeClaims.state, "pending"))).limit(1);
|
||||
if (pending) await notifyDeliveryWork(db, DELIVERY_QUEUES.executionContinuation);
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
import { markBudgetCancelledWorkForResume } from "./budget-resume-continuation.js";
|
||||
import { notifyMuseContactVerification } from "../modules/external-agents/index.js";
|
||||
import { promoteDeferredWorkIntents } from "./work-intent-outbox.js";
|
||||
import { notifyDeferredExecutionEligibility } from "./deferred-execution-notifications.js";
|
||||
@@ -70,6 +71,8 @@ export type BudgetEnforcementScope = {
|
||||
|
||||
export type BudgetServiceHooks = {
|
||||
cancelWorkForScope?: (scope: BudgetEnforcementScope) => Promise<void>;
|
||||
/** Deliver budget-resume continuations; the enforcement worker calls it after Stop delivery. */
|
||||
deliverResumeContinuations?: () => Promise<void>;
|
||||
};
|
||||
|
||||
function currentUtcMonthWindow(now = new Date()) {
|
||||
@@ -1007,6 +1010,8 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
|
||||
incidentId: string,
|
||||
input: BudgetIncidentResolutionInput,
|
||||
actorUserId: string,
|
||||
/** The resolving actor's session source, kept for the continuation Retry. */
|
||||
actorSource?: string | null,
|
||||
): Promise<BudgetIncident> => {
|
||||
const incident = await db
|
||||
.select()
|
||||
@@ -1066,6 +1071,24 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
|
||||
}
|
||||
|
||||
await resumeScopeFromBudget(policy);
|
||||
// Claim the work this stop cancelled for one continuation each, before
|
||||
// the incident closes: a crash before here leaves the incident open,
|
||||
// and resolving it again keeps existing claims. The continuation
|
||||
// worker delivers each claim once its stop has settled.
|
||||
const claimedRunIds = await markBudgetCancelledWorkForResume(db, {
|
||||
companyId, scopeType: policy.scopeType as BudgetScopeType, scopeId: policy.scopeId,
|
||||
incidentId: incident.id, incidentOpenedAt: incident.createdAt, actorUserId, actorSource,
|
||||
});
|
||||
if (claimedRunIds.length > 0) await recordActivity({
|
||||
companyId,
|
||||
actorType: "user",
|
||||
actorId: actorUserId,
|
||||
action: "budget.cancelled_work_resume_requested",
|
||||
entityType: "budget_incident",
|
||||
entityId: incident.id,
|
||||
details: { scopeType: policy.scopeType, scopeId: policy.scopeId, runCount: claimedRunIds.length,
|
||||
sampleRunIds: claimedRunIds.slice(0, 20) },
|
||||
});
|
||||
await db
|
||||
.update(budgetIncidents)
|
||||
.set({
|
||||
@@ -1355,11 +1378,12 @@ export function budgetService(db: Db, _hooks: BudgetServiceHooks = {}) {
|
||||
getInvocationBlock: (companyId: string, agentId: string | null, context?: AdmissionContext) =>
|
||||
withBudgetSpan("budget.admission.inspect", () => withAccountingReadSnapshot(db, companyId,
|
||||
tx => budgetServiceInTransaction(tx).getInvocationBlock(companyId, agentId, context))),
|
||||
resolveIncident: async (companyId: string, incidentId: string, input: BudgetIncidentResolutionInput, actorUserId: string) => {
|
||||
resolveIncident: async (companyId: string, incidentId: string, input: BudgetIncidentResolutionInput, actorUserId: string,
|
||||
actorSource?: string | null) => {
|
||||
const [scope] = await db.select({ scopeType: budgetIncidents.scopeType, scopeId: budgetIncidents.scopeId }).from(budgetIncidents)
|
||||
.where(and(eq(budgetIncidents.id, incidentId), eq(budgetIncidents.companyId, companyId)));
|
||||
if (!scope) throw notFound("Budget incident not found");
|
||||
return mutate(companyId, service => service.resolveIncident(companyId, incidentId, input, actorUserId), scope);
|
||||
return mutate(companyId, service => service.resolveIncident(companyId, incidentId, input, actorUserId, actorSource), scope);
|
||||
},
|
||||
};
|
||||
}
|
||||
@@ -8,6 +8,7 @@ import { obsoleteExecutionConversationHold, unresolvedExecutionDispositions } fr
|
||||
import { hasPendingNativeReplacements } from "./native-runtime/native-safe-replacement.js";
|
||||
import { nextAbandonedNativeOwnerReleaseAt } from "./native-runtime/native-failed-run-owner-release.js";
|
||||
import { hasPendingPauseResume } from "./issue-tree-resume.js";
|
||||
import { hasPendingBudgetResume } from "./budget-resume-continuation.js";
|
||||
|
||||
export function registerExecutionControlWork(
|
||||
coordinator: ReturnType<typeof createDeliveryWorkCoordinator>,
|
||||
@@ -94,7 +95,7 @@ export function executionControlSchedule(db: Db) {
|
||||
eq(issueRecoveryActions.status, "resolved"),
|
||||
sql`${issueRecoveryActions.evidence}->>'continuationDelivery' = 'pending'`,
|
||||
)).limit(1);
|
||||
return rows.length || await hasPendingPauseResume(db) ? retry() : null;
|
||||
return rows.length || await hasPendingPauseResume(db) || await hasPendingBudgetResume(db) ? retry() : null;
|
||||
},
|
||||
async status() {
|
||||
const rows = await db.select({ id: heartbeatRuns.id }).from(heartbeatRuns)
|
||||
|
||||
@@ -2,7 +2,7 @@ import { randomUUID } from "node:crypto";
|
||||
import { contextCompactionService } from "../context-compaction.js";
|
||||
import { claimNativeSessionResumptions } from "../native-runtime/native-finalization-reconciler.js";
|
||||
import type { Db } from "@paperclipai/db";
|
||||
import { authUsers, agents, agentWakeupRequests, caseHandoffs, cases, chatCompletionDeliveries, companies, companyMemberships, createDb, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS, environmentLeases, heartbeatRuns, issueComments, issueThreadInteractions, issues, nativeRunFinalizations } from "@paperclipai/db";
|
||||
import { authUsers, agents, agentWakeupRequests, budgetIncidents, budgetPolicies, caseHandoffs, cases, chatCompletionDeliveries, companies, companyMemberships, costEvents, createDb, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS, environmentLeases, heartbeatRuns, issueComments, issueThreadInteractions, issues, nativeRunFinalizations } from "@paperclipai/db";
|
||||
import { eq, inArray, sql } from "drizzle-orm";
|
||||
import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import { resolveExternalTestDatabaseUrl } from "../../__tests__/helpers/external-test-database.js";
|
||||
@@ -810,6 +810,51 @@ describe.skipIf(!support.supported)("heartbeat queue database wiring", () => {
|
||||
await expectSavedMessages([...firstMessages, ...secondMessages]);
|
||||
});
|
||||
|
||||
// H1 bug C review: a queued run admission cancels for a board-owned budget
|
||||
// block records which budget stopped it, so Raise budget and resume can
|
||||
// continue it.
|
||||
async function admissionBudgetBlock(options: { openIncident: boolean }) {
|
||||
const fixtureRows = await fixture("queued");
|
||||
const { company, agent } = fixtureRows;
|
||||
const [policy] = await db.insert(budgetPolicies).values({ companyId: company.id, scopeType: "company", scopeId: company.id,
|
||||
windowKind: "lifetime", amount: 1, hardStopEnabled: true, isActive: true }).returning();
|
||||
await db.insert(costEvents).values({ companyId: company.id, agentId: agent.id, provider: "test", model: "test",
|
||||
costCents: 5, occurredAt: new Date() });
|
||||
if (options.openIncident) await db.insert(budgetIncidents).values({ companyId: company.id, policyId: policy.id,
|
||||
scopeType: "company", scopeId: company.id, metric: "billed_cents", windowKind: "lifetime",
|
||||
windowStart: new Date(0), windowEnd: new Date(Date.now() + 86_400_000), thresholdType: "hard",
|
||||
amountLimit: 1, amountObserved: 5 });
|
||||
const deps = callbacks(db);
|
||||
deps.getAgent.mockResolvedValue(agent);
|
||||
deps.budgets.getInvocationBlock.mockResolvedValue({ scopeType: "company", scopeId: company.id,
|
||||
scopeName: company.name, reason: "Company cannot start work because its budget hard-stop is still exceeded." });
|
||||
deps.cancelRunInternal.mockImplementation(async (runId: string) => {
|
||||
await db.update(heartbeatRuns).set({ status: "cancelled", finishedAt: new Date() }).where(eq(heartbeatRuns.id, runId));
|
||||
});
|
||||
return { ...fixtureRows, deps };
|
||||
}
|
||||
|
||||
it("leaves a queued run queued when a resume closed its budget incident first", async () => {
|
||||
// Raise budget and resume closed the incident between the admission
|
||||
// check and the cancel: the re-read finds no block, so the task runs.
|
||||
const { run, deps } = await admissionBudgetBlock({ openIncident: false });
|
||||
deps.budgets.getInvocationBlock.mockResolvedValueOnce({ scopeType: "company", scopeId: run.companyId,
|
||||
scopeName: "Budget", reason: "Company cannot start work because its budget hard-stop is still exceeded." })
|
||||
.mockResolvedValueOnce(null);
|
||||
expect(await createHeartbeatQueue(db, deps).claimQueuedRun({ ...run, responsibleUserId: "queue-test-user" })).toBeNull();
|
||||
expect(deps.cancelRunInternal).not.toHaveBeenCalled();
|
||||
expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, run.id)))[0].status).toBe("queued");
|
||||
});
|
||||
|
||||
it("records the budget that cancelled a queued run at admission", async () => {
|
||||
const { company, agent, run, deps } = await admissionBudgetBlock({ openIncident: true });
|
||||
void agent;
|
||||
expect(await createHeartbeatQueue(db, deps).claimQueuedRun({ ...run, responsibleUserId: "queue-test-user" })).toBeNull();
|
||||
expect(deps.cancelRunInternal).toHaveBeenCalledWith(run.id, expect.stringContaining("budget"), expect.anything());
|
||||
expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, run.id)))[0].resultJson)
|
||||
.toMatchObject({ budgetAdmissionStop: { scopeType: "company", scopeId: company.id } });
|
||||
});
|
||||
|
||||
it("rechecks a handoff's assignment version under the issue claim lock", async () => {
|
||||
const { company, agent, issue, wake, run } = await fixture("queued");
|
||||
const [record] = await db.insert(cases).values({ companyId: company.id, caseNumber: 1, identifier: "QM-CASE-1",
|
||||
|
||||
@@ -103,6 +103,7 @@ import {
|
||||
cases,
|
||||
agentWakeupRequests,
|
||||
activityLog,
|
||||
budgetIncidents,
|
||||
chatConversations,
|
||||
chatEndpoints,
|
||||
companies,
|
||||
@@ -923,7 +924,39 @@ export function createHeartbeatQueue(db: Db, dependencies: HeartbeatQueueDepende
|
||||
// message, Retry, or recovery it carries; keep it queued. Settling
|
||||
// accounting re-notifies run dispatch. A reached hard stop is the board's.
|
||||
if (!run.startedAt && !(await isBoardOwnedBudgetBlock(db, run.companyId, budgetBlock))) return null;
|
||||
await cancelRunInternal(run.id, budgetBlock.reason);
|
||||
// Record which budget stops this queued run, so Raise budget and resume
|
||||
// continues it (budget-resume-continuation.ts). Saved before the cancel
|
||||
// so a crash between them cannot lose it; the resume matches it only
|
||||
// against a recorded cancellation with this same reason.
|
||||
const budgetAdmissionStop = { scopeType: budgetBlock.scopeType, scopeId: budgetBlock.scopeId,
|
||||
reason: budgetBlock.reason.slice(0, 512) };
|
||||
const stamped = await db.transaction(async tx => {
|
||||
// The same commit touches the scope's open incident. A Raise budget
|
||||
// and resume whose repeatable-read snapshot predates this commit then
|
||||
// conflicts on that row and retries, so its resume sees the stamp
|
||||
// (a still-queued stamped run is claimed and waits for this cancel).
|
||||
const open = await tx.update(budgetIncidents).set({ updatedAt: new Date() }).where(and(
|
||||
eq(budgetIncidents.companyId, run.companyId), eq(budgetIncidents.scopeType, budgetBlock.scopeType),
|
||||
eq(budgetIncidents.scopeId, budgetBlock.scopeId), eq(budgetIncidents.status, "open")))
|
||||
.returning({ id: budgetIncidents.id });
|
||||
if (!open.length) return false;
|
||||
await tx.update(heartbeatRuns).set({ resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb)
|
||||
|| jsonb_build_object('budgetAdmissionStop', jsonb_build_object('scopeType', ${budgetAdmissionStop.scopeType}::text,
|
||||
'scopeId', ${budgetAdmissionStop.scopeId}::text, 'reason', ${budgetAdmissionStop.reason}::text))` })
|
||||
.where(and(eq(heartbeatRuns.id, run.id), eq(heartbeatRuns.status, run.status)));
|
||||
return true;
|
||||
});
|
||||
if (!stamped) {
|
||||
// No open incident: a resume may have just closed it and raised the
|
||||
// budget. Read the block again; resumed work stays queued.
|
||||
const current = await budgets.getInvocationBlock(run.companyId, run.agentId, {
|
||||
issueId: readNonEmptyString(context.issueId), projectId: readNonEmptyString(context.projectId) });
|
||||
if (!current) return null;
|
||||
if (!run.startedAt && !(await isBoardOwnedBudgetBlock(db, run.companyId, current))) return null;
|
||||
await cancelRunInternal(run.id, current.reason);
|
||||
return null;
|
||||
}
|
||||
await cancelRunInternal(run.id, budgetBlock.reason, { resultJson: { budgetAdmissionStop } });
|
||||
return null;
|
||||
}
|
||||
|
||||
|
||||
@@ -73,6 +73,7 @@ import {
|
||||
} from "../adapter-execution-control.js";
|
||||
import { AdapterStopTimeoutError } from "../adapter-stop-timeout.js";
|
||||
import { budgetCancellationScopePredicate, pendingBudgetRunCancellationPredicate } from "./budget-cancellation.js";
|
||||
import { notifyBudgetStopsSettled } from "../budget-resume-continuation.js";
|
||||
import { withBudgetSpan } from "../budget-tracing.js";
|
||||
import { captureException } from "../../sentry.js";
|
||||
import { randomUUID } from "node:crypto";
|
||||
@@ -1842,6 +1843,9 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont
|
||||
}
|
||||
}
|
||||
}));
|
||||
// A resumed budget owes these runs a continuation once their Stop lands.
|
||||
await notifyBudgetStopsSettled(db, runIds).catch(error =>
|
||||
logger.warn({ err: error, companyId: scope.companyId }, "Budget resume continuation notification deferred"));
|
||||
// Drain the batch before rejecting; one failed provider must not prevent
|
||||
// the remaining runs from receiving their already-committed Stop.
|
||||
if (failures.length) throw new AggregateError(failures, "Budget cancellation delivery remains pending");
|
||||
|
||||
Reference in new issue
Block a user