mirror of
https://github.com/paperclipai/paperclip.git
synced 2026-11-04 02:32:07 +01:00
perf(budgets): scope accounting and deliver enforcement asynchronously (#17046)
## Thinking Path > - Paperclip manages AI agents and their paid work. > - Budgets must prevent new work and stop reported overruns. > - Cost settlement also ran serial lifecycle reconciliation for every eligible company agent. > - That made request latency grow with unrelated idle agents and delayed results on stop delivery. > - This change separates atomic admission and accounting from durable asynchronous enforcement. > - Scoped live usage detects overruns without writing provisional ledger charges. > - The result is constant query scope for ordinary settlement and recoverable stop delivery. ## Linked Issues or Issue Description **What happened?** Settlement and budget checks entered company-wide lifecycle reconciliation. A company with 1,000 unrelated idle agents performed 1,001 accounting transactions for one ordinary cost evaluation. User requests could also wait for provider termination. Running usage checkpoints did not enforce reported monetary exposure. **Expected behavior** Keep the final reservation gate authoritative. Commit charges, scope blocks, and enforcement intent before returning. Deliver stop requests after commit. Read only applicable policies and execution scopes for ordinary charges. Stop long-running work when supported reported or priced usage establishes an overrun. **Steps to reproduce** 1. Create a disposable company with one company budget and idle agents. 2. Evaluate a charge below the budget with 0, 17, and 1,000 unrelated idle agents. 3. Compare query counts, company accounting transactions, and duration. 4. Use the controlled partial-usage and delayed-stop fixtures to check enforcement and request responsiveness. **Paperclip version or commit** The measured baseline is `6c9b87d0922f1ebfc43e8b0c36889ac5ba7b9837` from `master`. Related public work: Refs #16976, #2006, #3539, #13763, #14296, #16853, #13075. These cover the affected service charge, cancellation, and accounting paths. This change removes their shared reconciliation overhead. ## What Changed - Batch applicable policy exposure under the existing company accounting lock. Preserve exact fractional money and zero-reservation dispatch fences. - Commit a durable final budget dispatch claim before provider I/O. Retain a separate final issue/run ownership transaction through synchronous handoff, without holding the company accounting lock or awaiting provider work. Preserve native recovery lease checks and duplicate handoff fencing. Browser handoff rechecks current policy; it reuses a consumed rate slot only for the identical policy generation. - Separate ledger spend, live observed spend, and effective held capacity. Use the maximum of reservation and live observation for each unfinished operation. - Deliver budget cancellation through the existing coordinator. Revalidate generations and persist queued/running stop fences before external termination. Bound stop concurrency to four. - Apply agent holds only to the affected agent. Project budgets leave other projects running. Project real company transitions in bounded lifecycle batches. - Replace periodic broad reconciliation with durable pending work, startup repair, receipt replay, and concrete retry/window deadlines. Retain drain and idle-safety checks. - Enforce accepted cumulative usage checkpoints without provisional charges. Final accounting atomically replaces live exposure and settles reservations. - Add operator-gated budget spans, regression coverage, map updates, and a real-PostgreSQL comparison harness. No schema or first-party telemetry change. ## Verification - Full workspace `pnpm -r typecheck` passes on the final ownership fix, rebased onto `bbcf5dfe98`. Production `pnpm build` passes on final head `4c5f4de8b1`. Map, module-boundary, token, and Storybook catalog checks pass. - Focused real-PostgreSQL coverage includes the 292-case rebased accounting/dispatch run, 63 browser cases, 52 run-control/cancellation cases, 42 policy-generation cases, 77 lifecycle/commentary cases, and the exact Muse confirmation regression. Native executor unit coverage passes all 768 cases. Final ownership coverage passes: 74 adapter/use-case cases, nine public regression cases, and the additional deleted-project race. These exercise competing ownership, budget changes and charges between phases, zero/nonzero duplicate claims, rollback, recovery contention, read-only capacity rechecks, and controlled HTTP startup after commit. Another 74 pure/startup/tracing cases pass. The complete CI matrix and both 5/5 reviews passed on `1ac9b58`; its external security scan identified a stale-owner handoff race. The final fix restores synchronous issue/run ownership fencing after the accounting transaction commits. Final head `8e84f55ce8` has fresh Greptile 5/5 with zero new comments, Paperclip review 5/5, and a passing security scan. The complete latest-head [CI matrix](https://github.com/paperclipai/paperclip/actions/runs/39535189808) passes, including budget, browser, Runner, build, and canary jobs. The unchanged Runner close fixture passed on its sole same-head job retry. - Broad local `pnpm test:run` invocations were stopped before required source updates (rebase and review fixes). They are not claimed as full passes. Final full-suite coverage comes from the complete latest-head CI matrix; local focused coverage and the incomplete broad invocations are reported separately. - Reproduce performance with `pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-budget-enforcement.ts --baseline-ref=6c9b87d0922f1ebfc43e8b0c36889ac5ba7b9837`; repeat with `--reverse`. The harness uses disposable real PostgreSQL, 20 warm samples per variant/count, and records raw timings, SQL/lock counts, agent reads, unchanged agent state, and query plans. Final production-source measurement revision: `eff713a4f72c7f4b7a4002a701cb481d4bfccd95`; the extended harness is committed in `4c5f4de8b1`. | Unrelated idle agents | Baseline median / p95 | New median / p95 | Baseline SQL / locks | New SQL / locks | New agent reads | | --- | ---: | ---: | ---: | ---: | ---: | | 0 | 1.80 / 2.33 ms | 2.99 / 3.88 ms | 5 / 1 | 7 / 1 | 0 | | 17 | 15.53 / 19.55 ms | 2.24 / 2.60 ms | 90 / 18 | 7 / 1 | 0 | | 1,000 | 978.35 / 1,790.43 ms | 2.16 / 2.52 ms | 5,005 / 1,001 | 7 / 1 | 0 | - These measurements use local PostgreSQL 18.1 and Node 26.4.0 arm64. They isolate the cost-evaluation core and complete admission path, without hosted-region or provider latency. Prior-revision reversed-order controls also preserved the constant query counts; the final measurements above use the final production source. - Complete reserve → final ownership gate → immediate offline provider result is measured separately, with one responsible agent and one active operation at each unrelated-agent count. Admission p95 baseline/new: 0 agents **11.54 / 29.77 ms**; 17 agents **9.31 / 27.62 ms**; 1,000 agents **24.51 / 74.05 ms**. Baseline admission was already scoped; the stronger final fence adds bounded overhead. New counts stay constant at 53 Drizzle statements, three transactions, two company-accounting locks, and five agent-reading statements bound to the responsible agent. Zero unrelated-agent reads, exactly one provider callback per sample, and no leftover runs or holds. All admission scenarios stay below 250 ms. No external provider calls are included. - The partial-usage/delayed-termination fixture passed 20/20: receipt persistence p95 13 ms, receipt return to observed durable Stop fence p95 28 ms, complete receipt-start-to-fence p95 50 ms, unrelated database request p95 1 ms while termination was blocked. After releasing the controlled termination fixture, acknowledgement p95 was 14 ms. These are separately measured phases; fixture acknowledgement is not live-provider termination latency. - Models and provider deadlines are unchanged. No deployment, live inference, or staging acceptance is claimed. A quiet 30-minute follow-up waits for the merged revision on staging, then repeats classifier and controlled budget-stop acceptance in staging, recording hosted bookkeeping, queue delay, and actual provider acknowledgement separately. - Development failures retained: finalizer notification regressions (fixed; 27/27 affected cases passed), invalid browser admission fixture assumptions (fixed; browser coverage now passes 63/63), and first-head lifecycle/dispatch CI regressions (reproduced and repaired). A prior-head Cases E2E test hit its unchanged announcement-dismiss handler race while opening a filter; the same journey passed on `1ac9b58`. The external security review then found a stale-owner race between accounting commit and provider entry; final verification covers a competing run in that gap while retaining provider I/O outside the accounting lock. Early PostgreSQL startup attempts hit local semaphore pressure; serial reruns succeeded. The final-head Runner lane had one failure among 1,823 tests: the unchanged `timeout-attachment` fixture failed at its first process close after 6.28 seconds. The sole same-head job retry passed all 1,819 active tests (four skipped), including that fixture in 1.85 seconds. Three isolated local probes could not run correctly because the Rust fake-provider fixture binary was absent; cleanup masked the original startup failure. CI builds that fixture through its normal lane. These probes are inconclusive and are not counted as passes. ## Risks - This changes accounting, lifecycle, and cancellation ordering. Durable generations, source/sequence checks, ownership leases, and company locking remain required for correctness. - Provider stop acknowledgement can be slow or fail. Stop intent remains pending until existing termination proof succeeds; cancelled row status alone does not acknowledge delivery. - Mid-run limits depend on provider usage reports or supported price calculations. Providers that report only at completion cannot provide exact mid-call detection. - Signals are process-local wake hints for the supported single-active-writer model. Durable state and startup recovery remain authoritative. - There are no migrations, new queue tables, or cached balances. Repeat the classifier and controlled budget-stop acceptance after normal rollout. ## Model Used OpenAI Codex, based on GPT-6, with reasoning, tool use, code execution, and parallel coding agents. The host does not expose the exact backend model revision or context-window size, so those values are not asserted. ## Checklist - [x] I have included a thinking path that traces from project context to this change - [x] I have specified the model used (with version and capability details) - [x] I have checked ROADMAP.md and confirmed this PR does not duplicate planned core work - [x] I have searched GitHub for duplicate or related PRs and linked them above - [x] I have either (a) linked existing issues with `Fixes: #` / `Closes #` / `Refs #` OR (b) described the issue in-PR following the relevant issue template - [x] I have not referenced internal/instance-local Paperclip issues or links (only public GitHub `#NNN` / `github.com/paperclipai/paperclip` URLs) - [x] My branch name describes the change (e.g. `docs/...`, `fix/...`) and contains no internal Paperclip ticket id or instance-derived details - [x] I have run tests locally and they pass - [x] I have added or updated tests where applicable - [x] I have updated relevant documentation to reflect my changes - [x] I have considered and documented any risks above - [x] All Paperclip CI gates are green - [x] Greptile is 5/5 with no open P2s, recommendations, or follow-ups - [x] I will address all Greptile and reviewer comments before requesting merge --------- Co-authored-by: Paperclip <noreply@paperclip.ing>
This commit is contained in:
1 parent
dff63551fe
commit
30b753f614
85 files changed
+5204
-1111
No files matched your search
@@ -3,19 +3,20 @@ import { companies, type Db } from "@paperclipai/db";
|
||||
import { publishActivity, type ActivityPublication } from "./activity-log.js";
|
||||
import { notFound } from "../errors.js";
|
||||
import { logger } from "../middleware/logger.js";
|
||||
import { withBudgetSpan } from "./budget-tracing.js";
|
||||
|
||||
/** Serialize ledger and budget mutations in one company. Always acquire this
|
||||
* before locking a run or a policy, so concurrent writers use one lock order.
|
||||
* NO KEY UPDATE permits FK KEY SHARE locks held by recovery callers. */
|
||||
export async function withAccountingTransaction<T>(db: Db, companyId: string, work: (tx: Db, publications: ActivityPublication[]) => Promise<T>): Promise<T> {
|
||||
const publications: ActivityPublication[] = [];
|
||||
const result = await db.transaction(async (transaction) => {
|
||||
const result = await withBudgetSpan("budget.accounting.transaction", () => db.transaction(async (transaction) => {
|
||||
const tx = transaction as unknown as Db;
|
||||
const [company] = await tx.select({ id: companies.id }).from(companies)
|
||||
.where(eq(companies.id, companyId)).for("no key update");
|
||||
const [company] = await withBudgetSpan("budget.accounting.lock_wait", () => tx.select({ id: companies.id }).from(companies)
|
||||
.where(eq(companies.id, companyId)).for("no key update"));
|
||||
if (!company) throw notFound("Company not found");
|
||||
return work(tx, publications);
|
||||
});
|
||||
}));
|
||||
publishAccountingActivities(companyId, publications);
|
||||
return result;
|
||||
}
|
||||
|
||||
@@ -45,10 +45,12 @@ export function createLifecycleDriver(db: Db, manager: PluginWorkerManager): Lif
|
||||
async runHost(agent) {
|
||||
if (agent.lifecycleState === "verifying") return verify(agent);
|
||||
if (["pausing", "terminating", "cleaning_up"].includes(agent.lifecycleState)) {
|
||||
const companyAgents = await db.select().from(agents).where(eq(agents.companyId, agent.companyId));
|
||||
const ids = agent.lifecycleState === "pausing" ? [agent.id]
|
||||
: [agent.id, ...listInvalidOrgChainDescendantIds(agent.id, companyAgents)];
|
||||
if (!await heartbeat.stopInvocationsForAgents(ids, "Agent lifecycle stop requested")) return "pending";
|
||||
: [agent.id, ...listInvalidOrgChainDescendantIds(agent.id,
|
||||
await db.select().from(agents).where(eq(agents.companyId, agent.companyId)))];
|
||||
const reason = agent.lifecycleHolds?.includes("company_archived")
|
||||
? "Cancelled because the company was archived" : "Agent lifecycle stop requested";
|
||||
if (!await heartbeat.stopInvocationsForAgents(ids, reason)) return "pending";
|
||||
}
|
||||
return "complete";
|
||||
},
|
||||
|
||||
@@ -2,19 +2,23 @@ import { logger } from "../middleware/logger.js";
|
||||
import type { Db } from "@paperclipai/db";
|
||||
import * as lifecycle from "../modules/agent-lifecycle/index.js";
|
||||
import { withAccountingTransaction } from "./accounting-transaction.js";
|
||||
import { budgetServiceInTransaction, deliverBudgetEnforcement, policyBlocks, type BudgetServiceHooks } from "./budgets.js";
|
||||
import { budgetServiceInTransaction, policyBlocks, type BudgetServiceHooks } from "./budgets.js";
|
||||
import { recordAgentStatusEvent, recordResourceCreationEvent } from "./resource-lifecycle-events.js";
|
||||
import { clearPrimaryAgent } from "./primary-agent.js";
|
||||
import { deleteAgentDependencies } from "./agent-deletion.js";
|
||||
import { prepareAgentConfiguration, completeAgentConfiguration } from "./agent-configuration.js";
|
||||
import { prepareAgentHire, initializeAgentHire, prepareAgentHireApproval, completeAgentHireApproval } from "./agent-hiring.js";
|
||||
import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notifications.js";
|
||||
import { fenceActiveCompanyAgentStopsInTransaction as fenceActiveCompanyStops } from "../modules/agent-lifecycle/index.js";
|
||||
import { applyAgentPolicyHoldsInTransaction } from "../modules/agent-lifecycle/index.js";
|
||||
import { trackIdleWork } from "./task-admission.js";
|
||||
|
||||
export { scheduleAgentLifecycle } from "../modules/agent-lifecycle/index.js";
|
||||
export { scheduleAgentLifecycle, processAgentLifecycle } from "../modules/agent-lifecycle/index.js";
|
||||
export type { LifecycleDriver, LifecycleAgent } from "../modules/agent-lifecycle/index.js";
|
||||
|
||||
export function createAgentLifecycleEffects(hooks: BudgetServiceHooks = {}): lifecycle.LifecycleEffects {
|
||||
export function createAgentLifecycleEffects(_hooks: BudgetServiceHooks = {}): lifecycle.LifecycleEffects {
|
||||
return {
|
||||
notifyWork: tx => notifyDeliveryWork(tx, DELIVERY_QUEUES.agentLifecycle),
|
||||
reportFailure(context, error) {
|
||||
const code = lifecycle.lifecycleFailureCode(error);
|
||||
logger.warn({ ...context, failureCode: code ?? "unclassified" }, "Agent lifecycle operation failed");
|
||||
@@ -26,7 +30,6 @@ export function createAgentLifecycleEffects(hooks: BudgetServiceHooks = {}): lif
|
||||
budgetServiceInTransaction(db, publications).upsertPolicy(companyId, {
|
||||
scopeType: "agent", scopeId: agentId, amount, isActive, windowKind: "calendar_month_utc",
|
||||
}, userId),
|
||||
enforceBudget: (db, companyId, agentId) => deliverBudgetEnforcement(db, hooks, companyId, agentId),
|
||||
recordCreation: (db, companyId, agentId) => recordResourceCreationEvent(db, companyId, "agent", agentId),
|
||||
recordStatus: recordAgentStatusEvent,
|
||||
clearPrimary: clearPrimaryAgent,
|
||||
@@ -55,3 +58,12 @@ export function configureAgentLifecycle(db: Db, driver: lifecycle.LifecycleDrive
|
||||
export function startAgentLifecycle(db: Db, driver: lifecycle.LifecycleDriver, canRun: () => boolean) {
|
||||
return lifecycle.startAgentLifecycle(db, createAgentLifecycleEffects(), trackedDriver(driver), canRun);
|
||||
}
|
||||
|
||||
/** Budget evaluation owns the transaction and supplies the aggregate agent decision. */
|
||||
export function applyAgentBudgetHoldInTransaction(tx: Db, companyId: string, agentId: string, blocked: boolean) {
|
||||
return applyAgentPolicyHoldsInTransaction(tx, createAgentLifecycleEffects(), companyId, agentId, { budget: blocked });
|
||||
}
|
||||
|
||||
export function fenceActiveCompanyAgentStopsInTransaction(tx: Db, companyId: string, status: "paused" | "archived") {
|
||||
return fenceActiveCompanyStops(tx, createAgentLifecycleEffects(), companyId, status);
|
||||
}
|
||||
@@ -1,5 +1,4 @@
|
||||
import { approvalResolutionRecords, resolvableApprovalStatuses as resolvableStatuses } from "../lib/approval-records.js";
|
||||
import { budgetService } from "./budgets.js";
|
||||
import { notifyHireApproved } from "./hire-hook.js";
|
||||
|
||||
import { and, asc, eq, inArray, sql, type SQL } from "drizzle-orm";
|
||||
@@ -177,7 +176,6 @@ export function approvalService(db: Db) {
|
||||
const result = await (status === "approved" ? lifecycle.approveHire : lifecycle.rejectHire)(target, userId, note);
|
||||
if (result?.hireApprovedAgentId && result.approval) {
|
||||
await reconcileApprovedBuiltInAgent(result.approval.companyId, result.approval.payload, db);
|
||||
await budgetService(db).deliverPendingEnforcement(result.approval.companyId);
|
||||
void notifyHireApproved(db, { companyId: result.approval.companyId, agentId: result.hireApprovedAgentId,
|
||||
source: "approval", sourceId: result.approval.id, approvedAt: result.approval.decidedAt ?? new Date() }).catch(() => {});
|
||||
}
|
||||
|
||||
@@ -7,7 +7,7 @@ import { withAccountingReadSnapshot, withAccountingTransaction } from "./account
|
||||
import { receiptFingerprint } from "./receipt-fingerprint.js";
|
||||
import { logActivity } from "./activity-log.js";
|
||||
import { updateMonthlySpendProjections } from "./costs.js";
|
||||
import { budgetServiceInTransaction, deliverBudgetEnforcement, type BudgetServiceHooks } from "./budgets.js";
|
||||
import { budgetServiceInTransaction, type BudgetServiceHooks } from "./budgets.js";
|
||||
|
||||
async function invoiceCostCandidates(tx: Db, companyId: string, biller: string, line: {
|
||||
costEventId?: string | null; runId?: string | null; providerRequestId?: string | null; model?: string | null;
|
||||
@@ -180,7 +180,6 @@ export function billingReconciliationService(db: Db, hooks: BudgetServiceHooks =
|
||||
details: { adjustmentId: adjustment.id, previousCents: input.expectedCents, correctedCents: input.correctedCents, reason: input.reason, pricing: input.pricing } }, publications);
|
||||
return adjustment;
|
||||
});
|
||||
await deliverBudgetEnforcement(db, hooks, companyId);
|
||||
return result;
|
||||
},
|
||||
};
|
||||
|
||||
@@ -117,6 +117,7 @@ export function browserUseClient(
|
||||
path: string,
|
||||
method = "GET",
|
||||
body?: unknown,
|
||||
beforeDispatch?: () => Promise<void>,
|
||||
): Promise<unknown> {
|
||||
const until = backoff.get(scope) ?? 0;
|
||||
if (until > Date.now())
|
||||
@@ -126,6 +127,9 @@ export function browserUseClient(
|
||||
until - Date.now(),
|
||||
true,
|
||||
);
|
||||
// The owning service commits its final admission fence immediately before
|
||||
// this one paid POST. Gate failures never become uncertain provider work.
|
||||
if (beforeDispatch) await beforeDispatch();
|
||||
let response: Response;
|
||||
try {
|
||||
const init: RequestInit = {
|
||||
@@ -207,8 +211,8 @@ export function browserUseClient(
|
||||
allowResizing?: true;
|
||||
};
|
||||
maxCostUsd?: number;
|
||||
}) {
|
||||
return created.parse(await call("/runs", "POST", body));
|
||||
}, beforeDispatch?: () => Promise<void>) {
|
||||
return created.parse(await call("/runs", "POST", body, beforeDispatch));
|
||||
},
|
||||
async status(runId: string) {
|
||||
return z
|
||||
|
||||
+208
-103
@@ -16,7 +16,8 @@ import {
|
||||
issues,
|
||||
heartbeatRuns,
|
||||
agents,
|
||||
costEvents,
|
||||
companies,
|
||||
toolCatalogEntries,
|
||||
} from "@paperclipai/db";
|
||||
import {
|
||||
BROWSER_USE_IDLE_MS,
|
||||
@@ -24,10 +25,13 @@ import {
|
||||
type BrowserUseSettings,
|
||||
type BrowserUseViewportRequest,
|
||||
type TaskBrowser,
|
||||
centsToUnits,
|
||||
unitsToCents,
|
||||
centsToUsd,
|
||||
usdToUnits,
|
||||
} from "@paperclipai/shared";
|
||||
import {
|
||||
browserUseClient,
|
||||
browserUseCostCap,
|
||||
BrowserUseError,
|
||||
browserUseTerminal,
|
||||
browserUseViewerUrl,
|
||||
@@ -36,15 +40,17 @@ import {
|
||||
type BrowserUseRequest,
|
||||
} from "./browser-use-client.js";
|
||||
import { secretService } from "./secrets.js";
|
||||
import { budgetService, type BudgetServiceHooks } from "./budgets.js";
|
||||
import { budgetService, budgetServiceInTransaction, type BudgetServiceHooks } from "./budgets.js";
|
||||
import { issueTreeControlService } from "./issue-tree-control.js";
|
||||
import { accessService } from "./access.js";
|
||||
import { toolAccessPolicyService } from "./tool-access-policy.js";
|
||||
import { createCostEventInTransaction } from "./costs.js";
|
||||
import { createFinanceEventInTransaction } from "./finance.js";
|
||||
import { withAccountingTransaction } from "./accounting-transaction.js";
|
||||
import { logActivity } from "./activity-log.js";
|
||||
import { logActivity, type ActivityPublication } from "./activity-log.js";
|
||||
import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notifications.js";
|
||||
import { browserUseViewports } from "./browser-use-viewport.js";
|
||||
import { withBudgetSpan } from "./budget-tracing.js";
|
||||
import { forbidden, notFound, conflict } from "../errors.js";
|
||||
|
||||
type Session = typeof sessions.$inferSelect;
|
||||
@@ -54,6 +60,7 @@ type Binding = {
|
||||
runId: string | null;
|
||||
issueId: string | null;
|
||||
approvedInvocationId?: string;
|
||||
gatewayId?: string | null;
|
||||
};
|
||||
type Grant = typeof connectionGrants.$inferSelect;
|
||||
const argsSchema = z
|
||||
@@ -67,6 +74,17 @@ const argsSchema = z
|
||||
const terminalSessions = ["closed", "failed"];
|
||||
const liveSessions = ["starting", "running", "idle", "stopping"];
|
||||
const requestMarker = (invocationId: string) => `\n\n[Paperclip request: ${invocationId}]`;
|
||||
type BrowserBudgetHold = { invocationId: string; amountCents: string; projectId: string | null;
|
||||
state: "held" | "released" | "settled"; uncapped?: boolean;
|
||||
dispatchCommittedAt?: string; accountedCentsExact?: string; billingUnknown?: boolean };
|
||||
const resultObject = (value: unknown): Record<string, unknown> =>
|
||||
value && typeof value === "object" && !Array.isArray(value) ? value as Record<string, unknown> : {};
|
||||
const budgetHold = (value: unknown) => resultObject(value).paperclipBudgetHold as BrowserBudgetHold | undefined;
|
||||
function publicResult(value: unknown) {
|
||||
if (!value || typeof value !== "object" || Array.isArray(value)) return value;
|
||||
const { paperclipBudgetHold: _internal, ...visible } = resultObject(value);
|
||||
return Object.keys(visible).length ? visible : null;
|
||||
}
|
||||
|
||||
export function browserUseService(
|
||||
db: Db,
|
||||
@@ -114,6 +132,10 @@ export function browserUseService(
|
||||
r.configPath.toLowerCase() === "headers.x-browser-use-api-key",
|
||||
);
|
||||
if (!ref) throw forbidden("Reconnect Browser Use to restore its API key.");
|
||||
const [credentialSnapshot] = await db.select().from(companySecrets).where(and(
|
||||
eq(companySecrets.companyId, s.companyId), eq(companySecrets.id, ref.secretId)));
|
||||
if (!credentialSnapshot || (!cleanup && credentialSnapshot.status !== "active"))
|
||||
throw forbidden("Browser credential is unavailable.");
|
||||
if (row.grant.kind === "user") {
|
||||
const [secret] = await db
|
||||
.select()
|
||||
@@ -154,6 +176,7 @@ export function browserUseService(
|
||||
);
|
||||
return {
|
||||
...row,
|
||||
credentialSnapshot,
|
||||
client: browserUseClient(
|
||||
{ "X-Browser-Use-API-Key": key },
|
||||
request,
|
||||
@@ -271,7 +294,7 @@ export function browserUseService(
|
||||
status: s.status,
|
||||
runStatus: run?.status ?? null,
|
||||
progress: run?.progress ?? null,
|
||||
result: run?.result ?? null,
|
||||
result: publicResult(run?.result ?? null),
|
||||
costCents: run?.accountedCents ?? 0,
|
||||
browsers: ownedBrowsers,
|
||||
error: s.error,
|
||||
@@ -307,7 +330,8 @@ export function browserUseService(
|
||||
connectionId: grant.connectionId,
|
||||
grantId: grant.id,
|
||||
};
|
||||
const { client } = await credentials(scope);
|
||||
const selected = await credentials(scope);
|
||||
const { client } = selected;
|
||||
if (name === "browser_profiles") {
|
||||
// A safe Test action that creates no paid work and exposes no unapproved profile.
|
||||
if (!config.allowedProfileIds.length) return { profiles: [] };
|
||||
@@ -403,27 +427,89 @@ export function browserUseService(
|
||||
.limit(1);
|
||||
if (existing[0])
|
||||
return snapshot(await owned(binding, grant, existing[0].sessionId));
|
||||
const issue = await assertBinding(binding, approved);
|
||||
const { policies } = await budgets.overview(binding.companyId);
|
||||
const caps = policies
|
||||
.filter(
|
||||
(p) =>
|
||||
p.isActive &&
|
||||
p.hardStopEnabled &&
|
||||
p.metric === "billed_cents" &&
|
||||
p.amount > 0 &&
|
||||
(p.scopeType === "company" ||
|
||||
(p.scopeType === "agent" && p.scopeId === binding.agentId) ||
|
||||
(p.scopeType === "project" && p.scopeId === issue.projectId)),
|
||||
)
|
||||
.map((p) => p.remainingAmount / 100);
|
||||
const maxCostUsd = browserUseCostCap(
|
||||
config.maxCostUsd,
|
||||
args.maxCostUsd,
|
||||
...caps,
|
||||
);
|
||||
const admissionCheck = async (tx: Db, publications: ActivityPublication[], hold?: BrowserBudgetHold) => {
|
||||
const [parent] = await tx.select().from(heartbeatRuns).where(and(
|
||||
eq(heartbeatRuns.id, binding.runId!), eq(heartbeatRuns.companyId, binding.companyId),
|
||||
eq(heartbeatRuns.agentId, binding.agentId!))).for("update");
|
||||
const [issue] = await tx.select().from(issues).where(and(eq(issues.id, binding.issueId!), eq(issues.companyId, binding.companyId)));
|
||||
const [currentAgent] = await tx.select().from(agents).where(and(eq(agents.id, binding.agentId!), eq(agents.companyId, binding.companyId)));
|
||||
const [company] = await tx.select({ status: companies.status }).from(companies).where(eq(companies.id, binding.companyId));
|
||||
if (!parent || (!approved && parent.status !== "running") || parent.resultJson?.cancellation || parent.resultJson?.startupCancellation
|
||||
|| !issue || issue.assigneeAgentId !== binding.agentId || issue.status === "cancelled"
|
||||
|| !currentAgent || currentAgent.lifecycleState !== "ready" || ["paused", "terminated", "pending_approval"].includes(currentAgent.status)
|
||||
|| company?.status !== "active") return { reason: "The browser's owning task, company or agent run is no longer active." };
|
||||
const [invocation] = await tx.select().from(toolInvocations).where(and(
|
||||
eq(toolInvocations.id, invocationId), eq(toolInvocations.companyId, binding.companyId)));
|
||||
if (invocation && (invocation.agentId !== binding.agentId || invocation.runId !== binding.runId
|
||||
|| invocation.issueId !== binding.issueId || invocation.connectionId !== grant.connectionId
|
||||
|| invocation.status !== "executing" || !["agent", "user", "system", "plugin"].includes(invocation.actorType)))
|
||||
return { reason: "The browser invocation is no longer authorized." };
|
||||
if (approved && invocation?.approvalState !== "approved") return { reason: "The browser approval is no longer authorized." };
|
||||
if (await issueTreeControlService(tx).getActivePauseHoldGate(binding.companyId, issue.id)) return { reason: "This task is paused." };
|
||||
const [current] = await tx.select({ grant: connectionGrants, connection: toolConnections }).from(connectionGrants)
|
||||
.innerJoin(toolConnections, eq(toolConnections.id, connectionGrants.connectionId)).where(and(
|
||||
eq(connectionGrants.id, grant.id), eq(connectionGrants.companyId, binding.companyId),
|
||||
eq(toolConnections.id, grant.connectionId), eq(toolConnections.companyId, binding.companyId))).for("share");
|
||||
const [secret] = await tx.select().from(companySecrets).where(and(eq(companySecrets.id, selected.credentialSnapshot.id),
|
||||
eq(companySecrets.companyId, binding.companyId))).for("share");
|
||||
// Reuse the gateway's policy evaluator, including profile precedence,
|
||||
// exclusions and explicit denies. REST gallery tools do not need MCP
|
||||
// installation rows, and a matching binding alone is not authorization.
|
||||
const [entry] = await tx.select().from(toolCatalogEntries).where(and(
|
||||
eq(toolCatalogEntries.companyId, binding.companyId), eq(toolCatalogEntries.connectionId, grant.connectionId),
|
||||
eq(toolCatalogEntries.toolName, name))).limit(1);
|
||||
if (!entry || (invocation?.catalogEntryId && invocation.catalogEntryId !== entry.id)
|
||||
|| (invocation?.upstreamToolName && invocation.upstreamToolName !== name))
|
||||
return { reason: "Browser tool is no longer available." };
|
||||
const access = await toolAccessPolicyService(tx).decide({
|
||||
companyId: binding.companyId, actor: { actorType: (invocation?.actorType ?? "agent") as "agent" | "user" | "system" | "plugin",
|
||||
actorId: invocation?.actorId ?? binding.agentId!, agentId: binding.agentId },
|
||||
runContext: { heartbeatRunId: binding.runId, issueId: issue.id, projectId: issue.projectId, gatewayId: binding.gatewayId ?? invocation?.gatewayId },
|
||||
request: { applicationId: entry.applicationId, connectionId: grant.connectionId, catalogEntryId: entry.id,
|
||||
toolName: invocation?.toolName ?? entry.name, upstreamToolName: name, riskLevel: entry.riskLevel,
|
||||
applicationKey: invocation?.applicationKey, providerType: invocation?.providerType ?? "rest_api", arguments: value, sideEffecting: true },
|
||||
consumeRateLimit: false, writeAuditEvent: false,
|
||||
}, { admittedInvocationId: invocation?.id });
|
||||
if (!access.allowed && !(approved && access.decision === "require_approval")) return { reason: access.explanation };
|
||||
if (!current || current.grant.status !== "active" || current.connection.status !== "active" || !current.connection.enabled
|
||||
|| !isBrowserUseConnection(current.connection)
|
||||
|| current.grant.updatedAt.getTime() !== selected.grant.updatedAt.getTime()
|
||||
|| current.connection.updatedAt.getTime() !== selected.connection.updatedAt.getTime()
|
||||
|| JSON.stringify(current.grant.credentialSecretRefs) !== JSON.stringify(selected.grant.credentialSecretRefs)
|
||||
|| !secret || secret.status !== "active" || secret.latestVersion !== selected.credentialSnapshot.latestVersion)
|
||||
return { reason: "Browser access or its credential changed before dispatch." };
|
||||
const admission = await budgetServiceInTransaction(tx, publications).evaluateAdmission(binding.companyId, binding.agentId, { projectId: issue.projectId });
|
||||
if (admission.block) return { reason: admission.block.reason };
|
||||
if (hold && centsToUnits(hold.amountCents) === 0n && admission.policies.length)
|
||||
return { reason: "Browser budget changed before dispatch; try again with a bounded cost cap." };
|
||||
const [currentConfig] = await tx.select().from(settings).where(and(eq(settings.companyId, binding.companyId), eq(settings.grantId, grant.id)));
|
||||
if (args.profileId && !currentConfig?.allowedProfileIds.includes(args.profileId)) return { reason: "This profile is not allowed for this credential." };
|
||||
let amount: bigint | undefined;
|
||||
for (const cap of [currentConfig?.maxCostUsd, args.maxCostUsd]) {
|
||||
if (cap === null || cap === undefined) continue;
|
||||
const value = usdToUnits(cap);
|
||||
amount = amount === undefined || value < amount ? value : amount;
|
||||
}
|
||||
for (const [index, policy] of admission.policies.entries()) {
|
||||
const observed = admission.observations[index];
|
||||
const committed = centsToUnits(observed.ledgerCents) + centsToUnits(observed.effectiveHeldCents);
|
||||
const available = centsToUnits(policy.amount) - committed;
|
||||
if (hold ? available < 0n || (centsToUnits(hold.amountCents) === 0n && available === 0n) : available <= 0n)
|
||||
return { reason: "No browser budget remains; available capacity is reserved by unfinished work." };
|
||||
if (!hold) amount = amount === undefined || available < amount ? available : amount;
|
||||
}
|
||||
if (amount !== undefined && amount <= 0n) return { reason: "No browser budget remains." };
|
||||
if (hold && amount !== undefined && (centsToUnits(hold.amountCents) === 0n || amount < centsToUnits(hold.amountCents)))
|
||||
return { reason: "Browser cost settings changed before dispatch; try again." };
|
||||
if (hold && issue.projectId !== hold.projectId) return { reason: "The browser task's project changed before dispatch." };
|
||||
return { amountCents: unitsToCents(amount ?? 0n), maxCostUsd: amount === undefined ? undefined : Number(centsToUsd(unitsToCents(amount))), projectId: issue.projectId };
|
||||
};
|
||||
const sid = args.sessionId;
|
||||
const s = await db.transaction(async (tx) => {
|
||||
const admitted = await withAccountingTransaction(db, binding.companyId, async (tx, publications) => {
|
||||
const [duplicate] = await tx.select().from(runs).where(and(eq(runs.companyId, binding.companyId), eq(runs.invocationId, invocationId)));
|
||||
if (duplicate) return { existing: duplicate.sessionId };
|
||||
const check = await admissionCheck(tx, publications);
|
||||
if (check.reason) return { reason: check.reason };
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.browser);
|
||||
let row: Session;
|
||||
if (sid) {
|
||||
@@ -463,7 +549,10 @@ export function browserUseService(
|
||||
})
|
||||
.returning();
|
||||
}
|
||||
const hold: BrowserBudgetHold = { invocationId, amountCents: check.amountCents!, projectId: check.projectId!,
|
||||
state: "held", uncapped: check.maxCostUsd === undefined };
|
||||
await tx.insert(runs).values({
|
||||
result: { paperclipBudgetHold: hold },
|
||||
companyId: binding.companyId,
|
||||
sessionId: row.id,
|
||||
heartbeatRunId: binding.runId!,
|
||||
@@ -472,8 +561,13 @@ export function browserUseService(
|
||||
? new Date(Date.now() + BROWSER_USE_IDLE_MS)
|
||||
: null,
|
||||
});
|
||||
return row;
|
||||
return { session: row, hold, maxCostUsd: check.maxCostUsd };
|
||||
});
|
||||
if (admitted.reason) throw forbidden(admitted.reason);
|
||||
if (admitted.existing) return snapshot(await owned(binding, grant, admitted.existing));
|
||||
const s = admitted.session!;
|
||||
const maxCostUsd = admitted.maxCostUsd;
|
||||
let providerAttempted = false;
|
||||
let result: Awaited<ReturnType<typeof client.start>>;
|
||||
try {
|
||||
result = await client.start({
|
||||
@@ -488,6 +582,21 @@ export function browserUseService(
|
||||
...(args.profileId ? { profileId: args.profileId } : {}),
|
||||
},
|
||||
...(maxCostUsd === undefined ? {} : { maxCostUsd }),
|
||||
}, async () => {
|
||||
const reason = await withAccountingTransaction(db, binding.companyId, async (tx, publications) => {
|
||||
const [current] = await tx.select().from(runs).where(and(eq(runs.invocationId, invocationId), eq(runs.companyId, binding.companyId))).for("update");
|
||||
const [currentSession] = await tx.select().from(sessions).where(and(eq(sessions.id, s.id), eq(sessions.companyId, binding.companyId))).for("update");
|
||||
const hold = budgetHold(current?.result);
|
||||
if (!currentSession || currentSession.status !== "starting" || currentSession.stopRequested) return "Browser session was stopped before dispatch.";
|
||||
if (!current || current.status !== "creating" || current.providerRunId || hold?.state !== "held" || hold.dispatchCommittedAt)
|
||||
return "Browser provider dispatch no longer owns this invocation.";
|
||||
const check = await admissionCheck(tx, publications, hold);
|
||||
if (check.reason) return check.reason;
|
||||
await tx.update(runs).set({ result: { ...resultObject(current.result), paperclipBudgetHold: { ...hold, dispatchCommittedAt: new Date().toISOString() } } }).where(eq(runs.id, current.id));
|
||||
return null;
|
||||
});
|
||||
if (reason) throw new BrowserUseError(403, reason, 0, true);
|
||||
providerAttempted = true;
|
||||
});
|
||||
await db.transaction(async (tx) => {
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.browser);
|
||||
@@ -508,12 +617,19 @@ export function browserUseService(
|
||||
});
|
||||
} catch (error) {
|
||||
// A lost create response can mean paid work exists. Never submit it again.
|
||||
const rejected = error instanceof BrowserUseError && error.requestRejected;
|
||||
const message = rejected ? error.message :
|
||||
const rejected = !providerAttempted || (error instanceof BrowserUseError && error.requestRejected);
|
||||
const message = rejected ? (error instanceof Error ? error.message : "Browser dispatch was rejected.") :
|
||||
"Browser run creation could not be confirmed. Paperclip is locating and stopping possible provider work. Do not start another run yet.";
|
||||
await db.transaction(async tx => {
|
||||
await withAccountingTransaction(db, binding.companyId, async (tx, publications) => {
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.browser);
|
||||
const [current] = await tx.select().from(sessions).where(eq(sessions.id, s.id)).for("update");
|
||||
if (rejected) {
|
||||
await tx.update(runs).set({ result: sql`jsonb_set(${runs.result}, '{paperclipBudgetHold,state}', '"released"'::jsonb)` }).where(eq(runs.invocationId, invocationId));
|
||||
await budgetServiceInTransaction(tx, publications).evaluateAffectedScopes(binding.companyId, binding.agentId, admitted.hold!.projectId);
|
||||
} else {
|
||||
await tx.update(runs).set({ result: sql`jsonb_set(${runs.result}, '{paperclipBudgetHold,billingUnknown}', 'true'::jsonb)` }).where(eq(runs.invocationId, invocationId));
|
||||
await budgetServiceInTransaction(tx, publications).evaluateAffectedScopes(binding.companyId, binding.agentId, admitted.hold!.projectId);
|
||||
}
|
||||
// A restart must not leave a drained rejection attached to a starting
|
||||
// session: that combination has neither recovery work nor an idle timer.
|
||||
await tx.update(runs)
|
||||
@@ -917,75 +1033,47 @@ export function browserUseService(
|
||||
ReturnType<ReturnType<typeof browserUseClient>["summary"]>
|
||||
>,
|
||||
) {
|
||||
const cents = Math.max(0, Math.round(Number(summary.totalCostUsd) * 100));
|
||||
if (!Number.isSafeInteger(cents))
|
||||
throw new BrowserUseError(502, "Browser Use returned invalid usage.");
|
||||
const event = await withAccountingTransaction(db, s.companyId, async (tx, publications) => {
|
||||
const [issue] = await tx
|
||||
.select({ projectId: issues.projectId })
|
||||
.from(issues)
|
||||
let amount: bigint;
|
||||
try { amount = usdToUnits(summary.totalCostUsd); }
|
||||
catch { throw new BrowserUseError(502, "Browser Use returned invalid usage."); }
|
||||
if (amount < 0n || !browserUseTerminal(summary.status))
|
||||
throw new BrowserUseError(502, "Browser Use returned invalid terminal usage.");
|
||||
const cents = unitsToCents(amount);
|
||||
await withBudgetSpan("budget.accounting.settlement", () => withAccountingTransaction(db, s.companyId, async (tx, publications) => {
|
||||
const [current] = await tx.select().from(runs).where(and(eq(runs.id, run.id), eq(runs.companyId, s.companyId))).for("update");
|
||||
const hold = budgetHold(current.result);
|
||||
const previous = centsToUnits(hold?.accountedCentsExact ?? current.accountedCents);
|
||||
const [issue] = await tx.select({ projectId: issues.projectId }).from(issues)
|
||||
.where(and(eq(issues.id, s.issueId), eq(issues.companyId, s.companyId)));
|
||||
const [current] = await tx
|
||||
.select()
|
||||
.from(runs)
|
||||
.where(eq(runs.id, run.id))
|
||||
.for("update");
|
||||
if (cents <= current.accountedCents) return null;
|
||||
const projectId = hold ? hold.projectId : issue?.projectId ?? null;
|
||||
// Replace the separate browser hold with its ledger charge atomically.
|
||||
// Provider output never gets to overwrite the private accounting state.
|
||||
const settled: BrowserBudgetHold = { ...(hold ?? { invocationId: current.invocationId, amountCents: "0.0000000", projectId }), state: "settled",
|
||||
accountedCentsExact: unitsToCents(amount > previous ? amount : previous) };
|
||||
await tx.update(runs).set({
|
||||
result: { ...resultObject(sanitizeBrowserUse({ result: summary.result, output: summary.output })),
|
||||
paperclipBudgetHold: settled },
|
||||
eventsDrained: 1,
|
||||
accountedCents: Math.min(2_147_483_647, Math.round(Number(unitsToCents(amount > previous ? amount : previous)))),
|
||||
}).where(eq(runs.id, current.id));
|
||||
if (amount <= previous) {
|
||||
if (hold?.state === "held") await budgetServiceInTransaction(tx, publications).evaluateAffectedScopes(s.companyId, s.agentId, projectId);
|
||||
return;
|
||||
}
|
||||
const e = await createCostEventInTransaction(tx, s.companyId, {
|
||||
agentId: s.agentId,
|
||||
issueId: s.issueId,
|
||||
projectId: issue?.projectId,
|
||||
agentId: s.agentId, issueId: s.issueId, projectId,
|
||||
heartbeatRunId: run.heartbeatRunId,
|
||||
provider: "browser-use-cloud",
|
||||
biller: "browser-use-cloud",
|
||||
billingType: "metered_api",
|
||||
model: summary.model,
|
||||
costCents: cents - current.accountedCents,
|
||||
occurredAt: new Date(),
|
||||
billingCode: `browser-use-cloud:${run.id}`,
|
||||
idempotencyKey: `browser-use-cloud:${run.id}:${cents}`,
|
||||
provider: "browser-use-cloud", biller: "browser-use-cloud", billingType: "metered_api", model: summary.model,
|
||||
costCents: unitsToCents(amount - previous), occurredAt: new Date(),
|
||||
billingCode: `browser-use-cloud:${run.id}`, idempotencyKey: `browser-use-cloud:${run.id}:${cents}`,
|
||||
}, publications);
|
||||
await createFinanceEventInTransaction(tx, publications, s.companyId, {
|
||||
agentId: e.agentId,
|
||||
issueId: e.issueId,
|
||||
projectId: e.projectId,
|
||||
heartbeatRunId: e.heartbeatRunId,
|
||||
costEventId: e.id,
|
||||
billingCode: e.billingCode,
|
||||
eventKind: "inference_charge",
|
||||
direction: "debit",
|
||||
biller: "browser-use-cloud",
|
||||
provider: "browser-use-cloud",
|
||||
model: e.model,
|
||||
amountCents: e.costCents,
|
||||
currency: "USD",
|
||||
description: "Browser Use agent run",
|
||||
occurredAt: e.occurredAt,
|
||||
agentId: e.agentId, issueId: e.issueId, projectId: e.projectId, heartbeatRunId: e.heartbeatRunId,
|
||||
costEventId: e.id, billingCode: e.billingCode, eventKind: "inference_charge", direction: "debit",
|
||||
biller: "browser-use-cloud", provider: "browser-use-cloud", model: e.model,
|
||||
amountCents: e.costCentsExact, currency: "USD", description: "Browser Use agent run", occurredAt: e.occurredAt,
|
||||
});
|
||||
await tx
|
||||
.update(runs)
|
||||
.set({ accountedCents: cents })
|
||||
.where(eq(runs.id, run.id));
|
||||
return e;
|
||||
});
|
||||
// Re-evaluate on retries too: accounting may have committed immediately
|
||||
// before a crash interrupted budget enforcement.
|
||||
const lastEvent =
|
||||
event ??
|
||||
(
|
||||
await db
|
||||
.select()
|
||||
.from(costEvents)
|
||||
.where(
|
||||
and(
|
||||
eq(costEvents.companyId, s.companyId),
|
||||
eq(costEvents.billingCode, `browser-use-cloud:${run.id}`),
|
||||
),
|
||||
)
|
||||
.orderBy(desc(costEvents.createdAt))
|
||||
.limit(1)
|
||||
)[0];
|
||||
if (lastEvent) await budgets.evaluateCostEvent(lastEvent);
|
||||
}));
|
||||
}
|
||||
async function reconcile(s: Session) {
|
||||
const { grant, connection, client } = await credentials(s, true);
|
||||
@@ -1046,6 +1134,33 @@ export function browserUseService(
|
||||
active = true;
|
||||
continue;
|
||||
}
|
||||
if (budgetHold(run.result)?.state === "held" && !budgetHold(run.result)?.dispatchCommittedAt) {
|
||||
const repaired = await withAccountingTransaction(db, s.companyId, async (tx, publications) => {
|
||||
const [current] = await tx.select().from(runs).where(and(eq(runs.id, run.id), eq(runs.companyId, s.companyId))).for("update");
|
||||
const hold = budgetHold(current?.result);
|
||||
if (!current || current.providerRunId || current.status !== "creating" || hold?.state !== "held" || hold.dispatchCommittedAt
|
||||
|| Date.now() - current.createdAt.getTime() <= 60_000) return null;
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.browser);
|
||||
await tx.update(runs).set({ status: "failed", eventsDrained: 1,
|
||||
result: { ...resultObject(current.result), paperclipBudgetHold: { ...hold, state: "released" } } }).where(eq(runs.id, run.id));
|
||||
const [session] = await tx.select().from(sessions).where(eq(sessions.id, s.id)).for("update");
|
||||
const stopping = session.stopRequested === "end";
|
||||
const [updated] = await tx.update(sessions).set({
|
||||
status: stopping ? "stopping" : session.providerSessionId ? "idle" : "failed",
|
||||
idleDeadline: !stopping && session.providerSessionId ? new Date(Date.now() + BROWSER_USE_IDLE_MS) : null,
|
||||
updatedAt: new Date(),
|
||||
}).where(eq(sessions.id, s.id)).returning();
|
||||
await budgetServiceInTransaction(tx, publications).evaluateAffectedScopes(s.companyId, s.agentId, hold.projectId);
|
||||
return updated;
|
||||
});
|
||||
if (repaired) {
|
||||
Object.assign(s, repaired);
|
||||
run.eventsDrained = 1;
|
||||
run.status = "failed";
|
||||
if (terminalSessions.includes(s.status)) return;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
end = true;
|
||||
await db.update(sessions).set({ status: "stopping", stopRequested: "end" })
|
||||
.where(eq(sessions.id, s.id));
|
||||
@@ -1154,16 +1269,6 @@ export function browserUseService(
|
||||
if (browserUseTerminal(state.status) && drained) {
|
||||
const summary = await client.summary(run.providerRunId);
|
||||
await account(s, run, summary);
|
||||
await db
|
||||
.update(runs)
|
||||
.set({
|
||||
result: sanitizeBrowserUse({
|
||||
result: summary.result,
|
||||
output: summary.output,
|
||||
}),
|
||||
eventsDrained: 1,
|
||||
})
|
||||
.where(eq(runs.id, run.id));
|
||||
await signal(s, "browser.run_finished");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,86 @@
|
||||
import type { Db } from "@paperclipai/db";
|
||||
import { withBudgetSpan } from "./budget-tracing.js";
|
||||
import type { createDeliveryWorkCoordinator } from "./delivery-work-coordinator.js";
|
||||
import { DELIVERY_QUEUES } from "./delivery-work-notifications.js";
|
||||
import {
|
||||
deliverBudgetEnforcement,
|
||||
hasPendingBudgetEnforcement,
|
||||
nextBudgetReconciliationAt,
|
||||
reconcileDueBudgetPolicies,
|
||||
type BudgetServiceHooks,
|
||||
} from "./budgets.js";
|
||||
import { readCostAccountingWork } from "./run-cost-accounting.js";
|
||||
import { subscribeAllCompanyLiveEvents } from "./live-events.js";
|
||||
|
||||
type Coordinator = Pick<ReturnType<typeof createDeliveryWorkCoordinator>, "register">;
|
||||
const accountingWakes = new WeakMap<Db, () => void>();
|
||||
|
||||
/** A durable filesystem receipt can need recovery even when its database
|
||||
* attempt failed before a transaction identity existed. Startup also replays it. */
|
||||
export function wakeCostAccountingWork(db: Db): void {
|
||||
accountingWakes.get(db)?.();
|
||||
}
|
||||
|
||||
/** The version pair is durable cancellation intent. Signals only request an
|
||||
* inspection; startup and concrete deadlines recover work after interruption. */
|
||||
export function registerBudgetEnforcementWork(db: Db, coordinator: Coordinator, hooks: BudgetServiceHooks) {
|
||||
let startup = true;
|
||||
let reconciliationCursor: string | null = null;
|
||||
let reconciling = false;
|
||||
let nextDeadline: number | null = null;
|
||||
return coordinator.register(DELIVERY_QUEUES.budgetEnforcement, {
|
||||
retryMs: 1_000,
|
||||
runDuringDrain: true,
|
||||
async run(signal) {
|
||||
if (signal.aborted) return;
|
||||
if (startup || reconciling || (nextDeadline !== null && nextDeadline <= Date.now())) {
|
||||
const page = await reconcileDueBudgetPolicies(db, { startup, ...(reconciliationCursor ? { cursor: reconciliationCursor } : {}) });
|
||||
reconciliationCursor = page.nextCursor;
|
||||
reconciling = reconciliationCursor !== null;
|
||||
if (!reconciling) startup = false;
|
||||
}
|
||||
if (signal.aborted) return;
|
||||
await withBudgetSpan("budget.enforcement.deliver", () => deliverBudgetEnforcement(db, hooks));
|
||||
},
|
||||
async hasPending() {
|
||||
const [pending, next] = await Promise.all([
|
||||
hasPendingBudgetEnforcement(db), nextBudgetReconciliationAt(db),
|
||||
]);
|
||||
nextDeadline = next?.getTime() ?? null;
|
||||
return pending || reconciling;
|
||||
},
|
||||
nextRunAt: () => nextDeadline,
|
||||
});
|
||||
}
|
||||
|
||||
/** Accounting repairs admitted work and never dispatches another paid turn.
|
||||
* A waiting partial receipt stays held without polling unchanged debt. */
|
||||
export function registerCostAccountingWork(db: Db, coordinator: Coordinator, reconcile: () => Promise<unknown>) {
|
||||
let nextDeadline: number | null = null;
|
||||
const worker = coordinator.register(DELIVERY_QUEUES.accounting, {
|
||||
retryMs: 5_000,
|
||||
runDuringDrain: true,
|
||||
async run(signal) {
|
||||
if (signal.aborted) return;
|
||||
await withBudgetSpan("budget.accounting.recover", reconcile);
|
||||
},
|
||||
async hasPending() {
|
||||
const work = await readCostAccountingWork(db);
|
||||
nextDeadline = work.nextRunAt;
|
||||
return work.pending;
|
||||
},
|
||||
nextRunAt: () => nextDeadline,
|
||||
});
|
||||
accountingWakes.set(db, worker.wake);
|
||||
// The existing heartbeat terminal event is published after its persisted
|
||||
// transition. Receipt writes have transaction-aware accounting signals too.
|
||||
const unsubscribe = subscribeAllCompanyLiveEvents(event => {
|
||||
if (event.type === "heartbeat.run.status" &&
|
||||
["succeeded", "failed", "timed_out", "cancelled", "interrupted"].includes(String(event.payload.status))) worker.wake();
|
||||
});
|
||||
const stop = () => {
|
||||
unsubscribe();
|
||||
if (accountingWakes.get(db) === worker.wake) accountingWakes.delete(db);
|
||||
};
|
||||
return { ...worker, stop };
|
||||
}
|
||||
@@ -0,0 +1,167 @@
|
||||
import { addCents, normalizeCents, usdToUnits, unitsToCents } from "@paperclipai/shared";
|
||||
import { and, eq, gte, inArray, lt, or, sql, type SQL } from "drizzle-orm";
|
||||
import { budgetPolicies, budgetReservations, browserUseRuns, browserUseSessions, costEvents, heartbeatRuns, issues, nativeRunFinalizations, type Db } from "@paperclipai/db";
|
||||
import { withBudgetSpan } from "./budget-tracing.js";
|
||||
|
||||
export interface BudgetExposure {
|
||||
ledgerCents: string;
|
||||
liveObservedCents: string;
|
||||
effectiveHeldCents: string;
|
||||
pendingTerminalRunCount: number;
|
||||
unpricedEventCount: number;
|
||||
/** Unconfirmed POSTs or unfinished uncapped browser work, separate from ledger data quality. */
|
||||
unknownBrowserRunCount: number;
|
||||
}
|
||||
|
||||
export type SpendPolicy = Pick<typeof budgetPolicies.$inferSelect, "companyId" | "scopeType" | "scopeId" | "windowKind" | "metric">;
|
||||
export type ObservedSpend = BudgetExposure & {
|
||||
windowStart: Date; windowEnd: Date; total: number; totalExact: string;
|
||||
pendingRunCount: number; recoveringRunCount: number;
|
||||
};
|
||||
|
||||
export function budgetWindow(kind: string, now = new Date()) {
|
||||
return kind === "lifetime" ? { start: new Date(0), end: new Date(Date.UTC(9999, 0, 1)) }
|
||||
: { start: new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth(), 1)), end: new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth() + 1, 1)) };
|
||||
}
|
||||
|
||||
/** Only normalized per-run prices are exposure. An incomplete receipt can
|
||||
* contain a known partial price while still being unpriced for settlement. */
|
||||
export function knownRunCostCents(usage: Record<string, unknown>): string | null {
|
||||
if (usage.billingType === "subscription_included" || usage.billingType === "subscription") return "0";
|
||||
if (usage.usageSource === "session_cumulative" || usage.usageBasis === "session_cumulative") return null;
|
||||
for (const value of [usage.cacheAdjustedCostUsd, usage.costUsdExact, usage.costUsd]) {
|
||||
if ((typeof value !== "number" || !Number.isFinite(value) || value < 0) && typeof value !== "string") continue;
|
||||
try {
|
||||
const amount = usdToUnits(value);
|
||||
if (amount >= 0n) return unitsToCents(amount);
|
||||
} catch { /* An older malformed field cannot erase a reliable later field. */ }
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
/** SQL compatibility for active runs written before the exposure projection.
|
||||
* Numeric text is bounded and checked before casting legacy JSON fields. */
|
||||
function normalizedNumericCents(value: SQL, usd = false) {
|
||||
const scale = usd ? 1_000_000_000 : 10_000_000;
|
||||
return sql`case when length(${value}) <= 160 and ${value} ~ '^[+]?[0-9]+([.][0-9]+)?([eE][+-]?[0-9]{1,3})?$'
|
||||
then case when abs(coalesce(substring(${value} from '[eE]([+-]?[0-9]+)$'), '0')::int) <= 100
|
||||
then case when round((${value})::numeric * ${scale}) <= 999999999999999999999999::numeric
|
||||
then round((${value})::numeric * ${scale}) / 10000000::numeric else null end
|
||||
else null end else null end`;
|
||||
}
|
||||
|
||||
const usage = heartbeatRuns.usageJson;
|
||||
const projected = normalizedNumericCents(sql`${usage}->'budgetExposure'->>'costCents'`);
|
||||
const hasProjection = sql`jsonb_typeof(${usage}->'budgetExposure') = 'object'
|
||||
and jsonb_typeof(${usage}->'budgetExposure'->'sourceId') = 'string'
|
||||
and jsonb_typeof(${usage}->'budgetExposure'->'sequence') = 'number'`;
|
||||
const legacy = sql`case when coalesce(${usage}->>'usageSource', 'per_run') <> 'session_cumulative'
|
||||
and coalesce(${usage}->>'usageBasis', 'per_run') <> 'session_cumulative' then
|
||||
coalesce(${normalizedNumericCents(sql`${usage}->>'cacheAdjustedCostUsd'`, true)},
|
||||
${normalizedNumericCents(sql`${usage}->>'costUsdExact'`, true)}, ${normalizedNumericCents(sql`${usage}->>'costUsd'`, true)}) else null end`;
|
||||
// The projection already separates attempts and cumulative/session counters.
|
||||
// Raw compatibility fields are a fallback, never stronger monetary evidence.
|
||||
// A canonical unknown projection must remain unknown, not inherit an unsafe
|
||||
// cumulative aggregate retained only for old accounting compatibility.
|
||||
const liveCost = sql`case when ${heartbeatRuns.costAccountedAt} is null and ${heartbeatRuns.id} is not null
|
||||
and coalesce(${usage}->>'billingType', '') not in ('subscription_included', 'subscription')
|
||||
then case when ${hasProjection} then coalesce(${projected}, 0)
|
||||
else coalesce(${projected}, ${legacy}, 0) end else 0 end`;
|
||||
|
||||
/** Batched ledger, held-operation and unfinished-run aggregates. Never scans
|
||||
* agents or receipt history; each policy retains its own scope and window. */
|
||||
async function readObservedSpends(db: Db, policies: SpendPolicy[], now: Date): Promise<ObservedSpend[]> {
|
||||
const windows = policies.map(policy => budgetWindow(policy.windowKind, now));
|
||||
const billable = policies.map((policy, index) => ({ policy, index })).filter(({ policy }) => policy.metric === "billed_cents");
|
||||
const ledgerColumns: Record<string, SQL> = {}, heldColumns: Record<string, SQL> = {}, runColumns: Record<string, SQL> = {}, browserColumns: Record<string, SQL> = {};
|
||||
const ledgerPredicates: SQL[] = [], heldPredicates: SQL[] = [], runPredicates: SQL[] = [], browserPredicates: SQL[] = [];
|
||||
for (const { policy, index } of billable) {
|
||||
const { start, end } = windows[index];
|
||||
const ledger = and(eq(costEvents.companyId, policy.companyId),
|
||||
policy.scopeType === "agent" ? eq(costEvents.agentId, policy.scopeId) : undefined,
|
||||
policy.scopeType === "project" ? eq(costEvents.projectId, policy.scopeId) : undefined,
|
||||
policy.windowKind === "calendar_month_utc" ? and(gte(costEvents.occurredAt, start), lt(costEvents.occurredAt, end)) : undefined)!;
|
||||
const held = and(eq(budgetReservations.companyId, policy.companyId), eq(budgetReservations.state, "held"),
|
||||
policy.scopeType === "agent" ? eq(budgetReservations.agentId, policy.scopeId) : undefined,
|
||||
policy.scopeType === "project" ? eq(budgetReservations.projectId, policy.scopeId) : undefined)!;
|
||||
// Hosted browser work is a separate paid operation from its parent run.
|
||||
// Its durable hold survives unknown POST outcomes and monthly rollover.
|
||||
const browserHold = sql`${browserUseRuns.result}->'paperclipBudgetHold'`;
|
||||
const hasBrowserHold = sql`coalesce(jsonb_typeof(${browserHold}) = 'object'
|
||||
and ${browserHold}->>'invocationId' = ${browserUseRuns.invocationId}::text, false)`;
|
||||
// Earlier browser rows did not persist a cap/dispatch projection. Their
|
||||
// unfinished work remains unknown until a terminal receipt is accounted.
|
||||
const legacyBrowserPending = sql`not ${hasBrowserHold} and ${browserUseRuns.eventsDrained} = 0
|
||||
and (${browserUseRuns.providerRunId} is not null or ${browserUseRuns.status} in ('creating', 'unknown', 'queued', 'dispatching', 'running'))`;
|
||||
const browser = and(eq(browserUseRuns.companyId, policy.companyId), eq(browserUseSessions.companyId, policy.companyId),
|
||||
sql`((${hasBrowserHold} and ${browserHold}->>'state' = 'held') or (${legacyBrowserPending}))`,
|
||||
policy.scopeType === "agent" ? eq(browserUseSessions.agentId, policy.scopeId) : undefined,
|
||||
policy.scopeType === "project" ? sql`case when ${hasBrowserHold} then ${browserHold}->>'projectId'
|
||||
else ${issues.projectId}::text end = ${policy.scopeId}` : undefined)!;
|
||||
browserPredicates.push(browser);
|
||||
browserColumns[`held${index}`] = sql`coalesce(sum(${normalizedNumericCents(sql`${browserHold}->>'amountCents'`)}) filter (where ${browser}), 0)::text`;
|
||||
// A successful uncapped POST still has no bounded monetary exposure until
|
||||
// terminal accounting. Introducing a budget cannot treat that work as free.
|
||||
// Older holds had no explicit flag: zero meant an uncapped dispatch. New
|
||||
// capped zero holds are distinct, and unstarted reservations are not spend.
|
||||
const uncapped = sql`case when ${browserHold}->>'uncapped' = 'false' then false
|
||||
when ${browserHold}->>'uncapped' = 'true' then true
|
||||
else coalesce(${normalizedNumericCents(sql`${browserHold}->>'amountCents'`)}, 0) = 0 end`;
|
||||
const dispatched = sql`(${browserUseRuns.providerRunId} is not null or ${browserHold}->>'dispatchCommittedAt' is not null)`;
|
||||
browserColumns[`unknown${index}`] = sql`count(*) filter (where ${browser} and
|
||||
((${legacyBrowserPending}) or ${browserHold}->>'billingUnknown' = 'true' or (${uncapped} and ${dispatched})))::int`;
|
||||
const inWindow = policy.windowKind === "calendar_month_utc"
|
||||
? sql`coalesce(${heartbeatRuns.finishedAt}, ${now.toISOString()}::timestamptz) >= ${start.toISOString()}::timestamptz
|
||||
and coalesce(${heartbeatRuns.finishedAt}, ${now.toISOString()}::timestamptz) < ${end.toISOString()}::timestamptz` : sql`true`;
|
||||
const run = and(eq(heartbeatRuns.companyId, policy.companyId), eq(heartbeatRuns.costAccountingPending, true),
|
||||
policy.scopeType === "agent" ? eq(heartbeatRuns.agentId, policy.scopeId) : undefined,
|
||||
policy.scopeType === "project" ? sql`${usage}->'ledgerScope'->>'projectId' = ${policy.scopeId}` : undefined)!;
|
||||
const terminal = and(run, inArray(heartbeatRuns.status, ["succeeded", "failed", "timed_out", "cancelled", "interrupted"]), inWindow)!;
|
||||
// Keep correlated subqueries in nested SQL chunks. Drizzle strips column
|
||||
// qualifiers from a single-table selection's direct column chunks.
|
||||
const hasReservation = sql`exists (select 1 from ${budgetReservations}
|
||||
where ${budgetReservations.companyId} = ${heartbeatRuns.companyId}
|
||||
and ${budgetReservations.runId} = ${heartbeatRuns.id} and ${budgetReservations.state} = 'held')`;
|
||||
const recovering = sql`exists (select 1 from ${nativeRunFinalizations}
|
||||
where ${nativeRunFinalizations.runId} = ${heartbeatRuns.id}
|
||||
and ${nativeRunFinalizations.companyId} = ${heartbeatRuns.companyId}
|
||||
and ${nativeRunFinalizations.resultId} is null and ${nativeRunFinalizations.phase} <> 'terminal_failure')`;
|
||||
ledgerPredicates.push(ledger); heldPredicates.push(held); runPredicates.push(run);
|
||||
ledgerColumns[`total${index}`] = sql`coalesce(sum(${costEvents.costCents}) filter (where ${ledger}), 0)::text`;
|
||||
ledgerColumns[`unpriced${index}`] = sql`count(*) filter (where ${ledger} and ${costEvents.costStatus} = 'unpriced' and ${costEvents.billingType} <> 'subscription_included')::int`;
|
||||
runColumns[`live${index}`] = sql`coalesce(sum(${liveCost}) filter (where ${run} and ${inWindow}), 0)::text`;
|
||||
runColumns[`unreserved${index}`] = sql`coalesce(sum(${liveCost}) filter (where ${run} and not ${hasReservation}), 0)::text`;
|
||||
// Reservations carry over UTC windows. A reservation and its running
|
||||
// observation describe one operation, so count their maximum, not both.
|
||||
heldColumns[`held${index}`] = sql`coalesce(sum(greatest(${budgetReservations.amountCents}, ${liveCost})) filter (where ${held}), 0)::text`;
|
||||
runColumns[`pending${index}`] = sql`count(*) filter (where ${terminal})::int`;
|
||||
runColumns[`recovering${index}`] = sql`count(*) filter (where ${terminal} and ${heartbeatRuns.runtimeMode} = 'native' and ${recovering})::int`;
|
||||
}
|
||||
const [ledger = {}] = billable.length ? await db.select(ledgerColumns).from(costEvents).where(or(...ledgerPredicates)) : [];
|
||||
const [held = {}] = billable.length ? await db.select(heldColumns).from(budgetReservations)
|
||||
.leftJoin(heartbeatRuns, and(eq(heartbeatRuns.id, budgetReservations.runId), eq(heartbeatRuns.companyId, budgetReservations.companyId)))
|
||||
.where(or(...heldPredicates)) : [];
|
||||
const [runs = {}] = billable.length ? await db.select(runColumns).from(heartbeatRuns).where(or(...runPredicates)) : [];
|
||||
const [browser = {}] = billable.length ? await db.select(browserColumns).from(browserUseRuns)
|
||||
.innerJoin(browserUseSessions, and(eq(browserUseSessions.id, browserUseRuns.sessionId), eq(browserUseSessions.companyId, browserUseRuns.companyId)))
|
||||
.leftJoin(issues, and(eq(issues.id, browserUseSessions.issueId), eq(issues.companyId, browserUseSessions.companyId)))
|
||||
.where(or(...browserPredicates)) : [];
|
||||
return policies.map((_, index) => {
|
||||
const ledgerCents = normalizeCents(String(ledger[`total${index}`] ?? 0));
|
||||
const liveObservedCents = normalizeCents(String(runs[`live${index}`] ?? 0));
|
||||
const totalExact = addCents(ledgerCents, liveObservedCents);
|
||||
const pendingRunCount = Number(runs[`pending${index}`] ?? 0);
|
||||
return { windowStart: windows[index].start, windowEnd: windows[index].end,
|
||||
ledgerCents, liveObservedCents, effectiveHeldCents: addCents(String(held[`held${index}`] ?? 0), String(runs[`unreserved${index}`] ?? 0), String(browser[`held${index}`] ?? 0)),
|
||||
total: Number(totalExact), totalExact, pendingRunCount, pendingTerminalRunCount: pendingRunCount,
|
||||
unpricedEventCount: Number(ledger[`unpriced${index}`] ?? 0), unknownBrowserRunCount: Number(browser[`unknown${index}`] ?? 0), recoveringRunCount: Number(runs[`recovering${index}`] ?? 0) };
|
||||
});
|
||||
}
|
||||
|
||||
export async function computeObservedSpend(db: Db, policy: SpendPolicy) {
|
||||
return (await computeObservedSpends(db, [policy]))[0];
|
||||
}
|
||||
|
||||
export function computeObservedSpends(db: Db, policies: SpendPolicy[], now = new Date()): Promise<ObservedSpend[]> {
|
||||
return withBudgetSpan("budget.exposure.evaluate", () => readObservedSpends(db, policies, now), { "budget.policy_count": policies.length });
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
import type { AdapterUsageCheckpoint } from "@paperclipai/adapter-utils";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { createLiveRunUsageProjection, liveRunUsageCostCents } from "./budget-live-usage.js";
|
||||
|
||||
const openai: AdapterUsageCheckpoint = {
|
||||
complete: false, usageBasis: "per_run", provider: "openai", biller: "openai",
|
||||
billingType: "metered_api", model: "gpt-6.1-sol", costStatus: "unpriced",
|
||||
usage: { inputTokens: 1000, cachedInputTokens: 1000, outputTokens: 1000 },
|
||||
};
|
||||
const anthropic: AdapterUsageCheckpoint = {
|
||||
complete: false, usageBasis: "per_run", provider: "anthropic", biller: "anthropic",
|
||||
billingType: "metered_api", model: "claude-sonnet-5", costStatus: "unpriced",
|
||||
usage: { inputTokens: 22462, cachedInputTokens: 208218, cacheWriteTokens: 22450, outputTokens: 2450 },
|
||||
};
|
||||
|
||||
describe("live run budget pricing", () => {
|
||||
it("projects partial OpenAI token usage while preserving final-accounting uncertainty", () => {
|
||||
const receipt = Object.freeze({ ...openai, usage: Object.freeze({ ...openai.usage! }) });
|
||||
expect(liveRunUsageCostCents(receipt)).toBe("1.2100000");
|
||||
expect(receipt).toMatchObject({ complete: false, costStatus: "unpriced" });
|
||||
expect(receipt.costUsdExact).toBeUndefined();
|
||||
});
|
||||
|
||||
it("projects partial Anthropic token usage using the reviewed cache-write bound", () => {
|
||||
expect(liveRunUsageCostCents(anthropic)).toBe("15.5967600");
|
||||
expect(anthropic).toMatchObject({ complete: false, costStatus: "unpriced" });
|
||||
});
|
||||
|
||||
it("uses explicit provider cost before estimates even on incomplete unpriced receipts", () => {
|
||||
expect(liveRunUsageCostCents({ ...openai, costUsdExact: "0.000000001" })).toBe("0.0000001");
|
||||
expect(liveRunUsageCostCents({ ...openai, costUsd: 0 })).toBe("0.0000000");
|
||||
expect(liveRunUsageCostCents({ ...openai, costUsd: 5, cacheAdjustedCostUsd: 0.123456789 })).toBe("12.3456789");
|
||||
expect(liveRunUsageCostCents({ ...openai, model: "unsupported", costUsd: 0.05 })).toBe("5.0000000");
|
||||
});
|
||||
|
||||
it("applies supported OpenAI context and processing-tier guards", () => {
|
||||
expect(liveRunUsageCostCents({ ...openai, pricingContext: { serviceTier: "priority", contextTier: "long" } }))
|
||||
.toBe("3.8400000");
|
||||
expect(liveRunUsageCostCents({ ...openai, pricingContext: { serviceTier: "unrecognized" } })).toBeNull();
|
||||
});
|
||||
|
||||
it.each([
|
||||
{ usageBasis: "session_cumulative" }, { usageBasis: undefined }, { provider: "unknown" }, { biller: "openrouter" },
|
||||
{ billingType: "unknown" }, { model: "unsupported" }, { model: "constructor" }, { usage: undefined },
|
||||
{ usage: { inputTokens: 1, cacheWriteTokens: 2, outputTokens: 1 } },
|
||||
{ usageByModel: [{ model: "gpt-6.1-sol", usage: { inputTokens: 1, outputTokens: 1 }, costUsd: 1 }] },
|
||||
] as Partial<AdapterUsageCheckpoint>[])("does not invent a price for unsupported OpenAI usage: %j", override => {
|
||||
expect(liveRunUsageCostCents({ ...openai, ...override })).toBeNull();
|
||||
});
|
||||
|
||||
it.each([
|
||||
{ biller: "aws_bedrock" }, { model: "claude-sonnet-latest" }, { pricingContext: { serviceTier: "priority" } },
|
||||
{ usage: { inputTokens: 1, outputTokens: 1 } },
|
||||
{ usage: { inputTokens: 1, cacheWriteTokens: 2, cachedInputTokens: 0, outputTokens: 1 } },
|
||||
] as Partial<AdapterUsageCheckpoint>[])("retains Anthropic rate-card restrictions: %j", override => {
|
||||
expect(liveRunUsageCostCents({ ...anthropic, ...override })).toBeNull();
|
||||
});
|
||||
|
||||
it.each(["subscription", "subscription_included"] as const)("keeps %s usage outside billed exposure", billingType => {
|
||||
expect(liveRunUsageCostCents({ ...openai, billingType, costUsd: 10 })).toBe("0");
|
||||
});
|
||||
|
||||
it("does not use cumulative explicit cost as a per-run price", () => {
|
||||
expect(liveRunUsageCostCents({ ...openai, usageBasis: "session_cumulative", costUsd: 10 })).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe("live exposure across provider attempts", () => {
|
||||
it("adds a token-priced retry after an explicitly free rejected attempt", () => {
|
||||
const observe = createLiveRunUsageProjection();
|
||||
expect(observe("rejected", { ...openai, complete: true, costUsd: 0, usage: { inputTokens: 0, outputTokens: 0 } }))
|
||||
.toBe("0.0000000");
|
||||
expect(observe("retry", openai)).toBe("1.2100000");
|
||||
});
|
||||
|
||||
it("prices different provider attempts separately instead of applying one model to aggregate counters", () => {
|
||||
const observe = createLiveRunUsageProjection();
|
||||
expect(observe("openai", openai)).toBe("1.2100000");
|
||||
expect(observe("anthropic", anthropic)).toBe("16.8067600");
|
||||
});
|
||||
|
||||
it("preserves a known attempt price when its next partial checkpoint omits it", () => {
|
||||
const observe = createLiveRunUsageProjection();
|
||||
observe("first", { ...openai, model: "unsupported", costUsd: 0.8 });
|
||||
observe("first", { ...openai, model: "unsupported" });
|
||||
expect(observe("second", { ...openai, model: "unsupported", costUsd: 0.6 })).toBe("140.0000000");
|
||||
expect(observe("second", { ...openai, model: "unsupported", costUsd: 0.4 })).toBe("140.0000000");
|
||||
});
|
||||
|
||||
it("accepts a complete attempt correction without changing other attempts", () => {
|
||||
const observe = createLiveRunUsageProjection();
|
||||
observe("first", { ...openai, costUsd: 0.8 });
|
||||
observe("second", { ...openai, costUsd: 0.6 });
|
||||
expect(observe("first", { ...openai, complete: true, costUsd: 0.5 })).toBe("110.0000000");
|
||||
});
|
||||
|
||||
it("restores the durable aggregate projection after same-run recorder replacement", () => {
|
||||
const observe = createLiveRunUsageProjection({ receipt: { ...openai, costUsd: 0 }, costCents: "140" });
|
||||
expect(observe("default", { ...openai, model: "unsupported" })).toBe("140.0000000");
|
||||
expect(observe("next", { ...openai, costUsd: 0.3 })).toBe("170.0000000");
|
||||
});
|
||||
|
||||
it("restores explicit unknown exposure without pricing unsafe compatibility aggregates", () => {
|
||||
const observe = createLiveRunUsageProjection({ receipt: { ...openai, costUsd: 100 }, costCents: null });
|
||||
expect(observe("default", { ...openai, model: "unsupported" })).toBeNull();
|
||||
});
|
||||
|
||||
it("keeps unpriced attempts unknown while retaining known portions", () => {
|
||||
const observe = createLiveRunUsageProjection();
|
||||
expect(observe("unknown", { ...openai, usageBasis: "session_cumulative", costUsd: 100 })).toBeNull();
|
||||
expect(observe("known", { ...openai, costUsd: 0.1 })).toBe("10.0000000");
|
||||
expect(observe("subscription", { ...anthropic, billingType: "subscription_included", costUsd: 100 }))
|
||||
.toBe("10.0000000");
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,36 @@
|
||||
import type { AdapterUsageCheckpoint } from "@paperclipai/adapter-utils";
|
||||
import { priceAnthropicReceipt } from "./anthropic-pricing.js";
|
||||
import { knownRunCostCents } from "./budget-exposure.js";
|
||||
import { priceCodexReceipt } from "./codex-pricing.js";
|
||||
import { addCents, compareCents, normalizeCents } from "@paperclipai/shared";
|
||||
|
||||
/** Live admission can use a reviewed rate-card estimate before a receipt is
|
||||
* final. The original receipt stays incomplete/unpriced for ledger settlement. */
|
||||
export function liveRunUsageCostCents(receipt: AdapterUsageCheckpoint): string | null {
|
||||
const reported = knownRunCostCents({ ...receipt, usageSource: receipt.usageBasis });
|
||||
if (reported !== null) return reported;
|
||||
if (receipt.usageBasis !== "per_run") return null;
|
||||
const projection = priceAnthropicReceipt(priceCodexReceipt({ ...receipt, complete: true }));
|
||||
return knownRunCostCents({ ...projection, usageSource: projection.usageBasis });
|
||||
}
|
||||
|
||||
/** Attempt receipts can omit a later price without retracting earlier spend.
|
||||
* Keep that knowledge per attempt before combining different models/billers. */
|
||||
export function createLiveRunUsageProjection(initial?: {
|
||||
receipt: AdapterUsageCheckpoint;
|
||||
costCents?: string | null;
|
||||
}) {
|
||||
const costs = new Map<string, string>();
|
||||
if (initial) {
|
||||
const cost = initial.costCents !== undefined ? initial.costCents : liveRunUsageCostCents(initial.receipt);
|
||||
if (cost !== null) costs.set("default", normalizeCents(cost));
|
||||
}
|
||||
return (attemptId: string, receipt: AdapterUsageCheckpoint): string | null => {
|
||||
const reported = liveRunUsageCostCents(receipt);
|
||||
const previous = costs.get(attemptId);
|
||||
if (reported !== null && (receipt.complete || previous === undefined || compareCents(reported, previous) > 0)) {
|
||||
costs.set(attemptId, normalizeCents(reported));
|
||||
}
|
||||
return costs.size ? addCents(...costs.values()) : null;
|
||||
};
|
||||
}
|
||||
@@ -1,15 +1,78 @@
|
||||
import { and, eq, sql } from "drizzle-orm";
|
||||
import { agents, budgetPolicies, budgetReservations, heartbeatRuns, nativeRunFinalizations, type Db } from "@paperclipai/db";
|
||||
import { agents, budgetReservations, heartbeatRuns, nativeRunFinalizations, type Db } from "@paperclipai/db";
|
||||
import { centsToUnits, unitsToCents } from "@paperclipai/shared";
|
||||
import { conflict, notFound } from "../errors.js";
|
||||
import { withAccountingTransaction } from "./accounting-transaction.js";
|
||||
import { budgetServiceInTransaction, budgetPoliciesForRun, computeObservedSpend } from "./budgets.js";
|
||||
import { budgetServiceInTransaction } from "./budgets.js";
|
||||
import { withBudgetSpan } from "./budget-tracing.js";
|
||||
import type { ActivityPublication } from "./activity-log.js";
|
||||
|
||||
/** Native finalization takes its coordinator before issue/run locks. Recovery
|
||||
* admission uses the same order after the company accounting lock, so its
|
||||
* later ownership checks never wait for a coordinator while holding a run. */
|
||||
export async function lockRunBudgetRecoveryOwnerInTransaction(
|
||||
tx: Db, companyId: string, runId: string, recoveryLeaseOwner?: string,
|
||||
) {
|
||||
if (!recoveryLeaseOwner) return;
|
||||
const [coordinator] = await tx.select({ leaseOwner: nativeRunFinalizations.leaseOwner, leaseExpiresAt: nativeRunFinalizations.leaseExpiresAt }).from(nativeRunFinalizations)
|
||||
.where(and(eq(nativeRunFinalizations.companyId, companyId), eq(nativeRunFinalizations.runId, runId))).for("update");
|
||||
if (!coordinator || coordinator.leaseOwner !== recoveryLeaseOwner || !coordinator.leaseExpiresAt || coordinator.leaseExpiresAt.getTime() <= Date.now())
|
||||
throw conflict("Native recovery no longer owns provider dispatch");
|
||||
}
|
||||
|
||||
/** Called by the final dispatch gate with the company and run already locked.
|
||||
* Admission may have waited through preparation; that earlier reservation is
|
||||
* evidence of ownership, not permission to bypass a newer charge or pause.
|
||||
* Recovery callers prelock their coordinator before taking issue/run locks. */
|
||||
export async function recheckRunBudgetReservationInTransaction(
|
||||
tx: Db,
|
||||
publications: ActivityPublication[],
|
||||
run: typeof heartbeatRuns.$inferSelect,
|
||||
reservationId: string,
|
||||
recoveryLeaseOwner?: string,
|
||||
): Promise<string | null> {
|
||||
if (run.resultJson?.cancellation || run.resultJson?.startupCancellation)
|
||||
return "Run cancellation already requested";
|
||||
if (run.costAccountedAt || !["queued", "running"].includes(run.status))
|
||||
return "Run can no longer start provider work";
|
||||
const [reservation] = await tx.select().from(budgetReservations).where(and(
|
||||
eq(budgetReservations.id, reservationId), eq(budgetReservations.companyId, run.companyId),
|
||||
eq(budgetReservations.runId, run.id), eq(budgetReservations.agentId, run.agentId),
|
||||
eq(budgetReservations.state, "held"),
|
||||
)).for("update");
|
||||
if (!reservation) return "Provider dispatch no longer owns a held budget reservation";
|
||||
if (recoveryLeaseOwner) {
|
||||
const [owner] = run.runtimeMode === "native" ? await tx.select({ runId: nativeRunFinalizations.runId })
|
||||
.from(nativeRunFinalizations).where(and(eq(nativeRunFinalizations.companyId, run.companyId),
|
||||
eq(nativeRunFinalizations.runId, run.id), eq(nativeRunFinalizations.leaseOwner, recoveryLeaseOwner),
|
||||
sql`${nativeRunFinalizations.leaseExpiresAt} > clock_timestamp()`)) : [];
|
||||
if (!owner) return "Native recovery no longer owns provider dispatch";
|
||||
}
|
||||
const { block, policies, observations } = await budgetServiceInTransaction(tx, publications)
|
||||
.evaluateAdmission(run.companyId, run.agentId, { projectId: reservation.projectId });
|
||||
if (block) return block.reason;
|
||||
const [agent] = await tx.select({ status: agents.status, lifecycleState: agents.lifecycleState }).from(agents)
|
||||
.where(and(eq(agents.companyId, run.companyId), eq(agents.id, run.agentId)));
|
||||
if (!agent || agent.lifecycleState !== "ready" || ["paused", "terminated", "pending_approval"].includes(agent.status))
|
||||
return "Agent is paused or its lifecycle is not ready for new work";
|
||||
for (const [index, policy] of policies.entries()) {
|
||||
const observed = observations[index];
|
||||
// This operation's hold is already included in exposure. Do not propose it
|
||||
// twice, or reject a nonzero reservation that exactly fills the budget.
|
||||
const committed = centsToUnits(observed.ledgerCents) + centsToUnits(observed.effectiveHeldCents);
|
||||
const limit = centsToUnits(policy.amount);
|
||||
if (committed > limit || (centsToUnits(reservation.amountCents) === 0n && committed >= limit))
|
||||
return "Available budget is reserved by unfinished runs";
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
/** Reserve before dispatch, under the same company lock as charges and policy
|
||||
* changes. Estimates constrain admission; they cannot cap a provider's bill.
|
||||
* A reservation survives timeouts and restarts until accounting proves closure. */
|
||||
export async function reserveRunBudget(db: Db, companyId: string, runId: string, projectId: string | null, ledgerScope: Record<string, unknown> = {}, recoveryLeaseOwner?: string) {
|
||||
return withAccountingTransaction(db, companyId, async (tx, publications) => {
|
||||
const result = await withBudgetSpan("budget.admission.reserve", () => withAccountingTransaction(db, companyId, async (tx, publications) => {
|
||||
await lockRunBudgetRecoveryOwnerInTransaction(tx, companyId, runId, recoveryLeaseOwner);
|
||||
const [run] = await tx.select().from(heartbeatRuns).where(and(eq(heartbeatRuns.id, runId), eq(heartbeatRuns.companyId, companyId))).for("update");
|
||||
if (!run) throw notFound("Run not found");
|
||||
if (run.resultJson?.cancellation || run.resultJson?.startupCancellation) throw conflict("Run cancellation already requested");
|
||||
@@ -22,40 +85,37 @@ export async function reserveRunBudget(db: Db, companyId: string, runId: string,
|
||||
&& existing.projectId === projectId ? await tx.select({ runId: nativeRunFinalizations.runId })
|
||||
.from(nativeRunFinalizations).where(and(eq(nativeRunFinalizations.companyId, companyId),
|
||||
eq(nativeRunFinalizations.runId, runId), eq(nativeRunFinalizations.leaseOwner, recoveryLeaseOwner),
|
||||
sql`${nativeRunFinalizations.leaseExpiresAt} > now()`)).for("update") : [];
|
||||
sql`${nativeRunFinalizations.leaseExpiresAt} > clock_timestamp()`)) : [];
|
||||
if (owner) {
|
||||
// Recovery has already made this run queued/running, so its own pending
|
||||
// accounting is excluded. Current overruns, other unfinished runs and
|
||||
// scope pauses still apply before any more paid work starts.
|
||||
const block = await budgetServiceInTransaction(tx, publications).getInvocationBlock(companyId, run.agentId, { projectId });
|
||||
if (block) throw conflict(block.reason);
|
||||
const [agent] = await tx.select({ status: agents.status }).from(agents)
|
||||
if (block) return { admissionBlockReason: block.reason };
|
||||
const [agent] = await tx.select({ status: agents.status, lifecycleState: agents.lifecycleState }).from(agents)
|
||||
.where(and(eq(agents.companyId, companyId), eq(agents.id, run.agentId)));
|
||||
if (agent?.status === "paused") throw conflict("Agent is paused and cannot start new work");
|
||||
if (!agent || agent.lifecycleState !== "ready" || ["paused", "terminated", "pending_approval"].includes(agent.status))
|
||||
return { admissionBlockReason: "Agent is paused and cannot start new work" };
|
||||
return { ...existing, reused: true };
|
||||
}
|
||||
throw conflict("Provider dispatch has already reserved this run");
|
||||
}
|
||||
const block = await budgetServiceInTransaction(tx, publications).getInvocationBlock(companyId, run.agentId, { projectId });
|
||||
if (block) throw conflict(block.reason);
|
||||
const policies = await tx.select().from(budgetPolicies).where(and(
|
||||
eq(budgetPolicies.companyId, companyId), eq(budgetPolicies.isActive, true), eq(budgetPolicies.hardStopEnabled, true),
|
||||
eq(budgetPolicies.metric, "billed_cents"), sql`${budgetPolicies.amount} > 0`,
|
||||
budgetPoliciesForRun(companyId, run.agentId, projectId),
|
||||
));
|
||||
const { block, policies, observations } = await budgetServiceInTransaction(tx, publications).evaluateAdmission(companyId, run.agentId, { projectId });
|
||||
if (block) return { admissionBlockReason: block.reason };
|
||||
// A budget release can start asynchronous resume hooks. Capacity alone
|
||||
// does not authorize dispatch before that lifecycle operation is ready.
|
||||
const [agent] = await tx.select({ status: agents.status, lifecycleState: agents.lifecycleState }).from(agents)
|
||||
.where(and(eq(agents.companyId, companyId), eq(agents.id, run.agentId)));
|
||||
if (!agent || agent.lifecycleState !== "ready" || ["paused", "terminated", "pending_approval"].includes(agent.status))
|
||||
return { admissionBlockReason: "Agent is paused or its lifecycle is not ready for new work" };
|
||||
const amount = policies.reduce((max, policy) => { const next = centsToUnits(policy.reservationCents); return next > max ? next : max; }, 0n);
|
||||
for (const policy of policies) {
|
||||
const [held] = await tx.select({ amount: sql<string>`coalesce(sum(${budgetReservations.amountCents}), 0)::text` }).from(budgetReservations).where(and(
|
||||
eq(budgetReservations.companyId, companyId), eq(budgetReservations.state, "held"),
|
||||
policy.scopeType === "agent" ? eq(budgetReservations.agentId, run.agentId) : undefined,
|
||||
policy.scopeType === "project" ? eq(budgetReservations.projectId, projectId!) : undefined,
|
||||
));
|
||||
const observed = await computeObservedSpend(tx, policy);
|
||||
for (const [index, policy] of policies.entries()) {
|
||||
const observed = observations[index];
|
||||
// Even zero-estimate runs cannot enter a scope whose available capacity
|
||||
// is completely reserved. Outstanding reservations carry across months.
|
||||
const committed = centsToUnits(observed.totalExact) + centsToUnits(held.amount);
|
||||
const committed = centsToUnits(observed.ledgerCents) + centsToUnits(observed.effectiveHeldCents);
|
||||
const limit = centsToUnits(policy.amount);
|
||||
if (committed >= limit || committed + amount > limit) throw conflict("Available budget is reserved by unfinished runs");
|
||||
if (committed >= limit || committed + amount > limit) return { admissionBlockReason: "Available budget is reserved by unfinished runs" };
|
||||
}
|
||||
// A zero-valued row also fences duplicate dispatch when estimates are off.
|
||||
const [reservation] = await tx.insert(budgetReservations).values({ companyId, runId, agentId: run.agentId, projectId,
|
||||
@@ -64,5 +124,9 @@ export async function reserveRunBudget(db: Db, companyId: string, runId: string,
|
||||
usageJson: sql`coalesce(${heartbeatRuns.usageJson}, '{}'::jsonb) || ${JSON.stringify({ accountingReceiptReady: false, ledgerScope: { ...(run.usageJson?.ledgerScope as object ?? {}), ...ledgerScope, projectId } })}::jsonb`,
|
||||
}).where(eq(heartbeatRuns.id, runId));
|
||||
return { ...reservation, reused: false };
|
||||
});
|
||||
}));
|
||||
// Commit required block/resume transitions even when dispatch is denied.
|
||||
// Throwing inside the transaction would discard those durable obligations.
|
||||
if ("admissionBlockReason" in result) throw conflict(result.admissionBlockReason ?? "Budget admission blocked");
|
||||
return result;
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
import { afterEach, expect, it, vi } from "vitest";
|
||||
import { withBudgetSpan } from "./budget-tracing.js";
|
||||
|
||||
const start = vi.hoisted(() => vi.fn());
|
||||
vi.mock("@opentelemetry/api", () => ({ trace: { getTracer: () => ({ startActiveSpan: start }) } }));
|
||||
afterEach(() => { vi.unstubAllEnvs(); vi.resetAllMocks(); });
|
||||
|
||||
it("leaves tracing unused without an operator endpoint", async () => {
|
||||
vi.stubEnv("OTEL_EXPORTER_OTLP_ENDPOINT", "");
|
||||
await expect(withBudgetSpan("budget.accounting.transaction", async () => 42)).resolves.toBe(42);
|
||||
expect(start).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it.each(["before", "after", "end", "rejected", "missing"])("keeps accounting authoritative when tracing fails %s", async failure => {
|
||||
vi.stubEnv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://collector.test");
|
||||
start.mockImplementation((_name, _options, callback) => {
|
||||
if (failure === "before") throw new Error("before callback");
|
||||
if (failure === "rejected") return Promise.reject(new Error("tracer rejected"));
|
||||
if (failure === "missing") return undefined;
|
||||
const result = callback({ end: () => { if (failure === "end") throw new Error("end failed"); } });
|
||||
if (failure === "after") throw new Error("after callback");
|
||||
return result;
|
||||
});
|
||||
const work = vi.fn(async () => 42);
|
||||
await expect(withBudgetSpan("budget.accounting.transaction", work)).resolves.toBe(42);
|
||||
expect(work).toHaveBeenCalledOnce();
|
||||
const error = new Error("original accounting failure");
|
||||
const failed = vi.fn(async () => { throw error; });
|
||||
await expect(withBudgetSpan("budget.accounting.transaction", failed)).rejects.toBe(error);
|
||||
expect(failed).toHaveBeenCalledOnce();
|
||||
});
|
||||
@@ -0,0 +1,18 @@
|
||||
import { trace, type Span } from "@opentelemetry/api";
|
||||
|
||||
/** Operator-gated OpenTelemetry only; these spans contain no prompt, receipt,
|
||||
* credential, or provider error text. Diagnostics never change accounting. */
|
||||
export async function withBudgetSpan<T>(name: string, work: () => Promise<T>, attributes?: Record<string, number | boolean>) {
|
||||
if (!process.env.OTEL_EXPORTER_OTLP_ENDPOINT?.trim()) return work();
|
||||
let operation: Promise<T> | undefined;
|
||||
const run = (span?: Span) => operation ??= Promise.resolve().then(work).finally(() => {
|
||||
try { span?.end(); } catch { /* Diagnostics cannot change accounting. */ }
|
||||
});
|
||||
try {
|
||||
const traced = trace.getTracer("paperclip.budgets").startActiveSpan(name, { attributes }, span => run(span));
|
||||
// Some external SDK implementations can fail before or after the callback.
|
||||
// Own the accounting promise so either failure still executes work once.
|
||||
void Promise.resolve(traced).catch(() => undefined);
|
||||
} catch { /* Preserve the accounting operation if span creation fails. */ }
|
||||
return operation ?? run();
|
||||
}
|
||||
+303
-202
@@ -1,6 +1,5 @@
|
||||
import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notifications.js";
|
||||
import { compareCents, normalizeCents } from "@paperclipai/shared";
|
||||
import { and, desc, eq, gte, inArray, lt, ne, sql, or, isNull, type SQL } from "drizzle-orm";
|
||||
import { addCents, compareCents, normalizeCents } from "@paperclipai/shared";
|
||||
import { and, desc, eq, inArray, sql, or, isNull } from "drizzle-orm";
|
||||
import type { Db } from "@paperclipai/db";
|
||||
|
||||
import {
|
||||
@@ -8,10 +7,10 @@ import {
|
||||
approvals,
|
||||
budgetIncidents,
|
||||
budgetPolicies,
|
||||
budgetReservations,
|
||||
companies,
|
||||
costEvents,
|
||||
heartbeatRuns,
|
||||
nativeRunFinalizations,
|
||||
projects,
|
||||
} from "@paperclipai/db";
|
||||
import type {
|
||||
@@ -31,12 +30,19 @@ import { HttpError, notFound, unprocessable } from "../errors.js";
|
||||
import { logger } from "../middleware/logger.js";
|
||||
import { withAccountingTransaction } from "./accounting-transaction.js";
|
||||
import { logActivity, type LogActivityInput, type ActivityPublication } from "./activity-log.js";
|
||||
import { computeObservedSpend, computeObservedSpends, type ObservedSpend } from "./budget-exposure.js";
|
||||
import { notifyDeliveryWork, DELIVERY_QUEUES } from "./delivery-work-notifications.js";
|
||||
import { listPendingBudgetCancellationScopes } from "./heartbeat/budget-cancellation.js";
|
||||
import { withBudgetSpan } from "./budget-tracing.js";
|
||||
export { computeObservedSpend, computeObservedSpends } from "./budget-exposure.js";
|
||||
export type { BudgetExposure } from "./budget-exposure.js";
|
||||
|
||||
type ScopeRecord = {
|
||||
companyId: string;
|
||||
name: string;
|
||||
paused: boolean;
|
||||
pauseReason: PauseReason | null;
|
||||
budgetHeld: boolean;
|
||||
};
|
||||
|
||||
type PolicyRow = typeof budgetPolicies.$inferSelect;
|
||||
@@ -109,6 +115,7 @@ async function resolveScopeRecord(db: Db, scopeType: BudgetScopeType, scopeId: s
|
||||
name: row.name,
|
||||
paused: row.status === "paused" || Boolean(row.pausedAt),
|
||||
pauseReason: (row.pauseReason as ScopeRecord["pauseReason"]) ?? null,
|
||||
budgetHeld: row.pauseReason === "budget",
|
||||
};
|
||||
}
|
||||
|
||||
@@ -119,6 +126,7 @@ async function resolveScopeRecord(db: Db, scopeType: BudgetScopeType, scopeId: s
|
||||
name: agents.name,
|
||||
status: agents.status,
|
||||
pauseReason: agents.pauseReason,
|
||||
lifecycleHolds: agents.lifecycleHolds,
|
||||
})
|
||||
.from(agents)
|
||||
.where(eq(agents.id, scopeId))
|
||||
@@ -129,6 +137,7 @@ async function resolveScopeRecord(db: Db, scopeType: BudgetScopeType, scopeId: s
|
||||
name: row.name,
|
||||
paused: row.status === "paused",
|
||||
pauseReason: (row.pauseReason as ScopeRecord["pauseReason"]) ?? null,
|
||||
budgetHeld: row.lifecycleHolds.includes("budget"),
|
||||
};
|
||||
}
|
||||
|
||||
@@ -148,6 +157,7 @@ async function resolveScopeRecord(db: Db, scopeType: BudgetScopeType, scopeId: s
|
||||
name: row.name,
|
||||
paused: Boolean(row.pausedAt),
|
||||
pauseReason: (row.pauseReason as ScopeRecord["pauseReason"]) ?? null,
|
||||
budgetHeld: row.pauseReason === "budget",
|
||||
};
|
||||
}
|
||||
|
||||
@@ -169,68 +179,50 @@ export function budgetPoliciesForRun(companyId: string, agentId: string | null,
|
||||
);
|
||||
}
|
||||
|
||||
type SpendPolicy = Pick<PolicyRow, "companyId" | "scopeType" | "scopeId" | "windowKind" | "metric">;
|
||||
type ObservedSpend = { windowStart: Date; windowEnd: Date; total: number; totalExact: string; unpricedEventCount: number; pendingRunCount: number; recoveringRunCount: number };
|
||||
|
||||
/** One ledger scan and one pending-run scan for the affected policies. The
|
||||
* company lock still protects decisions; nothing is cached across mutations.
|
||||
* Each aggregate keeps its own scope/window, including mixed lifetime and
|
||||
* monthly policies. An outer predicate excludes rows no policy needs. */
|
||||
async function computeObservedSpends(db: Db, policies: SpendPolicy[]): Promise<ObservedSpend[]> {
|
||||
const now = new Date();
|
||||
const windows = policies.map(policy => resolveWindow(policy.windowKind as BudgetWindowKind, now));
|
||||
const billable = policies.map((policy, index) => ({ policy, index })).filter(({ policy }) => policy.metric === "billed_cents");
|
||||
const ledgerColumns: Record<string, SQL> = {};
|
||||
const pendingColumns: Record<string, SQL> = {};
|
||||
const ledgerPredicates: SQL[] = [], pendingPredicates: SQL[] = [];
|
||||
for (const { policy, index } of billable) {
|
||||
const { start, end } = windows[index];
|
||||
const ledger = and(
|
||||
eq(costEvents.companyId, policy.companyId),
|
||||
policy.scopeType === "agent" ? eq(costEvents.agentId, policy.scopeId) : undefined,
|
||||
policy.scopeType === "project" ? eq(costEvents.projectId, policy.scopeId) : undefined,
|
||||
policy.windowKind === "calendar_month_utc" ? and(gte(costEvents.occurredAt, start), lt(costEvents.occurredAt, end)) : undefined,
|
||||
)!;
|
||||
const pending = and(
|
||||
eq(heartbeatRuns.companyId, policy.companyId), eq(heartbeatRuns.costAccountingPending, true),
|
||||
inArray(heartbeatRuns.status, ["succeeded", "failed", "timed_out", "cancelled", "interrupted"]),
|
||||
policy.scopeType === "agent" ? eq(heartbeatRuns.agentId, policy.scopeId) : undefined,
|
||||
policy.scopeType === "project" ? sql`${heartbeatRuns.usageJson}->'ledgerScope'->>'projectId' = ${policy.scopeId}` : undefined,
|
||||
policy.windowKind === "calendar_month_utc" ? sql`coalesce(${heartbeatRuns.finishedAt}, ${heartbeatRuns.createdAt}) >= ${start.toISOString()}::timestamptz
|
||||
and coalesce(${heartbeatRuns.finishedAt}, ${heartbeatRuns.createdAt}) < ${end.toISOString()}::timestamptz` : undefined,
|
||||
)!;
|
||||
ledgerPredicates.push(ledger); pendingPredicates.push(pending);
|
||||
ledgerColumns[`total${index}`] = sql`coalesce(sum(${costEvents.costCents}) filter (where ${ledger}), 0)::text`;
|
||||
ledgerColumns[`unpriced${index}`] = sql`count(*) filter (where ${ledger} and ${costEvents.costStatus} = 'unpriced' and ${costEvents.billingType} <> 'subscription_included')::int`;
|
||||
pendingColumns[`pending${index}`] = sql`count(*) filter (where ${pending})::int`;
|
||||
pendingColumns[`recovering${index}`] = sql`count(*) filter (where ${pending} and ${heartbeatRuns.runtimeMode} = 'native'
|
||||
and exists (select 1 from ${nativeRunFinalizations} where ${nativeRunFinalizations.runId} = ${heartbeatRuns.id}
|
||||
and ${nativeRunFinalizations.companyId} = ${heartbeatRuns.companyId}
|
||||
and ${nativeRunFinalizations.resultId} is null and ${nativeRunFinalizations.phase} <> 'terminal_failure'))::int`;
|
||||
}
|
||||
const [ledger = {}] = billable.length ? await db.select(ledgerColumns).from(costEvents).where(or(...ledgerPredicates)) : [];
|
||||
const [pending = {}] = billable.length ? await db.select(pendingColumns).from(heartbeatRuns).where(or(...pendingPredicates)) : [];
|
||||
return policies.map((_, index) => ({
|
||||
windowStart: windows[index].start, windowEnd: windows[index].end,
|
||||
total: Number(ledger[`total${index}`] ?? 0), totalExact: normalizeCents(String(ledger[`total${index}`] ?? 0)),
|
||||
unpricedEventCount: Number(ledger[`unpriced${index}`] ?? 0), pendingRunCount: Number(pending[`pending${index}`] ?? 0),
|
||||
recoveringRunCount: Number(pending[`recovering${index}`] ?? 0),
|
||||
}));
|
||||
}
|
||||
|
||||
export async function computeObservedSpend(db: Db, policy: SpendPolicy) {
|
||||
return (await computeObservedSpends(db, [policy]))[0];
|
||||
function unpricedUsageBlocks(policy: PolicyRow, observed: ObservedSpend) {
|
||||
return policy.unpricedUsagePolicy !== "allow"
|
||||
&& (observed.unpricedEventCount > 0 || observed.unknownBrowserRunCount > 0);
|
||||
}
|
||||
|
||||
function observedBlocks(policy: PolicyRow, observed: Awaited<ReturnType<typeof computeObservedSpend>>) {
|
||||
if (!policy.isActive || !policy.hardStopEnabled || policy.amount <= 0) return false;
|
||||
return compareCents(observed.totalExact, policy.amount) >= 0 || (observed.pendingRunCount - observed.recoveringRunCount) > 0 || (policy.unpricedUsagePolicy !== "allow" && observed.unpricedEventCount > 0);
|
||||
return compareCents(observed.totalExact, policy.amount) >= 0 || (observed.pendingRunCount - observed.recoveringRunCount) > 0 || unpricedUsageBlocks(policy, observed);
|
||||
}
|
||||
export async function policyBlocks(db: Db, policy: PolicyRow) {
|
||||
if (!policy.isActive || !policy.hardStopEnabled || policy.amount <= 0) return false;
|
||||
return observedBlocks(policy, await computeObservedSpend(db, policy));
|
||||
}
|
||||
|
||||
/** Supplemental handoff check after authoritative accounting admission commits.
|
||||
* Read only: the caller can hold issue/run ownership, so this must never take
|
||||
* the company lock, reconcile holds, or deliver effects. A later concurrent
|
||||
* charge targets already-admitted work through the normal durable Stop path. */
|
||||
export async function recheckAdmittedRunBudgetReadOnly(db: Db, input: {
|
||||
companyId: string; agentId: string; runId: string; projectId: string | null; reservationId: string;
|
||||
}): Promise<string | null> {
|
||||
return withBudgetSpan("budget.admission.evaluate", async () => {
|
||||
const [reservation] = await db.select({ amountCents: budgetReservations.amountCents }).from(budgetReservations).where(and(
|
||||
eq(budgetReservations.id, input.reservationId), eq(budgetReservations.companyId, input.companyId),
|
||||
eq(budgetReservations.runId, input.runId), eq(budgetReservations.agentId, input.agentId),
|
||||
sql`${budgetReservations.projectId} is not distinct from ${input.projectId}::uuid`,
|
||||
eq(budgetReservations.state, "held"),
|
||||
));
|
||||
if (!reservation) return "Provider dispatch no longer owns a held budget reservation";
|
||||
const policies = await db.select().from(budgetPolicies).where(and(eq(budgetPolicies.companyId, input.companyId),
|
||||
eq(budgetPolicies.isActive, true), eq(budgetPolicies.hardStopEnabled, true), eq(budgetPolicies.metric, "billed_cents"),
|
||||
sql`${budgetPolicies.amount} > 0`, budgetPoliciesForRun(input.companyId, input.agentId, input.projectId)));
|
||||
const observations = await computeObservedSpends(db, policies);
|
||||
for (const [index, policy] of policies.entries()) {
|
||||
const observed = observations[index];
|
||||
if (observedBlocks(policy, observed)) return "Applicable budget no longer permits provider work";
|
||||
const capacity = compareCents(addCents(observed.ledgerCents, observed.effectiveHeldCents), policy.amount);
|
||||
if (capacity > 0 || (compareCents(reservation.amountCents, 0) === 0 && capacity >= 0))
|
||||
return "Available budget is reserved by unfinished runs";
|
||||
}
|
||||
return null;
|
||||
});
|
||||
}
|
||||
|
||||
function buildApprovalPayload(input: {
|
||||
policy: PolicyRow;
|
||||
scopeName: string;
|
||||
@@ -283,7 +275,11 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
|
||||
const recordActivity = (input: LogActivityInput) => logActivity(db, input, publications);
|
||||
async function pauseScopeForBudget(policy: PolicyRow) {
|
||||
const now = new Date();
|
||||
if (policy.scopeType === "agent") return;
|
||||
if (policy.scopeType === "agent") {
|
||||
const { applyAgentBudgetHoldInTransaction } = await import("./agent-lifecycle.js");
|
||||
await applyAgentBudgetHoldInTransaction(db, policy.companyId, policy.scopeId, true);
|
||||
return;
|
||||
}
|
||||
|
||||
if (policy.scopeType === "project") {
|
||||
await db
|
||||
@@ -293,10 +289,13 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
|
||||
pausedAt: now,
|
||||
updatedAt: now,
|
||||
})
|
||||
.where(and(eq(projects.id, policy.scopeId), or(isNull(projects.pausedAt), eq(projects.pauseReason, "budget"))));
|
||||
.where(and(eq(projects.id, policy.scopeId), isNull(projects.pausedAt)));
|
||||
return;
|
||||
}
|
||||
|
||||
const scope = await resolveScopeRecord(db, "company", policy.scopeId);
|
||||
if (scope.paused) return;
|
||||
await notifyDeliveryWork(db, DELIVERY_QUEUES.agentLifecycle, policy.companyId);
|
||||
await notifyDeliveryWork(db, DELIVERY_QUEUES.skillSourceSync);
|
||||
await db
|
||||
.update(companies)
|
||||
@@ -306,13 +305,20 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
|
||||
pausedAt: now,
|
||||
updatedAt: now,
|
||||
})
|
||||
.where(and(eq(companies.id, policy.scopeId), or(eq(companies.status, "active"), and(eq(companies.status, "paused"), eq(companies.pauseReason, "budget")))));
|
||||
.where(and(eq(companies.id, policy.scopeId), eq(companies.status, "active")));
|
||||
}
|
||||
|
||||
async function pauseAndCancelScopeForBudget(policy: PolicyRow) {
|
||||
async function pauseAndCancelScopeForBudget(policy: PolicyRow, requestCancellation = true) {
|
||||
await pauseScopeForBudget(policy);
|
||||
if (!requestCancellation) return;
|
||||
// A pending generation already covers all admitted work in this scope.
|
||||
// Rechecks must not invalidate a worker that is delivering that intent.
|
||||
await notifyDeliveryWork(db, DELIVERY_QUEUES.budgetEnforcement, policy.companyId);
|
||||
await db.update(budgetPolicies).set({
|
||||
enforcementVersion: sql`${budgetPolicies.enforcementVersion} + 1`,
|
||||
enforcementVersion: sql`case when ${budgetPolicies.enforcementVersion} > ${budgetPolicies.enforcementDeliveredVersion}
|
||||
then ${budgetPolicies.enforcementVersion} else ${budgetPolicies.enforcementVersion} + 1 end`,
|
||||
updatedAt: sql`case when ${budgetPolicies.enforcementVersion} > ${budgetPolicies.enforcementDeliveredVersion}
|
||||
then ${budgetPolicies.updatedAt} else now() end`,
|
||||
}).where(eq(budgetPolicies.id, policy.id));
|
||||
}
|
||||
|
||||
@@ -321,11 +327,14 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
|
||||
eq(budgetPolicies.companyId, policy.companyId), eq(budgetPolicies.scopeType, policy.scopeType),
|
||||
eq(budgetPolicies.scopeId, policy.scopeId), eq(budgetPolicies.isActive, true), eq(budgetPolicies.hardStopEnabled, true),
|
||||
));
|
||||
for (const candidate of policies) {
|
||||
if (await policyBlocks(db, candidate)) return;
|
||||
}
|
||||
const observations = await computeObservedSpends(db, policies);
|
||||
if (policies.some((candidate, index) => observedBlocks(candidate, observations[index]))) return;
|
||||
const now = new Date();
|
||||
if (policy.scopeType === "agent") return;
|
||||
if (policy.scopeType === "agent") {
|
||||
const { applyAgentBudgetHoldInTransaction } = await import("./agent-lifecycle.js");
|
||||
await applyAgentBudgetHoldInTransaction(db, policy.companyId, policy.scopeId, false);
|
||||
return;
|
||||
}
|
||||
|
||||
if (policy.scopeType === "project") {
|
||||
await db
|
||||
@@ -340,6 +349,9 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
|
||||
}
|
||||
|
||||
await notifyDeliveryWork(db, DELIVERY_QUEUES.skillSourceSync);
|
||||
const scope = await resolveScopeRecord(db, "company", policy.scopeId);
|
||||
if (scope.pauseReason !== "budget") return;
|
||||
await notifyDeliveryWork(db, DELIVERY_QUEUES.agentLifecycle, policy.companyId);
|
||||
await db
|
||||
.update(companies)
|
||||
.set({
|
||||
@@ -559,23 +571,25 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
|
||||
});
|
||||
}
|
||||
|
||||
async function reconcileScope(companyId: string, scopeType: BudgetScopeType, scopeId: string) {
|
||||
const scope = await resolveScopeRecord(db, scopeType, scopeId);
|
||||
async function reconcileScope(companyId: string, scopeType: BudgetScopeType, scopeId: string,
|
||||
prepared?: { scope: ScopeRecord; policies: PolicyRow[]; observations: ObservedSpend[] }) {
|
||||
const scope = prepared?.scope ?? await resolveScopeRecord(db, scopeType, scopeId);
|
||||
if (scope.companyId !== companyId) throw notFound("Budget scope not found");
|
||||
const policies = await db.select().from(budgetPolicies).where(and(
|
||||
const policies = prepared?.policies ?? await db.select().from(budgetPolicies).where(and(
|
||||
eq(budgetPolicies.companyId, companyId), eq(budgetPolicies.scopeType, scopeType), eq(budgetPolicies.scopeId, scopeId),
|
||||
));
|
||||
let recoveringPolicy: PolicyRow | undefined;
|
||||
for (const policy of policies) {
|
||||
const observed = await computeObservedSpend(db, policy);
|
||||
const observations = prepared?.observations ?? await computeObservedSpends(db, policies);
|
||||
for (const [index, policy] of policies.entries()) {
|
||||
const observed = observations[index];
|
||||
if (observedBlocks(policy, observed)) {
|
||||
await createIncidentIfNeeded(policy, "hard", observed);
|
||||
await pauseAndCancelScopeForBudget(policy);
|
||||
await pauseAndCancelScopeForBudget(policy, !scope.budgetHeld);
|
||||
const title = scopeType === "company" ? "Company" : scopeType === "agent" ? "Agent" : "Project";
|
||||
if (observed.pendingRunCount > 0) {
|
||||
return { scopeType, scopeId, scopeName: scope.name, reason: `${title} cannot start work while completed runs await accounting.` };
|
||||
}
|
||||
if (policy.unpricedUsagePolicy !== "allow" && observed.unpricedEventCount > 0) {
|
||||
if (unpricedUsageBlocks(policy, observed)) {
|
||||
return { scopeType, scopeId, scopeName: scope.name, reason: `${title} cannot start work because recorded usage has no reliable price.` };
|
||||
}
|
||||
return { scopeType, scopeId, scopeName: scope.name, reason: scopeType === "project"
|
||||
@@ -600,12 +614,112 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
|
||||
if (scope.paused && scope.pauseReason === "budget" && !policies.length) {
|
||||
return { scopeType, scopeId, scopeName: scope.name, reason: "Budget pause requires a policy or an explicit operator resume." };
|
||||
}
|
||||
if (scope.pauseReason === "budget" && policies[0]) await resumeScopeFromBudget(policies[0]);
|
||||
if (scope.budgetHeld && policies[0]) await resumeScopeFromBudget(policies[0]);
|
||||
if (recoveringPolicy) return { scopeType, scopeId, scopeName: scope.name,
|
||||
reason: "New work must wait while a native run recovers its unfinished accounting." };
|
||||
return null;
|
||||
}
|
||||
|
||||
async function evaluateAffectedScopes(companyId: string, agentId: string | null, projectId: string | null) {
|
||||
const candidatePolicies = await db
|
||||
.select()
|
||||
.from(budgetPolicies)
|
||||
.where(
|
||||
and(
|
||||
eq(budgetPolicies.companyId, companyId),
|
||||
eq(budgetPolicies.isActive, true),
|
||||
budgetPoliciesForRun(companyId, agentId, projectId),
|
||||
),
|
||||
);
|
||||
|
||||
const policies = candidatePolicies.filter(policy => policy.metric === "billed_cents" && policy.amount > 0);
|
||||
const observations = await computeObservedSpends(db, policies);
|
||||
for (const [index, policy] of policies.entries()) {
|
||||
const observed = observations[index];
|
||||
const observedAmount = observed.total;
|
||||
const softThreshold = Math.ceil((policy.amount * policy.warnPercent) / 100);
|
||||
|
||||
if (policy.notifyEnabled && compareCents(observed.totalExact, softThreshold) >= 0 && (!policy.hardStopEnabled || compareCents(observed.totalExact, policy.amount) < 0)) {
|
||||
const softIncident = await createIncidentIfNeeded(policy, "soft", observed);
|
||||
if (softIncident?.created) {
|
||||
await recordActivity({
|
||||
companyId: policy.companyId,
|
||||
actorType: "system",
|
||||
actorId: "budget_service",
|
||||
action: "budget.soft_threshold_crossed",
|
||||
entityType: "budget_incident",
|
||||
entityId: softIncident.incident.id,
|
||||
details: {
|
||||
scopeType: policy.scopeType,
|
||||
scopeId: policy.scopeId,
|
||||
amountObserved: observedAmount,
|
||||
amountLimit: policy.amount,
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if (observedBlocks(policy, observed)) {
|
||||
await resolveOpenSoftIncidents(policy.id);
|
||||
const hardIncident = await createIncidentIfNeeded(policy, "hard", observed);
|
||||
await pauseAndCancelScopeForBudget(policy);
|
||||
if (hardIncident?.created) {
|
||||
await recordActivity({
|
||||
companyId: policy.companyId,
|
||||
actorType: "system",
|
||||
actorId: "budget_service",
|
||||
action: "budget.hard_threshold_crossed",
|
||||
entityType: "budget_incident",
|
||||
entityId: hardIncident.incident.id,
|
||||
details: {
|
||||
scopeType: policy.scopeType,
|
||||
scopeId: policy.scopeId,
|
||||
amountObserved: observedAmount,
|
||||
amountLimit: policy.amount,
|
||||
approvalId: hardIncident.incident.approvalId ?? null,
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
// A final receipt or price correction can reduce a partial observation.
|
||||
// Release only this scope's budget-owned hold after all its policies agree.
|
||||
const healthyScopes = new Map<string, PolicyRow>();
|
||||
for (const [index, policy] of policies.entries()) {
|
||||
if (!observedBlocks(policy, observations[index])) healthyScopes.set(`${policy.scopeType}:${policy.scopeId}`, policy);
|
||||
}
|
||||
for (const policy of healthyScopes.values()) {
|
||||
const scope = await resolveScopeRecord(db, policy.scopeType as BudgetScopeType, policy.scopeId);
|
||||
if (scope.budgetHeld) await resumeScopeFromBudget(policy);
|
||||
}
|
||||
}
|
||||
|
||||
async function evaluateAdmissionInTransaction(companyId: string, agentId: string | null,
|
||||
context?: { issueId?: string | null; projectId?: string | null }) {
|
||||
const scopes: Array<{ scopeType: BudgetScopeType; scopeId: string }> = [
|
||||
{ scopeType: "company", scopeId: companyId }, ...(agentId ? [{ scopeType: "agent" as const, scopeId: agentId }] : []),
|
||||
...(context?.projectId ? [{ scopeType: "project" as const, scopeId: context.projectId }] : []),
|
||||
];
|
||||
const records = await Promise.all(scopes.map(scope => resolveScopeRecord(db, scope.scopeType, scope.scopeId)));
|
||||
if (records.some(scope => scope.companyId !== companyId)) throw notFound("Budget scope not found");
|
||||
const policies = await db.select().from(budgetPolicies).where(and(eq(budgetPolicies.companyId, companyId),
|
||||
budgetPoliciesForRun(companyId, agentId, context?.projectId ?? null)));
|
||||
const observations = await computeObservedSpends(db, policies);
|
||||
const reservationIndexes = policies.map((policy, i) => policy.isActive && policy.hardStopEnabled && policy.metric === "billed_cents" && policy.amount > 0 ? i : -1).filter(i => i >= 0);
|
||||
const admission = (block: Awaited<ReturnType<typeof reconcileScope>>) => ({ block,
|
||||
policies: reservationIndexes.map(i => policies[i]), observations: reservationIndexes.map(i => observations[i]) });
|
||||
for (const [index, { scopeType, scopeId }] of scopes.entries()) {
|
||||
const indexes = policies.map((policy, i) => policy.scopeType === scopeType && policy.scopeId === scopeId ? i : -1).filter(i => i >= 0);
|
||||
const block = await reconcileScope(companyId, scopeType, scopeId, { scope: records[index],
|
||||
policies: indexes.map(i => policies[i]), observations: indexes.map(i => observations[i]) });
|
||||
if (block) return admission(block);
|
||||
}
|
||||
return admission(null);
|
||||
}
|
||||
|
||||
const evaluateAdmission = (...args: Parameters<typeof evaluateAdmissionInTransaction>) =>
|
||||
withBudgetSpan("budget.admission.evaluate", () => evaluateAdmissionInTransaction(...args));
|
||||
|
||||
return {
|
||||
reconcileScope,
|
||||
listPolicies: async (companyId: string): Promise<BudgetPolicy[]> => {
|
||||
@@ -652,6 +766,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
|
||||
// setting; editing a disabled positive policy must preserve that choice.
|
||||
const nextIsActive = amount > 0 && (input.isActive ?? (existing?.amount === 0 ? true : existing?.isActive) ?? true);
|
||||
const now = new Date();
|
||||
await notifyDeliveryWork(db, DELIVERY_QUEUES.budgetEnforcement, companyId);
|
||||
const row = existing
|
||||
? await db
|
||||
.update(budgetPolicies)
|
||||
@@ -665,6 +780,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
|
||||
isActive: nextIsActive,
|
||||
updatedByUserId: actorUserId,
|
||||
updatedAt: now,
|
||||
enforcementVersion: sql`${budgetPolicies.enforcementVersion} + 1`,
|
||||
})
|
||||
.where(eq(budgetPolicies.id, existing.id))
|
||||
.returning()
|
||||
@@ -773,88 +889,12 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
|
||||
};
|
||||
},
|
||||
|
||||
evaluateCostEvent: async (event: typeof costEvents.$inferSelect) => {
|
||||
const candidatePolicies = await db
|
||||
.select()
|
||||
.from(budgetPolicies)
|
||||
.where(
|
||||
and(
|
||||
eq(budgetPolicies.companyId, event.companyId),
|
||||
eq(budgetPolicies.isActive, true),
|
||||
budgetPoliciesForRun(event.companyId, event.agentId, event.projectId),
|
||||
),
|
||||
);
|
||||
evaluateAffectedScopes,
|
||||
evaluateCostEvent: (event: typeof costEvents.$inferSelect) => evaluateAffectedScopes(event.companyId, event.agentId, event.projectId),
|
||||
|
||||
const policies = candidatePolicies.filter(policy => policy.metric === "billed_cents" && policy.amount > 0);
|
||||
const observations = await computeObservedSpends(db, policies);
|
||||
for (const [index, policy] of policies.entries()) {
|
||||
const observed = observations[index];
|
||||
const observedAmount = observed.total;
|
||||
const softThreshold = Math.ceil((policy.amount * policy.warnPercent) / 100);
|
||||
|
||||
if (policy.notifyEnabled && compareCents(observed.totalExact, softThreshold) >= 0 && (!policy.hardStopEnabled || compareCents(observed.totalExact, policy.amount) < 0)) {
|
||||
const softIncident = await createIncidentIfNeeded(policy, "soft", observed);
|
||||
if (softIncident?.created) {
|
||||
await recordActivity({
|
||||
companyId: policy.companyId,
|
||||
actorType: "system",
|
||||
actorId: "budget_service",
|
||||
action: "budget.soft_threshold_crossed",
|
||||
entityType: "budget_incident",
|
||||
entityId: softIncident.incident.id,
|
||||
details: {
|
||||
scopeType: policy.scopeType,
|
||||
scopeId: policy.scopeId,
|
||||
amountObserved: observedAmount,
|
||||
amountLimit: policy.amount,
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if (observedBlocks(policy, observed)) {
|
||||
await resolveOpenSoftIncidents(policy.id);
|
||||
const hardIncident = await createIncidentIfNeeded(policy, "hard", observed);
|
||||
await pauseAndCancelScopeForBudget(policy);
|
||||
if (hardIncident?.created) {
|
||||
await recordActivity({
|
||||
companyId: policy.companyId,
|
||||
actorType: "system",
|
||||
actorId: "budget_service",
|
||||
action: "budget.hard_threshold_crossed",
|
||||
entityType: "budget_incident",
|
||||
entityId: hardIncident.incident.id,
|
||||
details: {
|
||||
scopeType: policy.scopeType,
|
||||
scopeId: policy.scopeId,
|
||||
amountObserved: observedAmount,
|
||||
amountLimit: policy.amount,
|
||||
approvalId: hardIncident.incident.approvalId ?? null,
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
|
||||
getInvocationBlock: async (
|
||||
companyId: string,
|
||||
agentId: string | null,
|
||||
context?: { issueId?: string | null; projectId?: string | null },
|
||||
) => {
|
||||
const scopes: Array<{ scopeType: BudgetScopeType; scopeId: string }> = [
|
||||
{ scopeType: "company", scopeId: companyId }, ...(agentId ? [{ scopeType: "agent" as const, scopeId: agentId }] : []),
|
||||
...(context?.projectId ? [{ scopeType: "project" as const, scopeId: context.projectId }] : []),
|
||||
];
|
||||
for (const { scopeType, scopeId } of scopes) {
|
||||
if ((await resolveScopeRecord(db, scopeType, scopeId)).companyId !== companyId) throw notFound("Budget scope not found");
|
||||
}
|
||||
for (const { scopeType, scopeId } of scopes) {
|
||||
const block = await reconcileScope(companyId, scopeType, scopeId);
|
||||
if (block) return block;
|
||||
}
|
||||
return null;
|
||||
},
|
||||
evaluateAdmission,
|
||||
getInvocationBlock: async (companyId: string, agentId: string | null, context?: { issueId?: string | null; projectId?: string | null }) =>
|
||||
(await evaluateAdmission(companyId, agentId, context)).block,
|
||||
|
||||
resolveIncident: async (
|
||||
companyId: string,
|
||||
@@ -874,7 +914,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
|
||||
if (input.action === "raise_budget_and_resume") {
|
||||
const nextAmount = Math.max(0, Math.floor(input.amount ?? 0));
|
||||
const observed = await computeObservedSpend(db, policy);
|
||||
if (policy.unpricedUsagePolicy !== "allow" && observed.unpricedEventCount > 0) {
|
||||
if (unpricedUsageBlocks(policy, observed)) {
|
||||
throw unprocessable("Resolve unpriced usage or explicitly allow it in the budget policy before resuming");
|
||||
}
|
||||
if (observed.pendingRunCount > 0) throw unprocessable("Completed runs must finish accounting before resuming");
|
||||
@@ -883,6 +923,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
|
||||
}
|
||||
|
||||
const now = new Date();
|
||||
await notifyDeliveryWork(db, DELIVERY_QUEUES.budgetEnforcement, companyId);
|
||||
await db
|
||||
.update(budgetPolicies)
|
||||
.set({
|
||||
@@ -890,6 +931,7 @@ export function budgetServiceInTransaction(db: Db, publications: ActivityPublica
|
||||
isActive: true,
|
||||
updatedByUserId: actorUserId,
|
||||
updatedAt: now,
|
||||
enforcementVersion: sql`${budgetPolicies.enforcementVersion} + 1`,
|
||||
})
|
||||
.where(eq(budgetPolicies.id, policy.id));
|
||||
|
||||
@@ -972,64 +1014,123 @@ export async function withCurrentBudgetEnforcement<T>(db: Db, scope: BudgetEnfor
|
||||
});
|
||||
}
|
||||
|
||||
/** Cancellation is an at-least-once external effect. A failed delivery never
|
||||
* rolls back committed spend, and its version remains pending for recovery. */
|
||||
export async function deliverBudgetEnforcement(db: Db, hooks: BudgetServiceHooks, companyId?: string, agentId?: string | null) {
|
||||
const { createAgentLifecycle } = await import("./agent-lifecycle.js");
|
||||
// Global recovery uses the bounded lifecycle sweep for policy holds.
|
||||
if (companyId) await createAgentLifecycle(db).reconcilePolicyHolds(companyId, agentId);
|
||||
async function deliverScopes<T>(items: T[], deliver: (item: T) => Promise<void>) {
|
||||
let next = 0;
|
||||
await Promise.all(Array.from({ length: Math.min(4, items.length) }, async () => {
|
||||
while (next < items.length) await deliver(items[next++]);
|
||||
}));
|
||||
}
|
||||
|
||||
/** Cancellation is an at-least-once external effect. Only the delivery worker
|
||||
* invokes this: request paths commit intent and return without awaiting it. */
|
||||
export async function deliverBudgetEnforcement(db: Db, hooks: BudgetServiceHooks, companyId?: string, _agentId?: string | null) {
|
||||
if (!hooks.cancelWorkForScope) return;
|
||||
const attempted = new Set<string>();
|
||||
const scopeKey = (scope: BudgetEnforcementScope) => `${scope.companyId}:${scope.scopeType}:${scope.scopeId}`;
|
||||
const pending = await db.select().from(budgetPolicies).where(and(
|
||||
sql`${budgetPolicies.enforcementVersion} > ${budgetPolicies.enforcementDeliveredVersion}`,
|
||||
companyId ? eq(budgetPolicies.companyId, companyId) : undefined,
|
||||
));
|
||||
for (const policy of pending) {
|
||||
try {
|
||||
const cancellation = await withAccountingTransaction(db, policy.companyId, async (tx) => {
|
||||
const [current] = await tx.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.id));
|
||||
if (!current || current.enforcementVersion !== policy.enforcementVersion || !(await policyBlocks(tx, current))) return null;
|
||||
return { companyId: policy.companyId, scopeType: policy.scopeType as BudgetScopeType, scopeId: policy.scopeId, createdBefore: new Date(), enforcement: { policyId: policy.id, version: policy.enforcementVersion } };
|
||||
});
|
||||
if (cancellation) await hooks.cancelWorkForScope(cancellation);
|
||||
await db.update(budgetPolicies).set({ enforcementDeliveredVersion: policy.enforcementVersion })
|
||||
.where(and(eq(budgetPolicies.id, policy.id), eq(budgetPolicies.enforcementVersion, policy.enforcementVersion)));
|
||||
} catch (error) {
|
||||
logger.warn({ err: error, policyId: policy.id }, "Budget cancellation delivery pending; accounting is committed");
|
||||
}
|
||||
}
|
||||
)).orderBy(budgetPolicies.updatedAt, budgetPolicies.id).limit(100);
|
||||
await deliverScopes(pending, policy => withBudgetSpan("budget.enforcement.scope", async () => {
|
||||
try {
|
||||
const cancellation = await withAccountingTransaction(db, policy.companyId, async tx => {
|
||||
const [current] = await tx.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.id));
|
||||
if (!current || current.enforcementVersion !== policy.enforcementVersion || !(await policyBlocks(tx, current))) return null;
|
||||
return { companyId: policy.companyId, scopeType: policy.scopeType as BudgetScopeType, scopeId: policy.scopeId,
|
||||
createdBefore: new Date(), enforcement: { policyId: policy.id, version: policy.enforcementVersion } };
|
||||
});
|
||||
if (cancellation) {
|
||||
attempted.add(scopeKey(cancellation));
|
||||
await hooks.cancelWorkForScope!(cancellation);
|
||||
}
|
||||
await db.update(budgetPolicies).set({ enforcementDeliveredVersion: policy.enforcementVersion })
|
||||
.where(and(eq(budgetPolicies.id, policy.id), eq(budgetPolicies.enforcementVersion, policy.enforcementVersion)));
|
||||
} catch (error) {
|
||||
logger.warn({ err: error, policyId: policy.id }, "Budget cancellation delivery pending; accounting is committed");
|
||||
// Rotate a failed scope behind other pending work. A poisoned delivery
|
||||
// must not monopolize the bounded batch on every retry.
|
||||
await db.update(budgetPolicies).set({ updatedAt: new Date() }).where(and(eq(budgetPolicies.id, policy.id),
|
||||
eq(budgetPolicies.enforcementVersion, policy.enforcementVersion))).catch(() => undefined);
|
||||
}
|
||||
}, { "budget.enforcement.queue_delay_ms": Math.max(0, Date.now() - policy.updatedAt.getTime()) }));
|
||||
// A committed run Stop outlives a policy edit or deletion. It remains owed
|
||||
// even if no policy generation now authorizes a fresh cancellation.
|
||||
const committed = (await listPendingBudgetCancellationScopes(db))
|
||||
.filter(scope => (!companyId || scope.companyId === companyId) && !attempted.has(scopeKey(scope)));
|
||||
await deliverScopes(committed, async scope => {
|
||||
try { await hooks.cancelWorkForScope!(scope); }
|
||||
catch (error) { logger.warn({ err: error, companyId: scope.companyId, scopeType: scope.scopeType, scopeId: scope.scopeId }, "Committed budget Stop remains pending"); }
|
||||
});
|
||||
}
|
||||
|
||||
export function budgetService(db: Db, hooks: BudgetServiceHooks = {}) {
|
||||
const reads = budgetServiceInTransaction(db);
|
||||
async function mutate<T>(companyId: string, work: (service: ReturnType<typeof budgetServiceInTransaction>) => Promise<T>, agentId?: string | null) {
|
||||
const result = await withAccountingTransaction(db, companyId, (tx, publications) => work(budgetServiceInTransaction(tx, publications)));
|
||||
await deliverBudgetEnforcement(db, hooks, companyId, agentId);
|
||||
return result;
|
||||
export async function hasPendingBudgetEnforcement(db: Db): Promise<boolean> {
|
||||
const [pending] = await db.select({ id: budgetPolicies.id }).from(budgetPolicies)
|
||||
.where(sql`${budgetPolicies.enforcementVersion} > ${budgetPolicies.enforcementDeliveredVersion}`).limit(1);
|
||||
return Boolean(pending) || (await listPendingBudgetCancellationScopes(db, 1)).length > 0;
|
||||
}
|
||||
|
||||
const budgetOwnedScope = or(
|
||||
and(eq(budgetPolicies.scopeType, "company"), sql`exists (select 1 from ${companies} where ${companies.id} = ${budgetPolicies.scopeId} and ${companies.pauseReason} = 'budget')`),
|
||||
and(eq(budgetPolicies.scopeType, "agent"), sql`exists (select 1 from ${agents} where ${agents.id} = ${budgetPolicies.scopeId} and ${agents.lifecycleHolds} @> '["budget"]'::jsonb)`),
|
||||
and(eq(budgetPolicies.scopeType, "project"), sql`exists (select 1 from ${projects} where ${projects.id} = ${budgetPolicies.scopeId} and ${projects.pauseReason} = 'budget')`),
|
||||
)!;
|
||||
|
||||
const openIncidentScope = sql`exists (select 1 from ${budgetIncidents}
|
||||
where ${budgetIncidents.policyId} = ${budgetPolicies.id} and ${budgetIncidents.status} = 'open')`;
|
||||
const unfinishedRunScope = sql`exists (select 1 from ${heartbeatRuns}
|
||||
where ${heartbeatRuns.companyId} = ${budgetPolicies.companyId} and ${heartbeatRuns.costAccountingPending} = true
|
||||
and ${heartbeatRuns.costAccountedAt} is null and (${budgetPolicies.scopeType} = 'company'
|
||||
or (${budgetPolicies.scopeType} = 'agent' and ${heartbeatRuns.agentId} = ${budgetPolicies.scopeId})
|
||||
or (${budgetPolicies.scopeType} = 'project' and ${heartbeatRuns.usageJson}->'ledgerScope'->>'projectId' = ${budgetPolicies.scopeId}::text)))`;
|
||||
|
||||
export async function nextBudgetReconciliationAt(db: Db): Promise<Date | null> {
|
||||
const [held] = await db.select({ id: budgetPolicies.id }).from(budgetPolicies)
|
||||
.where(and(eq(budgetPolicies.windowKind, "calendar_month_utc"), or(budgetOwnedScope, openIncidentScope))).limit(1);
|
||||
return held ? currentUtcMonthWindow().end : null;
|
||||
}
|
||||
|
||||
/** Startup repairs configured policy scopes, not every agent. Afterwards only
|
||||
* held monthly scopes owe a timed reevaluation; writes reconcile their own scopes. */
|
||||
export async function reconcileDueBudgetPolicies(db: Db, options: { startup?: boolean; cursor?: string | null } = {}): Promise<{ nextCursor: string | null }> {
|
||||
const actionable = options.startup ? or(budgetOwnedScope, openIncidentScope,
|
||||
sql`${budgetPolicies.enforcementVersion} > ${budgetPolicies.enforcementDeliveredVersion}`,
|
||||
and(eq(budgetPolicies.isActive, true), eq(budgetPolicies.hardStopEnabled, true), sql`${budgetPolicies.amount} > 0`, unfinishedRunScope))
|
||||
: and(eq(budgetPolicies.windowKind, "calendar_month_utc"), or(budgetOwnedScope, openIncidentScope));
|
||||
const rows = await db.select({ id: budgetPolicies.id, companyId: budgetPolicies.companyId, scopeType: budgetPolicies.scopeType, scopeId: budgetPolicies.scopeId })
|
||||
.from(budgetPolicies).where(and(actionable, options.cursor ? sql`${budgetPolicies.id} > ${options.cursor}::uuid` : undefined))
|
||||
.orderBy(budgetPolicies.id).limit(100);
|
||||
const scopes = new Map(rows.map(scope => [`${scope.companyId}:${scope.scopeType}:${scope.scopeId}`, scope]));
|
||||
for (const scope of scopes.values()) {
|
||||
try {
|
||||
await withAccountingTransaction(db, scope.companyId, (tx, publications) => budgetServiceInTransaction(tx, publications)
|
||||
.reconcileScope(scope.companyId, scope.scopeType as BudgetScopeType, scope.scopeId));
|
||||
} catch (error) {
|
||||
if (error instanceof HttpError && error.status === 404) continue;
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
return { nextCursor: rows.length === 100 ? rows[rows.length - 1].id : null };
|
||||
}
|
||||
|
||||
export function budgetService(db: Db, _hooks: BudgetServiceHooks = {}) {
|
||||
const reads = budgetServiceInTransaction(db);
|
||||
const mutate = <T>(companyId: string, work: (service: ReturnType<typeof budgetServiceInTransaction>) => Promise<T>) =>
|
||||
withAccountingTransaction(db, companyId, (tx, publications) => work(budgetServiceInTransaction(tx, publications)));
|
||||
return {
|
||||
...reads,
|
||||
deliverPendingEnforcement: (companyId?: string) => deliverBudgetEnforcement(db, hooks, companyId),
|
||||
reconcilePolicies: async () => {
|
||||
const scopes = await db.selectDistinct({ companyId: budgetPolicies.companyId, scopeType: budgetPolicies.scopeType, scopeId: budgetPolicies.scopeId })
|
||||
.from(budgetPolicies).orderBy(budgetPolicies.companyId, budgetPolicies.scopeType, budgetPolicies.scopeId);
|
||||
// Each scope commits independently: deleted targets and a failed scope
|
||||
// must not roll back recovery or cancellation delivery for other work.
|
||||
for (const scope of scopes) {
|
||||
try {
|
||||
await mutate(scope.companyId, (service) => service.reconcileScope(scope.companyId, scope.scopeType as BudgetScopeType, scope.scopeId));
|
||||
} catch (error) {
|
||||
if (error instanceof HttpError && error.status === 404) continue;
|
||||
logger.warn({ err: error, ...scope }, "Budget scope recovery pending; continuing remaining scopes");
|
||||
}
|
||||
}
|
||||
},
|
||||
evaluateAdmission: (companyId: string, agentId: string | null, context?: { issueId?: string | null; projectId?: string | null }) =>
|
||||
mutate(companyId, service => service.evaluateAdmission(companyId, agentId, context)),
|
||||
evaluateAffectedScopes: (companyId: string, agentId: string | null, projectId: string | null) =>
|
||||
mutate(companyId, service => service.evaluateAffectedScopes(companyId, agentId, projectId)),
|
||||
reconcileScope: (companyId: string, scopeType: BudgetScopeType, scopeId: string) =>
|
||||
mutate(companyId, service => service.reconcileScope(companyId, scopeType, scopeId)),
|
||||
reconcilePolicies: () => reconcileDueBudgetPolicies(db, { startup: true }),
|
||||
upsertPolicy: (companyId: string, input: BudgetPolicyUpsertInput, actorUserId: string | null) =>
|
||||
mutate(companyId, (service) => service.upsertPolicy(companyId, input, actorUserId)),
|
||||
evaluateCostEvent: (event: typeof costEvents.$inferSelect) =>
|
||||
mutate(event.companyId, (service) => service.evaluateCostEvent(event)),
|
||||
mutate(companyId, service => service.upsertPolicy(companyId, input, actorUserId)),
|
||||
evaluateCostEvent: (event: typeof costEvents.$inferSelect) => mutate(event.companyId, service => service.evaluateCostEvent(event)),
|
||||
getInvocationBlock: (companyId: string, agentId: string | null, context?: { issueId?: string | null; projectId?: string | null }) =>
|
||||
mutate(companyId, (service) => service.getInvocationBlock(companyId, agentId, context), agentId),
|
||||
mutate(companyId, service => service.getInvocationBlock(companyId, agentId, context)),
|
||||
resolveIncident: (companyId: string, incidentId: string, input: BudgetIncidentResolutionInput, actorUserId: string) =>
|
||||
mutate(companyId, (service) => service.resolveIncident(companyId, incidentId, input, actorUserId)),
|
||||
mutate(companyId, service => service.resolveIncident(companyId, incidentId, input, actorUserId)),
|
||||
};
|
||||
}
|
||||
@@ -1,9 +1,9 @@
|
||||
import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notifications.js";
|
||||
import { isRunnerApiResponseAsset } from "./asset-content-access.js";
|
||||
import { deleteCompany } from "./company-deletion.js";
|
||||
import { createAgentLifecycle } from "./agent-lifecycle.js";
|
||||
import { fenceActiveCompanyAgentStopsInTransaction } from "./agent-lifecycle.js";
|
||||
import { publishAccountingActivities } from "./accounting-transaction.js";
|
||||
import { budgetServiceInTransaction, deliverBudgetEnforcement, type BudgetServiceHooks } from "./budgets.js";
|
||||
import { budgetServiceInTransaction, type BudgetServiceHooks } from "./budgets.js";
|
||||
import { and, count, eq, gte, inArray, isNull, lt, notInArray, sql } from "drizzle-orm";
|
||||
import type { Db } from "@paperclipai/db";
|
||||
import {
|
||||
@@ -28,7 +28,6 @@ import {
|
||||
rekeyCompanyIssueIdentifiers,
|
||||
} from "./issue-prefix.js";
|
||||
import { environmentService } from "./environments.js";
|
||||
import { heartbeatService } from "./heartbeat.js";
|
||||
import { logActivity, type ActivityPublication } from "./activity-log.js";
|
||||
import { builtInAgentService } from "./built-in-agents.js";
|
||||
|
||||
@@ -48,24 +47,22 @@ const SYSTEM_COMPANY_ACTOR: CompanyActivityActor = {
|
||||
runId: null,
|
||||
};
|
||||
|
||||
export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
|
||||
export function companyService(db: Db, _budgetHooks: BudgetServiceHooks = {}) {
|
||||
const environmentsSvc = environmentService(db);
|
||||
const heartbeat = heartbeatService(db);
|
||||
const builtInAgents = builtInAgentService(db);
|
||||
|
||||
type CompanyTx = Parameters<Parameters<typeof db.transaction>[0]>[0];
|
||||
|
||||
async function applyArchiveCascadeInTx(tx: CompanyTx, id: string) {
|
||||
const pausedAgentRows = await tx.select({ id: agents.id }).from(agents).where(and(eq(agents.companyId, id), notInArray(agents.status, ["paused", "terminated", "pending_approval"])));
|
||||
const [agentCount] = await tx.select({ value: count() }).from(agents).where(and(eq(agents.companyId, id), notInArray(agents.status, ["paused", "terminated", "pending_approval"])));
|
||||
|
||||
const activeRunIds = await tx
|
||||
.select({ id: heartbeatRuns.id })
|
||||
const [runCount] = await tx
|
||||
.select({ value: count() })
|
||||
.from(heartbeatRuns)
|
||||
.where(and(
|
||||
eq(heartbeatRuns.companyId, id),
|
||||
inArray(heartbeatRuns.status, ["queued", "running"]),
|
||||
))
|
||||
.then((rows) => rows.map((row) => row.id));
|
||||
inArray(heartbeatRuns.status, ["queued", "running", "scheduled_retry"]),
|
||||
));
|
||||
|
||||
await tx
|
||||
.update(agentWakeupRequests)
|
||||
@@ -81,18 +78,16 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
|
||||
isNull(agentWakeupRequests.runId),
|
||||
));
|
||||
|
||||
return { agentsPaused: pausedAgentRows.length, activeRunIds };
|
||||
return { agentsPauseRequested: agentCount.value, runsCancellationRequested: runCount.value };
|
||||
}
|
||||
|
||||
async function finalizeArchive(
|
||||
id: string,
|
||||
actor: CompanyActivityActor,
|
||||
cascade: { agentsPaused: number; activeRunIds: string[] },
|
||||
cascade: { agentsPauseRequested: number; runsCancellationRequested: number },
|
||||
) {
|
||||
for (const runId of cascade.activeRunIds) {
|
||||
await heartbeat.cancelRun(runId, "Cancelled because the company was archived");
|
||||
}
|
||||
|
||||
// The committed company status blocks dispatch immediately. Lifecycle
|
||||
// projection owns provider termination, including paused agents with work.
|
||||
await logActivity(db, {
|
||||
companyId: id,
|
||||
actorType: actor.actorType,
|
||||
@@ -103,8 +98,8 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
|
||||
entityType: "company",
|
||||
entityId: id,
|
||||
details: {
|
||||
agentsPaused: cascade.agentsPaused,
|
||||
runsCancelled: cascade.activeRunIds.length,
|
||||
agentsPauseRequested: cascade.agentsPauseRequested,
|
||||
runsCancellationRequested: cascade.runsCancellationRequested,
|
||||
},
|
||||
});
|
||||
}
|
||||
@@ -304,14 +299,14 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
|
||||
) => {
|
||||
const budgetPublications: ActivityPublication[] = [];
|
||||
const result = await db.transaction(async (tx) => {
|
||||
if (data.status !== undefined || data.budgetMonthlyCents !== undefined) {
|
||||
await tx.select({ id: companies.id }).from(companies).where(eq(companies.id, id)).for("no key update");
|
||||
}
|
||||
const existing = await getCompanyQuery(tx)
|
||||
.where(eq(companies.id, id))
|
||||
.then((rows) => rows[0] ?? null);
|
||||
if (!existing) return null;
|
||||
|
||||
if (data.budgetMonthlyCents !== undefined) {
|
||||
await tx.select({ id: companies.id }).from(companies).where(eq(companies.id, id)).for("no key update");
|
||||
}
|
||||
const { logoAssetId, ...companyPatch } = data;
|
||||
const willReactivate = existing.status !== "active" && companyPatch.status === "active";
|
||||
const willArchive = existing.status !== "archived" && companyPatch.status === "archived";
|
||||
@@ -335,6 +330,7 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
|
||||
|
||||
if (companyPatch.status !== undefined && companyPatch.status !== existing.status) {
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.skillSourceSync);
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.agentLifecycle, id);
|
||||
}
|
||||
const updated = await tx
|
||||
.update(companies)
|
||||
@@ -375,6 +371,9 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
|
||||
}
|
||||
|
||||
const archiveCascade = willArchive ? await applyArchiveCascadeInTx(tx, id) : null;
|
||||
if (companyPatch.status !== existing.status && (companyPatch.status === "paused" || companyPatch.status === "archived")) {
|
||||
await fenceActiveCompanyAgentStopsInTransaction(tx as unknown as Db, id, companyPatch.status);
|
||||
}
|
||||
|
||||
if (logoAssetId === null) {
|
||||
await tx.delete(companyLogos).where(eq(companyLogos.companyId, id));
|
||||
@@ -422,9 +421,7 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
|
||||
};
|
||||
});
|
||||
if (!result) return null;
|
||||
await createAgentLifecycle(db).reconcilePolicyHolds(id);
|
||||
publishAccountingActivities(id, budgetPublications);
|
||||
if (data.budgetMonthlyCents !== undefined) await deliverBudgetEnforcement(db, budgetHooks, id);
|
||||
// Post-commit, fire-and-forget, and BEFORE any finalization that
|
||||
// could throw: a Cloud-pinned primary company that crossed the
|
||||
// archived boundary (either direction) rings the harness so the
|
||||
@@ -476,6 +473,7 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
|
||||
.select({ status: companies.status })
|
||||
.from(companies)
|
||||
.where(eq(companies.id, id))
|
||||
.for("no key update")
|
||||
.then((rows) => rows[0] ?? null);
|
||||
if (!existing) return null;
|
||||
|
||||
@@ -483,6 +481,7 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
|
||||
|
||||
if (!wasAlreadyArchived) {
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.skillSourceSync);
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.agentLifecycle, id);
|
||||
await tx
|
||||
.update(companies)
|
||||
.set({ status: "archived", updatedAt: new Date() })
|
||||
@@ -490,6 +489,7 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
|
||||
}
|
||||
|
||||
const cascade = wasAlreadyArchived ? null : await applyArchiveCascadeInTx(tx, id);
|
||||
if (!wasAlreadyArchived) await fenceActiveCompanyAgentStopsInTransaction(tx as unknown as Db, id, "archived");
|
||||
|
||||
const row = await getCompanyQuery(tx)
|
||||
.where(eq(companies.id, id))
|
||||
@@ -502,7 +502,6 @@ export function companyService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
|
||||
};
|
||||
});
|
||||
if (!result) return null;
|
||||
await createAgentLifecycle(db).reconcilePolicyHolds(id);
|
||||
|
||||
// Same doorbell rule as update(): the archive is committed, so ring
|
||||
// before finalization, which can throw without undoing it.
|
||||
|
||||
@@ -3,7 +3,7 @@ import { agentObservabilityService } from "./agent-observability.js";
|
||||
import { withAccountHomeSecretMutationLock } from "@paperclipai/adapter-codex-local/server";
|
||||
import type { Db } from "@paperclipai/db";
|
||||
import { eq, inArray } from "drizzle-orm";
|
||||
import { createAgentLifecycle, scheduleAgentLifecycle } from "./agent-lifecycle.js";
|
||||
import { createAgentLifecycle, processAgentLifecycle } from "./agent-lifecycle.js";
|
||||
import { agentLifecycleCompanyDeletion } from "../modules/agent-lifecycle/company-deletion.js";
|
||||
import {
|
||||
companies,
|
||||
@@ -50,7 +50,7 @@ export async function deleteCompany(db: Db, id: string) {
|
||||
const companyAgents = await db.select({ id: agents.id }).from(agents).where(eq(agents.companyId, id));
|
||||
for (const agent of companyAgents) {
|
||||
await lifecycle.terminateAgent(agent.id);
|
||||
await scheduleAgentLifecycle(db, agent.id);
|
||||
await processAgentLifecycle(db, agent.id);
|
||||
}
|
||||
return withAccountHomeSecretMutationLock(undefined, id, () => db.transaction(async (tx) => {
|
||||
// Exclude accounting writers before taking child locks. KEY SHARE must
|
||||
|
||||
@@ -5,7 +5,7 @@ import { alias } from "drizzle-orm/pg-core";
|
||||
import type { Db } from "@paperclipai/db";
|
||||
import { activityLog, agents, agentRuntimeState, authUsers, companies, companyMemberships, costEvents, heartbeatRuns, issues, projects, goals, aiSubscriptions } from "@paperclipai/db";
|
||||
import { notFound, unprocessable, conflict } from "../errors.js";
|
||||
import { budgetService, budgetServiceInTransaction, type BudgetServiceHooks } from "./budgets.js";
|
||||
import { budgetServiceInTransaction, type BudgetServiceHooks } from "./budgets.js";
|
||||
import { logActivity, type LogActivityInput, type ActivityPublication } from "./activity-log.js";
|
||||
import { withAccountingReadSnapshot, withAccountingTransaction } from "./accounting-transaction.js";
|
||||
import { visibleIssueCondition } from "./issue-visibility.js";
|
||||
@@ -82,7 +82,8 @@ export async function getMonthlySpendTotal(
|
||||
return Number(row?.total ?? 0);
|
||||
}
|
||||
|
||||
export async function createCostEventInTransaction(db: Db, companyId: string, data: Omit<typeof costEvents.$inferInsert, "companyId" | "receiptHash" | "costCents"> & { costCents: MoneyInput }, publications: ActivityPublication[] = [], actor?: Pick<LogActivityInput, "actorType" | "actorId" | "agentId">) {
|
||||
export async function createCostEventInTransaction(db: Db, companyId: string, data: Omit<typeof costEvents.$inferInsert, "companyId" | "receiptHash" | "costCents"> & { costCents: MoneyInput }, publications: ActivityPublication[] = [], actor?: Pick<LogActivityInput, "actorType" | "actorId" | "agentId">,
|
||||
options: { deferBudgetEvaluation?: boolean } = {}) {
|
||||
const parsed = ((data.usageKind === "decision" || data.usageKind === "fast_response" || data.usageKind === "prompt_injection") ? createServiceCostEventSchema : createCostEventSchema).safeParse({ ...data, occurredAt: data.occurredAt.toISOString() });
|
||||
if (!parsed.success) throw unprocessable("Invalid cost receipt", parsed.error.flatten());
|
||||
const values = {
|
||||
@@ -152,7 +153,9 @@ export async function createCostEventInTransaction(db: Db, companyId: string, da
|
||||
updatedAt: new Date(),
|
||||
}).where(and(eq(agentRuntimeState.agentId, event.agentId), eq(agentRuntimeState.companyId, companyId)));
|
||||
}
|
||||
await budgetServiceInTransaction(db, publications).evaluateCostEvent(event);
|
||||
// A multi-model run receipt inserts all parts before one scope evaluation.
|
||||
// The owning accounting transaction must perform that deferred evaluation.
|
||||
if (!options.deferBudgetEvaluation) await budgetServiceInTransaction(db, publications).evaluateCostEvent(event);
|
||||
await logActivity(db, {
|
||||
companyId, actorType: actor?.actorType ?? "system", actorId: actor?.actorId ?? "cost_accounting", agentId: actor?.agentId ?? event.agentId,
|
||||
...((event.usageKind === "decision" || event.usageKind === "fast_response" || event.usageKind === "prompt_injection") ? { responsibleUserIdOverride: event.responsibleUserId } : {}),
|
||||
@@ -162,13 +165,11 @@ export async function createCostEventInTransaction(db: Db, companyId: string, da
|
||||
return { ...event, costCentsExact: values.costCents };
|
||||
}
|
||||
|
||||
export function costService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
|
||||
const budgets = budgetService(db, budgetHooks);
|
||||
export function costService(db: Db, _budgetHooks: BudgetServiceHooks = {}) {
|
||||
return {
|
||||
createEvent: async (companyId: string, data: Omit<typeof costEvents.$inferInsert, "companyId" | "receiptHash" | "costCents"> & { costCents: MoneyInput }, actor?: Pick<LogActivityInput, "actorType" | "actorId" | "agentId">) => {
|
||||
if (data.idempotencyKey?.startsWith("heartbeat:")) throw unprocessable("The heartbeat idempotency namespace is reserved for run receipts");
|
||||
const event = await withAccountingTransaction(db, companyId, (tx, publications) => createCostEventInTransaction(tx, companyId, data, publications, actor));
|
||||
await budgets.deliverPendingEnforcement(companyId);
|
||||
return event;
|
||||
},
|
||||
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notifications.js";
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { and, desc, eq, gte, lt, lte, or, sql } from "drizzle-orm";
|
||||
import {
|
||||
agents, authUsers, budgetPolicies, budgetReservations, companies, companyMemberships, companyDecisionModels, costEvents,
|
||||
agents, authUsers, budgetReservations, companies, companyMemberships, companyDecisionModels, costEvents,
|
||||
decisionInvocations, heartbeatRuns, issues, projects, companySecrets, type Db,
|
||||
} from "@paperclipai/db";
|
||||
import {
|
||||
@@ -14,7 +15,8 @@ import { aiConnectionService } from "./ai-connections.js";
|
||||
import { accessService } from "./access.js";
|
||||
import { type AuthorizationActor } from "./authorization.js";
|
||||
import { withAccountingTransaction } from "./accounting-transaction.js";
|
||||
import { budgetPoliciesForRun, budgetService, budgetServiceInTransaction, computeObservedSpend, type BudgetServiceHooks } from "./budgets.js";
|
||||
import { withBudgetSpan } from "./budget-tracing.js";
|
||||
import { budgetServiceInTransaction, type BudgetServiceHooks } from "./budgets.js";
|
||||
import { createCostEventInTransaction } from "./costs.js";
|
||||
import { logActivity } from "./activity-log.js";
|
||||
import { decisionReceipt, runDecisionProvider, type DecisionProviderOutcome } from "./decision-model-provider.js";
|
||||
@@ -138,7 +140,7 @@ export function decisionModelService(db: Db, options: { provider?: typeof runDec
|
||||
return (context: DecisionContext) => { let result = cache.get(context); if (!result) { result = availability(context); cache.set(context, result); } return result; };
|
||||
}
|
||||
async function settle(invocation: typeof decisionInvocations.$inferSelect, outcome: DecisionProviderOutcome) {
|
||||
await withAccountingTransaction(db, invocation.companyId, async (tx, publications) => {
|
||||
await withBudgetSpan("budget.accounting.settlement", () => withAccountingTransaction(db, invocation.companyId, async (tx, publications) => {
|
||||
const [current] = await tx.select().from(decisionInvocations).where(and(eq(decisionInvocations.companyId, invocation.companyId), eq(decisionInvocations.id, invocation.id))).for("update");
|
||||
if (!current || current.status !== "running") return;
|
||||
const receipt = outcome.receipt;
|
||||
@@ -161,6 +163,7 @@ export function decisionModelService(db: Db, options: { provider?: typeof runDec
|
||||
providerRequestId: receipt.providerRequestId, pricingProvenance: receipt.pricingProvenance, occurredAt: invocation.startedAt,
|
||||
}, publications, { actorType: invocation.actorType as "system" | "user" | "agent", actorId: invocation.actorId, agentId: links.agentId });
|
||||
const finishedAt = new Date();
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, invocation.companyId);
|
||||
await tx.update(decisionInvocations).set({ status: outcome.answers ? "succeeded" : receipt.costStatus === "unpriced" ? "unknown" : "failed",
|
||||
errorCode: outcome.errorCode ?? null, costEventId: event.id, providerRequestId: receipt.providerRequestId,
|
||||
inputTokens: receipt.inputTokens, outputTokens: receipt.outputTokens, finishedAt, durationMs: Math.max(0, finishedAt.getTime() - invocation.startedAt.getTime()),
|
||||
@@ -168,8 +171,7 @@ export function decisionModelService(db: Db, options: { provider?: typeof runDec
|
||||
if (receipt.costStatus !== "unpriced") await tx.update(budgetReservations).set({ state: outcome.noProviderWork ? "released" : "settled", settledAt: finishedAt }).where(and(
|
||||
eq(budgetReservations.companyId, invocation.companyId), eq(budgetReservations.decisionInvocationId, invocation.id), eq(budgetReservations.state, "held"),
|
||||
));
|
||||
});
|
||||
await budgetService(db, options.budgetHooks).deliverPendingEnforcement(invocation.companyId);
|
||||
}));
|
||||
}
|
||||
async function decide(context: DecisionContext, request: DecisionRequest, signal?: AbortSignal): Promise<DecisionResult> {
|
||||
const parsed = decisionRequestSchema.safeParse(request);
|
||||
@@ -202,23 +204,17 @@ export function decisionModelService(db: Db, options: { provider?: typeof runDec
|
||||
if (!currentSecret || currentSecret.status !== "active" || currentSecret.version !== credentialSnapshot.version)
|
||||
return { reason: "connection_unavailable" as const };
|
||||
const budgets = budgetServiceInTransaction(tx, publications);
|
||||
const block = await budgets.getInvocationBlock(context.companyId, current.agentId, { projectId: current.projectId });
|
||||
const { block, policies, observations } = await budgets.evaluateAdmission(context.companyId, current.agentId, { projectId: current.projectId });
|
||||
const [company] = await tx.select({ status: companies.status }).from(companies).where(eq(companies.id, context.companyId));
|
||||
const [agent] = current.agentId ? await tx.select({ status: agents.status }).from(agents).where(and(eq(agents.id, current.agentId), eq(agents.companyId, context.companyId))) : [];
|
||||
if (block || company.status !== "active" || (current.agentId && (!agent || agent.status === "paused" || agent.status === "terminated"))) return { reason: "budget_blocked" as const };
|
||||
const policies = await tx.select().from(budgetPolicies).where(and(eq(budgetPolicies.companyId, context.companyId), eq(budgetPolicies.isActive, true),
|
||||
eq(budgetPolicies.hardStopEnabled, true), eq(budgetPolicies.metric, "billed_cents"), sql`${budgetPolicies.amount} > 0`, budgetPoliciesForRun(context.companyId, current.agentId, current.projectId)));
|
||||
const amount = policies.reduce((max, p) => { const n = centsToUnits(p.reservationCents); return n > max ? n : max; }, 0n);
|
||||
for (const policy of policies) {
|
||||
const [held] = await tx.select({ total: sql<string>`coalesce(sum(${budgetReservations.amountCents}),0)::text` }).from(budgetReservations).where(and(
|
||||
eq(budgetReservations.companyId, context.companyId), eq(budgetReservations.state, "held"),
|
||||
policy.scopeType === "agent" ? eq(budgetReservations.agentId, current.agentId!) : undefined,
|
||||
policy.scopeType === "project" ? eq(budgetReservations.projectId, current.projectId!) : undefined,
|
||||
));
|
||||
const committed = centsToUnits((await computeObservedSpend(tx, policy)).totalExact) + centsToUnits(held.total);
|
||||
for (const [index, policy] of policies.entries()) {
|
||||
const committed = centsToUnits(observations[index].ledgerCents) + centsToUnits(observations[index].effectiveHeldCents);
|
||||
if (committed >= centsToUnits(policy.amount) || committed + amount > centsToUnits(policy.amount)) return { reason: "budget_blocked" as const };
|
||||
}
|
||||
signal?.throwIfAborted();
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, context.companyId);
|
||||
const [invocation] = await tx.insert(decisionInvocations).values({ id: randomUUID(), companyId: context.companyId, feature: context.feature.id,
|
||||
actorType: current.actorType, actorId: current.actorId, responsibleUserId: current.userId, agentId: current.agentId,
|
||||
issueId: current.issueId, projectId: current.projectId, runId: current.runId, identityContextId: current.identityContextId,
|
||||
|
||||
@@ -29,6 +29,9 @@ export function createDeliveryWorkCoordinator(input: {
|
||||
// An empty queue's maintenance may wait until dispatch admission reopens.
|
||||
// Notifications and unresolved producer writes still reconcile in drain.
|
||||
pausePassiveMaintenance?: boolean;
|
||||
// Accounting and stop delivery finish already-admitted obligations;
|
||||
// they must continue during admission drain, never warm standby.
|
||||
runDuringDrain?: boolean;
|
||||
}) {
|
||||
if (stopped) throw new Error("Delivery coordinator is stopped");
|
||||
if (workers.has(queue)) throw new Error(`Delivery worker already registered: ${queue}`);
|
||||
@@ -64,7 +67,7 @@ export function createDeliveryWorkCoordinator(input: {
|
||||
const attempt = new AbortController();
|
||||
controller = attempt;
|
||||
running = Promise.resolve().then(() => reconcileDatabaseWork(input.owner, queue))
|
||||
.then(() => workerStopped || !input.canRun() ? undefined : task.run(attempt.signal)).then(() => workerStopped ? false : task.hasPending()).then(pending => {
|
||||
.then(() => workerStopped || !(input.canRun() || (task.runDuringDrain && input.canReconcile?.())) ? undefined : task.run(attempt.signal)).then(() => workerStopped ? false : task.hasPending()).then(pending => {
|
||||
pendingWork = pending || databaseWorkPending(input.owner, queue);
|
||||
if (pendingWork) schedule(task.retryMs);
|
||||
else { finishIntent?.(); finishIntent = undefined; }
|
||||
|
||||
@@ -2,6 +2,9 @@ import { signalDatabaseWork, subscribeDatabaseWork } from "@paperclipai/db";
|
||||
|
||||
/** Topics for existing durable queues; register intent before their writes. */
|
||||
export const DELIVERY_QUEUES = {
|
||||
budgetEnforcement: "budget-enforcement",
|
||||
accounting: "cost-accounting",
|
||||
agentLifecycle: "agent-lifecycle",
|
||||
chatBoardMessages: "chat-board-messages",
|
||||
chatTaskStarts: "chat-task-starts",
|
||||
chatReceiptReactions: "chat-receipt-reactions",
|
||||
|
||||
@@ -20,7 +20,6 @@ import {
|
||||
issueAttachments,
|
||||
projects,
|
||||
authUsers,
|
||||
budgetPolicies,
|
||||
budgetReservations,
|
||||
companies,
|
||||
companyFastResponses,
|
||||
@@ -60,11 +59,9 @@ import {
|
||||
type ActivityPublication,
|
||||
} from "./activity-log.js";
|
||||
import { withAccountingTransaction } from "./accounting-transaction.js";
|
||||
import { withBudgetSpan } from "./budget-tracing.js";
|
||||
import {
|
||||
budgetPoliciesForRun,
|
||||
budgetService,
|
||||
budgetServiceInTransaction,
|
||||
computeObservedSpend,
|
||||
type BudgetServiceHooks,
|
||||
} from "./budgets.js";
|
||||
import { createCostEventInTransaction } from "./costs.js";
|
||||
@@ -567,7 +564,7 @@ export function fastResponseService(
|
||||
);
|
||||
}
|
||||
async function settle(request: Request, outcome: FastResponseOutcome) {
|
||||
await withAccountingTransaction(
|
||||
await withBudgetSpan("budget.accounting.settlement", () => withAccountingTransaction(
|
||||
db,
|
||||
request.companyId,
|
||||
async (tx, publications) => {
|
||||
@@ -660,10 +657,8 @@ export function fastResponseService(
|
||||
),
|
||||
);
|
||||
},
|
||||
);
|
||||
await budgetService(db, options.budgetHooks).deliverPendingEnforcement(
|
||||
request.companyId,
|
||||
);
|
||||
));
|
||||
|
||||
}
|
||||
async function process(
|
||||
request: Request,
|
||||
@@ -853,60 +848,15 @@ export function fastResponseService(
|
||||
(!agent || ["paused", "terminated"].includes(agent.status)))
|
||||
)
|
||||
return null;
|
||||
if (
|
||||
await budgetServiceInTransaction(tx, []).getInvocationBlock(
|
||||
request.companyId,
|
||||
request.agentId,
|
||||
{ projectId: current.issue?.projectId },
|
||||
)
|
||||
)
|
||||
return null;
|
||||
const policies = await tx
|
||||
.select()
|
||||
.from(budgetPolicies)
|
||||
.where(
|
||||
and(
|
||||
eq(budgetPolicies.companyId, request.companyId),
|
||||
eq(budgetPolicies.isActive, true),
|
||||
eq(budgetPolicies.hardStopEnabled, true),
|
||||
eq(budgetPolicies.metric, "billed_cents"),
|
||||
sql`${budgetPolicies.amount} > 0`,
|
||||
budgetPoliciesForRun(
|
||||
request.companyId,
|
||||
request.agentId,
|
||||
current.issue?.projectId ?? null,
|
||||
),
|
||||
),
|
||||
);
|
||||
const { block, policies, observations } = await budgetServiceInTransaction(tx, []).evaluateAdmission(
|
||||
request.companyId, request.agentId, { projectId: current.issue?.projectId });
|
||||
if (block) return null;
|
||||
const amount = policies.reduce((max, p) => {
|
||||
const n = centsToUnits(p.reservationCents);
|
||||
return n > max ? n : max;
|
||||
}, 0n);
|
||||
for (const policy of policies) {
|
||||
const [held] = await tx
|
||||
.select({
|
||||
total: sql<string>`coalesce(sum(${budgetReservations.amountCents}),0)::text`,
|
||||
})
|
||||
.from(budgetReservations)
|
||||
.where(
|
||||
and(
|
||||
eq(budgetReservations.companyId, request.companyId),
|
||||
eq(budgetReservations.state, "held"),
|
||||
policy.scopeType === "agent"
|
||||
? eq(budgetReservations.agentId, request.agentId!)
|
||||
: undefined,
|
||||
policy.scopeType === "project"
|
||||
? eq(
|
||||
budgetReservations.projectId,
|
||||
current.issue!.projectId!,
|
||||
)
|
||||
: undefined,
|
||||
),
|
||||
);
|
||||
const committed =
|
||||
centsToUnits(
|
||||
(await computeObservedSpend(tx, policy)).totalExact,
|
||||
) + centsToUnits(held.total);
|
||||
for (const [index, policy] of policies.entries()) {
|
||||
const committed = centsToUnits(observations[index].ledgerCents) + centsToUnits(observations[index].effectiveHeldCents);
|
||||
if (
|
||||
committed >= centsToUnits(policy.amount) ||
|
||||
committed + amount > centsToUnits(policy.amount)
|
||||
|
||||
@@ -330,7 +330,7 @@ import {
|
||||
projects,
|
||||
projectWorkspaces,
|
||||
} from "@paperclipai/db";
|
||||
import { HttpError, notFound } from "../errors.js";
|
||||
import { conflict, HttpError, notFound } from "../errors.js";
|
||||
|
||||
import { logger } from "../middleware/logger.js";
|
||||
import { publishLiveEvent } from "./live-events.js";
|
||||
@@ -929,6 +929,14 @@ export function heartbeatService(
|
||||
const budgetHooks = {
|
||||
cancelWorkForScope: cancelBudgetScopeWork,
|
||||
};
|
||||
const reconcileCostAccounting = createCostAccountingReconciler(db, budgetHooks, async () => {
|
||||
const outcomes = await Promise.allSettled([
|
||||
decisionModelService(db, { budgetHooks }).recoverInterrupted(),
|
||||
promptInjectionPreventionService(db, { budgetHooks }).recoverInterrupted(),
|
||||
]);
|
||||
const failed = outcomes.find(outcome => outcome.status === "rejected");
|
||||
if (failed?.status === "rejected") throw failed.reason;
|
||||
});
|
||||
const budgets = budgetService(db, budgetHooks);
|
||||
const recovery = recoveryService(db, {
|
||||
enqueueWakeup,
|
||||
@@ -2740,12 +2748,13 @@ export function heartbeatService(
|
||||
}
|
||||
const dispatchResolvedInteractionContinuationAfterAdmission = async <T>(
|
||||
dispatch: (markDispatchStarted: () => void) => Promise<T>,
|
||||
reservation: { id: string; reused: boolean },
|
||||
): Promise<
|
||||
{ dispatched: true; resultPromise: Promise<T> } | { dispatched: false }
|
||||
> => {
|
||||
await controllerLease.assertOwned("dispatching");
|
||||
// Recheck after workspace/credential preparation, immediately before the
|
||||
// provider handoff. Never hold validation locks while adapter code runs.
|
||||
// Recheck after preparation. The final locked gate commits a dispatch
|
||||
// claim before entering adapter bootstrap or contacting a provider.
|
||||
await authorizeFailedChatRetryExecution();
|
||||
if (
|
||||
!(await withChatControlRecoveryGate(
|
||||
@@ -2771,33 +2780,27 @@ export function heartbeatService(
|
||||
}
|
||||
return { dispatched: false };
|
||||
}
|
||||
if (
|
||||
!issueId ||
|
||||
(!isResolvedInteractionContinuationWakeContext(context) &&
|
||||
typeof context.workspaceReconstructionActionId !== "string" &&
|
||||
!["native_safe_replacement", "native_provider_overloaded"].includes(run.scheduledRetryReason ?? ""))
|
||||
) {
|
||||
return { dispatched: true, resultPromise: dispatch(() => {}) };
|
||||
const enforceContinuationOwnership = Boolean(issueId && (isResolvedInteractionContinuationWakeContext(context) ||
|
||||
typeof context.workspaceReconstructionActionId === "string" ||
|
||||
["native_safe_replacement", "native_provider_overloaded"].includes(run.scheduledRetryReason ?? "")));
|
||||
if (issueId && enforceContinuationOwnership) {
|
||||
await options.beforeResolvedInteractionContinuationDispatchCheck?.({ runId: run.id, issueId });
|
||||
await options.afterResolvedInteractionContinuationDispatchCheck?.({ runId: run.id, issueId });
|
||||
}
|
||||
await options.beforeResolvedInteractionContinuationDispatchCheck?.({
|
||||
runId: run.id,
|
||||
issueId,
|
||||
});
|
||||
|
||||
await options.afterResolvedInteractionContinuationDispatchCheck?.({
|
||||
runId: run.id,
|
||||
issueId,
|
||||
});
|
||||
const gate = await runDispatch.dispatchResolvedInteractionIfCurrent({
|
||||
runId: run.id,
|
||||
companyId: run.companyId,
|
||||
expectedStatus: "running",
|
||||
// Synchronous handoff under the ownership lock; the gate commits
|
||||
// without awaiting the adapter's asynchronous bootstrap or finalizer.
|
||||
reservationId: reservation.id,
|
||||
enforceContinuationOwnership,
|
||||
...(runOptions.nativeLeaseOwner ? { recoveryLeaseOwner: runOptions.nativeLeaseOwner } : {}),
|
||||
// The gate persists admission under the ownership lock, then invokes
|
||||
// this callback after commit so provider I/O holds no accounting lock.
|
||||
dispatch,
|
||||
});
|
||||
|
||||
if (gate.dispatched) return gate;
|
||||
if ("admissionBlockReason" in gate) throw conflict(gate.admissionBlockReason);
|
||||
if (gate.cancellation.outcome === "cancelled") {
|
||||
applyRunDispatchPostCommitEffects(
|
||||
gate.cancellation.postCommitEffects,
|
||||
@@ -2818,7 +2821,7 @@ export function heartbeatService(
|
||||
return await dispatchResolvedInteractionContinuationAfterAdmission((markDispatchStarted) => {
|
||||
entered = true;
|
||||
return dispatch(markDispatchStarted);
|
||||
});
|
||||
}, reservation);
|
||||
} finally {
|
||||
if (!entered && !reservation.reused) {
|
||||
await db.update(heartbeatRuns).set({ costAccountingPending: true,
|
||||
@@ -4092,11 +4095,7 @@ export function heartbeatService(
|
||||
scanSilentActiveRuns,
|
||||
|
||||
reconcileTaskWatchdogs,
|
||||
reconcileCostAccounting: async () => {
|
||||
await decisionModelService(db, { budgetHooks }).recoverInterrupted();
|
||||
await promptInjectionPreventionService(db, { budgetHooks }).recoverInterrupted();
|
||||
return createCostAccountingReconciler(db, budgetHooks)();
|
||||
},
|
||||
reconcileCostAccounting,
|
||||
|
||||
buildRunOutputSilence,
|
||||
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
import type { Db } from "@paperclipai/db";
|
||||
import { agents, companies, createDb, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS, heartbeatRuns } from "@paperclipai/db";
|
||||
import { sql } from "drizzle-orm";
|
||||
import { afterAll, afterEach, beforeAll, beforeEach, vi, type Mock } from "vitest";
|
||||
import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../../../__tests__/helpers/embedded-postgres.js";
|
||||
import { runningProcesses } from "../../../adapters/index.js";
|
||||
import { terminateLocalService } from "../../local-service-supervisor.js";
|
||||
import { createHeartbeatRunState } from "../run-state.js";
|
||||
import { createHeartbeatRunPreparation } from "../run-preparation.js";
|
||||
import { createHeartbeatLifecycle, type HeartbeatLifecycleDependencies } from "../run-lifecycle.js";
|
||||
import type { HeartbeatRunControlDependencies } from "../run-control.js";
|
||||
|
||||
type TestCallbacks = HeartbeatRunControlDependencies & ReturnType<typeof createHeartbeatLifecycle> & {
|
||||
getRun: Mock<HeartbeatRunControlDependencies["getRun"]>;
|
||||
getAgent: Mock<HeartbeatRunControlDependencies["getAgent"]>;
|
||||
setRunStatusFromLive: Mock<HeartbeatRunControlDependencies["setRunStatusFromLive"]>;
|
||||
getSchedulingSuppression: Mock<HeartbeatRunControlDependencies["getSchedulingSuppression"]>;
|
||||
options: { closeWarmNativeSessionsForRun: Mock<NonNullable<HeartbeatRunControlDependencies["options"]["closeWarmNativeSessionsForRun"]>> };
|
||||
};
|
||||
|
||||
/** Shared real lifecycle callbacks; each focused suite controls its provider mocks. */
|
||||
export function callbacks(db: Db): TestCallbacks {
|
||||
const state = createHeartbeatRunState(db);
|
||||
const lifecycleDependencies = {
|
||||
...state,
|
||||
...createHeartbeatRunPreparation(db),
|
||||
issuesSvc: {
|
||||
addComment: vi.fn<HeartbeatLifecycleDependencies["issuesSvc"]["addComment"]>(),
|
||||
listReviewAttention: vi.fn<HeartbeatLifecycleDependencies["issuesSvc"]["listReviewAttention"]>(),
|
||||
},
|
||||
recovery: { escalateStrandedAssignedIssue: vi.fn<HeartbeatLifecycleDependencies["recovery"]["escalateStrandedAssignedIssue"]>() },
|
||||
companySkills: { completeTestRunForIssue: vi.fn<HeartbeatLifecycleDependencies["companySkills"]["completeTestRunForIssue"]>() },
|
||||
budgets: { getInvocationBlock: vi.fn(async () => null) },
|
||||
treeControlSvc: { getActivePauseHoldGate: vi.fn(async () => null) },
|
||||
enqueueWakeup: vi.fn<HeartbeatLifecycleDependencies["enqueueWakeup"]>(),
|
||||
getCurrentUserRedactionOptions: vi.fn(async () => ({})),
|
||||
getAgentInvokability: vi.fn<HeartbeatLifecycleDependencies["getAgentInvokability"]>(async () => ({ invokable: true })),
|
||||
emitTerminalAgentTaskRun: vi.fn<HeartbeatLifecycleDependencies["emitTerminalAgentTaskRun"]>(),
|
||||
budgetHooks: {},
|
||||
} satisfies HeartbeatLifecycleDependencies;
|
||||
const lifecycle = createHeartbeatLifecycle(db, lifecycleDependencies);
|
||||
return {
|
||||
...state,
|
||||
...lifecycle,
|
||||
getRun: vi.fn(state.getRun),
|
||||
getAgent: vi.fn(state.getAgent),
|
||||
setRunStatusFromLive: vi.fn(lifecycle.setRunStatusFromLive),
|
||||
setRunStatus: vi.fn(lifecycle.setRunStatus),
|
||||
setWakeupStatus: vi.fn(lifecycle.setWakeupStatus),
|
||||
appendRunEvent: vi.fn(lifecycle.appendRunEvent),
|
||||
finalizeAgentStatus: vi.fn<HeartbeatRunControlDependencies["finalizeAgentStatus"]>(),
|
||||
options: { closeWarmNativeSessionsForRun: vi.fn(async () => ({ closed: 1, busy: 0, failed: 0 })) },
|
||||
envOrchestrator: { releaseForRun: vi.fn<HeartbeatRunControlDependencies["envOrchestrator"]["releaseForRun"]>(async () => ({ released: [], errors: [] })) },
|
||||
enqueueWakeup: lifecycleDependencies.enqueueWakeup,
|
||||
startNextQueuedRunForAgent: vi.fn<HeartbeatRunControlDependencies["startNextQueuedRunForAgent"]>(),
|
||||
sweepPendingCleanupLeases: vi.fn<HeartbeatRunControlDependencies["sweepPendingCleanupLeases"]>(),
|
||||
timerClaimWasFirstHeartbeat: vi.fn<HeartbeatRunControlDependencies["timerClaimWasFirstHeartbeat"]>(() => undefined),
|
||||
wakeQueue: { releaseIssueExecution: vi.fn<HeartbeatRunControlDependencies["wakeQueue"]["releaseIssueExecution"]>(async () => ({ outcome: { kind: "released" }, postCommitEffects: [] })) },
|
||||
applyWakeQueuePostCommitEffects: vi.fn(async () => undefined),
|
||||
getSchedulingSuppression: vi.fn(async () => ({ suppressed: false })),
|
||||
activeRunExecutions: new Set<string>(),
|
||||
processRunCancellationSettlements: new Map<string, { settled: Promise<void>; failed: boolean; error?: unknown }>(),
|
||||
failedProcessRunCancellations: new Map<string, { settled: Promise<void>; failed: boolean; error?: unknown }>(),
|
||||
} satisfies HeartbeatRunControlDependencies;
|
||||
}
|
||||
export const runControlPostgresSupport = await getEmbeddedPostgresTestSupport();
|
||||
if (!runControlPostgresSupport.supported) console.warn(`Skipping run-control database tests: ${runControlPostgresSupport.reason}`);
|
||||
|
||||
export function setupRunControlDatabase(name: string, ready: (db: Db) => void) {
|
||||
let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>>;
|
||||
let db: Db;
|
||||
beforeAll(async () => {
|
||||
database = await startEmbeddedPostgresTestDatabase(name);
|
||||
db = createDb(database.connectionString);
|
||||
ready(db);
|
||||
}, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS);
|
||||
afterAll(async () => {
|
||||
await db?.$client.end();
|
||||
await database?.cleanup();
|
||||
});
|
||||
beforeEach(async () => {
|
||||
vi.mocked(terminateLocalService).mockReset().mockResolvedValue(undefined);
|
||||
await db.execute(sql`truncate table companies restart identity cascade`);
|
||||
});
|
||||
afterEach(() => runningProcesses.clear());
|
||||
|
||||
async function fixture() {
|
||||
const [company] = await db.insert(companies).values({ name: "Run control", issuePrefix: "RC" }).returning();
|
||||
const [agent] = await db.insert(agents).values({ companyId: company.id, name: "Controlled agent", adapterType: "codex_local" }).returning();
|
||||
return { company, agent };
|
||||
}
|
||||
async function liveRun(companyId: string, agentId: string) {
|
||||
const [run] = await db.insert(heartbeatRuns).values({ companyId, agentId, status: "running", runtimeMode: "legacy", runtimeModeResolvedAt: new Date() }).returning();
|
||||
runningProcesses.set(run.id, { child: { pid: 424242 } as never, graceSec: 4, processGroupId: 424242 });
|
||||
return run;
|
||||
}
|
||||
return { fixture, liveRun };
|
||||
}
|
||||
@@ -0,0 +1,469 @@
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { spawn } from "node:child_process";
|
||||
import { once } from "node:events";
|
||||
import type { Db } from "@paperclipai/db";
|
||||
import { agents, agentWakeupRequests, browserUseRuns, browserUseSessions, budgetPolicies, companies, connectionGrants, costEvents, environmentLeases, heartbeatRuns, issues, projects, toolApplications, toolConnections } from "@paperclipai/db";
|
||||
import { eq, inArray } from "drizzle-orm";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { runningProcesses } from "../../adapters/index.js";
|
||||
import { terminateLocalService } from "../local-service-supervisor.js";
|
||||
import { persistUsageReceipt } from "../usage-receipts.js";
|
||||
import { registerBudgetEnforcementWork } from "../budget-enforcement-worker.js";
|
||||
import { createDeliveryWorkCoordinator } from "../delivery-work-coordinator.js";
|
||||
import { DELIVERY_QUEUES, subscribeDeliveryWork } from "../delivery-work-notifications.js";
|
||||
import { PROCESS_START_REQUESTED } from "../native-local-process-stop.js";
|
||||
import { createPostgresRunDispatchAdapter } from "../../modules/run-dispatch/adapters/postgres.js";
|
||||
import { createHeartbeatRunControl } from "./run-control.js";
|
||||
import { listPendingBudgetCancellationScopes } from "./budget-cancellation.js";
|
||||
import { callbacks, runControlPostgresSupport, setupRunControlDatabase } from "./__tests__/run-control-fixtures.js";
|
||||
|
||||
vi.mock("../live-events.js", () => ({ publishLiveEvent: vi.fn() }));
|
||||
vi.mock("../local-service-supervisor.js", async (importOriginal) => ({
|
||||
...await importOriginal<typeof import("../local-service-supervisor.js")>(),
|
||||
terminateLocalService: vi.fn(async () => undefined),
|
||||
}));
|
||||
|
||||
describe.skipIf(!runControlPostgresSupport.supported)("scoped budget cancellation", () => {
|
||||
let db: Db;
|
||||
const { fixture, liveRun } = setupRunControlDatabase("paperclip-heartbeat-budget-cancellation", value => { db = value; });
|
||||
|
||||
it("rechecks stale budget enforcement before cancelling runs or queued wakes", async () => {
|
||||
const { company, agent } = await fixture();
|
||||
const [run] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id, status: "queued", runtimeMode: "legacy" }).returning();
|
||||
const [wake] = await db.insert(agentWakeupRequests).values({ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued" }).returning();
|
||||
const deps = callbacks(db);
|
||||
await createHeartbeatRunControl(db, deps).cancelBudgetScopeWork({ companyId: company.id, scopeType: "agent", scopeId: agent.id, enforcement: { policyId: randomUUID(), version: 1 } });
|
||||
expect((await deps.getRun(run.id))?.status).toBe("queued");
|
||||
expect((await db.select().from(agentWakeupRequests).where(eq(agentWakeupRequests.id, wake.id)))[0].status).toBe("queued");
|
||||
expect(deps.setRunStatusFromLive).not.toHaveBeenCalled();
|
||||
expect(deps.finalizeAgentStatus).not.toHaveBeenCalled();
|
||||
expect(terminateLocalService).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("does not apply an earlier enforcement generation after a budget increase", async () => {
|
||||
const { company, agent } = await fixture();
|
||||
const [policy] = await db.insert(budgetPolicies).values({ companyId: company.id,
|
||||
scopeType: "agent", scopeId: agent.id, windowKind: "lifetime", amount: 10_000,
|
||||
enforcementVersion: 2, enforcementDeliveredVersion: 1 }).returning();
|
||||
const run = await liveRun(company.id, agent.id);
|
||||
const [wake] = await db.insert(agentWakeupRequests).values({ companyId: company.id, agentId: agent.id,
|
||||
source: "assignment", status: "queued" }).returning();
|
||||
await createHeartbeatRunControl(db, callbacks(db)).cancelBudgetScopeWork({ companyId: company.id,
|
||||
scopeType: "agent", scopeId: agent.id, enforcement: { policyId: policy.id, version: 1 } });
|
||||
expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, run.id)))[0]).toMatchObject({ status: "running", resultJson: null });
|
||||
expect((await db.select().from(agentWakeupRequests).where(eq(agentWakeupRequests.id, wake.id)))[0].status).toBe("queued");
|
||||
expect(terminateLocalService).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("fences every affected run and queued wake before bounded provider delivery", async () => {
|
||||
const { company, agent } = await fixture();
|
||||
const runs = await db.insert(heartbeatRuns).values(Array.from({ length: 5 }, () => ({
|
||||
companyId: company.id, agentId: agent.id, status: "running", runtimeMode: "legacy", runtimeModeResolvedAt: new Date(),
|
||||
}))).returning();
|
||||
const [unrelatedAgent] = await db.insert(agents).values({ companyId: company.id, name: "Unrelated" }).returning();
|
||||
const unrelated = await db.insert(heartbeatRuns).values({ companyId: company.id,
|
||||
agentId: unrelatedAgent.id, status: "running", runtimeMode: "legacy" }).returning().then(rows => rows[0]);
|
||||
const wakes = await db.insert(agentWakeupRequests).values([
|
||||
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued" },
|
||||
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "deferred_issue_execution" },
|
||||
{ companyId: company.id, agentId: unrelatedAgent.id, source: "assignment", status: "queued" },
|
||||
]).returning();
|
||||
runs.forEach((run, i) => runningProcesses.set(run.id, {
|
||||
child: { pid: 425000 + i } as never, graceSec: 4, processGroupId: 425000 + i,
|
||||
}));
|
||||
const releases = new Map<number, () => void>();
|
||||
let active = 0, maximumActive = 0, releaseAll = false;
|
||||
vi.mocked(terminateLocalService).mockImplementation(async input => {
|
||||
active += 1;
|
||||
maximumActive = Math.max(maximumActive, active);
|
||||
if (!releaseAll) await new Promise<void>(resolve => releases.set(input.pid, resolve));
|
||||
active -= 1;
|
||||
});
|
||||
const deps = callbacks(db);
|
||||
const delivery = createHeartbeatRunControl(db, deps).cancelBudgetScopeWork({ companyId: company.id,
|
||||
scopeType: "agent", scopeId: agent.id });
|
||||
try {
|
||||
await vi.waitFor(() => expect(releases.size).toBe(4));
|
||||
const persisted = await db.select().from(heartbeatRuns).where(inArray(heartbeatRuns.id, runs.map(run => run.id)));
|
||||
expect(persisted.every(run => run.status === "running" && Boolean(run.resultJson?.startupCancellation))).toBe(true);
|
||||
const pendingWakes = await db.select().from(agentWakeupRequests).where(inArray(agentWakeupRequests.id, wakes.map(wake => wake.id)));
|
||||
expect(wakes.map(wake => pendingWakes.find(row => row.id === wake.id)?.status)).toEqual(["cancelled", "cancelled", "queued"]);
|
||||
expect((await deps.getRun(unrelated.id))?.resultJson).toBeNull();
|
||||
// A slow first Stop must not stop another slot from making progress.
|
||||
const firstPid = releases.keys().next().value!;
|
||||
const otherPid = [...releases.keys()].find(pid => pid !== firstPid)!;
|
||||
releases.get(otherPid)!();
|
||||
await vi.waitFor(() => expect(releases.size).toBe(5));
|
||||
expect(maximumActive).toBe(4);
|
||||
expect(active).toBe(4);
|
||||
} finally {
|
||||
releaseAll = true;
|
||||
releases.forEach(release => release());
|
||||
await delivery;
|
||||
}
|
||||
expect((await db.select().from(heartbeatRuns).where(inArray(heartbeatRuns.id, runs.map(run => run.id))))
|
||||
.every(run => run.status === "cancelled" && (run.resultJson?.executionCancellation as { state?: string })?.state === "acknowledged")).toBe(true);
|
||||
expect((await deps.getRun(unrelated.id))?.status).toBe("running");
|
||||
});
|
||||
|
||||
it("wakes scoped Stop delivery from a partial usage overrun without waiting for adapter termination", async () => {
|
||||
const { company, agent } = await fixture();
|
||||
const [project, otherProject] = await db.insert(projects).values([
|
||||
{ companyId: company.id, name: "Over budget" }, { companyId: company.id, name: "Still admitted" },
|
||||
]).returning();
|
||||
const sourceId = randomUUID();
|
||||
const run = await liveRun(company.id, agent.id);
|
||||
await db.update(heartbeatRuns).set({ contextSnapshot: { projectId: project.id }, costAccountingPending: true,
|
||||
usageJson: { accountingReceiptSourceId: sourceId, accountingReceiptSequence: 0,
|
||||
accountingReceiptReady: false, ledgerScope: { projectId: project.id } },
|
||||
}).where(eq(heartbeatRuns.id, run.id));
|
||||
const [queued, unrelated] = await db.insert(heartbeatRuns).values([
|
||||
{ companyId: company.id, agentId: agent.id, status: "queued", runtimeMode: "legacy", contextSnapshot: { projectId: project.id } },
|
||||
{ companyId: company.id, agentId: agent.id, status: "running", runtimeMode: "legacy", contextSnapshot: { projectId: otherProject.id } },
|
||||
]).returning();
|
||||
const wakes = await db.insert(agentWakeupRequests).values([
|
||||
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued", payload: { projectId: project.id } },
|
||||
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "deferred_issue_execution", payload: { projectId: project.id } },
|
||||
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued", payload: { projectId: otherProject.id } },
|
||||
]).returning();
|
||||
const [policy] = await db.insert(budgetPolicies).values({ companyId: company.id, scopeType: "project",
|
||||
scopeId: project.id, windowKind: "lifetime", amount: 100, notifyEnabled: false }).returning();
|
||||
let releaseTermination!: () => void;
|
||||
const termination = new Promise<void>(resolve => { releaseTermination = resolve; });
|
||||
vi.mocked(terminateLocalService).mockImplementation(async () => termination);
|
||||
const deps = callbacks(db);
|
||||
const control = createHeartbeatRunControl(db, deps);
|
||||
const errors = vi.fn();
|
||||
const coordinator = createDeliveryWorkCoordinator({ owner: db, canRun: () => true, onError: errors });
|
||||
const worker = registerBudgetEnforcementWork(db, coordinator, { cancelWorkForScope: control.cancelBudgetScopeWork });
|
||||
try {
|
||||
// Drain startup recovery first: only the receipt's committed notification
|
||||
// may wake the idle worker below. Do not call wake or deliver explicitly.
|
||||
await worker.ready;
|
||||
expect(terminateLocalService).not.toHaveBeenCalled();
|
||||
expect(coordinator.nextWakeAt()).toBeNull();
|
||||
const startedAt = performance.now();
|
||||
await persistUsageReceipt(db, { schema: "paperclip/accounting-receipt/v1", id: randomUUID(),
|
||||
companyId: company.id, runId: run.id, sourceId, sequence: 1, receivedAt: new Date().toISOString(),
|
||||
adapterType: agent.adapterType, receipt: { complete: false, usageBasis: "per_run", costUsd: 1.1,
|
||||
provider: "test", biller: "test", model: "budget-stop-fixture", billingType: "metered_api", costStatus: "unpriced" },
|
||||
});
|
||||
const receiptReturnedAt = performance.now();
|
||||
await vi.waitFor(async () => {
|
||||
const persisted = await deps.getRun(run.id);
|
||||
expect(persisted).toMatchObject({ status: "running", usageJson: { budgetExposure: { sequence: 1, costCents: "110.0000000" } },
|
||||
resultJson: { startupCancellation: { budgetEnforcement: { policyId: policy.id, scopeType: "project", scopeId: project.id } } } });
|
||||
expect(terminateLocalService).toHaveBeenCalledTimes(1);
|
||||
}, { interval: 10, timeout: 3_000 });
|
||||
const fenceObservedAt = performance.now();
|
||||
const fenced = await deps.getRun(run.id);
|
||||
expect(fenced?.resultJson?.executionCancellation).toBeUndefined();
|
||||
expect((await deps.getRun(queued.id))?.resultJson?.startupCancellation).toBeTruthy();
|
||||
const persistedWakes = await db.select().from(agentWakeupRequests).where(inArray(agentWakeupRequests.id, wakes.map(wake => wake.id)));
|
||||
expect(wakes.map(wake => persistedWakes.find(row => row.id === wake.id)?.status)).toEqual(["cancelled", "cancelled", "queued"]);
|
||||
const [pending] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.id));
|
||||
expect(pending.enforcementVersion).toBeGreaterThan(pending.enforcementDeliveredVersion);
|
||||
|
||||
const requestStartedAt = performance.now();
|
||||
let responseTimeout: ReturnType<typeof setTimeout> | undefined;
|
||||
try {
|
||||
const rows = await Promise.race([
|
||||
db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, unrelated.id)),
|
||||
new Promise<never>((_, reject) => { responseTimeout = setTimeout(() => reject(new Error("Unrelated database request blocked by Stop delivery")), 1_000); }),
|
||||
]);
|
||||
expect(rows[0]).toMatchObject({ status: "running", resultJson: null });
|
||||
} finally { clearTimeout(responseTimeout); }
|
||||
const unrelatedRequestMs = performance.now() - requestStartedAt;
|
||||
const terminationReleasedAt = performance.now();
|
||||
releaseTermination();
|
||||
await vi.waitFor(async () => {
|
||||
expect(await deps.getRun(run.id)).toMatchObject({ status: "cancelled", resultJson: { executionCancellation: { state: "acknowledged" } } });
|
||||
const [delivered] = await db.select().from(budgetPolicies).where(eq(budgetPolicies.id, policy.id));
|
||||
expect(delivered.enforcementDeliveredVersion).toBe(delivered.enforcementVersion);
|
||||
}, { interval: 10, timeout: 3_000 });
|
||||
expect((await deps.getRun(unrelated.id))?.status).toBe("running");
|
||||
expect(errors).not.toHaveBeenCalled();
|
||||
console.info("Partial usage budget Stop timing", { receiptMs: Math.round(receiptReturnedAt - startedAt),
|
||||
receiptReturnToFenceMs: Math.round(fenceObservedAt - receiptReturnedAt), unrelatedRequestMs: Math.round(unrelatedRequestMs),
|
||||
controlledTerminationWaitMs: Math.round(terminationReleasedAt - fenceObservedAt),
|
||||
terminationReleaseToAcknowledgementMs: Math.round(performance.now() - terminationReleasedAt) });
|
||||
} finally {
|
||||
releaseTermination();
|
||||
await coordinator.stop();
|
||||
}
|
||||
});
|
||||
|
||||
it("drains other provider stops after a failed delivery and retries the unfinished run", async () => {
|
||||
const { company, agent } = await fixture();
|
||||
const failed = await liveRun(company.id, agent.id);
|
||||
const [other] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id,
|
||||
status: "running", runtimeMode: "legacy", runtimeModeResolvedAt: new Date() }).returning();
|
||||
runningProcesses.set(other.id, { child: { pid: 434343 } as never, graceSec: 4, processGroupId: 434343 });
|
||||
vi.mocked(terminateLocalService).mockImplementation(async input => {
|
||||
if (input.pid === 424242) throw new Error("provider unreachable");
|
||||
});
|
||||
const deps = callbacks(db);
|
||||
const control = createHeartbeatRunControl(db, deps);
|
||||
const scope = { companyId: company.id, scopeType: "agent" as const, scopeId: agent.id };
|
||||
await expect(control.cancelBudgetScopeWork(scope)).rejects.toThrow("Budget cancellation delivery remains pending");
|
||||
expect((await deps.getRun(other.id))?.status).toBe("cancelled");
|
||||
expect(await deps.getRun(failed.id)).toMatchObject({ status: "running", resultJson: { startupCancellation: expect.any(Object) } });
|
||||
expect(runningProcesses.has(failed.id)).toBe(true);
|
||||
vi.mocked(terminateLocalService).mockResolvedValue(undefined);
|
||||
await control.cancelBudgetScopeWork(scope);
|
||||
expect(await deps.getRun(failed.id)).toMatchObject({ status: "cancelled", resultJson: { executionCancellation: { state: "acknowledged" } } });
|
||||
});
|
||||
|
||||
it("shares four Stop slots across independent scopes and service instances", async () => {
|
||||
const { company, agent } = await fixture();
|
||||
const [otherAgent] = await db.insert(agents).values({ companyId: company.id, name: "Another scope" }).returning();
|
||||
const runs = await db.insert(heartbeatRuns).values([agent.id, otherAgent.id].flatMap(agentId =>
|
||||
Array.from({ length: 3 }, () => ({ companyId: company.id, agentId,
|
||||
status: "running", runtimeMode: "legacy", runtimeModeResolvedAt: new Date() })))).returning();
|
||||
runs.forEach((run, i) => runningProcesses.set(run.id, {
|
||||
child: { pid: 435000 + i } as never, graceSec: 4, processGroupId: 435000 + i,
|
||||
}));
|
||||
let active = 0, maximumActive = 0, releaseAll = false;
|
||||
const releases = new Map<number, () => void>();
|
||||
vi.mocked(terminateLocalService).mockImplementation(async input => {
|
||||
active += 1;
|
||||
maximumActive = Math.max(maximumActive, active);
|
||||
if (!releaseAll) await new Promise<void>(resolve => releases.set(input.pid, resolve));
|
||||
active -= 1;
|
||||
});
|
||||
const deliveries = [agent, otherAgent].map(selected => createHeartbeatRunControl(db, callbacks(db))
|
||||
.cancelBudgetScopeWork({ companyId: company.id, scopeType: "agent", scopeId: selected.id }));
|
||||
try {
|
||||
await vi.waitFor(() => expect(releases.size).toBe(4));
|
||||
expect(maximumActive).toBe(4);
|
||||
expect((await db.select().from(heartbeatRuns).where(inArray(heartbeatRuns.id, runs.map(run => run.id))))
|
||||
.every(run => Boolean(run.resultJson?.startupCancellation))).toBe(true);
|
||||
releases.values().next().value!();
|
||||
await vi.waitFor(() => expect(releases.size).toBe(5));
|
||||
expect(maximumActive).toBe(4);
|
||||
} finally {
|
||||
releaseAll = true;
|
||||
releases.forEach(release => release());
|
||||
await Promise.all(deliveries);
|
||||
}
|
||||
});
|
||||
|
||||
it("blocks retry promotion after the budget batch commits but before Stop delivery", async () => {
|
||||
const { company, agent } = await fixture();
|
||||
const [run] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id,
|
||||
status: "scheduled_retry", runtimeMode: "legacy", runtimeModeResolvedAt: new Date(),
|
||||
scheduledRetryAt: new Date(), scheduledRetryReason: "transient_failure" }).returning();
|
||||
const deps = callbacks(db);
|
||||
const getAgent = deps.getAgent.getMockImplementation()!;
|
||||
let entered!: () => void, release!: () => void;
|
||||
const deliveryEntered = new Promise<void>(resolve => { entered = resolve; });
|
||||
const continueDelivery = new Promise<void>(resolve => { release = resolve; });
|
||||
deps.getAgent.mockImplementationOnce(async id => { entered(); await continueDelivery; return getAgent(id); });
|
||||
const cancellation = createHeartbeatRunControl(db, deps).cancelBudgetScopeWork({ companyId: company.id,
|
||||
scopeType: "agent", scopeId: agent.id });
|
||||
try {
|
||||
await deliveryEntered;
|
||||
expect(await createPostgresRunDispatchAdapter(db).promoteOrCancelDueRetry({
|
||||
companyId: company.id, runId: run.id, now: new Date(),
|
||||
})).toMatchObject({ outcome: "not_promoted" });
|
||||
expect(await deps.getRun(run.id)).toMatchObject({ status: "scheduled_retry",
|
||||
resultJson: { startupCancellation: expect.any(Object) } });
|
||||
} finally {
|
||||
release();
|
||||
await cancellation;
|
||||
}
|
||||
expect(await deps.getRun(run.id)).toMatchObject({ status: "cancelled" });
|
||||
});
|
||||
|
||||
it("recovers a committed Stop after its policy was deleted without fencing newly admitted work", async () => {
|
||||
const { company, agent } = await fixture();
|
||||
const oldRun = await liveRun(company.id, agent.id);
|
||||
const [newRun] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id,
|
||||
status: "queued", runtimeMode: "legacy" }).returning();
|
||||
const intent = { companyId: company.id, scopeType: "agent", scopeId: agent.id,
|
||||
policyId: randomUUID(), version: 1 };
|
||||
await db.update(heartbeatRuns).set({ resultJson: { startupCancellation: {
|
||||
requestedAt: new Date().toISOString(), budgetEnforcement: intent,
|
||||
} } }).where(eq(heartbeatRuns.id, oldRun.id));
|
||||
const scopes = await listPendingBudgetCancellationScopes(db);
|
||||
expect(scopes).toHaveLength(1);
|
||||
const deps = callbacks(db);
|
||||
await createHeartbeatRunControl(db, deps).cancelBudgetScopeWork(scopes[0]);
|
||||
expect(await deps.getRun(oldRun.id)).toMatchObject({ status: "cancelled",
|
||||
resultJson: { startupCancellation: { budgetEnforcement: { deliveredAt: expect.any(String) } } } });
|
||||
expect(await deps.getRun(newRun.id)).toMatchObject({ status: "queued", resultJson: null });
|
||||
expect(await listPendingBudgetCancellationScopes(db)).toEqual([]);
|
||||
});
|
||||
|
||||
it("retains a terminal budget Stop until actual stop acknowledgement exists", async () => {
|
||||
const { company, agent } = await fixture();
|
||||
const [run] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id,
|
||||
status: "cancelled", runtimeMode: "legacy", startedAt: new Date(), resultJson: {
|
||||
startupCancellation: { budgetEnforcement: { companyId: company.id, scopeType: "agent", scopeId: agent.id,
|
||||
policyId: randomUUID(), version: 1 } },
|
||||
} }).returning();
|
||||
const control = createHeartbeatRunControl(db, callbacks(db));
|
||||
const [scope] = await listPendingBudgetCancellationScopes(db);
|
||||
await expect(control.cancelBudgetScopeWork(scope)).rejects.toThrow("Budget cancellation delivery remains pending");
|
||||
expect(await listPendingBudgetCancellationScopes(db)).toHaveLength(1);
|
||||
expect(terminateLocalService).not.toHaveBeenCalled();
|
||||
await db.update(heartbeatRuns).set({ resultJson: { ...run.resultJson,
|
||||
executionCancellation: { state: "acknowledged", acknowledgedAt: new Date().toISOString() },
|
||||
} }).where(eq(heartbeatRuns.id, run.id));
|
||||
await control.cancelBudgetScopeWork(scope);
|
||||
expect(await listPendingBudgetCancellationScopes(db)).toEqual([]);
|
||||
});
|
||||
|
||||
it.each(["local exit", "remote receipt", "unproven", "later launch"] as const)("reconciles normal completion before queued budget Stop using %s evidence", async evidence => {
|
||||
const { company, agent } = await fixture();
|
||||
const [run] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id,
|
||||
status: "running", runtimeMode: "native", startedAt: new Date() }).returning();
|
||||
const deps = callbacks(db);
|
||||
if (evidence === "local exit" || evidence === "later launch") {
|
||||
const child = spawn(process.execPath, ["-e", "setTimeout(() => process.exit(0), 50)"], { detached: true, stdio: "ignore" });
|
||||
const exited = once(child, "exit");
|
||||
try {
|
||||
expect(child.pid).toBeGreaterThan(1);
|
||||
await deps.persistRunProcessMetadata(run.id, { pid: child.pid!, processGroupId: child.pid!,
|
||||
startedAt: new Date().toISOString(), targetKind: "local" });
|
||||
await exited;
|
||||
} finally { if (child.exitCode === null) child.kill("SIGKILL"); }
|
||||
if (evidence === "later launch") await deps.appendRunEvent(run, {
|
||||
eventType: PROCESS_START_REQUESTED, stream: "system", level: "info", message: "A later launch invalidated the earlier exit proof",
|
||||
});
|
||||
}
|
||||
if (evidence === "remote receipt") {
|
||||
const leaseId = randomUUID(), providerLeaseId = randomUUID();
|
||||
await db.insert(environmentLeases).values({ id: leaseId, companyId: company.id, heartbeatRunId: run.id,
|
||||
provider: "daytona", providerLeaseId, status: "released", releasedAt: new Date(), cleanupStatus: "success",
|
||||
metadata: { remoteExecutionTermination: { schema: "paperclip.remote-termination.v1", companyId: company.id,
|
||||
runId: run.id, leaseId, provider: "daytona", providerLeaseId, state: "stopped", confirmedAt: new Date().toISOString() } },
|
||||
});
|
||||
}
|
||||
const getRun = deps.getRun.getMockImplementation()!;
|
||||
deps.getRun.mockImplementationOnce(async id => {
|
||||
// Budget delivery persisted its fence, then ordinary native completion
|
||||
// won the race before this run acquired a provider Stop delivery slot.
|
||||
expect((await getRun(id))?.resultJson?.startupCancellation).toBeTruthy();
|
||||
await db.update(heartbeatRuns).set({ status: "succeeded", finishedAt: new Date(), nativePhase: "committed",
|
||||
processPid: null, processGroupId: null, processStartedAt: null }).where(eq(heartbeatRuns.id, id));
|
||||
return getRun(id);
|
||||
});
|
||||
const policyId = randomUUID();
|
||||
await db.insert(budgetPolicies).values({ id: policyId, companyId: company.id, scopeType: "agent",
|
||||
scopeId: agent.id, windowKind: "lifetime", amount: 1, enforcementVersion: 1 });
|
||||
await db.insert(costEvents).values({ companyId: company.id, agentId: agent.id, provider: "test",
|
||||
model: "overrun", inputTokens: 1, outputTokens: 1, costCents: 2, occurredAt: new Date() });
|
||||
const scope = { companyId: company.id, scopeType: "agent" as const, scopeId: agent.id,
|
||||
enforcement: { policyId, version: 1 } };
|
||||
const control = createHeartbeatRunControl(db, deps);
|
||||
const succeeded = evidence === "local exit" || evidence === "remote receipt";
|
||||
const delivery = control.cancelBudgetScopeWork(scope);
|
||||
if (succeeded) await delivery;
|
||||
else await expect(delivery).rejects.toThrow("Budget cancellation delivery remains pending");
|
||||
const persisted = await getRun(run.id);
|
||||
expect(persisted).toMatchObject({ status: "succeeded", nativePhase: "committed" });
|
||||
expect(persisted?.resultJson?.executionCancellation).toBeUndefined();
|
||||
const intent = (persisted?.resultJson?.startupCancellation as { budgetEnforcement?: Record<string, unknown> })?.budgetEnforcement;
|
||||
if (succeeded) expect(intent).toMatchObject({ deliveredAt: expect.any(String),
|
||||
terminationProof: evidence === "local exit" ? "local_process_exit" : "provider_termination_receipt" });
|
||||
else expect(intent?.deliveredAt).toBeUndefined();
|
||||
expect(await listPendingBudgetCancellationScopes(db)).toHaveLength(succeeded ? 0 : 1);
|
||||
expect(terminateLocalService).not.toHaveBeenCalled();
|
||||
expect(deps.setRunStatusFromLive).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("moves an attempted failed Stop behind another pending scope", async () => {
|
||||
const { company, agent } = await fixture();
|
||||
const [otherAgent] = await db.insert(agents).values({ companyId: company.id, name: "Later recovery" }).returning();
|
||||
await db.insert(heartbeatRuns).values([agent, otherAgent].map((selected, index) => ({
|
||||
companyId: company.id, agentId: selected.id, status: "cancelled", runtimeMode: "legacy",
|
||||
resultJson: { startupCancellation: { requestedAt: `2020-01-0${index + 1}T00:00:00.000Z`,
|
||||
budgetEnforcement: { companyId: company.id, scopeType: "agent", scopeId: selected.id,
|
||||
policyId: randomUUID(), version: 1 } } },
|
||||
})));
|
||||
const [first] = await listPendingBudgetCancellationScopes(db, 1);
|
||||
expect(first.scopeId).toBe(agent.id);
|
||||
await expect(createHeartbeatRunControl(db, callbacks(db)).cancelBudgetScopeWork(first))
|
||||
.rejects.toThrow("Budget cancellation delivery remains pending");
|
||||
const [next] = await listPendingBudgetCancellationScopes(db, 1);
|
||||
expect(next.scopeId).toBe(otherAgent.id);
|
||||
const [attempted] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.agentId, agent.id));
|
||||
expect((attempted.resultJson?.startupCancellation as { budgetEnforcement?: { lastAttemptAt?: string } })
|
||||
?.budgetEnforcement?.lastAttemptAt).toEqual(expect.any(String));
|
||||
});
|
||||
|
||||
it("limits project budget cancellation by company, effective project, and admission cutoff", async () => {
|
||||
const { company, agent } = await fixture();
|
||||
const [otherCompany] = await db.insert(companies).values({ name: "Other company", issuePrefix: "OC" }).returning();
|
||||
const [otherAgent] = await db.insert(agents).values({ companyId: otherCompany.id, name: "Other company agent" }).returning();
|
||||
const projectRows = await db.insert(projects).values([
|
||||
{ companyId: company.id, name: "Target" }, { companyId: company.id, name: "Override" },
|
||||
]).returning();
|
||||
const [issue] = await db.insert(issues).values({ companyId: company.id, projectId: projectRows[0].id, title: "Bound task", issueNumber: 1, identifier: "RC-1" }).returning();
|
||||
const old = new Date("2026-01-01T00:00:00Z");
|
||||
const cutoff = new Date("2026-01-02T00:00:00Z");
|
||||
const wakes = await db.insert(agentWakeupRequests).values([
|
||||
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued", createdAt: old, payload: { issueId: issue.id } },
|
||||
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued", createdAt: old, payload: { issueId: issue.id, projectId: projectRows[1].id } },
|
||||
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued", createdAt: cutoff, payload: { projectId: projectRows[0].id } },
|
||||
{ companyId: otherCompany.id, agentId: otherAgent.id, source: "assignment", status: "queued", createdAt: old, payload: { projectId: projectRows[0].id } },
|
||||
]).returning();
|
||||
await createHeartbeatRunControl(db, callbacks(db)).cancelBudgetScopeWork({ companyId: company.id, scopeType: "project", scopeId: projectRows[0].id, createdBefore: cutoff });
|
||||
const persisted = await db.select().from(agentWakeupRequests);
|
||||
expect(wakes.map(wake => persisted.find(row => row.id === wake.id)?.status)).toEqual(["cancelled", "queued", "queued", "queued"]);
|
||||
});
|
||||
|
||||
it("fences detached browsers after their parent run finishes using admitted project scope and cutoff", async () => {
|
||||
const { company, agent } = await fixture();
|
||||
const [project, otherProject] = await db.insert(projects).values([
|
||||
{ companyId: company.id, name: "Stopped project" }, { companyId: company.id, name: "Other project" },
|
||||
]).returning();
|
||||
const [movedIssue, targetIssue] = await db.insert(issues).values([
|
||||
{ companyId: company.id, projectId: otherProject.id, title: "Moved after browser admission" },
|
||||
{ companyId: company.id, projectId: project.id, title: "Current project" },
|
||||
]).returning();
|
||||
const [parent] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id,
|
||||
status: "succeeded", finishedAt: new Date(), runtimeMode: "legacy" }).returning();
|
||||
const [application] = await db.insert(toolApplications).values({ companyId: company.id, name: "Browser fixture", type: "rest_api" }).returning();
|
||||
const [connection] = await db.insert(toolConnections).values({ companyId: company.id, applicationId: application.id,
|
||||
name: "Browser fixture", uid: "budget-browser-fixture", transport: "rest_api" }).returning();
|
||||
const [grant] = await db.insert(connectionGrants).values({ companyId: company.id, connectionId: connection.id, kind: "organization" }).returning();
|
||||
const cutoff = new Date();
|
||||
const old = new Date(cutoff.getTime() - 1_000);
|
||||
const later = new Date(cutoff.getTime() + 60_000);
|
||||
const [admitted, other, legacy, newer, noProject] = await db.insert(browserUseSessions).values([
|
||||
{ issueId: movedIssue.id, stopRequested: "cancel" }, { issueId: targetIssue.id },
|
||||
{ issueId: targetIssue.id }, { issueId: targetIssue.id, createdAt: later }, { issueId: targetIssue.id },
|
||||
].map(values => ({ companyId: company.id, agentId: agent.id, connectionId: connection.id, grantId: grant.id,
|
||||
providerSessionId: randomUUID(), status: "running", createdAt: old, nextPollAt: later,
|
||||
leaseToken: randomUUID(), leaseUntil: later, error: "retained cleanup context", ...values }))).returning();
|
||||
await db.insert(browserUseRuns).values([admitted, other, legacy, noProject].map((session, index) => {
|
||||
const invocationId = randomUUID();
|
||||
return { companyId: company.id, sessionId: session.id, heartbeatRunId: parent.id, invocationId,
|
||||
status: "running", detachedUntil: later, createdAt: old,
|
||||
result: index === 2 ? null : { paperclipBudgetHold: { invocationId, amountCents: "10", state: "held",
|
||||
projectId: index === 0 ? project.id : index === 1 ? otherProject.id : null } },
|
||||
};
|
||||
}));
|
||||
const wake = vi.fn();
|
||||
const unsubscribe = subscribeDeliveryWork(db, DELIVERY_QUEUES.browser, wake);
|
||||
const deps = callbacks(db);
|
||||
try {
|
||||
await createHeartbeatRunControl(db, deps).cancelBudgetScopeWork({ companyId: company.id,
|
||||
scopeType: "project", scopeId: project.id, createdBefore: cutoff });
|
||||
const persisted = await db.select().from(browserUseSessions).where(eq(browserUseSessions.companyId, company.id));
|
||||
for (const original of [admitted, legacy]) {
|
||||
const stopped = persisted.find(session => session.id === original.id)!;
|
||||
expect(stopped).toMatchObject({ status: "stopping", stopRequested: "end", providerSessionId: original.providerSessionId,
|
||||
leaseToken: original.leaseToken, leaseUntil: original.leaseUntil, error: original.error });
|
||||
expect(stopped.nextPollAt.getTime()).toBeLessThan(later.getTime());
|
||||
}
|
||||
for (const original of [other, newer, noProject]) expect(persisted.find(session => session.id === original.id)).toEqual(original);
|
||||
expect(await deps.getRun(parent.id)).toMatchObject({ status: "succeeded", resultJson: null });
|
||||
expect(wake).toHaveBeenCalledTimes(1);
|
||||
expect(terminateLocalService).not.toHaveBeenCalled();
|
||||
expect(deps.setRunStatusFromLive).not.toHaveBeenCalled();
|
||||
} finally { unsubscribe(); }
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,43 @@
|
||||
import { heartbeatRuns, type Db } from "@paperclipai/db";
|
||||
import { isUuidLike } from "@paperclipai/shared";
|
||||
import { and, eq, sql } from "drizzle-orm";
|
||||
import type { BudgetEnforcementScope } from "../budgets.js";
|
||||
|
||||
/** Committed Stop intent survives policy edits, deletion, and terminal status.
|
||||
* Run control writes deliveredAt only after its ordinary stop proof succeeds. */
|
||||
export function pendingBudgetRunCancellationPredicate() {
|
||||
return and(
|
||||
sql`jsonb_typeof(${heartbeatRuns.resultJson}->'startupCancellation'->'budgetEnforcement') = 'object'`,
|
||||
sql`${heartbeatRuns.resultJson}->'startupCancellation'->'budgetEnforcement'->>'deliveredAt' is null`,
|
||||
);
|
||||
}
|
||||
|
||||
export function budgetCancellationScopePredicate(scope: Pick<BudgetEnforcementScope, "companyId" | "scopeType" | "scopeId">) {
|
||||
return and(eq(heartbeatRuns.companyId, scope.companyId),
|
||||
sql`${heartbeatRuns.resultJson}->'startupCancellation'->'budgetEnforcement'->>'companyId' = ${scope.companyId}`,
|
||||
sql`${heartbeatRuns.resultJson}->'startupCancellation'->'budgetEnforcement'->>'scopeType' = ${scope.scopeType}`,
|
||||
sql`${heartbeatRuns.resultJson}->'startupCancellation'->'budgetEnforcement'->>'scopeId' = ${scope.scopeId}`);
|
||||
}
|
||||
|
||||
export async function listPendingBudgetCancellationScopes(db: Db, limit = 100): Promise<BudgetEnforcementScope[]> {
|
||||
const intent = sql`${heartbeatRuns.resultJson}->'startupCancellation'->'budgetEnforcement'`;
|
||||
const scopeType = sql<string>`${intent}->>'scopeType'`, scopeId = sql<string>`${intent}->>'scopeId'`;
|
||||
const attemptAt = sql`coalesce((${intent}->>'lastAttemptAt')::timestamptz,
|
||||
(${heartbeatRuns.resultJson}->'startupCancellation'->>'requestedAt')::timestamptz, ${heartbeatRuns.updatedAt})`;
|
||||
const rows = await db.select({ companyId: heartbeatRuns.companyId,
|
||||
intent: sql<Record<string, unknown>>`(array_agg(${intent} order by ${attemptAt}, ${heartbeatRuns.id}))[1]`,
|
||||
}).from(heartbeatRuns).where(pendingBudgetRunCancellationPredicate())
|
||||
.groupBy(heartbeatRuns.companyId, scopeType, scopeId)
|
||||
.orderBy(sql`min(${attemptAt})`, heartbeatRuns.companyId, scopeType, scopeId)
|
||||
.limit(limit);
|
||||
return rows.flatMap(({ companyId, intent }) => {
|
||||
if (intent.companyId !== companyId || !["company", "agent", "project"].includes(String(intent.scopeType))
|
||||
|| typeof intent.scopeId !== "string" || !isUuidLike(intent.scopeId)
|
||||
|| typeof intent.policyId !== "string" || !isUuidLike(intent.policyId)
|
||||
|| !Number.isSafeInteger(intent.version) || Number(intent.version) < 0) return [];
|
||||
const cutoff = typeof intent.createdBefore === "string" ? new Date(intent.createdBefore) : null;
|
||||
return [{ companyId, scopeType: intent.scopeType as BudgetEnforcementScope["scopeType"], scopeId: intent.scopeId,
|
||||
enforcement: { policyId: intent.policyId, version: Number(intent.version) },
|
||||
...(cutoff && Number.isFinite(cutoff.getTime()) ? { createdBefore: cutoff } : {}) }];
|
||||
});
|
||||
}
|
||||
@@ -1,9 +1,8 @@
|
||||
import { randomUUID } from "node:crypto";
|
||||
import type { Db } from "@paperclipai/db";
|
||||
import { activityLog, agents, agentWakeupRequests, caseEvents, caseHandoffs, cases, companies, createDb, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS, environmentLeases, heartbeatRuns, issues, nativeRunFinalizations, projects } from "@paperclipai/db";
|
||||
import { eq, inArray, sql } from "drizzle-orm";
|
||||
import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../../__tests__/helpers/embedded-postgres.js";
|
||||
import { activityLog, agents, agentWakeupRequests, caseEvents, caseHandoffs, cases, companies, environmentLeases, heartbeatRuns, issues, nativeRunFinalizations } from "@paperclipai/db";
|
||||
import { eq, inArray } from "drizzle-orm";
|
||||
import { afterEach, describe, expect, it, vi } from "vitest";
|
||||
import { runningProcesses } from "../../adapters/index.js";
|
||||
import {
|
||||
adapterExecutionControls,
|
||||
@@ -16,14 +15,13 @@ import { HttpError } from "../../errors.js";
|
||||
import { captureException } from "../../sentry.js";
|
||||
import { terminateLocalService } from "../local-service-supervisor.js";
|
||||
import { createHeartbeatRunState } from "./run-state.js";
|
||||
import { createHeartbeatRunPreparation } from "./run-preparation.js";
|
||||
import { createHeartbeatLifecycle, type HeartbeatLifecycleDependencies } from "./run-lifecycle.js";
|
||||
import { createHeartbeatRunControl, type HeartbeatRunControlDependencies } from "./run-control.js";
|
||||
import { createHeartbeatRunControl } from "./run-control.js";
|
||||
import { preserveCancelledNativeWorkspaces } from "../native-runtime/native-failed-workspace-recovery.js";
|
||||
import { createPostgresRunDispatchAdapter } from "../../modules/run-dispatch/adapters/postgres.js";
|
||||
import { caseHandoffService } from "../case-handoffs.js";
|
||||
import { instanceSettingsService } from "../instance-settings.js";
|
||||
import { issueService } from "../issues.js";
|
||||
import { callbacks, runControlPostgresSupport, setupRunControlDatabase } from "./__tests__/run-control-fixtures.js";
|
||||
|
||||
vi.mock("../live-events.js", () => ({ publishLiveEvent: vi.fn() }));
|
||||
vi.mock("../adapter-execution-control.js", async (importOriginal) => {
|
||||
@@ -40,50 +38,6 @@ vi.mock("../local-service-supervisor.js", async (importOriginal) => ({
|
||||
}));
|
||||
|
||||
type Run = typeof heartbeatRuns.$inferSelect;
|
||||
function callbacks(db: Db) {
|
||||
const state = createHeartbeatRunState(db);
|
||||
const lifecycleDependencies = {
|
||||
...state,
|
||||
...createHeartbeatRunPreparation(db),
|
||||
issuesSvc: {
|
||||
addComment: vi.fn<HeartbeatLifecycleDependencies["issuesSvc"]["addComment"]>(),
|
||||
listReviewAttention: vi.fn<HeartbeatLifecycleDependencies["issuesSvc"]["listReviewAttention"]>(),
|
||||
},
|
||||
recovery: { escalateStrandedAssignedIssue: vi.fn<HeartbeatLifecycleDependencies["recovery"]["escalateStrandedAssignedIssue"]>() },
|
||||
companySkills: { completeTestRunForIssue: vi.fn<HeartbeatLifecycleDependencies["companySkills"]["completeTestRunForIssue"]>() },
|
||||
budgets: { getInvocationBlock: vi.fn(async () => null) },
|
||||
treeControlSvc: { getActivePauseHoldGate: vi.fn(async () => null) },
|
||||
enqueueWakeup: vi.fn<HeartbeatLifecycleDependencies["enqueueWakeup"]>(),
|
||||
getCurrentUserRedactionOptions: vi.fn(async () => ({})),
|
||||
getAgentInvokability: vi.fn<HeartbeatLifecycleDependencies["getAgentInvokability"]>(async () => ({ invokable: true })),
|
||||
emitTerminalAgentTaskRun: vi.fn<HeartbeatLifecycleDependencies["emitTerminalAgentTaskRun"]>(),
|
||||
budgetHooks: {},
|
||||
} satisfies HeartbeatLifecycleDependencies;
|
||||
const lifecycle = createHeartbeatLifecycle(db, lifecycleDependencies);
|
||||
return {
|
||||
...state,
|
||||
...lifecycle,
|
||||
getRun: vi.fn(state.getRun),
|
||||
getAgent: vi.fn(state.getAgent),
|
||||
setRunStatusFromLive: vi.fn(lifecycle.setRunStatusFromLive),
|
||||
setRunStatus: vi.fn(lifecycle.setRunStatus),
|
||||
setWakeupStatus: vi.fn(lifecycle.setWakeupStatus),
|
||||
appendRunEvent: vi.fn(lifecycle.appendRunEvent),
|
||||
finalizeAgentStatus: vi.fn<HeartbeatRunControlDependencies["finalizeAgentStatus"]>(),
|
||||
options: { closeWarmNativeSessionsForRun: vi.fn(async () => ({ closed: 1, busy: 0, failed: 0 })) },
|
||||
envOrchestrator: { releaseForRun: vi.fn<HeartbeatRunControlDependencies["envOrchestrator"]["releaseForRun"]>(async () => ({ released: [], errors: [] })) },
|
||||
enqueueWakeup: lifecycleDependencies.enqueueWakeup,
|
||||
startNextQueuedRunForAgent: vi.fn<HeartbeatRunControlDependencies["startNextQueuedRunForAgent"]>(),
|
||||
sweepPendingCleanupLeases: vi.fn<HeartbeatRunControlDependencies["sweepPendingCleanupLeases"]>(),
|
||||
timerClaimWasFirstHeartbeat: vi.fn<HeartbeatRunControlDependencies["timerClaimWasFirstHeartbeat"]>(() => undefined),
|
||||
wakeQueue: { releaseIssueExecution: vi.fn<HeartbeatRunControlDependencies["wakeQueue"]["releaseIssueExecution"]>(async () => ({ outcome: { kind: "released" }, postCommitEffects: [] })) },
|
||||
applyWakeQueuePostCommitEffects: vi.fn(async () => undefined),
|
||||
getSchedulingSuppression: vi.fn(async () => ({ suppressed: false })),
|
||||
activeRunExecutions: new Set<string>(),
|
||||
processRunCancellationSettlements: new Map<string, { settled: Promise<void>; failed: boolean; error?: unknown }>(),
|
||||
failedProcessRunCancellations: new Map<string, { settled: Promise<void>; failed: boolean; error?: unknown }>(),
|
||||
} satisfies HeartbeatRunControlDependencies;
|
||||
}
|
||||
function guardedDatabase() {
|
||||
const access = vi.fn(() => { throw new Error("Unexpected run-control database access"); });
|
||||
return { db: new Proxy({}, { get: access }) as Db, access };
|
||||
@@ -171,35 +125,9 @@ describe("heartbeat run-control module boundary", () => {
|
||||
});
|
||||
});
|
||||
|
||||
const support = await getEmbeddedPostgresTestSupport();
|
||||
if (!support.supported) console.warn(`Skipping run-control database tests: ${support.reason}`);
|
||||
describe.skipIf(!support.supported)("heartbeat run-control database wiring", () => {
|
||||
let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>>;
|
||||
describe.skipIf(!runControlPostgresSupport.supported)("heartbeat run-control database wiring", () => {
|
||||
let db: Db;
|
||||
beforeAll(async () => {
|
||||
database = await startEmbeddedPostgresTestDatabase("paperclip-heartbeat-run-control-module");
|
||||
db = createDb(database.connectionString);
|
||||
}, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS);
|
||||
afterAll(async () => {
|
||||
await db?.$client.end();
|
||||
await database?.cleanup();
|
||||
});
|
||||
beforeEach(async () => {
|
||||
vi.mocked(terminateLocalService).mockReset().mockResolvedValue(undefined);
|
||||
await db.execute(sql`truncate table companies restart identity cascade`);
|
||||
});
|
||||
afterEach(() => runningProcesses.clear());
|
||||
|
||||
async function fixture() {
|
||||
const [company] = await db.insert(companies).values({ name: "Run control", issuePrefix: "RC" }).returning();
|
||||
const [agent] = await db.insert(agents).values({ companyId: company.id, name: "Controlled agent", adapterType: "codex_local" }).returning();
|
||||
return { company, agent };
|
||||
}
|
||||
async function liveRun(companyId: string, agentId: string) {
|
||||
const [run] = await db.insert(heartbeatRuns).values({ companyId, agentId, status: "running", runtimeMode: "legacy", runtimeModeResolvedAt: new Date() }).returning();
|
||||
runningProcesses.set(run.id, { child: { pid: 424242 } as never, graceSec: 4, processGroupId: 424242 });
|
||||
return run;
|
||||
}
|
||||
const { fixture, liveRun } = setupRunControlDatabase("paperclip-heartbeat-run-control-module", value => { db = value; });
|
||||
|
||||
it.each(["queued", "running", "cancelled", "succeeded"] as const)("does not supersede a retry that becomes %s between lookup and the cancellation fence", async status => {
|
||||
const { company, agent } = await fixture();
|
||||
@@ -704,37 +632,4 @@ describe.skipIf(!support.supported)("heartbeat run-control database wiring", ()
|
||||
expect(wakes.map(wake => persisted.find(row => row.id === wake.id)?.status)).toEqual(["cancelled", "cancelled", "claimed", "queued", "queued"]);
|
||||
});
|
||||
|
||||
it("rechecks stale budget enforcement before cancelling runs or queued wakes", async () => {
|
||||
const { company, agent } = await fixture();
|
||||
const [run] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id, status: "queued", runtimeMode: "legacy" }).returning();
|
||||
const [wake] = await db.insert(agentWakeupRequests).values({ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued" }).returning();
|
||||
const deps = callbacks(db);
|
||||
await createHeartbeatRunControl(db, deps).cancelBudgetScopeWork({ companyId: company.id, scopeType: "agent", scopeId: agent.id, enforcement: { policyId: randomUUID(), version: 1 } });
|
||||
expect((await deps.getRun(run.id))?.status).toBe("queued");
|
||||
expect((await db.select().from(agentWakeupRequests).where(eq(agentWakeupRequests.id, wake.id)))[0].status).toBe("queued");
|
||||
expect(deps.setRunStatusFromLive).not.toHaveBeenCalled();
|
||||
expect(deps.finalizeAgentStatus).not.toHaveBeenCalled();
|
||||
expect(terminateLocalService).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("limits project budget cancellation by company, effective project, and admission cutoff", async () => {
|
||||
const { company, agent } = await fixture();
|
||||
const [otherCompany] = await db.insert(companies).values({ name: "Other company", issuePrefix: "OC" }).returning();
|
||||
const [otherAgent] = await db.insert(agents).values({ companyId: otherCompany.id, name: "Other company agent" }).returning();
|
||||
const projectRows = await db.insert(projects).values([
|
||||
{ companyId: company.id, name: "Target" }, { companyId: company.id, name: "Override" },
|
||||
]).returning();
|
||||
const [issue] = await db.insert(issues).values({ companyId: company.id, projectId: projectRows[0].id, title: "Bound task", issueNumber: 1, identifier: "RC-1" }).returning();
|
||||
const old = new Date("2026-01-01T00:00:00Z");
|
||||
const cutoff = new Date("2026-01-02T00:00:00Z");
|
||||
const wakes = await db.insert(agentWakeupRequests).values([
|
||||
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued", createdAt: old, payload: { issueId: issue.id } },
|
||||
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued", createdAt: old, payload: { issueId: issue.id, projectId: projectRows[1].id } },
|
||||
{ companyId: company.id, agentId: agent.id, source: "assignment", status: "queued", createdAt: cutoff, payload: { projectId: projectRows[0].id } },
|
||||
{ companyId: otherCompany.id, agentId: otherAgent.id, source: "assignment", status: "queued", createdAt: old, payload: { projectId: projectRows[0].id } },
|
||||
]).returning();
|
||||
await createHeartbeatRunControl(db, callbacks(db)).cancelBudgetScopeWork({ companyId: company.id, scopeType: "project", scopeId: projectRows[0].id, createdBefore: cutoff });
|
||||
const persisted = await db.select().from(agentWakeupRequests);
|
||||
expect(wakes.map(wake => persisted.find(row => row.id === wake.id)?.status)).toEqual(["cancelled", "queued", "queued", "queued"]);
|
||||
});
|
||||
});
|
||||
@@ -1,5 +1,5 @@
|
||||
import { readTaskRetryIntent, reconcileTaskChatRetryReceipts } from "../task-retry.js";
|
||||
import { recordedLocalExecutionStopped } from "../native-local-process-stop.js";
|
||||
import { PROCESS_IDENTITY_RECORDED, PROCESS_START_REQUESTED, readNativeLocalProcessStop, recordedLocalExecutionStopped } from "../native-local-process-stop.js";
|
||||
import { isHeartbeatRunTerminalStatus } from "./run-lifecycle.js";
|
||||
import { pendingCaseHandoffForRun, cancelCaseHandoffsForIssue, isCaseHandoffSourceReleased } from "./case-handoff-lifecycle.js";
|
||||
import { publishActivity, type ActivityPublication } from "../activity-log.js";
|
||||
@@ -57,6 +57,8 @@ import {
|
||||
waitForAdapterStop,
|
||||
} from "../adapter-execution-control.js";
|
||||
import { AdapterStopTimeoutError } from "../adapter-stop-timeout.js";
|
||||
import { budgetCancellationScopePredicate, pendingBudgetRunCancellationPredicate } from "./budget-cancellation.js";
|
||||
import { withBudgetSpan } from "../budget-tracing.js";
|
||||
import { captureException } from "../../sentry.js";
|
||||
import { randomUUID } from "node:crypto";
|
||||
import {
|
||||
@@ -82,9 +84,12 @@ import {
|
||||
} from "@paperclipai/shared";
|
||||
import {
|
||||
agentWakeupRequests,
|
||||
browserUseRuns,
|
||||
browserUseSessions,
|
||||
companies,
|
||||
environmentLeases,
|
||||
heartbeatRuns,
|
||||
heartbeatRunEvents,
|
||||
issueComments,
|
||||
issueRecoveryActions,
|
||||
issues,
|
||||
@@ -96,6 +101,7 @@ import {
|
||||
notFound,
|
||||
} from "../../errors.js";
|
||||
import { logger } from "../../middleware/logger.js";
|
||||
import { DELIVERY_QUEUES, notifyDeliveryWork } from "../delivery-work-notifications.js";
|
||||
import {
|
||||
queuedCommentIdsFromWakePayload,
|
||||
withQueuedCommentIdsInWakePayload,
|
||||
@@ -134,6 +140,27 @@ type ProcessRunCancellationSettlement = {
|
||||
error?: unknown;
|
||||
};
|
||||
|
||||
// This bounds external Stop concurrency across scopes and heartbeat service
|
||||
// instances sharing a database. Durable cancellation rows own recovery.
|
||||
const budgetStopDeliveryLimits = new WeakMap<Db, { active: number; waiting: Array<() => void> }>();
|
||||
|
||||
async function withBudgetStopDelivery<T>(db: Db, work: () => Promise<T>): Promise<T> {
|
||||
let limit = budgetStopDeliveryLimits.get(db);
|
||||
if (!limit) {
|
||||
limit = { active: 0, waiting: [] };
|
||||
budgetStopDeliveryLimits.set(db, limit);
|
||||
}
|
||||
if (limit.active >= 4) await new Promise<void>(resolve => limit.waiting.push(resolve));
|
||||
else limit.active += 1;
|
||||
try {
|
||||
return await withBudgetSpan("budget.cancellation.deliver", work);
|
||||
} finally {
|
||||
const next = limit.waiting.shift();
|
||||
if (next) next();
|
||||
else limit.active -= 1;
|
||||
}
|
||||
}
|
||||
|
||||
/** Cancellation shares executor ownership and delegates lifecycle and admission effects. */
|
||||
export interface HeartbeatRunControlDependencies extends Pick<ReturnType<typeof createHeartbeatRunState>,
|
||||
"getRun" | "getAgent">,
|
||||
@@ -255,6 +282,43 @@ export type CancelRunOptions = {
|
||||
suppressImmediateRecovery?: boolean;
|
||||
};
|
||||
|
||||
/** The same durable fence protects both individual Stop and budget fan-out.
|
||||
* It records intent only; provider termination still needs its own receipt. */
|
||||
function runCancellationFenceSql(input: {
|
||||
cancellation: ReturnType<typeof requestedRunCancellation>;
|
||||
errorCode: string;
|
||||
reassignmentStopRequested?: boolean;
|
||||
preserveCancellation?: boolean;
|
||||
budgetEnforcement?: BudgetEnforcementScope;
|
||||
}) {
|
||||
const cancellation = JSON.stringify(input.cancellation);
|
||||
const startupCancellation = sql`jsonb_build_object(
|
||||
'requestedAt', ${new Date().toISOString()}::text,
|
||||
'errorCode', ${input.errorCode}::text,
|
||||
'beforeNativeSelection', ${heartbeatRuns.runtimeMode} = 'legacy'
|
||||
and ${heartbeatRuns.runtimeModeResolvedAt} is null
|
||||
and ${heartbeatRuns.executionStage} = 'preparing'
|
||||
and coalesce(${heartbeatRuns.runnerProfileJson}->'adapterDispatch'->>'adapterType' = 'paperclip_runner', false)
|
||||
)`;
|
||||
let startupFence = input.preserveCancellation
|
||||
? sql`coalesce(${heartbeatRuns.resultJson}->'startupCancellation', ${startupCancellation})`
|
||||
: startupCancellation;
|
||||
if (input.budgetEnforcement) {
|
||||
const intent = JSON.stringify({ companyId: input.budgetEnforcement.companyId,
|
||||
scopeType: input.budgetEnforcement.scopeType, scopeId: input.budgetEnforcement.scopeId,
|
||||
policyId: input.budgetEnforcement.enforcement?.policyId, version: input.budgetEnforcement.enforcement?.version,
|
||||
createdBefore: input.budgetEnforcement.createdBefore?.toISOString() });
|
||||
startupFence = sql`${startupFence} || jsonb_build_object('budgetEnforcement',
|
||||
coalesce(nullif(${heartbeatRuns.resultJson}->'startupCancellation'->'budgetEnforcement', 'null'::jsonb), ${intent}::jsonb))`;
|
||||
}
|
||||
return sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) ||
|
||||
${JSON.stringify(input.reassignmentStopRequested ? { reassignmentStopRequested: true } : {})}::jsonb ||
|
||||
jsonb_build_object('cancellation', ${input.preserveCancellation
|
||||
? sql`coalesce(${heartbeatRuns.resultJson}->'cancellation', ${cancellation}::jsonb)`
|
||||
: sql`${cancellation}::jsonb`}) ||
|
||||
jsonb_build_object('startupCancellation', ${startupFence})`;
|
||||
}
|
||||
|
||||
/** Run control binds cancellation, resource release, and saved-input resumption. */
|
||||
export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunControlDependencies) {
|
||||
const {
|
||||
@@ -793,7 +857,7 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont
|
||||
}
|
||||
}
|
||||
|
||||
async function listProjectScopedRunIds(companyId: string, projectId: string, createdBefore?: Date) {
|
||||
async function listProjectScopedRunIds(companyId: string, projectId: string, createdBefore?: Date, database: Db = db) {
|
||||
const runIssueId = sql<
|
||||
string | null
|
||||
>`${heartbeatRuns.contextSnapshot} ->> 'issueId'`;
|
||||
@@ -801,7 +865,7 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont
|
||||
string | null
|
||||
>`coalesce(${heartbeatRuns.contextSnapshot} ->> 'projectId', ${issues.projectId}::text)`;
|
||||
|
||||
const rows = await db
|
||||
const rows = await database
|
||||
.selectDistinctOn([heartbeatRuns.id], { id: heartbeatRuns.id })
|
||||
.from(heartbeatRuns)
|
||||
.leftJoin(
|
||||
@@ -866,69 +930,68 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont
|
||||
|
||||
async function cancelPendingWakeupsForBudgetScope(
|
||||
scope: BudgetEnforcementScope,
|
||||
tx: Db,
|
||||
) {
|
||||
return withCurrentBudgetEnforcement(db, scope, async (tx) => {
|
||||
const now = new Date();
|
||||
let wakeupIds: string[] = [];
|
||||
const now = new Date();
|
||||
let wakeupIds: string[] = [];
|
||||
|
||||
if (scope.scopeType === "company") {
|
||||
wakeupIds = await tx
|
||||
.select({ id: agentWakeupRequests.id })
|
||||
.from(agentWakeupRequests)
|
||||
.where(
|
||||
and(
|
||||
eq(agentWakeupRequests.companyId, scope.companyId),
|
||||
inArray(agentWakeupRequests.status, [
|
||||
"queued",
|
||||
"deferred_issue_execution",
|
||||
]),
|
||||
sql`${agentWakeupRequests.runId} is null`,
|
||||
),
|
||||
)
|
||||
.then((rows) => rows.map((row) => row.id));
|
||||
} else if (scope.scopeType === "agent") {
|
||||
wakeupIds = await tx
|
||||
.select({ id: agentWakeupRequests.id })
|
||||
.from(agentWakeupRequests)
|
||||
.where(
|
||||
and(
|
||||
eq(agentWakeupRequests.companyId, scope.companyId),
|
||||
eq(agentWakeupRequests.agentId, scope.scopeId),
|
||||
inArray(agentWakeupRequests.status, [
|
||||
"queued",
|
||||
"deferred_issue_execution",
|
||||
]),
|
||||
sql`${agentWakeupRequests.runId} is null`,
|
||||
),
|
||||
)
|
||||
.then((rows) => rows.map((row) => row.id));
|
||||
} else {
|
||||
wakeupIds = await listProjectScopedWakeupIds(
|
||||
scope.companyId,
|
||||
scope.scopeId,
|
||||
tx,
|
||||
);
|
||||
}
|
||||
if (scope.scopeType === "company") {
|
||||
wakeupIds = await tx
|
||||
.select({ id: agentWakeupRequests.id })
|
||||
.from(agentWakeupRequests)
|
||||
.where(
|
||||
and(
|
||||
eq(agentWakeupRequests.companyId, scope.companyId),
|
||||
inArray(agentWakeupRequests.status, [
|
||||
"queued",
|
||||
"deferred_issue_execution",
|
||||
]),
|
||||
sql`${agentWakeupRequests.runId} is null`,
|
||||
),
|
||||
)
|
||||
.then((rows) => rows.map((row) => row.id));
|
||||
} else if (scope.scopeType === "agent") {
|
||||
wakeupIds = await tx
|
||||
.select({ id: agentWakeupRequests.id })
|
||||
.from(agentWakeupRequests)
|
||||
.where(
|
||||
and(
|
||||
eq(agentWakeupRequests.companyId, scope.companyId),
|
||||
eq(agentWakeupRequests.agentId, scope.scopeId),
|
||||
inArray(agentWakeupRequests.status, [
|
||||
"queued",
|
||||
"deferred_issue_execution",
|
||||
]),
|
||||
sql`${agentWakeupRequests.runId} is null`,
|
||||
),
|
||||
)
|
||||
.then((rows) => rows.map((row) => row.id));
|
||||
} else {
|
||||
wakeupIds = await listProjectScopedWakeupIds(
|
||||
scope.companyId,
|
||||
scope.scopeId,
|
||||
tx,
|
||||
);
|
||||
}
|
||||
|
||||
if (wakeupIds.length === 0) return 0;
|
||||
if (wakeupIds.length === 0) return 0;
|
||||
|
||||
await tx
|
||||
.update(agentWakeupRequests)
|
||||
.set({
|
||||
status: "cancelled",
|
||||
finishedAt: now,
|
||||
error: "Cancelled due to budget pause",
|
||||
updatedAt: now,
|
||||
})
|
||||
.where(and(
|
||||
inArray(agentWakeupRequests.id, wakeupIds),
|
||||
inArray(agentWakeupRequests.status, ["queued", "deferred_issue_execution"]),
|
||||
isNull(agentWakeupRequests.runId),
|
||||
scope.createdBefore ? lt(agentWakeupRequests.createdAt, scope.createdBefore) : undefined,
|
||||
));
|
||||
await tx
|
||||
.update(agentWakeupRequests)
|
||||
.set({
|
||||
status: "cancelled",
|
||||
finishedAt: now,
|
||||
error: "Cancelled due to budget pause",
|
||||
updatedAt: now,
|
||||
})
|
||||
.where(and(
|
||||
inArray(agentWakeupRequests.id, wakeupIds),
|
||||
inArray(agentWakeupRequests.status, ["queued", "deferred_issue_execution"]),
|
||||
isNull(agentWakeupRequests.runId),
|
||||
scope.createdBefore ? lt(agentWakeupRequests.createdAt, scope.createdBefore) : undefined,
|
||||
));
|
||||
|
||||
return wakeupIds.length;
|
||||
});
|
||||
return wakeupIds.length;
|
||||
}
|
||||
|
||||
function cancellationTerminationGraceMs(
|
||||
@@ -946,6 +1009,81 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont
|
||||
runId: string,
|
||||
reason = "Cancelled by control plane",
|
||||
options: CancelRunOptions = {},
|
||||
) {
|
||||
return cancelRun(runId, reason, options);
|
||||
}
|
||||
|
||||
async function acknowledgeBudgetCancellation(run: typeof heartbeatRuns.$inferSelect, terminationProof?: string) {
|
||||
if (!parseObject(parseObject(run.resultJson?.startupCancellation).budgetEnforcement).scopeType) return;
|
||||
const acknowledged = await db.update(heartbeatRuns).set({ resultJson: sql`jsonb_set(${heartbeatRuns.resultJson},
|
||||
'{startupCancellation,budgetEnforcement}', ${heartbeatRuns.resultJson}->'startupCancellation'->'budgetEnforcement'
|
||||
|| ${JSON.stringify({ deliveredAt: new Date().toISOString(), ...(terminationProof ? { terminationProof } : {}) })}::jsonb)` })
|
||||
.where(and(eq(heartbeatRuns.companyId, run.companyId), eq(heartbeatRuns.id, run.id),
|
||||
pendingBudgetRunCancellationPredicate(),
|
||||
terminationProof ? and(eq(heartbeatRuns.status, run.status), eq(heartbeatRuns.nextEventSeq, run.nextEventSeq),
|
||||
sql`${heartbeatRuns.processPid} is not distinct from ${run.processPid}`,
|
||||
sql`${heartbeatRuns.processGroupId} is not distinct from ${run.processGroupId}`,
|
||||
sql`${heartbeatRuns.processStartedAt} is not distinct from ${run.processStartedAt?.toISOString() ?? null}::timestamptz`) : undefined,
|
||||
)).returning({ id: heartbeatRuns.id });
|
||||
if (terminationProof && !acknowledged.length)
|
||||
throw conflict("Execution evidence changed before acknowledging budget Stop", { runId: run.id });
|
||||
}
|
||||
|
||||
function hasRecordedStopAcknowledgement(run: typeof heartbeatRuns.$inferSelect) {
|
||||
if (parseObject(run.resultJson?.executionCancellation).state === "acknowledged") return true;
|
||||
const intent = parseObject(run.resultJson?.nativeCancellation);
|
||||
return intent.schema === "paperclip.native-cancellation.v1" && intent.runId === run.id
|
||||
&& intent.companyId === run.companyId && intent.issueId === run.nativeIssueId
|
||||
&& intent.scope === "run" && intent.reasonCode === "cancellation_run_only"
|
||||
&& intent.dispatchState === "acknowledged" && intent.dispatched === true
|
||||
&& typeof intent.intentAuditId === "string" && typeof intent.acknowledgementAuditId === "string";
|
||||
}
|
||||
|
||||
async function completedRunTerminationProof(run: typeof heartbeatRuns.$inferSelect): Promise<string | null> {
|
||||
// Logical completion can win the race with queued Stop delivery. It does
|
||||
// not prove a provider exited, and an in-process owner must still settle.
|
||||
if (!isHeartbeatRunTerminalStatus(run.status) || activeRunExecutions.has(run.id) ||
|
||||
runningProcesses.has(run.id) || adapterExecutionControls.has(run.id) || processRunCancellationSettlements.has(run.id)) return null;
|
||||
const leases = await db.select().from(environmentLeases).where(and(
|
||||
eq(environmentLeases.companyId, run.companyId), eq(environmentLeases.heartbeatRunId, run.id),
|
||||
));
|
||||
if (leases.some(lease => lease.provider !== "local"))
|
||||
return leases.every(hasRemoteTerminationReceipt) ? "provider_termination_receipt" : null;
|
||||
if (run.processPid || run.processGroupId || run.processStartedAt)
|
||||
return await recordedLocalExecutionStopped(db, run) ? "local_process_exit" : null;
|
||||
if (await readNativeLocalProcessStop(db, run.companyId, run.id)) return "local_process_stop_receipt";
|
||||
// Normal native finalization clears mutable process fields. Rebind the
|
||||
// latest server-owned launch identity before checking that exact local
|
||||
// process and group. Any newer start invalidates the older observation.
|
||||
const [identity] = await db.select({ eventType: heartbeatRunEvents.eventType, payload: heartbeatRunEvents.payload })
|
||||
.from(heartbeatRunEvents).where(and(eq(heartbeatRunEvents.companyId, run.companyId),
|
||||
eq(heartbeatRunEvents.runId, run.id), isNull(heartbeatRunEvents.sourceEventId),
|
||||
inArray(heartbeatRunEvents.eventType, [PROCESS_START_REQUESTED, PROCESS_IDENTITY_RECORDED, "native.local_process_stopped"])))
|
||||
.orderBy(desc(heartbeatRunEvents.seq)).limit(1);
|
||||
const recorded = identity?.payload;
|
||||
if (identity?.eventType !== PROCESS_IDENTITY_RECORDED || typeof recorded?.processPid !== "number" ||
|
||||
!Number.isSafeInteger(recorded.processPid) || recorded.processPid <= 1 ||
|
||||
(recorded.processGroupId !== null && (typeof recorded.processGroupId !== "number" ||
|
||||
!Number.isSafeInteger(recorded.processGroupId) || recorded.processGroupId <= 1)) ||
|
||||
typeof recorded.processStartedAt !== "string") return null;
|
||||
const processStartedAt = new Date(recorded.processStartedAt);
|
||||
if (Number.isNaN(processStartedAt.getTime())) return null;
|
||||
return await recordedLocalExecutionStopped(db, { ...run, processPid: recorded.processPid,
|
||||
processGroupId: recorded.processGroupId, processStartedAt }) ? "local_process_exit" : null;
|
||||
}
|
||||
|
||||
function deliveredBudgetCancellationResult(run: typeof heartbeatRuns.$inferSelect, deliveredAt: Date) {
|
||||
const startup = parseObject(run.resultJson?.startupCancellation);
|
||||
const intent = parseObject(startup.budgetEnforcement);
|
||||
return intent.scopeType ? { startupCancellation: { ...startup,
|
||||
budgetEnforcement: { ...intent, deliveredAt: deliveredAt.toISOString() } } } : {};
|
||||
}
|
||||
|
||||
async function cancelRun(
|
||||
runId: string,
|
||||
reason: string,
|
||||
options: CancelRunOptions,
|
||||
committedBudgetCancellation = false,
|
||||
) {
|
||||
let run = await getRun(runId);
|
||||
if (!run) throw notFound("Heartbeat run not found");
|
||||
@@ -1041,11 +1179,19 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont
|
||||
!CANCELLABLE_HEARTBEAT_RUN_STATUSES.includes(
|
||||
run.status as (typeof CANCELLABLE_HEARTBEAT_RUN_STATUSES)[number],
|
||||
)
|
||||
)
|
||||
) {
|
||||
if (committedBudgetCancellation) {
|
||||
const terminationProof = hasRecordedStopAcknowledgement(run) ? undefined : await completedRunTerminationProof(run);
|
||||
if (terminationProof === null) {
|
||||
throw conflict("Budget Stop is committed, but provider termination has not been acknowledged", { runId: run.id });
|
||||
}
|
||||
await acknowledgeBudgetCancellation(run, terminationProof);
|
||||
}
|
||||
return run;
|
||||
}
|
||||
const agent = await getAgent(run.agentId);
|
||||
const errorCode = options.errorCode ?? "cancelled";
|
||||
const cancellation = (options.expectedStatus && run.resultJson?.startupCancellation
|
||||
const cancellation = ((options.expectedStatus || committedBudgetCancellation) && run.resultJson?.startupCancellation
|
||||
? readRunCancellation(run.resultJson) : null) ?? requestedRunCancellation(options.resultJson ?? {}, reason);
|
||||
|
||||
const pendingProcessCancellation = processRunCancellationSettlements.get(
|
||||
@@ -1069,23 +1215,13 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont
|
||||
// Established legacy processes must still be stopped if the database is
|
||||
// unavailable. Only native or not-yet-dispatched preparation needs this
|
||||
// additional durable fence before its existing cancellation path.
|
||||
if (options.expectedStatus || options.budgetEnforcement || (!options.cancellationRequestId && (run.runtimeMode === "native" || (!run.runtimeModeResolvedAt && !running && !control)))) {
|
||||
if (!committedBudgetCancellation && (options.expectedStatus || options.budgetEnforcement || (!options.cancellationRequestId && (run.runtimeMode === "native" || (!run.runtimeModeResolvedAt && !running && !control))))) {
|
||||
const fence = async (tx: Db) => {
|
||||
const [fenced] = await tx.update(heartbeatRuns).set({
|
||||
// Record handoff intent before native cancellation can finalize and release
|
||||
// the run. Only its audited stop acknowledgement suppresses recovery.
|
||||
resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) ||
|
||||
${JSON.stringify(options.errorCode === "issue_reassigned" && options.resultJson?.reassignmentStopConfirmed === true
|
||||
? { reassignmentStopRequested: true } : {})}::jsonb ||
|
||||
${JSON.stringify({ cancellation })}::jsonb ||
|
||||
jsonb_build_object('startupCancellation', jsonb_build_object(
|
||||
'requestedAt', ${new Date().toISOString()}::text,
|
||||
'errorCode', ${errorCode}::text,
|
||||
'beforeNativeSelection', ${heartbeatRuns.runtimeMode} = 'legacy'
|
||||
and ${heartbeatRuns.runtimeModeResolvedAt} is null
|
||||
and ${heartbeatRuns.executionStage} = 'preparing'
|
||||
and coalesce(${heartbeatRuns.runnerProfileJson}->'adapterDispatch'->>'adapterType' = 'paperclip_runner', false)
|
||||
))`,
|
||||
resultJson: runCancellationFenceSql({ cancellation, errorCode,
|
||||
reassignmentStopRequested: options.errorCode === "issue_reassigned" && options.resultJson?.reassignmentStopConfirmed === true }),
|
||||
}).where(and(eq(heartbeatRuns.id, runId), inArray(heartbeatRuns.status,
|
||||
options.expectedStatus ? [options.expectedStatus] : pendingNativeRetry ? [...CANCELLABLE_HEARTBEAT_RUN_STATUSES, "failed"] : [...CANCELLABLE_HEARTBEAT_RUN_STATUSES],
|
||||
))).returning();
|
||||
@@ -1229,6 +1365,7 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont
|
||||
}
|
||||
// The owned adapter already finalized this run and its lifecycle.
|
||||
// Do not replay the process cancellation side effects below.
|
||||
await acknowledgeBudgetCancellation(stopped);
|
||||
return { run: stopped, updated: false };
|
||||
}
|
||||
}
|
||||
@@ -1271,6 +1408,10 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont
|
||||
errorCode, errorMessage: reason,
|
||||
})
|
||||
: {}),
|
||||
// Commit this receipt with the terminal transition. A
|
||||
// crash after provider Stop cannot strand an undispatched
|
||||
// queued run whose normal path needs no process receipt.
|
||||
...deliveredBudgetCancellationResult(run, finishedAt),
|
||||
// The native cancellation helper may have advanced a durable
|
||||
// pending intent to its acknowledged state after `run` was
|
||||
// first read. Never let that stale snapshot overwrite the
|
||||
@@ -1428,17 +1569,86 @@ export function createHeartbeatRunControl(db: Db, dependencies: HeartbeatRunCont
|
||||
};
|
||||
}
|
||||
|
||||
async function stopBrowserSessionsForBudgetScope(tx: Db, scope: BudgetEnforcementScope) {
|
||||
// Browser work can outlive its terminal parent run. Its admitted project
|
||||
// snapshot remains authoritative even if the task moves to another project.
|
||||
const activeHold = sql`(select ${browserUseRuns.result}->'paperclipBudgetHold'
|
||||
from ${browserUseRuns} where ${browserUseRuns.companyId} = ${browserUseSessions.companyId}
|
||||
and ${browserUseRuns.sessionId} = ${browserUseSessions.id}
|
||||
and ${browserUseRuns.result}->'paperclipBudgetHold'->>'state' = 'held'
|
||||
and ${browserUseRuns.result}->'paperclipBudgetHold'->>'invocationId' = ${browserUseRuns.invocationId}::text
|
||||
order by ${browserUseRuns.createdAt} desc, ${browserUseRuns.id} desc limit 1)`;
|
||||
const project = sql`case when ${activeHold} ? 'projectId' then ${activeHold}->>'projectId'
|
||||
else (select ${issues.projectId}::text from ${issues}
|
||||
where ${issues.id} = ${browserUseSessions.issueId} and ${issues.companyId} = ${scope.companyId}) end`;
|
||||
const sessions = await tx.select({ id: browserUseSessions.id }).from(browserUseSessions).where(and(
|
||||
eq(browserUseSessions.companyId, scope.companyId),
|
||||
inArray(browserUseSessions.status, ["starting", "running", "idle", "stopping"]),
|
||||
scope.scopeType === "agent" ? eq(browserUseSessions.agentId, scope.scopeId) : undefined,
|
||||
scope.scopeType === "project" ? sql`${project} = ${scope.scopeId}` : undefined,
|
||||
scope.createdBefore ? lt(browserUseSessions.createdAt, scope.createdBefore) : undefined,
|
||||
)).for("update");
|
||||
if (!sessions.length) return;
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.browser, scope.companyId);
|
||||
await tx.update(browserUseSessions).set({ status: "stopping", stopRequested: "end",
|
||||
nextPollAt: new Date(), updatedAt: new Date(),
|
||||
}).where(and(eq(browserUseSessions.companyId, scope.companyId),
|
||||
inArray(browserUseSessions.id, sessions.map(session => session.id))));
|
||||
}
|
||||
|
||||
async function cancelBudgetScopeWork(scope: BudgetEnforcementScope) {
|
||||
const runIds = scope.scopeType === "project"
|
||||
? await listProjectScopedRunIds(scope.companyId, scope.scopeId, scope.createdBefore)
|
||||
: await db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(and(
|
||||
eq(heartbeatRuns.companyId, scope.companyId),
|
||||
scope.scopeType === "agent" ? eq(heartbeatRuns.agentId, scope.scopeId) : undefined,
|
||||
scope.createdBefore ? lt(heartbeatRuns.createdAt, scope.createdBefore) : undefined,
|
||||
const reason = "Cancelled due to budget pause";
|
||||
const newlyFenced = await withCurrentBudgetEnforcement(db, scope, async (tx) => {
|
||||
// Queue admission and every selected run lose authority in this commit,
|
||||
// before any provider can delay cancellation delivery.
|
||||
await cancelPendingWakeupsForBudgetScope(scope, tx);
|
||||
await stopBrowserSessionsForBudgetScope(tx, scope);
|
||||
const candidates = scope.scopeType === "project"
|
||||
? await listProjectScopedRunIds(scope.companyId, scope.scopeId, scope.createdBefore, tx)
|
||||
: await tx.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(and(
|
||||
eq(heartbeatRuns.companyId, scope.companyId),
|
||||
scope.scopeType === "agent" ? eq(heartbeatRuns.agentId, scope.scopeId) : undefined,
|
||||
scope.createdBefore ? lt(heartbeatRuns.createdAt, scope.createdBefore) : undefined,
|
||||
inArray(heartbeatRuns.status, [...CANCELLABLE_HEARTBEAT_RUN_STATUSES]),
|
||||
)).then((rows) => rows.map((row) => row.id));
|
||||
if (!candidates.length) return [];
|
||||
return tx.update(heartbeatRuns).set({ resultJson: runCancellationFenceSql({
|
||||
cancellation: requestedRunCancellation({}, reason), errorCode: "cancelled", preserveCancellation: true, budgetEnforcement: scope,
|
||||
}) }).where(and(eq(heartbeatRuns.companyId, scope.companyId), inArray(heartbeatRuns.id, candidates),
|
||||
inArray(heartbeatRuns.status, [...CANCELLABLE_HEARTBEAT_RUN_STATUSES]),
|
||||
)).then((rows) => rows.map((row) => row.id));
|
||||
for (const runId of runIds) await cancelRunInternal(runId, "Cancelled due to budget pause", { budgetEnforcement: scope });
|
||||
await cancelPendingWakeupsForBudgetScope(scope);
|
||||
scope.createdBefore ? lt(heartbeatRuns.createdAt, scope.createdBefore) : undefined,
|
||||
)).returning({ id: heartbeatRuns.id }).then((rows) => rows.map((row) => row.id));
|
||||
});
|
||||
// A later policy change revokes permission to fence new work, not the Stop
|
||||
// already committed for older work. Recovery must deliver those intents.
|
||||
const pending = await db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(and(
|
||||
pendingBudgetRunCancellationPredicate(), budgetCancellationScopePredicate(scope),
|
||||
));
|
||||
const runIds = [...new Set([...(newlyFenced ?? []), ...pending.map(run => run.id)])];
|
||||
if (!runIds.length) return;
|
||||
const failures: unknown[] = [];
|
||||
let nextRun = 0;
|
||||
await Promise.all(Array.from({ length: Math.min(4, runIds.length) }, async () => {
|
||||
while (nextRun < runIds.length) {
|
||||
const runId = runIds[nextRun++]!;
|
||||
try {
|
||||
// The committed intent owns these old runs even if a later budget
|
||||
// edit admits new work. Do not revoke or repeat that earlier fence.
|
||||
await withBudgetStopDelivery(db, async () => {
|
||||
const [attempted] = await db.update(heartbeatRuns).set({ resultJson: sql`jsonb_set(${heartbeatRuns.resultJson},
|
||||
'{startupCancellation,budgetEnforcement,lastAttemptAt}', ${JSON.stringify(new Date().toISOString())}::jsonb)` })
|
||||
.where(and(eq(heartbeatRuns.companyId, scope.companyId), eq(heartbeatRuns.id, runId),
|
||||
pendingBudgetRunCancellationPredicate())).returning({ id: heartbeatRuns.id });
|
||||
if (attempted) await cancelRun(runId, reason, {}, true);
|
||||
});
|
||||
} catch (error) {
|
||||
failures.push(error);
|
||||
}
|
||||
}
|
||||
}));
|
||||
// Drain the batch before rejecting; one failed provider must not prevent
|
||||
// the remaining runs from receiving their already-committed Stop.
|
||||
if (failures.length) throw new AggregateError(failures, "Budget cancellation delivery remains pending");
|
||||
}
|
||||
|
||||
return {
|
||||
|
||||
@@ -9,6 +9,7 @@ import { getHeartbeatRunRuntimeStatus } from "../heartbeat-run-runtime-status.js
|
||||
import { createHeartbeatRunState } from "./run-state.js";
|
||||
import { createHeartbeatRunPreparation } from "./run-preparation.js";
|
||||
import { createHeartbeatLifecycle, persistHeartbeatRunProcessMetadata, type HeartbeatLifecycleDependencies } from "./run-lifecycle.js";
|
||||
import { DELIVERY_QUEUES, subscribeDeliveryWork } from "../delivery-work-notifications.js";
|
||||
|
||||
vi.mock("../live-events.js", () => ({ publishLiveEvent: vi.fn() }));
|
||||
|
||||
@@ -161,6 +162,43 @@ describe.skipIf(!support.supported)("heartbeat lifecycle database wiring", () =>
|
||||
expect(persisted.usageJson).toEqual({ accountingReceiptReady: false, usageSource: "per_run", inputTokens: 8, outputTokens: 3 });
|
||||
});
|
||||
|
||||
it.each(["setRunStatus", "setRunStatusIfRunning"] as const)("signals committed terminal accounting before %s publishes a failed live event", async operation => {
|
||||
const { run } = await fixture({ costAccountingPending: true });
|
||||
const settled = vi.fn();
|
||||
const unsubscribe = subscribeDeliveryWork(db, DELIVERY_QUEUES.accounting, settled);
|
||||
vi.mocked(publishLiveEvent).mockImplementationOnce(() => { throw new Error("Live publication unavailable"); });
|
||||
try {
|
||||
await expect(createHeartbeatLifecycle(db, callbacks(db))[operation](run.id, "succeeded", { finishedAt: new Date() }))
|
||||
.rejects.toThrow("Live publication unavailable");
|
||||
const [persisted] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, run.id));
|
||||
expect(persisted.status).toBe("succeeded");
|
||||
expect(persisted.costAccountingPending).toBe(true);
|
||||
expect(settled).toHaveBeenCalledOnce();
|
||||
} finally { unsubscribe(); }
|
||||
});
|
||||
|
||||
it("does not request accounting for an unchanged running status", async () => {
|
||||
const { run } = await fixture();
|
||||
const settled = vi.fn();
|
||||
const unsubscribe = subscribeDeliveryWork(db, DELIVERY_QUEUES.accounting, settled);
|
||||
try {
|
||||
await createHeartbeatLifecycle(db, callbacks(db)).setRunStatus(run.id, "running");
|
||||
expect(settled).not.toHaveBeenCalled();
|
||||
} finally { unsubscribe(); }
|
||||
});
|
||||
|
||||
it("signals legacy terminal accounting after the nested terminalizer transaction commits", async () => {
|
||||
const { run } = await fixture({ runtimeMode: "legacy", costAccountingPending: true });
|
||||
const settled = vi.fn();
|
||||
const unsubscribe = subscribeDeliveryWork(db, DELIVERY_QUEUES.accounting, settled);
|
||||
try {
|
||||
const result = await createHeartbeatLifecycle(db, callbacks(db)).setRunStatusIfRunning(run.id, "failed", { finishedAt: new Date() });
|
||||
expect(result.updated).toBe(true);
|
||||
expect(result.run?.status).toBe("failed");
|
||||
expect(settled).toHaveBeenCalledOnce();
|
||||
} finally { unsubscribe(); }
|
||||
});
|
||||
|
||||
it.each(["claimed", "cancelled"] as const)("settles a %s wake without reviving a revoked request", async (status) => {
|
||||
const { wake } = await fixture();
|
||||
await db.update(agentWakeupRequests).set({ status }).where(eq(agentWakeupRequests.id, wake.id));
|
||||
|
||||
@@ -26,7 +26,6 @@ import {
|
||||
} from "../native-runtime/native-cancellation-request.js";
|
||||
import { cancellationResultJson } from "../run-cancellation.js";
|
||||
import { CONTROLLER_LEASE_LOST_FAILURE_CODE } from "../legacy-controller-lease.js";
|
||||
import { accountRunCost } from "../run-cost-accounting.js";
|
||||
import { isWaitingConversation } from "../agent-conversations.js";
|
||||
import { PROCESS_IDENTITY_RECORDED } from "../native-local-process-stop.js";
|
||||
import {
|
||||
@@ -39,6 +38,7 @@ import {
|
||||
isConversationAdapter,
|
||||
} from "../conversation-continuation.js";
|
||||
import { terminalizeLegacyExecution } from "../legacy-execution-recovery.js";
|
||||
import { DELIVERY_QUEUES, notifyDeliveryWork } from "../delivery-work-notifications.js";
|
||||
import { buildHeartbeatRunStatusLiveEventPayload } from "../heartbeat-run-status-payload.js";
|
||||
import { randomUUID } from "node:crypto";
|
||||
import {
|
||||
@@ -1010,6 +1010,17 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl
|
||||
return recordHeartbeatRunRuntimeProgress(currentRun, update, issueId);
|
||||
}
|
||||
|
||||
async function writeRunStatus<T>(status: string, companyId: string | undefined, write: (tx: Db) => Promise<T>) {
|
||||
if (!isHeartbeatRunTerminalStatus(status)) return write(db);
|
||||
// Register on the outer transaction before the terminal write. An unknown
|
||||
// commit cannot strand newly-actionable accounting after the live event.
|
||||
return db.transaction(async transaction => {
|
||||
const tx = transaction as unknown as Db;
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, companyId);
|
||||
return write(tx);
|
||||
});
|
||||
}
|
||||
|
||||
async function setRunStatus(
|
||||
runId: string,
|
||||
status: string,
|
||||
@@ -1050,16 +1061,16 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl
|
||||
},
|
||||
};
|
||||
}
|
||||
const updated =
|
||||
const updated = await writeRunStatus(status, previousStatus?.companyId, writer =>
|
||||
previousStatus && previousStatus.runtimeMode === "legacy" && isHeartbeatRunTerminalStatus(status)
|
||||
? await terminalizeLegacyExecution({
|
||||
db,
|
||||
? terminalizeLegacyExecution({
|
||||
db: writer,
|
||||
run: previousStatus,
|
||||
status,
|
||||
patch,
|
||||
reconcileIfNeeded: true,
|
||||
})
|
||||
: await db
|
||||
: writer
|
||||
.update(heartbeatRuns)
|
||||
.set({
|
||||
status,
|
||||
@@ -1069,7 +1080,7 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl
|
||||
})
|
||||
.where(eq(heartbeatRuns.id, runId))
|
||||
.returning()
|
||||
.then((rows) => rows[0] ?? null);
|
||||
.then((rows) => rows[0] ?? null));
|
||||
|
||||
if (updated) {
|
||||
publishLiveEvent({
|
||||
@@ -1146,10 +1157,10 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl
|
||||
},
|
||||
};
|
||||
}
|
||||
const updated =
|
||||
const updated = await writeRunStatus(status, previousStatus?.companyId, writer =>
|
||||
previousStatus && previousStatus.runtimeMode === "legacy" && isHeartbeatRunTerminalStatus(status)
|
||||
? await terminalizeLegacyExecution({
|
||||
db,
|
||||
? terminalizeLegacyExecution({
|
||||
db: writer,
|
||||
run: previousStatus,
|
||||
status,
|
||||
patch,
|
||||
@@ -1157,7 +1168,7 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl
|
||||
reconcileIfNeeded: true,
|
||||
writeConditions: [nativeRunnerOwnershipNotHeldCondition(), ...(cancellationCondition ? [cancellationCondition] : [])],
|
||||
})
|
||||
: await db
|
||||
: writer
|
||||
.update(heartbeatRuns)
|
||||
.set({
|
||||
status,
|
||||
@@ -1174,7 +1185,7 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl
|
||||
// issue promotion, or release of the unfinished source lease.
|
||||
...(status === "failed" && patch?.errorCode === "process_lost"
|
||||
? [
|
||||
not(pendingNativeWorkspaceFinalizationCondition(db)),
|
||||
not(pendingNativeWorkspaceFinalizationCondition(writer)),
|
||||
// Result publication locks and updates this heartbeat in
|
||||
// its transaction. If this UPDATE waited behind it, the
|
||||
// joined tables can still use the older statement snapshot;
|
||||
@@ -1193,7 +1204,7 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl
|
||||
.catch((error: unknown) => {
|
||||
if (cancellationCondition) rethrowNativeCancellationLockConflict(error);
|
||||
throw error;
|
||||
});
|
||||
}));
|
||||
|
||||
if (updated) {
|
||||
publishLiveEvent({
|
||||
@@ -2975,7 +2986,6 @@ export function createHeartbeatLifecycle(db: Db, dependencies: HeartbeatLifecycl
|
||||
normalizedUsage?: UsageTotals | null,
|
||||
) {
|
||||
await ensureRuntimeState(agent);
|
||||
await accountRunCost(db, run.id, budgetHooks);
|
||||
await db
|
||||
.update(agentRuntimeState)
|
||||
.set({
|
||||
|
||||
@@ -6,6 +6,8 @@ import { connectionDeliveryWaitingForResponse, finishedIdleRun as finishedRun, f
|
||||
import type { Db } from "@paperclipai/db";
|
||||
import { sql } from "drizzle-orm";
|
||||
import { verifiedResourceCleanup } from "./cleanup-verification-proof.js";
|
||||
import { companyHoldMismatch } from "../modules/agent-lifecycle/index.js";
|
||||
import { pendingBudgetRunCancellationPredicate } from "./heartbeat/budget-cancellation.js";
|
||||
|
||||
export interface IdleSleepDrainStatus {
|
||||
ownerId?: string;
|
||||
@@ -52,6 +54,11 @@ const WORK_CHECKS = [
|
||||
OR process_stopped_at IS NULL OR next_attempt_at IS NOT NULL
|
||||
OR receipt ? 'baseline' OR receipt->>'cleanupPending' = 'true'`,
|
||||
`SELECT 1 FROM budget_reservations WHERE state NOT IN ('settled', 'released')`,
|
||||
`SELECT 1 FROM budget_policies WHERE enforcement_version > enforcement_delivered_version`,
|
||||
`SELECT 1 FROM agents WHERE lifecycle_state IN
|
||||
('preparing', 'verifying', 'pausing', 'resuming', 'terminating', 'cleaning_up')`,
|
||||
`SELECT 1 FROM decision_invocations WHERE status = 'running'`,
|
||||
`SELECT 1 FROM prompt_injection_invocations WHERE status = 'running'`,
|
||||
// Persistent machine credentials can receive work without a human opening
|
||||
// the instance. They need an ingress wake contract before we can sleep.
|
||||
`SELECT 1 FROM agent_api_keys WHERE revoked_at IS NULL`,
|
||||
@@ -275,6 +282,9 @@ export async function readIdleSleepSafety(
|
||||
(!scheduledRoutines || query !== `SELECT 1 FROM routines WHERE status NOT IN ('paused', 'archived')`)
|
||||
&& (!scheduledIssueMonitors || query !== `SELECT 1 FROM issues WHERE status NOT IN ('done', 'cancelled')`))
|
||||
.map((query) => sql`EXISTS (${sql.raw(query)})`);
|
||||
checks.push(sql`EXISTS (SELECT 1 FROM agents JOIN companies ON companies.id = agents.company_id
|
||||
WHERE ${companyHoldMismatch()})`);
|
||||
checks.push(sql`EXISTS (SELECT 1 FROM heartbeat_runs WHERE ${pendingBudgetRunCancellationPredicate()})`);
|
||||
// Only the current worker's owned runtime drain can exempt a plugin.
|
||||
// Missing workers, old SDKs, and newly installed plugins still block.
|
||||
checks.push(quietPlugins.length === 0
|
||||
|
||||
@@ -2,6 +2,7 @@ import { recoverLegacyUnsafeWorkspaceExports } from "./native-workspace-export-r
|
||||
import { hasPendingNativeChildCompletion } from "./native-child-completion-delivery.js";
|
||||
import { dismissAutomaticCompletionReviews, decisionHasRetiredAutomaticReview } from "./automatic-completion-reviews.js";
|
||||
import { logger } from "../../middleware/logger.js";
|
||||
import { DELIVERY_QUEUES, notifyDeliveryWork } from "../delivery-work-notifications.js";
|
||||
import { createHash, randomUUID } from "node:crypto";
|
||||
import { and, asc, desc, eq, gt, inArray, isNotNull, isNull, lt, lte, ne, not, notInArray, or, sql, type SQL } from "drizzle-orm";
|
||||
import type { Db } from "@paperclipai/db";
|
||||
@@ -489,6 +490,7 @@ export async function claimNativeSessionResumptions(input: {
|
||||
|| (!!checkpoint && !checkpointHasProviderIdentity)
|
||||
) {
|
||||
const failureCode = originalFailureCode;
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, row.run.companyId);
|
||||
await tx.update(nativeRunFinalizations).set({
|
||||
phase: "terminal_failure",
|
||||
leaseOwner: null,
|
||||
|
||||
@@ -62,6 +62,7 @@ import { waitForPendingRunFailureReports } from "../run-failure-report.js";
|
||||
import { deliverExecutionStatuses } from "../execution-status-delivery.js";
|
||||
import { commitNativeStatusDecision } from "./status-decision-committer.js";
|
||||
import { NATIVE_STATUS_ARBITER_POLICY_VERSION, type NativeStatusDecision } from "./status-arbiter.js";
|
||||
import { DELIVERY_QUEUES, subscribeDeliveryWork } from "../delivery-work-notifications.js";
|
||||
|
||||
const embeddedPostgresSupport = await getEmbeddedPostgresTestSupport();
|
||||
const describeEmbeddedPostgres = embeddedPostgresSupport.supported ? describe : describe.skip;
|
||||
@@ -946,6 +947,26 @@ describeEmbeddedPostgres("native run finalizer / status decision committer — a
|
||||
expect(alreadyEmittedByCommittedDecision).toBe(true);
|
||||
expect(agentTaskRunCalls(callsBefore)).toHaveLength(1);
|
||||
});
|
||||
|
||||
it("signals durable accounting when native commitment or permanent failure makes a receipt actionable", async () => {
|
||||
const settled = vi.fn();
|
||||
const unsubscribe = subscribeDeliveryWork(db, DELIVERY_QUEUES.accounting, settled);
|
||||
try {
|
||||
const committed = await seedNativeRun();
|
||||
await driveToCompleteResult(committed);
|
||||
const beforeCommit = settled.mock.calls.length;
|
||||
await finalizeNativeRun({ db, runId: committed.runId, workspaceFinalizeStatus: "succeeded", projectRunStatus: true });
|
||||
expect((await db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, committed.runId)))[0].phase).toBe("committed");
|
||||
expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, committed.runId)))[0].status).toBe("succeeded");
|
||||
expect(settled.mock.calls.length).toBeGreaterThan(beforeCommit);
|
||||
const failed = await seedNativeRun();
|
||||
await driveToCompleteResult(failed);
|
||||
const beforeFailure = settled.mock.calls.length;
|
||||
await recordNativeFinalizationFailure({ db, runId: failed.runId, error: new Error("native_fixture_permanent_failure"), permanent: true });
|
||||
expect((await db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, failed.runId)))[0].phase).toBe("terminal_failure");
|
||||
expect(settled.mock.calls.length).toBeGreaterThan(beforeFailure);
|
||||
} finally { unsubscribe(); }
|
||||
});
|
||||
});
|
||||
|
||||
|
||||
|
||||
@@ -64,6 +64,7 @@ import {
|
||||
restoreNativeChatReviewPresentationInTransaction,
|
||||
} from "./native-chat-review-presentation.js";
|
||||
import { logger } from "../../middleware/logger.js";
|
||||
import { DELIVERY_QUEUES, notifyDeliveryWork } from "../delivery-work-notifications.js";
|
||||
import {
|
||||
CHAT_RUN_PRESENTATION_AUTHORIZATION_REASON,
|
||||
findHeartbeatRunCompletionComment,
|
||||
@@ -418,6 +419,7 @@ async function recordRetryableFailure(input: {
|
||||
? "native_workspace_sync_out_retry_exhausted"
|
||||
: "native_finalization_retry_exhausted"
|
||||
: input.failureCode;
|
||||
if (phase === "terminal_failure") await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, input.run.companyId);
|
||||
await tx
|
||||
.update(nativeRunFinalizations)
|
||||
.set({
|
||||
@@ -660,69 +662,74 @@ async function projectCommittedRun(input: {
|
||||
const projectedStatus = projectNativeTerminalRunStatus(
|
||||
terminalState as "succeeded" | "failed" | "cancelled",
|
||||
);
|
||||
const now = new Date();
|
||||
const [updatedRun] = await input.db
|
||||
.update(heartbeatRuns)
|
||||
.set({
|
||||
executionStatusDeliveryId: randomUUID(),
|
||||
status: projectedStatus,
|
||||
finishedAt: input.run.finishedAt ?? now,
|
||||
nativePhase: "committed",
|
||||
nativePhaseUpdatedAt: now,
|
||||
...(terminalState === "succeeded"
|
||||
? {
|
||||
error: null,
|
||||
errorCode: null,
|
||||
// Capture the row being updated, not the earlier admission read:
|
||||
// cleanup may have recorded a new diagnostic in the meantime.
|
||||
// Other result metadata and all physical-owner evidence stay put.
|
||||
resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb)
|
||||
|| ${recoveredExecutionFailureMetadata()}
|
||||
|| jsonb_build_object('finalizationPhase', 'committed', 'failureCode', null, 'originalFailureCode', null, 'nextAttemptAt', null)`,
|
||||
}
|
||||
: { resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) || jsonb_build_object('finalizationPhase', 'committed', 'failureCode', null, 'originalFailureCode', null, 'nextAttemptAt', null)` }),
|
||||
updatedAt: now,
|
||||
})
|
||||
.where(
|
||||
const needsRepair = and(
|
||||
eq(heartbeatRuns.id, input.run.id),
|
||||
eq(heartbeatRuns.runtimeMode, "native"),
|
||||
or(
|
||||
inArray(heartbeatRuns.status, ["queued", "running", "failed"]),
|
||||
and(
|
||||
eq(heartbeatRuns.id, input.run.id),
|
||||
eq(heartbeatRuns.runtimeMode, "native"),
|
||||
eq(heartbeatRuns.status, "succeeded"),
|
||||
or(
|
||||
inArray(heartbeatRuns.status, ["queued", "running", "failed"]),
|
||||
and(
|
||||
eq(heartbeatRuns.status, "succeeded"),
|
||||
or(
|
||||
isNotNull(heartbeatRuns.error),
|
||||
isNotNull(heartbeatRuns.errorCode),
|
||||
isNull(heartbeatRuns.finishedAt),
|
||||
sql`${heartbeatRuns.nativePhase} is distinct from 'committed'`,
|
||||
sql`${heartbeatRuns.resultJson}->>'finalizationPhase' is distinct from 'committed'`,
|
||||
sql`${heartbeatRuns.resultJson}->>'nextAttemptAt' is not null`,
|
||||
sql`${heartbeatRuns.resultJson}->>'failureCode' is not null`,
|
||||
sql`${heartbeatRuns.resultJson}->>'originalFailureCode' is not null`,
|
||||
),
|
||||
),
|
||||
),
|
||||
// Reconciliation revisits committed results periodically. Only repair
|
||||
// a changed projection; rewriting an unchanged failed run would mint a
|
||||
// fresh status delivery (and failure toast) on every sweep. Check the
|
||||
// current row so concurrent replays cannot both queue the same repair.
|
||||
or(
|
||||
sql`${heartbeatRuns.status} is distinct from ${projectedStatus}`,
|
||||
isNotNull(heartbeatRuns.error),
|
||||
isNotNull(heartbeatRuns.errorCode),
|
||||
isNull(heartbeatRuns.finishedAt),
|
||||
sql`${heartbeatRuns.nativePhase} is distinct from 'committed'`,
|
||||
sql`${heartbeatRuns.resultJson}->>'finalizationPhase' is distinct from 'committed'`,
|
||||
sql`${heartbeatRuns.resultJson}->>'nextAttemptAt' is not null`,
|
||||
sql`${heartbeatRuns.resultJson}->>'failureCode' is not null`,
|
||||
sql`${heartbeatRuns.resultJson}->>'originalFailureCode' is not null`,
|
||||
...(terminalState === "succeeded"
|
||||
? [isNotNull(heartbeatRuns.error), isNotNull(heartbeatRuns.errorCode)]
|
||||
: []),
|
||||
),
|
||||
nativeRunnerOwnershipNotHeldCondition(),
|
||||
),
|
||||
)
|
||||
.returning();
|
||||
),
|
||||
// Reconciliation revisits committed results periodically. Only repair
|
||||
// a changed projection; rewriting an unchanged failed run would mint a
|
||||
// fresh status delivery (and failure toast) on every sweep. Check the
|
||||
// current row so concurrent replays cannot both queue the same repair.
|
||||
or(
|
||||
sql`${heartbeatRuns.status} is distinct from ${projectedStatus}`,
|
||||
isNull(heartbeatRuns.finishedAt),
|
||||
sql`${heartbeatRuns.nativePhase} is distinct from 'committed'`,
|
||||
sql`${heartbeatRuns.resultJson}->>'finalizationPhase' is distinct from 'committed'`,
|
||||
sql`${heartbeatRuns.resultJson}->>'nextAttemptAt' is not null`,
|
||||
sql`${heartbeatRuns.resultJson}->>'failureCode' is not null`,
|
||||
sql`${heartbeatRuns.resultJson}->>'originalFailureCode' is not null`,
|
||||
...(terminalState === "succeeded"
|
||||
? [isNotNull(heartbeatRuns.error), isNotNull(heartbeatRuns.errorCode)]
|
||||
: []),
|
||||
),
|
||||
nativeRunnerOwnershipNotHeldCondition(),
|
||||
);
|
||||
const [candidate] = await input.db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(needsRepair).limit(1);
|
||||
if (!candidate) return;
|
||||
const now = new Date();
|
||||
const [updatedRun] = await input.db.transaction(async transaction => {
|
||||
const tx = transaction as unknown as Db;
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, input.run.companyId);
|
||||
return tx
|
||||
.update(heartbeatRuns)
|
||||
.set({
|
||||
executionStatusDeliveryId: randomUUID(),
|
||||
status: projectedStatus,
|
||||
finishedAt: input.run.finishedAt ?? now,
|
||||
nativePhase: "committed",
|
||||
nativePhaseUpdatedAt: now,
|
||||
...(terminalState === "succeeded"
|
||||
? {
|
||||
error: null,
|
||||
errorCode: null,
|
||||
// Capture the row being updated, not the earlier admission read:
|
||||
// cleanup may have recorded a new diagnostic in the meantime.
|
||||
// Other result metadata and all physical-owner evidence stay put.
|
||||
resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb)
|
||||
|| ${recoveredExecutionFailureMetadata()}
|
||||
|| jsonb_build_object('finalizationPhase', 'committed', 'failureCode', null, 'originalFailureCode', null, 'nextAttemptAt', null)`,
|
||||
}
|
||||
: { resultJson: sql`coalesce(${heartbeatRuns.resultJson}, '{}'::jsonb) || jsonb_build_object('finalizationPhase', 'committed', 'failureCode', null, 'originalFailureCode', null, 'nextAttemptAt', null)` }),
|
||||
updatedAt: now,
|
||||
})
|
||||
.where(needsRepair)
|
||||
.returning();
|
||||
});
|
||||
// Metadata repairs can preserve the terminal status. Only a genuine status
|
||||
// transition should emit another terminal event. A failed run that leaves a
|
||||
// pending retry still needs its first failure report.
|
||||
@@ -1579,7 +1586,7 @@ export async function finalizeNativeRun(input: {
|
||||
statusVersionAfter: Number(committed.issue.statusVersion),
|
||||
workspaceFinalizeStatus: input.workspaceFinalizeStatus,
|
||||
};
|
||||
const [updatedRun] = await input.db
|
||||
const projectRun = (tx: Db) => tx
|
||||
.update(heartbeatRuns)
|
||||
.set({
|
||||
executionStatusDeliveryId: randomUUID(),
|
||||
@@ -1610,6 +1617,13 @@ export async function finalizeNativeRun(input: {
|
||||
})
|
||||
.where(eq(heartbeatRuns.id, run.id))
|
||||
.returning();
|
||||
const [updatedRun] = input.projectRunStatus
|
||||
? await input.db.transaction(async transaction => {
|
||||
const tx = transaction as unknown as Db;
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, run.companyId);
|
||||
return projectRun(tx);
|
||||
})
|
||||
: await projectRun(input.db);
|
||||
if (
|
||||
input.projectRunStatus &&
|
||||
!alreadyEmittedByCommittedDecision &&
|
||||
|
||||
@@ -79,12 +79,22 @@ import * as nativeSessionResume from "./native-session-resume.js";
|
||||
import { buildNativeHeartbeatPreparationSpans } from "./native-run-trace.js";
|
||||
import { NativeRunnerOwnershipUnverifiedError } from "./native-runner-ownership.js";
|
||||
import type { AdapterRuntimeEvent } from "../../adapters/index.js";
|
||||
import type { AdapterUsageCheckpoint } from "@paperclipai/adapter-utils";
|
||||
|
||||
vi.mock("node:fs", async (importOriginal) => {
|
||||
const actual = await importOriginal<typeof import("node:fs")>();
|
||||
return { ...actual, readFileSync: vi.fn(actual.readFileSync) };
|
||||
});
|
||||
|
||||
const usageRecorder = vi.hoisted(() => ({
|
||||
create: vi.fn<typeof import("../usage-receipts.js")["createRunUsageRecorder"]>(),
|
||||
capture: vi.fn<(receipt: AdapterUsageCheckpoint) => Promise<AdapterUsageCheckpoint>>(),
|
||||
}));
|
||||
vi.mock("../usage-receipts.js", async (importOriginal) => ({
|
||||
...await importOriginal<typeof import("../usage-receipts.js")>(),
|
||||
createRunUsageRecorder: usageRecorder.create,
|
||||
}));
|
||||
|
||||
const githubAccess = vi.hoisted(() => ({
|
||||
activate: vi.fn((_binding: { runId: string }) => vi.fn()),
|
||||
stop: vi.fn(async () => undefined),
|
||||
@@ -376,6 +386,9 @@ import {
|
||||
} from "./native-session-executor.js";
|
||||
|
||||
beforeEach(() => {
|
||||
usageRecorder.capture.mockReset().mockImplementation(async receipt => receipt);
|
||||
usageRecorder.create.mockReset().mockResolvedValue({ capture: usageRecorder.capture,
|
||||
complete: vi.fn(async () => ({ complete: true })), persistFailure: vi.fn(async () => undefined) });
|
||||
state.bundledPackManifestPath = null;
|
||||
state.bundledRunnerBinary = null;
|
||||
state.createAssignedMcpTools.mockReset();
|
||||
@@ -6183,21 +6196,42 @@ describe("native terminal-turn accounting", () => {
|
||||
});
|
||||
await expect(executePaperclipNativeSession({ db: leaseDb(), execution: { ...execution, provider: { kind: "codex", model: "gpt-6-astra" } } as NativeExecutionInput, runnerInstanceId: "runner", onUsage })).rejects.toBe(failure);
|
||||
expect(onUsage).toHaveBeenLastCalledWith(expect.objectContaining({ complete: true }));
|
||||
expect(usageRecorder.create).not.toHaveBeenCalled();
|
||||
});
|
||||
it("clears provisional unpriced status when direct database persistence receives complete usage", async () => {
|
||||
it("uses the canonical recorder for partial and completed usage when no onUsage callback is supplied", async () => {
|
||||
const updates: Array<{ table: unknown; values: Record<string, unknown> }> = [];
|
||||
state.execute.mockReset().mockImplementationOnce(async () => {
|
||||
await accountingEvents.committed!({ eventType: "item.completed", turnId: "turn", emittedAt: new Date().toISOString(), payload: { kind: "usage", usage: { runDelta: { inputTokens: 10, outputTokens: 2, providerCostUsd: 0.25 } } } } as unknown as PrpEvent);
|
||||
return { result: { summary: "cancelled" }, terminal: { runTerminalState: "cancelled" },
|
||||
turnId: "turn", normalizedSessionId: "session", providerSessionId: null, driverKind: "test", driverVersion: "1", nativeEventCount: 1, highestContiguousSourceSeq: 1, usage: null };
|
||||
});
|
||||
await executePaperclipNativeSession({ db: leaseDb(execution, {}, {}, updates), execution, runnerInstanceId: "runner" });
|
||||
const dialect = new PgDialect();
|
||||
const writes = updates.filter(update => update.values.usageJson).map(update =>
|
||||
JSON.parse(dialect.sqlToQuery(update.values.usageJson as SQL).params[0] as string) as Record<string, unknown>);
|
||||
expect(writes[0]).toMatchObject({ accountingReceiptReady: false, costStatus: "unpriced" });
|
||||
// Apply the same JSON merge as the database: omitted keys retain the old value.
|
||||
expect(Object.assign({}, ...writes)).toMatchObject({ accountingReceiptReady: true, costStatus: null, costUsd: 0.25 });
|
||||
const database = leaseDb(execution, {}, {}, updates);
|
||||
await executePaperclipNativeSession({ db: database, execution, runnerInstanceId: "runner" });
|
||||
expect(usageRecorder.create).toHaveBeenCalledExactlyOnceWith(database, {
|
||||
companyId: execution.binding.companyId, runId: execution.binding.runId, adapterType: "paperclip_runner",
|
||||
});
|
||||
expect(usageRecorder.capture).toHaveBeenNthCalledWith(1, expect.objectContaining({
|
||||
complete: false, costStatus: "unpriced", costUsd: 0.25, usageBasis: "per_run",
|
||||
}));
|
||||
expect(usageRecorder.capture).toHaveBeenLastCalledWith(expect.objectContaining({ complete: true, costUsd: 0.25 }));
|
||||
expect(usageRecorder.capture.mock.calls.at(-1)![0].costStatus).toBeUndefined();
|
||||
expect(updates.filter(update => update.values.usageJson)).toEqual([]);
|
||||
});
|
||||
|
||||
it.each(["recorder_creation", "receipt_capture"] as const)("does not hide a failed canonical accounting %s", async stage => {
|
||||
const failure = new Error(`accounting_${stage}_failed`);
|
||||
if (stage === "recorder_creation") usageRecorder.create.mockRejectedValueOnce(failure);
|
||||
else usageRecorder.capture.mockRejectedValueOnce(failure);
|
||||
state.execute.mockReset().mockImplementationOnce(async () => {
|
||||
await accountingEvents.committed!({ eventType: "item.completed", turnId: "turn", emittedAt: new Date().toISOString(),
|
||||
payload: { kind: "usage", usage: { runDelta: { inputTokens: 10, outputTokens: 2, providerCostUsd: 0.25 } } },
|
||||
} as unknown as PrpEvent);
|
||||
throw new Error("Execution must not advance past failed accounting capture");
|
||||
});
|
||||
await expect(executePaperclipNativeSession({ db: leaseDb(), execution, runnerInstanceId: "runner" })).rejects.toBe(failure);
|
||||
expect(usageRecorder.create).toHaveBeenCalledOnce();
|
||||
if (stage === "recorder_creation") expect(usageRecorder.capture).not.toHaveBeenCalled();
|
||||
else expect(usageRecorder.capture).toHaveBeenCalledOnce();
|
||||
});
|
||||
|
||||
it.each(["stale_warm_zero", "partial", "fresh", "fresh_zero", "new_turn", "later_partial", "merged_partial"])("requires fresh complete usage at native completion (%s)", async scenario => {
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import { createRunUsageRecorder } from "../usage-receipts.js";
|
||||
import { applyObservabilityEnvironment, type ProviderObservabilityLaunch } from "@paperclipai/adapter-utils/agent-observability";
|
||||
import { museRunnerBroker } from "../muse-runner-broker.js";
|
||||
import { nativeCaseHandoffResult } from "./native-case-handoff.js";
|
||||
@@ -8217,6 +8218,7 @@ async function executePaperclipNativeSessionWithinScope(
|
||||
)).orderBy(heartbeatRunEvents.seq);
|
||||
for (const row of history) turnAccounting.observe(record(row.payload).prpEvent as PrpEvent);
|
||||
}
|
||||
let fallbackUsageRecorder: ReturnType<typeof createRunUsageRecorder> | undefined;
|
||||
const persistAccountingUsage = async (usage: Record<string, unknown> | null, complete: boolean) => {
|
||||
const accountingUsage = normalizeNativeUsage(usage, { inputIncludesCacheReads: input.execution.provider.kind === "codex" });
|
||||
const accountingCost = accountingUsage ? nativeUsageCostUsd(usage, input.execution.provider) : undefined;
|
||||
@@ -8226,15 +8228,14 @@ async function executePaperclipNativeSessionWithinScope(
|
||||
: "unpriced" as const;
|
||||
const billing = resolveNativeBilling(input.execution.provider, input.runnerEnvironment, input.billingIdentity);
|
||||
const costUsdExact = typeof usage?.accountingCostUsdExact === "string" ? usage.accountingCostUsdExact : undefined;
|
||||
await input.onUsage?.({ usage: accountingUsage, ...billing, complete, usageBasis: "per_run", costUsdExact,
|
||||
model: input.execution.provider.model ?? "unknown", costUsd: accountingCost ?? null, costStatus });
|
||||
if (!input.onUsage) await input.db.update(heartbeatRuns).set({
|
||||
costAccountingPending: true,
|
||||
usageJson: sql`coalesce(${heartbeatRuns.usageJson}, '{}'::jsonb) || ${JSON.stringify({
|
||||
...accountingUsage, ...billing, accountingReceiptReady: complete, model: input.execution.provider.model ?? "unknown",
|
||||
costUsd: accountingCost ?? null, costUsdExact: costUsdExact ?? null, costStatus: costStatus ?? null, usageSource: "per_run",
|
||||
})}::jsonb`,
|
||||
}).where(and(eq(heartbeatRuns.id, input.execution.binding.runId), eq(heartbeatRuns.companyId, input.execution.binding.companyId), isNull(heartbeatRuns.costAccountedAt)));
|
||||
const checkpoint: AdapterUsageCheckpoint = { usage: accountingUsage, ...billing, complete, usageBasis: "per_run", costUsdExact,
|
||||
model: input.execution.provider.model ?? "unknown", costUsd: accountingCost ?? null, costStatus };
|
||||
if (input.onUsage) await input.onUsage(checkpoint);
|
||||
else {
|
||||
fallbackUsageRecorder ??= createRunUsageRecorder(input.db, { companyId: input.execution.binding.companyId,
|
||||
runId: input.execution.binding.runId, adapterType: "paperclip_runner" });
|
||||
await (await fallbackUsageRecorder).capture(checkpoint);
|
||||
}
|
||||
};
|
||||
// appendEvent awaits this receipt before the runtime may commit completeRun.
|
||||
// Replayed events must rebuild it too after a stop between event and receipt.
|
||||
|
||||
@@ -17,6 +17,8 @@ import { preserveNativeFailedWorkspace } from "./native-failed-workspace-recover
|
||||
import { settleStopOnlyCleanup } from "../sandbox-stop-and-retain.js";
|
||||
import { createRunDispatch } from "../../modules/run-dispatch/index.js";
|
||||
import { drainHeartbeatRunsToQuiescence } from "../../__tests__/helpers/drain-heartbeat-runs.js";
|
||||
import { reserveRunBudget } from "../budget-reservations.js";
|
||||
import { accountRunCost } from "../run-cost-accounting.js";
|
||||
const providerExecute = vi.hoisted(() => vi.fn(async () => ({ exitCode: 0, signal: null, timedOut: false,
|
||||
errorMessage: null, summary: "Bounded reconstruction completed.", provider: "test", model: "test" })));
|
||||
vi.mock("../../adapters/index.js", async () => ({
|
||||
@@ -173,6 +175,7 @@ describe("reconstruction after confirmed native workspace loss", () => {
|
||||
expect(successorId).not.toBe("undefined");
|
||||
await db.update(heartbeatRuns).set({ status: "running" }).where(eq(heartbeatRuns.id, successorId));
|
||||
await db.update(issues).set({ status: "in_progress", executionRunId: successorId }).where(eq(issues.id, f.issueId));
|
||||
const reservation = await reserveRunBudget(db, f.companyId, successorId, null);
|
||||
if (kind === "changed_evidence") {
|
||||
const [action] = await db.select().from(issueRecoveryActions).where(eq(issueRecoveryActions.id, later.actionId!));
|
||||
await db.update(issueRecoveryActions).set({ evidence: { ...action.evidence, executionReconciliation: {
|
||||
@@ -182,7 +185,7 @@ describe("reconstruction after confirmed native workspace loss", () => {
|
||||
if (kind === "extra_reconciled") await interveningCancelledRun(f);
|
||||
const dispatch = vi.fn(async () => "provider entered");
|
||||
const result = await createRunDispatch(db).dispatchResolvedInteractionIfCurrent({ companyId: f.companyId,
|
||||
runId: successorId, expectedStatus: "running", dispatch });
|
||||
runId: successorId, reservationId: reservation.id, expectedStatus: "running", dispatch });
|
||||
expect(result.dispatched).toBe(kind === "unchanged");
|
||||
expect(dispatch).toHaveBeenCalledTimes(kind === "unchanged" ? 1 : 0);
|
||||
});
|
||||
@@ -286,16 +289,17 @@ describe("reconstruction after confirmed native workspace loss", () => {
|
||||
expect(successorId).not.toBe("undefined");
|
||||
await db.update(heartbeatRuns).set({ status: "running" }).where(eq(heartbeatRuns.id, successorId));
|
||||
await db.update(issues).set({ status: "in_progress", executionRunId: successorId }).where(eq(issues.id, f.issueId));
|
||||
const reservation = await reserveRunBudget(db, f.companyId, successorId, null);
|
||||
if (kind === "blocked") await db.update(issues).set({ status: "blocked" }).where(eq(issues.id, f.issueId));
|
||||
if (kind === "attachment") await db.delete(issueAttachments).where(eq(issueAttachments.id, f.attachmentId));
|
||||
if (kind === "paused") await db.update(agents).set({ status: "paused" }).where(eq(agents.id, f.agentId));
|
||||
if (kind === "forged_successor") await db.update(issueRecoveryActions).set({ evidence: { ...action.evidence, continuationRunId: randomUUID() } }).where(eq(issueRecoveryActions.id, f.actionId));
|
||||
const dispatch = vi.fn(async () => "provider entered");
|
||||
const outcome = await createRunDispatch(db).dispatchResolvedInteractionIfCurrent({ companyId: f.companyId,
|
||||
runId: successorId, expectedStatus: "running", dispatch });
|
||||
runId: successorId, reservationId: reservation.id, expectedStatus: "running", dispatch });
|
||||
expect(outcome.dispatched).toBe(kind === "unchanged");
|
||||
expect(dispatch).toHaveBeenCalledTimes(kind === "unchanged" ? 1 : 0);
|
||||
if (!outcome.dispatched) expect(outcome.cancellation).toMatchObject({ outcome: "cancelled" });
|
||||
if (!outcome.dispatched) expect(outcome).toMatchObject({ cancellation: { outcome: "cancelled" } });
|
||||
});
|
||||
it.each(["unchanged", "attachment", "blocked", "paused", "forged_wake", "foreign_parent", "unrelated_run", "uncertain_parent", "new_recovery"] as const)("keeps bounded reconstruction retries under current authority: %s", async kind => {
|
||||
const f = await seed();
|
||||
@@ -314,6 +318,7 @@ describe("reconstruction after confirmed native workspace loss", () => {
|
||||
if (retry.outcome !== "scheduled" || !retry.run) throw new Error("Expected a bounded bootstrap retry");
|
||||
await db.update(heartbeatRuns).set({ status: "running" }).where(eq(heartbeatRuns.id, retry.run.id));
|
||||
await db.update(issues).set({ executionRunId: retry.run.id }).where(eq(issues.id, f.issueId));
|
||||
const reservation = await reserveRunBudget(db, f.companyId, retry.run.id, null);
|
||||
if (kind === "attachment") await db.delete(issueAttachments).where(eq(issueAttachments.id, f.attachmentId));
|
||||
if (kind === "blocked") await db.update(issues).set({ status: "blocked" }).where(eq(issues.id, f.issueId));
|
||||
if (kind === "paused") await db.update(agents).set({ status: "paused" }).where(eq(agents.id, f.agentId));
|
||||
@@ -330,21 +335,23 @@ describe("reconstruction after confirmed native workspace loss", () => {
|
||||
fingerprint: randomUUID(), evidence: { runId: firstId }, nextAction: "Reconcile the new execution", status: "active" });
|
||||
const dispatch = vi.fn(async () => "provider entered");
|
||||
const outcome = await createRunDispatch(db).dispatchResolvedInteractionIfCurrent({ companyId: f.companyId,
|
||||
runId: retry.run.id, expectedStatus: "running", dispatch });
|
||||
runId: retry.run.id, reservationId: reservation.id, expectedStatus: "running", dispatch });
|
||||
expect(outcome.dispatched).toBe(kind === "unchanged");
|
||||
expect(dispatch).toHaveBeenCalledTimes(kind === "unchanged" ? 1 : 0);
|
||||
if (kind === "unchanged") {
|
||||
expect(retry.run.scheduledRetryAttempt).toBe(1);
|
||||
await db.update(heartbeatRuns).set({ status: "failed", finishedAt: new Date(), errorCode: "adapter_failed",
|
||||
resultJson: { executionRecovery: { kind: "bootstrap", providerWorkStarted: false } } }).where(eq(heartbeatRuns.id, retry.run.id));
|
||||
await accountRunCost(db, retry.run.id);
|
||||
const second = await heartbeat.scheduleBoundedRetry(retry.run.id, { now: new Date(), random: () => 0, delayMs: 0 });
|
||||
expect(second.outcome).toBe("scheduled");
|
||||
if (second.outcome !== "scheduled" || !second.run) throw new Error("Expected the second bounded retry");
|
||||
expect(second.run.scheduledRetryAttempt).toBe(2);
|
||||
await db.update(heartbeatRuns).set({ status: "running" }).where(eq(heartbeatRuns.id, second.run.id));
|
||||
await db.update(issues).set({ executionRunId: second.run.id }).where(eq(issues.id, f.issueId));
|
||||
const secondReservation = await reserveRunBudget(db, f.companyId, second.run.id, null);
|
||||
expect(await createRunDispatch(db).dispatchResolvedInteractionIfCurrent({ companyId: f.companyId,
|
||||
runId: second.run.id, expectedStatus: "running", dispatch })).toMatchObject({ dispatched: true });
|
||||
runId: second.run.id, reservationId: secondReservation.id, expectedStatus: "running", dispatch })).toMatchObject({ dispatched: true });
|
||||
await db.update(heartbeatRuns).set({ status: "failed", finishedAt: new Date(),
|
||||
resultJson: { executionRecovery: { kind: "bootstrap", providerWorkStarted: false } } }).where(eq(heartbeatRuns.id, second.run.id));
|
||||
expect(await heartbeat.scheduleBoundedRetry(second.run.id)).toMatchObject({ outcome: "retry_exhausted" });
|
||||
|
||||
@@ -57,6 +57,7 @@ import {
|
||||
import { emitAgentTaskRun } from "../agent-task-run-telemetry.js";
|
||||
import { materializeNativeProviderCapacityRetry } from "./native-provider-capacity-retry.js";
|
||||
import { NATIVE_PROVIDER_OVERLOADED_CODE } from "./native-provider-failure.js";
|
||||
import { DELIVERY_QUEUES, notifyDeliveryWork } from "../delivery-work-notifications.js";
|
||||
|
||||
export class NativeStatusRaceError extends Error {
|
||||
readonly code = "native_status_race" as const;
|
||||
@@ -1056,6 +1057,7 @@ async function materializeDecisionEffect(input: {
|
||||
),
|
||||
);
|
||||
}
|
||||
await notifyDeliveryWork(input.tx, DELIVERY_QUEUES.accounting, input.companyId);
|
||||
const [run] = await input.tx
|
||||
.update(heartbeatRuns)
|
||||
.set({
|
||||
@@ -1845,6 +1847,7 @@ export async function commitNativeStatusDecision(input: {
|
||||
const finalizationError = input.decision.effects.find(
|
||||
effect => effect.kind === "record_finalization_error",
|
||||
);
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, input.companyId);
|
||||
await tx
|
||||
.update(nativeRunFinalizations)
|
||||
.set({
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notifications.js";
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { and, eq, lt, sql } from "drizzle-orm";
|
||||
import { agents, budgetPolicies, budgetReservations, companies, companyPromptInjectionPrevention,
|
||||
import { agents, budgetReservations, companies, companyPromptInjectionPrevention,
|
||||
promptInjectionInvocations, companySecrets, heartbeatRuns, runIdentityContexts, issues, projects, type Db } from "@paperclipai/db";
|
||||
import { PROMPT_INJECTION_DEFAULT_MODEL, PROMPT_INJECTION_POLICY_VERSION, decisionProviderForConnection,
|
||||
updatePromptInjectionPreventionSchema, centsToUnits, unitsToCents,
|
||||
@@ -9,7 +10,8 @@ import { PROMPT_INJECTION_DEFAULT_MODEL, PROMPT_INJECTION_POLICY_VERSION, decisi
|
||||
import { HttpError, forbidden } from "../errors.js";
|
||||
import { aiConnectionService } from "./ai-connections.js";
|
||||
import { withAccountingTransaction } from "./accounting-transaction.js";
|
||||
import { budgetPoliciesForRun, budgetService, budgetServiceInTransaction, computeObservedSpend, type BudgetServiceHooks } from "./budgets.js";
|
||||
import { withBudgetSpan } from "./budget-tracing.js";
|
||||
import { budgetServiceInTransaction, type BudgetServiceHooks } from "./budgets.js";
|
||||
import { createCostEventInTransaction } from "./costs.js";
|
||||
import { logActivity } from "./activity-log.js";
|
||||
import { runPromptInjectionProvider, promptInjectionInput, promptInjectionRequestFits, type PromptInjectionProviderOutcome } from "./prompt-injection-provider.js";
|
||||
@@ -102,7 +104,7 @@ export function promptInjectionPreventionService(db: Db, options: {
|
||||
return { config, connection, userId, agentId, runId, identityContextId, issueId, projectId, actorType: context.actorType, actorId: context.actorId };
|
||||
}
|
||||
async function settle(invocation: typeof promptInjectionInvocations.$inferSelect, outcome: PromptInjectionProviderOutcome) {
|
||||
await withAccountingTransaction(db, invocation.companyId, async (tx, publications) => {
|
||||
await withBudgetSpan("budget.accounting.settlement", () => withAccountingTransaction(db, invocation.companyId, async (tx, publications) => {
|
||||
const [current] = await tx.select().from(promptInjectionInvocations).where(and(eq(promptInjectionInvocations.companyId, invocation.companyId), eq(promptInjectionInvocations.id, invocation.id))).for("update");
|
||||
if (!current || current.status !== "running") return;
|
||||
const receipt = outcome.receipt;
|
||||
@@ -125,6 +127,7 @@ export function promptInjectionPreventionService(db: Db, options: {
|
||||
providerRequestId: receipt.providerRequestId, pricingProvenance: receipt.pricingProvenance, occurredAt: invocation.startedAt,
|
||||
}, publications, { actorType: invocation.actorType as "system" | "user" | "agent", actorId: invocation.actorId, agentId: links.agentId });
|
||||
const finishedAt = new Date();
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, invocation.companyId);
|
||||
await tx.update(promptInjectionInvocations).set({ status: outcome.verdict ? "succeeded" : receipt.costStatus === "unpriced" ? "unknown" : "failed",
|
||||
violation: outcome.verdict?.violation ?? null, findings: outcome.verdict?.findings ?? null, errorCode: outcome.errorCode ?? null, costEventId: event.id, providerRequestId: receipt.providerRequestId,
|
||||
inputTokens: receipt.inputTokens, outputTokens: receipt.outputTokens, finishedAt, durationMs: Math.max(0, finishedAt.getTime() - invocation.startedAt.getTime()),
|
||||
@@ -132,8 +135,7 @@ export function promptInjectionPreventionService(db: Db, options: {
|
||||
if (receipt.costStatus !== "unpriced") await tx.update(budgetReservations).set({ state: outcome.noProviderWork ? "released" : "settled", settledAt: finishedAt }).where(and(
|
||||
eq(budgetReservations.companyId, invocation.companyId), eq(budgetReservations.promptInjectionInvocationId, invocation.id), eq(budgetReservations.state, "held"),
|
||||
));
|
||||
});
|
||||
await budgetService(db, options.budgetHooks).deliverPendingEnforcement(invocation.companyId);
|
||||
}));
|
||||
}
|
||||
async function classify(context: PromptInjectionContext, content: string, signal?: AbortSignal): Promise<PromptInjectionModelOutcome & { usage?: DecisionUsage }> {
|
||||
signal?.throwIfAborted();
|
||||
@@ -171,23 +173,17 @@ export function promptInjectionPreventionService(db: Db, options: {
|
||||
if (!currentSecret || currentSecret.status !== "active" || currentSecret.version !== credentialSnapshot.version)
|
||||
return { reason: "connection_unavailable" as const };
|
||||
const budgets = budgetServiceInTransaction(tx, publications);
|
||||
const block = await budgets.getInvocationBlock(context.companyId, current.agentId, { projectId: current.projectId });
|
||||
const { block, policies, observations } = await budgets.evaluateAdmission(context.companyId, current.agentId, { projectId: current.projectId });
|
||||
const [company] = await tx.select({ status: companies.status }).from(companies).where(eq(companies.id, context.companyId));
|
||||
const [agent] = current.agentId ? await tx.select({ status: agents.status }).from(agents).where(and(eq(agents.id, current.agentId), eq(agents.companyId, context.companyId))) : [];
|
||||
if (block || company.status !== "active" || (current.agentId && (!agent || agent.status === "paused" || agent.status === "terminated"))) return { reason: "budget_blocked" as const };
|
||||
const policies = await tx.select().from(budgetPolicies).where(and(eq(budgetPolicies.companyId, context.companyId), eq(budgetPolicies.isActive, true),
|
||||
eq(budgetPolicies.hardStopEnabled, true), eq(budgetPolicies.metric, "billed_cents"), sql`${budgetPolicies.amount} > 0`, budgetPoliciesForRun(context.companyId, current.agentId, current.projectId)));
|
||||
const amount = policies.reduce((max, p) => { const n = centsToUnits(p.reservationCents); return n > max ? n : max; }, 0n);
|
||||
for (const policy of policies) {
|
||||
const [held] = await tx.select({ total: sql<string>`coalesce(sum(${budgetReservations.amountCents}),0)::text` }).from(budgetReservations).where(and(
|
||||
eq(budgetReservations.companyId, context.companyId), eq(budgetReservations.state, "held"),
|
||||
policy.scopeType === "agent" ? eq(budgetReservations.agentId, current.agentId!) : undefined,
|
||||
policy.scopeType === "project" ? eq(budgetReservations.projectId, current.projectId!) : undefined,
|
||||
));
|
||||
const committed = centsToUnits((await computeObservedSpend(tx, policy)).totalExact) + centsToUnits(held.total);
|
||||
for (const [index, policy] of policies.entries()) {
|
||||
const committed = centsToUnits(observations[index].ledgerCents) + centsToUnits(observations[index].effectiveHeldCents);
|
||||
if (committed >= centsToUnits(policy.amount) || committed + amount > centsToUnits(policy.amount)) return { reason: "budget_blocked" as const };
|
||||
}
|
||||
signal?.throwIfAborted();
|
||||
await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, context.companyId);
|
||||
const [invocation] = await tx.insert(promptInjectionInvocations).values({ id: randomUUID(), companyId: context.companyId, source: context.source, toolInvocationId: context.toolInvocationId ?? null, toolName: context.toolName,
|
||||
actorType: current.actorType, actorId: current.actorId, responsibleUserId: current.userId, agentId: current.agentId,
|
||||
issueId: current.issueId, projectId: current.projectId, runId: current.runId, identityContextId: current.identityContextId,
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { promises as fs } from "node:fs";
|
||||
import os from "node:os";
|
||||
import path from "node:path";
|
||||
import { and, eq, sql } from "drizzle-orm";
|
||||
import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest";
|
||||
import { agents, companies, costEvents, createDb, heartbeatRuns, promptInjectionInvocations, type Db } from "@paperclipai/db";
|
||||
import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../__tests__/helpers/embedded-postgres.js";
|
||||
import { accountRunCost, createCostAccountingReconciler, readCostAccountingWork } from "./run-cost-accounting.js";
|
||||
|
||||
const support = await getEmbeddedPostgresTestSupport();
|
||||
(support.supported ? describe : describe.skip)("run cost accounting work", () => {
|
||||
let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>>;
|
||||
let db: Db;
|
||||
let directory: string;
|
||||
beforeAll(async () => {
|
||||
database = await startEmbeddedPostgresTestDatabase("run-cost-work-");
|
||||
db = createDb(database.connectionString);
|
||||
directory = await fs.mkdtemp(path.join(os.tmpdir(), "run-cost-work-"));
|
||||
}, 60_000);
|
||||
afterEach(async () => { await db.transaction(async tx => {
|
||||
await tx.execute(sql`SET LOCAL client_min_messages = warning`);
|
||||
await tx.execute(sql`TRUNCATE companies CASCADE`);
|
||||
}); });
|
||||
afterAll(async () => { if (directory) await fs.rm(directory, { force: true, recursive: true }); await database?.cleanup(); });
|
||||
|
||||
async function seed(change: Partial<typeof heartbeatRuns.$inferInsert> = {}) {
|
||||
const companyId = randomUUID(), agentId = randomUUID(), runId = randomUUID();
|
||||
await db.insert(companies).values({ id: companyId, name: "Accounting", issuePrefix: "ACCT" });
|
||||
await db.insert(agents).values({ id: agentId, companyId, name: "Accounting", role: "engineer", status: "idle", adapterType: "process" });
|
||||
await db.insert(heartbeatRuns).values({ id: runId, companyId, agentId, invocationSource: "on_demand", status: "succeeded",
|
||||
costAccountingPending: true, finishedAt: new Date(),
|
||||
usageJson: { accountingReceiptReady: true, provider: "test", model: "test", billingType: "metered_api", costUsd: 0.01 }, ...change });
|
||||
return { companyId, agentId, runId };
|
||||
}
|
||||
|
||||
it("returns from settled or running calls without opening an accounting transaction", async () => {
|
||||
const { runId } = await seed({ costAccountingPending: false });
|
||||
const transaction = vi.fn(db.transaction.bind(db));
|
||||
const observed = new Proxy(db, { get(target, property) {
|
||||
if (property === "transaction") return transaction;
|
||||
const value = Reflect.get(target, property);
|
||||
return typeof value === "function" ? value.bind(target) : value;
|
||||
} }) as Db;
|
||||
expect(await accountRunCost(observed, runId)).toBe(false);
|
||||
await db.update(heartbeatRuns).set({ status: "running", costAccountingPending: true }).where(eq(heartbeatRuns.id, runId));
|
||||
expect(await accountRunCost(observed, runId)).toBe(false);
|
||||
expect(transaction).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("settles once without delivering cancellation in the caller", async () => {
|
||||
const { companyId, runId } = await seed();
|
||||
const cancelWorkForScope = vi.fn(async () => { throw new Error("Cancellation must be asynchronous"); });
|
||||
expect(await accountRunCost(db, runId, { cancelWorkForScope })).toBe(true);
|
||||
expect(await accountRunCost(db, runId, { cancelWorkForScope })).toBe(false);
|
||||
const events = await db.select().from(costEvents).where(and(eq(costEvents.companyId, companyId), eq(costEvents.heartbeatRunId, runId)));
|
||||
expect(events).toHaveLength(1);
|
||||
expect(events[0].costCents).toBe(1);
|
||||
expect(cancelWorkForScope).not.toHaveBeenCalled();
|
||||
expect(await readCostAccountingWork(db, Date.now(), directory)).toEqual({ pending: false, nextRunAt: null });
|
||||
});
|
||||
|
||||
it("repairs a stale pending marker without creating another charge", async () => {
|
||||
const { runId } = await seed({ costAccountedAt: new Date(), usageJson: { accountingReceiptReady: false } });
|
||||
expect((await readCostAccountingWork(db, Date.now(), directory)).pending).toBe(true);
|
||||
expect(await accountRunCost(db, runId)).toBe(false);
|
||||
const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId));
|
||||
expect(run.costAccountingPending).toBe(false);
|
||||
expect(await db.select().from(costEvents)).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("waits for a missing final receipt without scheduling unchanged debt", async () => {
|
||||
const { runId } = await seed({ usageJson: { accountingReceiptReady: false, inputTokens: 5 } });
|
||||
expect(await readCostAccountingWork(db, Date.now(), directory)).toEqual({ pending: false, nextRunAt: null });
|
||||
const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId));
|
||||
expect(run.costAccountingPending).toBe(true);
|
||||
await db.update(heartbeatRuns).set({ usageJson: { accountingReceiptReady: true, costUsd: 0.01 } }).where(eq(heartbeatRuns.id, runId));
|
||||
expect((await readCostAccountingWork(db, Date.now(), directory)).pending).toBe(true);
|
||||
});
|
||||
|
||||
it("reconstructs the actual interrupted-service deadline and disarms after settlement", async () => {
|
||||
const { companyId } = await seed({ costAccountingPending: false });
|
||||
const startedAt = new Date();
|
||||
const [invocation] = await db.insert(promptInjectionInvocations).values({ companyId, source: "settings_test", actorType: "user", actorId: "test",
|
||||
connectionId: randomUUID(), grantId: randomUUID(), policyVersion: "tool-result-v1", model: "openai/gpt-oss-safeguard-20b", startedAt }).returning();
|
||||
const work = await readCostAccountingWork(db, startedAt.getTime(), directory);
|
||||
expect(work).toEqual({ pending: false, nextRunAt: startedAt.getTime() + 120_001 });
|
||||
expect((await readCostAccountingWork(db, startedAt.getTime() + 120_001, directory)).pending).toBe(true);
|
||||
await db.update(promptInjectionInvocations).set({ status: "unknown" }).where(eq(promptInjectionInvocations.id, invocation.id));
|
||||
expect(await readCostAccountingWork(db, startedAt.getTime() + 120_001, directory)).toEqual({ pending: false, nextRunAt: null });
|
||||
});
|
||||
|
||||
it("shares service and run recovery across overlapping signals", async () => {
|
||||
let release!: () => void;
|
||||
const recover = vi.fn(() => new Promise<void>(resolve => { release = resolve; }));
|
||||
const reconcile = createCostAccountingReconciler(db, {}, recover);
|
||||
const first = reconcile(), second = reconcile();
|
||||
expect(first).toBe(second);
|
||||
expect(recover).toHaveBeenCalledOnce();
|
||||
release();
|
||||
await first;
|
||||
expect(recover).toHaveBeenCalledOnce();
|
||||
});
|
||||
|
||||
it("recovers run receipts even when an unrelated service invocation cannot settle", async () => {
|
||||
const { runId } = await seed();
|
||||
const reconcile = createCostAccountingReconciler(db, {}, async () => { throw new Error("Unrelated service receipt failed"); });
|
||||
await expect(reconcile()).rejects.toThrow("Unrelated service receipt failed");
|
||||
const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, runId));
|
||||
expect(run.costAccountedAt).not.toBeNull();
|
||||
expect(run.costAccountingPending).toBe(false);
|
||||
expect(await db.select().from(costEvents)).toHaveLength(1);
|
||||
});
|
||||
});
|
||||
@@ -1,10 +1,11 @@
|
||||
import { usdToUnits, unitsToCents } from "@paperclipai/shared";
|
||||
import { and, asc, eq, inArray, sql } from "drizzle-orm";
|
||||
import { accountingRuntimeBaselines, budgetReservations, agentRuntimeState, agents, costEvents, heartbeatRuns, nativeRunFinalizations, issues, projects, type Db } from "@paperclipai/db";
|
||||
import { and, asc, eq, inArray, or, sql } from "drizzle-orm";
|
||||
import { accountingRuntimeBaselines, budgetReservations, agentRuntimeState, agents, costEvents, heartbeatRuns, nativeRunFinalizations, issues, projects, decisionInvocations, promptInjectionInvocations, type Db } from "@paperclipai/db";
|
||||
import { withAccountingTransaction } from "./accounting-transaction.js";
|
||||
import { withBudgetSpan } from "./budget-tracing.js";
|
||||
import { conflict } from "../errors.js";
|
||||
import { createCostEventInTransaction } from "./costs.js";
|
||||
import { budgetService, deliverBudgetEnforcement, type BudgetServiceHooks } from "./budgets.js";
|
||||
import { budgetServiceInTransaction, type BudgetServiceHooks } from "./budgets.js";
|
||||
import { logger } from "../middleware/logger.js";
|
||||
import { promises as fs } from "node:fs";
|
||||
import { replayUsageReceipts, indexPendingUsageReceipts, recoverPendingRunUsageReceipts, type UsageReceiptIndex } from "./usage-receipts.js";
|
||||
@@ -21,13 +22,13 @@ export async function accountRunCost(db: Db, runId: string, hooks: BudgetService
|
||||
pending: heartbeatRuns.costAccountingPending, accountedAt: heartbeatRuns.costAccountedAt,
|
||||
issueId: sql<string | null>`${heartbeatRuns.usageJson}->'ledgerScope'->>'issueId'`,
|
||||
}).from(heartbeatRuns).where(eq(heartbeatRuns.id, runId));
|
||||
if (!identity) return false;
|
||||
// Finalization can call twice. Avoid disk work for settled or active runs,
|
||||
// but retain the locked recheck, stale-marker cleanup, and stop delivery.
|
||||
if (!identity || !identity.pending || !terminalStatuses.includes(identity.status)) return false;
|
||||
// Finalization can call twice. Settled or active runs owe no accounting work.
|
||||
// Pending acknowledgements still receive the locked stale-marker cleanup.
|
||||
const index = identity.pending && !identity.accountedAt && terminalStatuses.includes(identity.status)
|
||||
? receiptIndex ?? await indexPendingUsageReceipts() : undefined;
|
||||
let recoveredFiles: string[] = [];
|
||||
const accounted = await withAccountingTransaction(db, identity.companyId, async (tx, publications) => {
|
||||
const accounted = await withBudgetSpan("budget.accounting.settlement", () => withAccountingTransaction(db, identity.companyId, async (tx, publications) => {
|
||||
// Native recovery owns the issue before updating its run. Acquire the
|
||||
// charge's issue FK lock first so inserting it cannot reverse that order.
|
||||
const issueId = text(identity.issueId);
|
||||
@@ -38,6 +39,8 @@ export async function accountRunCost(db: Db, runId: string, hooks: BudgetService
|
||||
if (!run?.costAccountingPending || !terminalStatuses.includes(run.status)) return false;
|
||||
if (run.costAccountedAt) {
|
||||
await tx.update(heartbeatRuns).set({ costAccountingPending: false }).where(eq(heartbeatRuns.id, run.id));
|
||||
await budgetServiceInTransaction(tx, publications).evaluateAffectedScopes(run.companyId, run.agentId,
|
||||
text(object(object(run.usageJson).ledgerScope).projectId));
|
||||
return false;
|
||||
}
|
||||
// The bounded global replay can leave a newer receipt behind an older
|
||||
@@ -69,7 +72,11 @@ export async function accountRunCost(db: Db, runId: string, hooks: BudgetService
|
||||
await tx.update(heartbeatRuns).set({ costAccountingPending: false, costAccountedAt: new Date() }).where(eq(heartbeatRuns.id, run.id));
|
||||
await tx.update(budgetReservations).set({ state: preProviderFailure ? "released" : "settled", settledAt: new Date() })
|
||||
.where(and(eq(budgetReservations.companyId, run.companyId), eq(budgetReservations.runId, run.id), eq(budgetReservations.state, "held")));
|
||||
if (preProviderFailure) return true;
|
||||
if (preProviderFailure) {
|
||||
await budgetServiceInTransaction(tx, publications).evaluateAffectedScopes(run.companyId, run.agentId,
|
||||
text(object(object(run.usageJson).ledgerScope).projectId));
|
||||
return true;
|
||||
}
|
||||
const usage = object(run.usageJson);
|
||||
const scope = object(usage.ledgerScope);
|
||||
// Deleted attribution targets become unallocated; a deletion must never
|
||||
@@ -121,11 +128,12 @@ export async function accountRunCost(db: Db, runId: string, hooks: BudgetService
|
||||
...receipt, idempotencyKey: `heartbeat:${run.id}:model:${index}`, model: text(part.model)!,
|
||||
inputTokens: amount(partUsage.inputTokens) ?? 0, cachedInputTokens: amount(partUsage.cachedInputTokens) ?? 0,
|
||||
outputTokens: amount(partUsage.outputTokens) ?? 0, costCents: partCents,
|
||||
}, publications);
|
||||
}, publications, undefined, { deferBudgetEvaluation: true });
|
||||
}
|
||||
} else {
|
||||
await createCostEventInTransaction(tx, run.companyId, receipt, publications);
|
||||
await createCostEventInTransaction(tx, run.companyId, receipt, publications, undefined, { deferBudgetEvaluation: true });
|
||||
}
|
||||
await budgetServiceInTransaction(tx, publications).evaluateAffectedScopes(run.companyId, run.agentId, text(scope.projectId));
|
||||
const [event] = await tx.select({
|
||||
inputTokens: sql<string>`coalesce(sum(${costEvents.inputTokens}),0)::text`,
|
||||
cachedInputTokens: sql<string>`coalesce(sum(${costEvents.cachedInputTokens}),0)::text`,
|
||||
@@ -145,26 +153,52 @@ export async function accountRunCost(db: Db, runId: string, hooks: BudgetService
|
||||
}).where(eq(agentRuntimeState.agentId, run.agentId));
|
||||
await tx.update(heartbeatRuns).set({ costAccountingPending: false, costAccountedAt: new Date(), accountingProjectionVersion: "v2", accountingLastError: null, accountingLastAttemptAt: new Date() }).where(eq(heartbeatRuns.id, run.id));
|
||||
return true;
|
||||
});
|
||||
}));
|
||||
for (const file of recoveredFiles) {
|
||||
try { await fs.rm(file, { force: true }); }
|
||||
catch (error) { logger.warn({ err: error, runId }, "Accounting committed; recovered receipt cleanup will retry"); }
|
||||
}
|
||||
await deliverBudgetEnforcement(db, hooks, identity.companyId);
|
||||
return accounted;
|
||||
}
|
||||
|
||||
export async function reconcileRunCosts(db: Db, hooks: BudgetServiceHooks = {}) {
|
||||
const pending = await db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(and(
|
||||
function actionableRunAccounting() {
|
||||
return and(
|
||||
eq(heartbeatRuns.costAccountingPending, true), inArray(heartbeatRuns.status, terminalStatuses),
|
||||
sql`${heartbeatRuns.usageJson}->>'accountingCaptureFailed' is distinct from 'true'`,
|
||||
sql`(${heartbeatRuns.resultJson}->'executionRecovery'->>'providerWorkStarted' = 'false'
|
||||
or ${heartbeatRuns.usageJson}->>'accountingProviderWorkStarted' = 'false'
|
||||
or not exists (select 1 from ${nativeRunFinalizations} where ${nativeRunFinalizations.runId} = ${heartbeatRuns.id}
|
||||
and ${heartbeatRuns.runtimeMode} = 'native' and ${nativeRunFinalizations.resultId} is null
|
||||
and ${nativeRunFinalizations.phase} <> 'terminal_failure'))`,
|
||||
sql`(${heartbeatRuns.usageJson}->>'accountingReceiptReady' is distinct from 'false' or ${heartbeatRuns.resultJson}->'executionRecovery'->>'providerWorkStarted' = 'false' or ${heartbeatRuns.usageJson}->>'accountingProviderWorkStarted' = 'false')`,
|
||||
)).orderBy(asc(heartbeatRuns.updatedAt), asc(heartbeatRuns.id)).limit(100);
|
||||
or(sql`${heartbeatRuns.costAccountedAt} is not null`, and(
|
||||
sql`${heartbeatRuns.usageJson}->>'accountingCaptureFailed' is distinct from 'true'`,
|
||||
sql`(${heartbeatRuns.resultJson}->'executionRecovery'->>'providerWorkStarted' = 'false'
|
||||
or ${heartbeatRuns.usageJson}->>'accountingProviderWorkStarted' = 'false'
|
||||
or not exists (select 1 from ${nativeRunFinalizations} where ${nativeRunFinalizations.runId} = ${heartbeatRuns.id}
|
||||
and ${heartbeatRuns.runtimeMode} = 'native' and ${nativeRunFinalizations.resultId} is null
|
||||
and ${nativeRunFinalizations.phase} <> 'terminal_failure'))`,
|
||||
sql`(${heartbeatRuns.usageJson}->>'accountingReceiptReady' is distinct from 'false' or ${heartbeatRuns.resultJson}->'executionRecovery'->>'providerWorkStarted' = 'false' or ${heartbeatRuns.usageJson}->>'accountingProviderWorkStarted' = 'false')`,
|
||||
)),
|
||||
);
|
||||
}
|
||||
|
||||
/** Only unfinished, actionable debt warrants retries. Incomplete provider
|
||||
* receipts retain their accounting/idle hold and wake when evidence arrives. */
|
||||
export async function readCostAccountingWork(db: Db, now = Date.now(), receiptDirectory?: string) {
|
||||
const [pending] = await db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(actionableRunAccounting()).limit(1);
|
||||
const [deadline] = await db.execute<{ at: Date | string | null }>(sql`
|
||||
select min(started_at) + interval '120 seconds 1 millisecond' as at from (
|
||||
select ${decisionInvocations.startedAt} as started_at from ${decisionInvocations}
|
||||
where ${decisionInvocations.status} = 'running'
|
||||
union all
|
||||
select ${promptInjectionInvocations.startedAt} from ${promptInjectionInvocations}
|
||||
where ${promptInjectionInvocations.status} = 'running'
|
||||
) running_service_invocations`);
|
||||
const nextRunAt = deadline?.at == null ? null : new Date(deadline.at).getTime();
|
||||
const receipts = await indexPendingUsageReceipts(receiptDirectory);
|
||||
return {
|
||||
pending: !!pending || receipts.size > 0 || (nextRunAt !== null && nextRunAt <= now),
|
||||
nextRunAt,
|
||||
};
|
||||
}
|
||||
|
||||
export async function reconcileRunCosts(db: Db, hooks: BudgetServiceHooks = {}) {
|
||||
const pending = await db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(actionableRunAccounting())
|
||||
.orderBy(asc(heartbeatRuns.updatedAt), asc(heartbeatRuns.id)).limit(100);
|
||||
let accounted = 0;
|
||||
const receiptIndex = pending.length > 0 ? await indexPendingUsageReceipts() : undefined;
|
||||
for (const run of pending) {
|
||||
@@ -176,20 +210,24 @@ export async function reconcileRunCosts(db: Db, hooks: BudgetServiceHooks = {})
|
||||
logger.error({ err: error, runId: run.id }, "Run accounting remains pending for retry");
|
||||
}
|
||||
}
|
||||
await budgetService(db, hooks).reconcilePolicies();
|
||||
return { scanned: pending.length, accounted };
|
||||
}
|
||||
|
||||
/** Share one full sweep across overlapping scheduler ticks. Do not enqueue
|
||||
* another sweep: the next tick retries after either success or failure. */
|
||||
export function createCostAccountingReconciler(db: Db, hooks: BudgetServiceHooks = {}) {
|
||||
/** Share recovery across overlapping signals. The owning delivery worker
|
||||
* schedules retries only while a durable receipt or actionable debt remains. */
|
||||
export function createCostAccountingReconciler(db: Db, hooks: BudgetServiceHooks = {}, recoverServices?: () => Promise<void>) {
|
||||
let inFlight: Promise<Awaited<ReturnType<typeof reconcileRunCosts>>> | undefined;
|
||||
return () => {
|
||||
if (!inFlight) {
|
||||
inFlight = (async () => {
|
||||
try {
|
||||
let serviceRecoveryFailure: { error: unknown } | undefined;
|
||||
try { await recoverServices?.(); }
|
||||
catch (error) { serviceRecoveryFailure = { error }; }
|
||||
await replayUsageReceipts(db);
|
||||
return await reconcileRunCosts(db, hooks);
|
||||
const result = await reconcileRunCosts(db, hooks);
|
||||
if (serviceRecoveryFailure) throw serviceRecoveryFailure.error;
|
||||
return result;
|
||||
} finally {
|
||||
inFlight = undefined;
|
||||
}
|
||||
|
||||
@@ -188,6 +188,11 @@ function sha256(value: unknown): string {
|
||||
return createHash("sha256").update(stableStringify(value)).digest("hex");
|
||||
}
|
||||
|
||||
type ConsumedRateAdmission = { policyId: string; versionHash: string };
|
||||
function ratePolicyVersionHash(policy: typeof toolPolicies.$inferSelect): string {
|
||||
return sha256({ ...policy, createdAt: policy.createdAt.toISOString(), updatedAt: policy.updatedAt.toISOString() });
|
||||
}
|
||||
|
||||
function isoDateOrNull(value: unknown): string | null {
|
||||
if (!value) return null;
|
||||
if (value instanceof Date) return Number.isNaN(value.getTime()) ? null : value.toISOString();
|
||||
@@ -1336,6 +1341,32 @@ export function toolAccessPolicyService(db: Db) {
|
||||
});
|
||||
}
|
||||
|
||||
async function admittedRateLimitAdmissions(ctx: ToolAccessContext, invocationId: string): Promise<ConsumedRateAdmission[]> {
|
||||
const [invocation] = await db.select({
|
||||
actorType: toolInvocations.actorType, actorId: toolInvocations.actorId, agentId: toolInvocations.agentId,
|
||||
runId: toolInvocations.runId, issueId: toolInvocations.issueId, gatewayId: toolInvocations.gatewayId,
|
||||
applicationId: toolInvocations.applicationId, connectionId: toolInvocations.connectionId,
|
||||
catalogEntryId: toolInvocations.catalogEntryId, toolName: toolInvocations.toolName,
|
||||
providerType: toolInvocations.providerType, applicationKey: toolInvocations.applicationKey,
|
||||
upstreamToolName: toolInvocations.upstreamToolName, argumentsHash: toolInvocations.argumentsHash,
|
||||
status: toolInvocations.status, policyDecision: toolInvocations.policyDecision, approvalState: toolInvocations.approvalState,
|
||||
policyExplanation: toolInvocations.policyExplanation,
|
||||
}).from(toolInvocations).where(and(eq(toolInvocations.companyId, ctx.companyId), eq(toolInvocations.id, invocationId)));
|
||||
const authorized = invocation && invocation.status === "executing"
|
||||
&& (invocation.policyDecision === "allow" || (invocation.policyDecision === "require_approval" && invocation.approvalState === "approved"))
|
||||
&& invocation.actorType === ctx.actorType && invocation.actorId === ctx.actorId && invocation.agentId === ctx.agentId
|
||||
&& invocation.runId === ctx.heartbeatRunId && invocation.issueId === ctx.issueId && invocation.gatewayId === ctx.gatewayId
|
||||
&& invocation.applicationId === ctx.applicationId && invocation.connectionId === ctx.connectionId
|
||||
&& invocation.catalogEntryId === ctx.catalogEntryId && invocation.toolName === ctx.toolName
|
||||
&& invocation.providerType === ctx.providerType && invocation.applicationKey === ctx.applicationKey
|
||||
&& invocation.upstreamToolName === ctx.upstreamToolName
|
||||
&& invocation.argumentsHash === ctx.argumentsHash;
|
||||
const stored = authorized ? invocation.policyExplanation?.consumedRateAdmissions : null;
|
||||
return Array.isArray(stored) ? stored.filter((item): item is ConsumedRateAdmission => isRecord(item)
|
||||
&& typeof item.policyId === "string" && CANONICAL_UUID_RE.test(item.policyId)
|
||||
&& typeof item.versionHash === "string" && /^[a-f0-9]{64}$/.test(item.versionHash)) : [];
|
||||
}
|
||||
|
||||
/**
|
||||
* Decides one tool access request. Pass the same `cache` to every call of a
|
||||
* batch that shares one actor and run context (for example, one tools/list
|
||||
@@ -1343,14 +1374,32 @@ export function toolAccessPolicyService(db: Db) {
|
||||
*/
|
||||
async function decide(
|
||||
input: ToolAccessDecisionInput,
|
||||
options: { cache?: ToolAccessDecisionCache } = {},
|
||||
options: { cache?: ToolAccessDecisionCache; admittedInvocationId?: string } = {},
|
||||
): Promise<ToolAccessDecision> {
|
||||
const consumedRateAdmissions: ConsumedRateAdmission[] = [];
|
||||
const result = await decideInternal(input, options, consumedRateAdmissions);
|
||||
return consumedRateAdmissions.length ? { ...result,
|
||||
policyExplanation: { ...result.policyExplanation, consumedRateAdmissions },
|
||||
} : result;
|
||||
}
|
||||
|
||||
async function decideInternal(
|
||||
input: ToolAccessDecisionInput,
|
||||
options: { cache?: ToolAccessDecisionCache; admittedInvocationId?: string },
|
||||
consumedRateAdmissions: ConsumedRateAdmission[],
|
||||
): Promise<ToolAccessDecision> {
|
||||
// A decision that consumes a rate limit can also record a trust-rule hit
|
||||
// from the policy row, so it always reads fresh rows.
|
||||
const cache = input.consumeRateLimit === true ? undefined : options.cache;
|
||||
const cache = input.consumeRateLimit === true || options.admittedInvocationId ? undefined : options.cache;
|
||||
const loaded = await loadContext(input, cache);
|
||||
if (!loaded.ok) return loaded.decision;
|
||||
const { ctx, redaction } = loaded;
|
||||
// A final handoff rechecks current permissions for work already admitted by
|
||||
// the gateway. Its own last consumed slot cannot reject that same work.
|
||||
// Only durable consumption of an unchanged rule retains that slot; new or
|
||||
// edited rules still apply. This option never approves or consumes a call.
|
||||
const admittedRates = input.consumeRateLimit === false && options.admittedInvocationId
|
||||
? await admittedRateLimitAdmissions(ctx, options.admittedInvocationId) : [];
|
||||
const profileState = await effectiveProfiles(ctx, cache);
|
||||
const effectiveProfileIds = profileState.profiles.map((profile) => profile.id);
|
||||
const permittedByProfile = profileState.profiles.some((profile) => {
|
||||
@@ -1422,10 +1471,13 @@ export function toolAccessPolicyService(db: Db) {
|
||||
{ redactionPlan: redaction.redactionPlan, policyExplanation },
|
||||
);
|
||||
}
|
||||
const versionHash = ratePolicyVersionHash(policy);
|
||||
if (admittedRates.some(admission => admission.policyId === policy.id && admission.versionHash === versionHash)) continue;
|
||||
const state = await enforceRateLimit(policy, ctx, input.consumeRateLimit === true);
|
||||
if (state?.limited) {
|
||||
return decision("rate_limited", "rate_limited", "Tool access rate limit exceeded.", effectiveProfileIds, [policy.id], { rateLimitState: state, redactionPlan: redaction.redactionPlan, policyExplanation });
|
||||
}
|
||||
if (state && input.consumeRateLimit === true) consumedRateAdmissions.push({ policyId: policy.id, versionHash });
|
||||
continue;
|
||||
}
|
||||
if (policy.policyType === "trust_rule") {
|
||||
@@ -1608,6 +1660,7 @@ export function toolAccessPolicyService(db: Db) {
|
||||
agentId: ctx.agentId,
|
||||
issueId: ctx.issueId,
|
||||
runId: ctx.heartbeatRunId,
|
||||
gatewayId: ctx.gatewayId,
|
||||
applicationId: ctx.applicationId,
|
||||
connectionId: ctx.connectionId,
|
||||
catalogEntryId: ctx.catalogEntryId,
|
||||
@@ -1621,6 +1674,7 @@ export function toolAccessPolicyService(db: Db) {
|
||||
argumentsHash,
|
||||
argumentsSummary: redaction.summary,
|
||||
policyDecision: accessDecision.decision,
|
||||
policyExplanation: accessDecision.policyExplanation ?? null,
|
||||
matchedPolicyIds: accessDecision.matchedPolicyIds,
|
||||
approvalState: accessDecision.decision === "require_approval" ? "pending" : "not_required",
|
||||
status,
|
||||
|
||||
@@ -8,13 +8,17 @@ import { z } from "zod";
|
||||
import { and, eq, isNull, sql } from "drizzle-orm";
|
||||
import { heartbeatRuns, runUsageReceipts, type Db } from "@paperclipai/db";
|
||||
import type { AdapterExecutionResult, AdapterUsageCheckpoint } from "@paperclipai/adapter-utils";
|
||||
import { pricingProvenanceSchema, usdToUnits } from "@paperclipai/shared";
|
||||
import { compareCents, normalizeCents, pricingProvenanceSchema, usdToUnits } from "@paperclipai/shared";
|
||||
import { idleAccountingSpoolPath, registerIdleSpoolDirectory } from "./idle-local-work.js";
|
||||
import { conflict, notFound } from "../errors.js";
|
||||
import { logger } from "../middleware/logger.js";
|
||||
import { receiptFingerprint } from "./receipt-fingerprint.js";
|
||||
import { withAccountingTransaction } from "./accounting-transaction.js";
|
||||
import { ensureDurableDirectory, syncDirectory } from "../lib/durable-directory.js";
|
||||
import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notifications.js";
|
||||
import { knownRunCostCents } from "./budget-exposure.js";
|
||||
import { createLiveRunUsageProjection, liveRunUsageCostCents } from "./budget-live-usage.js";
|
||||
import { budgetServiceInTransaction } from "./budgets.js";
|
||||
|
||||
const token = z.number().int().nonnegative().max(2_147_483_647);
|
||||
const price = z.number().finite().nonnegative().nullable().optional();
|
||||
@@ -36,6 +40,11 @@ const envelopeSchema = z.object({
|
||||
schema: z.literal("paperclip/accounting-receipt/v1"), id: z.string().uuid(), companyId: z.string().uuid(),
|
||||
runId: z.string().uuid(), sourceId: z.string().uuid(), sequence: z.number().int().positive(), receivedAt: z.iso.datetime(),
|
||||
adapterType: z.string().max(100), receipt: checkpointSchema,
|
||||
// Derived by the recorder, not a provider's final charge. Keeping it in the
|
||||
// fsynced envelope preserves known partial attempts through database replay.
|
||||
liveCostCents: z.string().max(160).refine(value => {
|
||||
try { return compareCents(value, 0) >= 0; } catch { return false; }
|
||||
}, "Invalid live cost").transform(normalizeCents).nullable().optional(),
|
||||
});
|
||||
export type UsageReceiptEnvelope = z.infer<typeof envelopeSchema>;
|
||||
export const usageReceiptSpoolPath = idleAccountingSpoolPath;
|
||||
@@ -58,7 +67,7 @@ export async function spoolUsageReceipt(envelope: UsageReceiptEnvelope, director
|
||||
export async function persistUsageReceipt(db: Db, raw: UsageReceiptEnvelope) {
|
||||
const envelope = envelopeSchema.parse(raw);
|
||||
const hash = receiptFingerprint(envelope);
|
||||
return withAccountingTransaction(db, envelope.companyId, async tx => {
|
||||
return withAccountingTransaction(db, envelope.companyId, async (tx, publications) => {
|
||||
const [run] = await tx.select().from(heartbeatRuns).where(and(eq(heartbeatRuns.id, envelope.runId), eq(heartbeatRuns.companyId, envelope.companyId))).for("update");
|
||||
if (!run) throw notFound("Accounting run not found");
|
||||
const [existing] = await tx.select().from(runUsageReceipts).where(and(eq(runUsageReceipts.companyId, envelope.companyId), eq(runUsageReceipts.runId, envelope.runId), eq(runUsageReceipts.sourceId, envelope.sourceId), eq(runUsageReceipts.sequence, envelope.sequence)));
|
||||
@@ -73,6 +82,23 @@ export async function persistUsageReceipt(db: Db, raw: UsageReceiptEnvelope) {
|
||||
if (run.costAccountedAt || previous.accountingReceiptSourceId !== envelope.sourceId || Number(previous.accountingReceiptSequence ?? 0) >= envelope.sequence) return;
|
||||
const receipt = envelope.receipt;
|
||||
const billingType = receipt.billingType === "api" ? "metered_api" : receipt.billingType === "subscription" ? "subscription_included" : receipt.billingType ?? "unknown";
|
||||
const oldProjection = previous.budgetExposure && typeof previous.budgetExposure === "object"
|
||||
? previous.budgetExposure as Record<string, unknown> : {};
|
||||
const canonicalProjection = typeof oldProjection.sourceId === "string" && typeof oldProjection.sequence === "number";
|
||||
let previousCost = canonicalProjection ? null : knownRunCostCents(previous);
|
||||
if (typeof oldProjection.costCents === "string") {
|
||||
try {
|
||||
const projected = normalizeCents(oldProjection.costCents);
|
||||
if (compareCents(projected, 0) >= 0 && (canonicalProjection || compareCents(projected, previousCost ?? "0") > 0)) previousCost = projected;
|
||||
} catch { /* Older malformed projections are not billing evidence. */ }
|
||||
}
|
||||
const reportedCost = envelope.liveCostCents !== undefined ? envelope.liveCostCents
|
||||
: liveRunUsageCostCents({ ...receipt, billingType } as AdapterUsageCheckpoint);
|
||||
const observedCost = reportedCost === null ? previousCost
|
||||
: !receipt.complete && previousCost !== null && compareCents(previousCost, reportedCost) > 0 ? previousCost : reportedCost;
|
||||
// Running checkpoints cannot settle yet. Their terminal transition owns
|
||||
// the accounting wake; only a late receipt needs to reactivate that worker.
|
||||
if (!["queued", "running"].includes(run.status)) await notifyDeliveryWork(tx, DELIVERY_QUEUES.accounting, envelope.companyId);
|
||||
await tx.update(heartbeatRuns).set({
|
||||
costAccountingPending: true,
|
||||
accountingLastError: previous.accountingCaptureFailed === true ? "usage_capture_failed" : null,
|
||||
@@ -83,9 +109,15 @@ export async function persistUsageReceipt(db: Db, raw: UsageReceiptEnvelope) {
|
||||
accountingReceiptReady: previous.accountingCaptureFailed !== true && receipt.complete && receipt.usageBasis !== "session_cumulative", accountingReceiptId: envelope.id,
|
||||
accountingReceiptSequence: envelope.sequence, accountingReceiptReceivedAt: envelope.receivedAt,
|
||||
costStatus: receipt.costStatus ?? null,
|
||||
budgetExposure: { sourceId: envelope.sourceId, sequence: envelope.sequence, costCents: observedCost },
|
||||
pricingProvenance: receipt.pricingProvenance ?? { source: receipt.costUsd != null || receipt.costUsdExact != null ? "provider_reported" : "unknown", version: "accounting-receipt/v1" },
|
||||
},
|
||||
}).where(and(eq(heartbeatRuns.id, run.id), isNull(heartbeatRuns.costAccountedAt)));
|
||||
if (observedCost !== null && compareCents(observedCost, previousCost ?? "0") !== 0) {
|
||||
const scope = previous.ledgerScope && typeof previous.ledgerScope === "object" ? previous.ledgerScope as Record<string, unknown> : {};
|
||||
await budgetServiceInTransaction(tx, publications).evaluateAffectedScopes(envelope.companyId, run.agentId,
|
||||
typeof scope.projectId === "string" ? scope.projectId : null);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
@@ -202,6 +234,7 @@ export async function createRunUsageRecorder(db: Db, input: { companyId: string;
|
||||
let sequence = 0;
|
||||
let currentAttempt = "default";
|
||||
const attempts = new Map<string, AdapterUsageCheckpoint>();
|
||||
let initialLiveUsage: Parameters<typeof createLiveRunUsageProjection>[0];
|
||||
// Native run deltas retain their baseline across same-run recovery. Restore
|
||||
// the previous snapshot as the same attempt, never add it a second time.
|
||||
// An empty final result must not erase the last durable observation.
|
||||
@@ -210,8 +243,14 @@ export async function createRunUsageRecorder(db: Db, input: { companyId: string;
|
||||
const [previous] = await db.select().from(runUsageReceipts).where(and(
|
||||
eq(runUsageReceipts.id, priorReceiptId), eq(runUsageReceipts.runId, input.runId), eq(runUsageReceipts.companyId, input.companyId),
|
||||
));
|
||||
if (previous) attempts.set("default", envelopeSchema.parse(previous.receiptJson).receipt as AdapterUsageCheckpoint);
|
||||
if (previous) {
|
||||
const restored = envelopeSchema.parse(previous.receiptJson);
|
||||
const receipt = restored.receipt as AdapterUsageCheckpoint;
|
||||
attempts.set("default", receipt);
|
||||
initialLiveUsage = { receipt, costCents: restored.liveCostCents };
|
||||
}
|
||||
}
|
||||
const projectLiveUsage = createLiveRunUsageProjection(initialLiveUsage);
|
||||
let chain = Promise.resolve();
|
||||
let captureFailed = bound[0].usageJson?.accountingCaptureFailed === true;
|
||||
let finishFailedCapture: (() => void) | undefined;
|
||||
@@ -253,6 +292,7 @@ export async function createRunUsageRecorder(db: Db, input: { companyId: string;
|
||||
version: "accounting-receipt/v1",
|
||||
};
|
||||
attempts.set(attempt, priced);
|
||||
const liveCostCents = projectLiveUsage(attempt, priced);
|
||||
const parts = [...attempts.values()];
|
||||
const latest = attempts.get(currentAttempt)!;
|
||||
const receipt: AdapterUsageCheckpoint = { ...latest, complete: latest.complete && !captureFailed, usage: {
|
||||
@@ -289,10 +329,16 @@ export async function createRunUsageRecorder(db: Db, input: { companyId: string;
|
||||
// can be considered complete, even if a later attempt succeeded.
|
||||
receipt.complete = !captureFailed && parts.every(part => part.complete);
|
||||
}
|
||||
const envelope: UsageReceiptEnvelope = { schema: "paperclip/accounting-receipt/v1", id: randomUUID(), ...input, sourceId, sequence: ++sequence, receivedAt: new Date().toISOString(), receipt };
|
||||
const envelope: UsageReceiptEnvelope = { schema: "paperclip/accounting-receipt/v1", id: randomUUID(), ...input, sourceId, sequence: ++sequence, receivedAt: new Date().toISOString(), receipt, liveCostCents };
|
||||
const file = await spoolUsageReceipt(envelope, directory);
|
||||
try { await persistUsageReceipt(db, envelope); await fs.rm(file, { force: true }); }
|
||||
catch (error) { logger.warn({ err: error, runId: input.runId, receiptId: envelope.id }, "Usage receipt is durable on disk; database replay pending"); }
|
||||
catch (error) {
|
||||
// The fsynced receipt is durable even when the DB transaction failed.
|
||||
// Wake the owned spool recovery path; never certify that failed write as free.
|
||||
const { wakeCostAccountingWork } = await import("./budget-enforcement-worker.js");
|
||||
wakeCostAccountingWork(db);
|
||||
logger.warn({ err: error, runId: input.runId, receiptId: envelope.id }, "Usage receipt is durable on disk; database replay pending");
|
||||
}
|
||||
return receipt;
|
||||
}
|
||||
return {
|
||||
|
||||
Reference in new issue
Block a user