aevgarik/paperclip
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Harden budget enforcement and migration startup

Browse Source
This commit is contained in:
Dotta
2026-03-16 08:12:50 -05:00
parent 411952573e
commit 5f2c2ee0e2
15 changed files with 9473 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
cli/src/__tests__/company-delete.test.ts
View File

@@ -8,6 +8,8 @@ function makeCompany(overrides: Partial<Company>): Company {
name: "Alpha",
description: null,
status: "active",
pauseReason: null,
pausedAt: null,
issuePrefix: "ALP",
issueCounter: 1,
budgetMonthlyCents: 0,

2
packages/db/src/migrations/0033_shiny_black_tarantula.sql Normal file
View File

@@ -0,0 +1,2 @@
ALTER TABLE "companies" ADD COLUMN "pause_reason" text;--> statement-breakpoint
ALTER TABLE "companies" ADD COLUMN "paused_at" timestamp with time zone;

8934
packages/db/src/migrations/meta/0033_snapshot.json Normal file
View File

File diff suppressed because it is too large Load Diff

7
packages/db/src/migrations/meta/_journal.json
View File

@@ -232,6 +232,13 @@
"when": 1773542934499,
"tag": "0032_pretty_doctor_octopus",
"breakpoints": true
},
{
"idx": 33,
"version": "7",
"when": 1773664961967,
"tag": "0033_shiny_black_tarantula",
"breakpoints": true
}
]
}

2
packages/db/src/schema/companies.ts
View File

@@ -7,6 +7,8 @@ export const companies = pgTable(
name: text("name").notNull(),
description: text("description"),
status: text("status").notNull().default("active"),
pauseReason: text("pause_reason"),
pausedAt: timestamp("paused_at", { withTimezone: true }),
issuePrefix: text("issue_prefix").notNull().default("PAP"),
issueCounter: integer("issue_counter").notNull().default(0),
budgetMonthlyCents: integer("budget_monthly_cents").notNull().default(0),

4
packages/shared/src/types/company.ts
View File

@@ -1,10 +1,12 @@
import type { CompanyStatus } from "../constants.js";
import type { CompanyStatus, PauseReason } from "../constants.js";
export interface Company {
id: string;
name: string;
description: string | null;
status: CompanyStatus;
pauseReason: PauseReason | null;
pausedAt: Date | null;
issuePrefix: string;
issueCounter: number;
budgetMonthlyCents: number;

20
scripts/dev-runner.mjs
View File

@@ -34,6 +34,11 @@ const env = {
PAPERCLIP_UI_DEV_MIDDLEWARE: "true",
};
if (mode === "watch") {
env.PAPERCLIP_MIGRATION_PROMPT ??= "never";
env.PAPERCLIP_MIGRATION_AUTO_APPLY ??= "true";
}
if (tailscaleAuth) {
env.PAPERCLIP_DEPLOYMENT_MODE = "authenticated";
env.PAPERCLIP_DEPLOYMENT_EXPOSURE = "private";
@@ -113,7 +118,6 @@ async function runPnpm(args, options = {}) {
async function maybePreflightMigrations() {
if (mode !== "watch") return;
if (process.env.PAPERCLIP_MIGRATION_PROMPT === "never") return;
const status = await runPnpm(
["--filter", "@paperclipai/db", "exec", "tsx", "src/migration-status.ts", "--json"],
@@ -144,7 +148,7 @@ async function maybePreflightMigrations() {
return;
}
const autoApply = process.env.PAPERCLIP_MIGRATION_AUTO_APPLY === "true";
const autoApply = env.PAPERCLIP_MIGRATION_AUTO_APPLY === "true";
let shouldApply = autoApply;
if (!autoApply) {
@@ -167,7 +171,13 @@ async function maybePreflightMigrations() {
}
}
if (!shouldApply) return;
if (!shouldApply) {
process.stderr.write(
`[paperclip] Pending migrations detected (${formatPendingMigrationSummary(payload.pendingMigrations)}). ` +
"Refusing to start watch mode against a stale schema.\n",
);
process.exit(1);
}
const migrate = spawn(pnpmBin, ["db:migrate"], {
stdio: "inherit",
@@ -206,10 +216,6 @@ async function buildPluginSdk() {
await buildPluginSdk();
if (mode === "watch") {
env.PAPERCLIP_MIGRATION_PROMPT = "never";
}
const serverScript = mode === "watch" ? "dev:watch" : "dev";
const child = spawn(
pnpmBin,

2
server/package.json
View File

@@ -23,7 +23,7 @@
],
"scripts": {
"dev": "tsx src/index.ts",
"dev:watch": "cross-env PAPERCLIP_MIGRATION_PROMPT=never tsx watch --ignore ../ui/node_modules --ignore ../ui/.vite --ignore ../ui/dist src/index.ts",
"dev:watch": "cross-env PAPERCLIP_MIGRATION_PROMPT=never PAPERCLIP_MIGRATION_AUTO_APPLY=true tsx watch --ignore ../ui/node_modules --ignore ../ui/.vite --ignore ../ui/dist src/index.ts",
"prepare:ui-dist": "bash ../scripts/prepare-server-ui-dist.sh",
"build": "tsc",
"prepack": "pnpm run prepare:ui-dist",

221
server/src/__tests__/budgets-service.test.ts Normal file
View File

@@ -0,0 +1,221 @@
import { beforeEach, describe, expect, it, vi } from "vitest";
import { budgetService } from "../services/budgets.ts";
const mockLogActivity = vi.hoisted(() => vi.fn());
vi.mock("../services/activity-log.js", () => ({
logActivity: mockLogActivity,
}));
type SelectResult = unknown[];
function createDbStub(selectResults: SelectResult[]) {
const pendingSelects = [...selectResults];
const selectWhere = vi.fn(async () => pendingSelects.shift() ?? []);
const selectThen = vi.fn((resolve: (value: unknown[]) => unknown) => Promise.resolve(resolve(pendingSelects.shift() ?? [])));
const selectOrderBy = vi.fn(async () => pendingSelects.shift() ?? []);
const selectFrom = vi.fn(() => ({
where: selectWhere,
then: selectThen,
orderBy: selectOrderBy,
}));
const select = vi.fn(() => ({
from: selectFrom,
}));
const insertValues = vi.fn();
const insertReturning = vi.fn(async () => pendingInserts.shift() ?? []);
const insert = vi.fn(() => ({
values: insertValues.mockImplementation(() => ({
returning: insertReturning,
})),
}));
const updateSet = vi.fn();
const updateWhere = vi.fn(async () => pendingUpdates.shift() ?? []);
const update = vi.fn(() => ({
set: updateSet.mockImplementation(() => ({
where: updateWhere,
})),
}));
const pendingInserts: unknown[][] = [];
const pendingUpdates: unknown[][] = [];
return {
db: {
select,
insert,
update,
},
queueInsert: (rows: unknown[]) => {
pendingInserts.push(rows);
},
queueUpdate: (rows: unknown[] = []) => {
pendingUpdates.push(rows);
},
selectWhere,
insertValues,
updateSet,
};
}
describe("budgetService", () => {
beforeEach(() => {
vi.clearAllMocks();
});
it("creates a hard-stop incident and pauses an agent when spend exceeds a budget", async () => {
const policy = {
id: "policy-1",
companyId: "company-1",
scopeType: "agent",
scopeId: "agent-1",
metric: "billed_cents",
windowKind: "calendar_month_utc",
amount: 100,
warnPercent: 80,
hardStopEnabled: true,
notifyEnabled: false,
isActive: true,
};
const dbStub = createDbStub([
[policy],
[{ total: 150 }],
[],
[{
companyId: "company-1",
name: "Budget Agent",
status: "running",
pauseReason: null,
}],
]);
dbStub.queueInsert([{
id: "approval-1",
companyId: "company-1",
status: "pending",
}]);
dbStub.queueInsert([{
id: "incident-1",
companyId: "company-1",
policyId: "policy-1",
approvalId: "approval-1",
}]);
dbStub.queueUpdate([]);
const cancelWorkForScope = vi.fn().mockResolvedValue(undefined);
const service = budgetService(dbStub.db as any, { cancelWorkForScope });
await service.evaluateCostEvent({
companyId: "company-1",
agentId: "agent-1",
projectId: null,
} as any);
expect(dbStub.insertValues).toHaveBeenCalledWith(
expect.objectContaining({
companyId: "company-1",
type: "budget_override_required",
status: "pending",
}),
);
expect(dbStub.insertValues).toHaveBeenCalledWith(
expect.objectContaining({
companyId: "company-1",
policyId: "policy-1",
thresholdType: "hard",
amountLimit: 100,
amountObserved: 150,
approvalId: "approval-1",
}),
);
expect(dbStub.updateSet).toHaveBeenCalledWith(
expect.objectContaining({
status: "paused",
pauseReason: "budget",
pausedAt: expect.any(Date),
}),
);
expect(mockLogActivity).toHaveBeenCalledWith(
expect.anything(),
expect.objectContaining({
action: "budget.hard_threshold_crossed",
entityId: "incident-1",
}),
);
expect(cancelWorkForScope).toHaveBeenCalledWith({
companyId: "company-1",
scopeType: "agent",
scopeId: "agent-1",
});
});
it("blocks new work when an agent hard-stop remains exceeded even if the agent is not paused yet", async () => {
const agentPolicy = {
id: "policy-agent-1",
companyId: "company-1",
scopeType: "agent",
scopeId: "agent-1",
metric: "billed_cents",
windowKind: "calendar_month_utc",
amount: 100,
warnPercent: 80,
hardStopEnabled: true,
notifyEnabled: true,
isActive: true,
};
const dbStub = createDbStub([
[{
status: "running",
pauseReason: null,
companyId: "company-1",
name: "Budget Agent",
}],
[{
status: "active",
name: "Paperclip",
}],
[],
[agentPolicy],
[{ total: 120 }],
]);
const service = budgetService(dbStub.db as any);
const block = await service.getInvocationBlock("company-1", "agent-1");
expect(block).toEqual({
scopeType: "agent",
scopeId: "agent-1",
scopeName: "Budget Agent",
reason: "Agent cannot start because its budget hard-stop is still exceeded.",
});
});
it("surfaces a budget-owned company pause distinctly from a manual pause", async () => {
const dbStub = createDbStub([
[{
status: "idle",
pauseReason: null,
companyId: "company-1",
name: "Budget Agent",
}],
[{
status: "paused",
pauseReason: "budget",
name: "Paperclip",
}],
]);
const service = budgetService(dbStub.db as any);
const block = await service.getInvocationBlock("company-1", "agent-1");
expect(block).toEqual({
scopeType: "company",
scopeId: "company-1",
scopeName: "Paperclip",
reason: "Company is paused because its budget hard-stop was reached.",
});
});
});

4
server/src/__tests__/costs-service.test.ts
View File

@@ -37,6 +37,9 @@ const mockAgentService = vi.hoisted(() => ({
getById: vi.fn(),
update: vi.fn(),
}));
const mockHeartbeatService = vi.hoisted(() => ({
cancelBudgetScopeWork: vi.fn().mockResolvedValue(undefined),
}));
const mockLogActivity = vi.hoisted(() => vi.fn());
const mockFetchAllQuotaWindows = vi.hoisted(() => vi.fn());
const mockCostService = vi.hoisted(() => ({
@@ -75,6 +78,7 @@ vi.mock("../services/index.js", () => ({
financeService: () => mockFinanceService,
companyService: () => mockCompanyService,
agentService: () => mockAgentService,
heartbeatService: () => mockHeartbeatService,
logActivity: mockLogActivity,
}));

17
server/src/index.ts
View File

@@ -83,8 +83,7 @@ export async function startServer(): Promise<StartedServer> {
| "skipped"
| "already applied"
| "applied (empty database)"
| "applied (pending migrations)"
| "pending migrations skipped";
| "applied (pending migrations)";
function formatPendingMigrationSummary(migrations: string[]): string {
if (migrations.length === 0) return "none";
@@ -139,11 +138,10 @@ export async function startServer(): Promise<StartedServer> {
);
const apply = autoApply ? true : await promptApplyMigrations(state.pendingMigrations);
if (!apply) {
logger.warn(
{ pendingMigrations: state.pendingMigrations },
`${label} has pending migrations; continuing without applying. Run pnpm db:migrate to apply before startup.`,
throw new Error(
`${label} has pending migrations (${formatPendingMigrationSummary(state.pendingMigrations)}). ` +
"Refusing to start against a stale schema. Run pnpm db:migrate or set PAPERCLIP_MIGRATION_AUTO_APPLY=true.",
);
return "pending migrations skipped";
}
logger.info({ pendingMigrations: state.pendingMigrations }, `Applying ${state.pendingMigrations.length} pending migrations for ${label}`);
@@ -153,11 +151,10 @@ export async function startServer(): Promise<StartedServer> {
const apply = autoApply ? true : await promptApplyMigrations(state.pendingMigrations);
if (!apply) {
logger.warn(
{ pendingMigrations: state.pendingMigrations },
`${label} has pending migrations; continuing without applying. Run pnpm db:migrate to apply before startup.`,
throw new Error(
`${label} has pending migrations (${formatPendingMigrationSummary(state.pendingMigrations)}). ` +
"Refusing to start against a stale schema. Run pnpm db:migrate or set PAPERCLIP_MIGRATION_AUTO_APPLY=true.",
);
return "pending migrations skipped";
}
logger.info({ pendingMigrations: state.pendingMigrations }, `Applying ${state.pendingMigrations.length} pending migrations for ${label}`);

9
server/src/routes/costs.ts
View File

@@ -14,6 +14,7 @@ import {
financeService,
companyService,
agentService,
heartbeatService,
logActivity,
} from "../services/index.js";
import { assertBoard, assertCompanyAccess, getActorInfo } from "./authz.js";
@@ -22,9 +23,13 @@ import { badRequest } from "../errors.js";
export function costRoutes(db: Db) {
const router = Router();
const costs = costService(db);
const heartbeat = heartbeatService(db);
const budgetHooks = {
cancelWorkForScope: heartbeat.cancelBudgetScopeWork,
};
const costs = costService(db, budgetHooks);
const finance = financeService(db);
const budgets = budgetService(db);
const budgets = budgetService(db, budgetHooks);
const companies = companyService(db);
const agents = agentService(db);

43
server/src/services/budgets.ts
View File

@@ -34,6 +34,16 @@ type ScopeRecord = {
type PolicyRow = typeof budgetPolicies.$inferSelect;
type IncidentRow = typeof budgetIncidents.$inferSelect;
export type BudgetEnforcementScope = {
companyId: string;
scopeType: BudgetScopeType;
scopeId: string;
};
export type BudgetServiceHooks = {
cancelWorkForScope?: (scope: BudgetEnforcementScope) => Promise<void>;
};
function currentUtcMonthWindow(now = new Date()) {
const year = now.getUTCFullYear();
const month = now.getUTCMonth();
@@ -75,6 +85,8 @@ async function resolveScopeRecord(db: Db, scopeType: BudgetScopeType, scopeId: s
companyId: companies.id,
name: companies.name,
status: companies.status,
pauseReason: companies.pauseReason,
pausedAt: companies.pausedAt,
})
.from(companies)
.where(eq(companies.id, scopeId))
@@ -83,8 +95,8 @@ async function resolveScopeRecord(db: Db, scopeType: BudgetScopeType, scopeId: s
return {
companyId: row.companyId,
name: row.name,
paused: row.status === "paused",
pauseReason: row.status === "paused" ? "budget" : null,
paused: row.status === "paused" || Boolean(row.pausedAt),
pauseReason: (row.pauseReason as ScopeRecord["pauseReason"]) ?? null,
};
}
@@ -197,7 +209,7 @@ async function markApprovalStatus(
.where(eq(approvals.id, approvalId));
}
export function budgetService(db: Db) {
export function budgetService(db: Db, hooks: BudgetServiceHooks = {}) {
async function pauseScopeForBudget(policy: PolicyRow) {
const now = new Date();
if (policy.scopeType === "agent") {
@@ -229,11 +241,22 @@ export function budgetService(db: Db) {
.update(companies)
.set({
status: "paused",
pauseReason: "budget",
pausedAt: now,
updatedAt: now,
})
.where(eq(companies.id, policy.scopeId));
}
async function pauseAndCancelScopeForBudget(policy: PolicyRow) {
await pauseScopeForBudget(policy);
await hooks.cancelWorkForScope?.({
companyId: policy.companyId,
scopeType: policy.scopeType as BudgetScopeType,
scopeId: policy.scopeId,
});
}
async function resumeScopeFromBudget(policy: PolicyRow) {
const now = new Date();
if (policy.scopeType === "agent") {
@@ -265,9 +288,11 @@ export function budgetService(db: Db) {
.update(companies)
.set({
status: "active",
pauseReason: null,
pausedAt: null,
updatedAt: now,
})
.where(eq(companies.id, policy.scopeId));
.where(and(eq(companies.id, policy.scopeId), eq(companies.pauseReason, "budget")));
}
async function getPolicyRow(policyId: string) {
@@ -573,7 +598,7 @@ export function budgetService(db: Db) {
if (row.hardStopEnabled && observedAmount >= row.amount) {
await resolveOpenSoftIncidents(row.id);
await createIncidentIfNeeded(row, "hard", observedAmount);
await pauseScopeForBudget(row);
await pauseAndCancelScopeForBudget(row);
}
}
} else {
@@ -665,7 +690,7 @@ export function budgetService(db: Db) {
if (policy.hardStopEnabled && observedAmount >= policy.amount) {
await resolveOpenSoftIncidents(policy.id);
const hardIncident = await createIncidentIfNeeded(policy, "hard", observedAmount);
await pauseScopeForBudget(policy);
await pauseAndCancelScopeForBudget(policy);
if (hardIncident) {
await logActivity(db, {
companyId: policy.companyId,
@@ -707,6 +732,7 @@ export function budgetService(db: Db) {
const company = await db
.select({
status: companies.status,
pauseReason: companies.pauseReason,
name: companies.name,
})
.from(companies)
@@ -718,7 +744,10 @@ export function budgetService(db: Db) {
scopeType: "company" as const,
scopeId: companyId,
scopeName: company.name,
reason: "Company is paused and cannot start new work.",
reason:
company.pauseReason === "budget"
? "Company is paused because its budget hard-stop was reached."
: "Company is paused and cannot start new work.",
};
}

25
server/src/services/costs.ts
View File

@@ -2,7 +2,7 @@ import { and, desc, eq, gte, isNotNull, lte, sql } from "drizzle-orm";
import type { Db } from "@paperclipai/db";
import { activityLog, agents, companies, costEvents, issues, projects } from "@paperclipai/db";
import { notFound, unprocessable } from "../errors.js";
import { budgetService } from "./budgets.js";
import { budgetService, type BudgetServiceHooks } from "./budgets.js";
export interface CostDateRange {
from?: Date;
@@ -12,8 +12,8 @@ export interface CostDateRange {
const METERED_BILLING_TYPE = "metered_api";
const SUBSCRIPTION_BILLING_TYPES = ["subscription_included", "subscription_overage"] as const;
export function costService(db: Db) {
const budgets = budgetService(db);
export function costService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
const budgets = budgetService(db, budgetHooks);
return {
createEvent: async (companyId: string, data: Omit<typeof costEvents.$inferInsert, "companyId">) => {
const agent = await db
@@ -55,25 +55,6 @@ export function costService(db: Db) {
})
.where(eq(companies.id, companyId));
const updatedAgent = await db
.select()
.from(agents)
.where(eq(agents.id, event.agentId))
.then((rows) => rows[0] ?? null);
if (
updatedAgent &&
updatedAgent.budgetMonthlyCents > 0 &&
updatedAgent.spentMonthlyCents >= updatedAgent.budgetMonthlyCents &&
updatedAgent.status !== "paused" &&
updatedAgent.status !== "terminated"
) {
await db
.update(agents)
.set({ status: "paused", updatedAt: new Date() })
.where(eq(agents.id, updatedAgent.id));
}
await budgets.evaluateCostEvent(event);
return event;

303
server/src/services/heartbeat.ts
View File

@@ -23,7 +23,7 @@ import type { AdapterExecutionResult, AdapterInvocationMeta, AdapterSessionCodec
import { createLocalAgentJwt } from "../agent-auth-jwt.js";
import { parseObject, asBoolean, asNumber, appendWithCap, MAX_EXCERPT_BYTES } from "../adapters/utils.js";
import { costService } from "./costs.js";
import { budgetService } from "./budgets.js";
import { budgetService, type BudgetEnforcementScope } from "./budgets.js";
import { secretService } from "./secrets.js";
import { resolveDefaultAgentWorkspaceDir } from "../home-paths.js";
import { summarizeHeartbeatRunResultJson } from "./heartbeat-run-summary.js";
@@ -617,10 +617,13 @@ function resolveNextSessionState(input: {
export function heartbeatService(db: Db) {
const runLogStore = getRunLogStore();
const budgets = budgetService(db);
const secretsSvc = secretService(db);
const issuesSvc = issueService(db);
const activeRunExecutions = new Set<string>();
const budgetHooks = {
cancelWorkForScope: cancelBudgetScopeWork,
};
const budgets = budgetService(db, budgetHooks);
async function getAgent(agentId: string) {
return db
@@ -1203,6 +1206,26 @@ export function heartbeatService(db: Db) {
async function claimQueuedRun(run: typeof heartbeatRuns.$inferSelect) {
if (run.status !== "queued") return run;
const agent = await getAgent(run.agentId);
if (!agent) {
await cancelRunInternal(run.id, "Cancelled because the agent no longer exists");
return null;
}
if (agent.status === "paused" || agent.status === "terminated" || agent.status === "pending_approval") {
await cancelRunInternal(run.id, "Cancelled because the agent is not invokable");
return null;
}
const context = parseObject(run.contextSnapshot);
const budgetBlock = await budgets.getInvocationBlock(run.companyId, run.agentId, {
issueId: readNonEmptyString(context.issueId),
projectId: readNonEmptyString(context.projectId),
});
if (budgetBlock) {
await cancelRunInternal(run.id, budgetBlock.reason);
return null;
}
const claimedAt = new Date();
const claimed = await db
.update(heartbeatRuns)
@@ -1382,7 +1405,7 @@ export function heartbeatService(db: Db) {
.where(eq(agentRuntimeState.agentId, agent.id));
if (additionalCostCents > 0 || hasTokenUsage) {
const costs = costService(db);
const costs = costService(db, budgetHooks);
await costs.createEvent(agent.companyId, {
heartbeatRunId: run.id,
agentId: agent.id,
@@ -1405,6 +1428,9 @@ export function heartbeatService(db: Db) {
return withAgentStartLock(agentId, async () => {
const agent = await getAgent(agentId);
if (!agent) return [];
if (agent.status === "paused" || agent.status === "terminated" || agent.status === "pending_approval") {
return [];
}
const policy = parseHeartbeatPolicy(agent);
const runningCount = await countRunningRunsForAgent(agentId);
const availableSlots = Math.max(0, policy.maxConcurrentRuns - runningCount);
@@ -2758,6 +2784,205 @@ export function heartbeatService(db: Db) {
return newRun;
}
async function listProjectScopedRunIds(companyId: string, projectId: string) {
const runIssueId = sql<string | null>`${heartbeatRuns.contextSnapshot} ->> 'issueId'`;
const effectiveProjectId = sql<string | null>`coalesce(${heartbeatRuns.contextSnapshot} ->> 'projectId', ${issues.projectId}::text)`;
const rows = await db
.selectDistinctOn([heartbeatRuns.id], { id: heartbeatRuns.id })
.from(heartbeatRuns)
.leftJoin(
issues,
and(
eq(issues.companyId, companyId),
sql`${issues.id}::text = ${runIssueId}`,
),
)
.where(
and(
eq(heartbeatRuns.companyId, companyId),
inArray(heartbeatRuns.status, ["queued", "running"]),
sql`${effectiveProjectId} = ${projectId}`,
),
);
return rows.map((row) => row.id);
}
async function listProjectScopedWakeupIds(companyId: string, projectId: string) {
const wakeIssueId = sql<string | null>`${agentWakeupRequests.payload} ->> 'issueId'`;
const effectiveProjectId = sql<string | null>`coalesce(${agentWakeupRequests.payload} ->> 'projectId', ${issues.projectId}::text)`;
const rows = await db
.selectDistinctOn([agentWakeupRequests.id], { id: agentWakeupRequests.id })
.from(agentWakeupRequests)
.leftJoin(
issues,
and(
eq(issues.companyId, companyId),
sql`${issues.id}::text = ${wakeIssueId}`,
),
)
.where(
and(
eq(agentWakeupRequests.companyId, companyId),
inArray(agentWakeupRequests.status, ["queued", "deferred_issue_execution"]),
sql`${agentWakeupRequests.runId} is null`,
sql`${effectiveProjectId} = ${projectId}`,
),
);
return rows.map((row) => row.id);
}
async function cancelPendingWakeupsForBudgetScope(scope: BudgetEnforcementScope) {
const now = new Date();
let wakeupIds: string[] = [];
if (scope.scopeType === "company") {
wakeupIds = await db
.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 db
.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);
}
if (wakeupIds.length === 0) return 0;
await db
.update(agentWakeupRequests)
.set({
status: "cancelled",
finishedAt: now,
error: "Cancelled due to budget pause",
updatedAt: now,
})
.where(inArray(agentWakeupRequests.id, wakeupIds));
return wakeupIds.length;
}
async function cancelRunInternal(runId: string, reason = "Cancelled by control plane") {
const run = await getRun(runId);
if (!run) throw notFound("Heartbeat run not found");
if (run.status !== "running" && run.status !== "queued") return run;
const running = runningProcesses.get(run.id);
if (running) {
running.child.kill("SIGTERM");
const graceMs = Math.max(1, running.graceSec) * 1000;
setTimeout(() => {
if (!running.child.killed) {
running.child.kill("SIGKILL");
}
}, graceMs);
}
const cancelled = await setRunStatus(run.id, "cancelled", {
finishedAt: new Date(),
error: reason,
errorCode: "cancelled",
});
await setWakeupStatus(run.wakeupRequestId, "cancelled", {
finishedAt: new Date(),
error: reason,
});
if (cancelled) {
await appendRunEvent(cancelled, 1, {
eventType: "lifecycle",
stream: "system",
level: "warn",
message: "run cancelled",
});
await releaseIssueExecutionAndPromote(cancelled);
}
runningProcesses.delete(run.id);
await finalizeAgentStatus(run.agentId, "cancelled");
await startNextQueuedRunForAgent(run.agentId);
return cancelled;
}
async function cancelActiveForAgentInternal(agentId: string, reason = "Cancelled due to agent pause") {
const runs = await db
.select()
.from(heartbeatRuns)
.where(and(eq(heartbeatRuns.agentId, agentId), inArray(heartbeatRuns.status, ["queued", "running"])));
for (const run of runs) {
await setRunStatus(run.id, "cancelled", {
finishedAt: new Date(),
error: reason,
errorCode: "cancelled",
});
await setWakeupStatus(run.wakeupRequestId, "cancelled", {
finishedAt: new Date(),
error: reason,
});
const running = runningProcesses.get(run.id);
if (running) {
running.child.kill("SIGTERM");
runningProcesses.delete(run.id);
}
await releaseIssueExecutionAndPromote(run);
}
return runs.length;
}
async function cancelBudgetScopeWork(scope: BudgetEnforcementScope) {
if (scope.scopeType === "agent") {
await cancelActiveForAgentInternal(scope.scopeId, "Cancelled due to budget pause");
await cancelPendingWakeupsForBudgetScope(scope);
return;
}
const runIds =
scope.scopeType === "company"
? await db
.select({ id: heartbeatRuns.id })
.from(heartbeatRuns)
.where(
and(
eq(heartbeatRuns.companyId, scope.companyId),
inArray(heartbeatRuns.status, ["queued", "running"]),
),
)
.then((rows) => rows.map((row) => row.id))
: await listProjectScopedRunIds(scope.companyId, scope.scopeId);
for (const runId of runIds) {
await cancelRunInternal(runId, "Cancelled due to budget pause");
}
await cancelPendingWakeupsForBudgetScope(scope);
}
return {
list: async (companyId: string, agentId?: string, limit?: number) => {
const query = db
@@ -2930,77 +3155,11 @@ export function heartbeatService(db: Db) {
return { checked, enqueued, skipped };
},
cancelRun: async (runId: string) => {
const run = await getRun(runId);
if (!run) throw notFound("Heartbeat run not found");
if (run.status !== "running" && run.status !== "queued") return run;
cancelRun: (runId: string) => cancelRunInternal(runId),
const running = runningProcesses.get(run.id);
if (running) {
running.child.kill("SIGTERM");
const graceMs = Math.max(1, running.graceSec) * 1000;
setTimeout(() => {
if (!running.child.killed) {
running.child.kill("SIGKILL");
}
}, graceMs);
}
cancelActiveForAgent: (agentId: string) => cancelActiveForAgentInternal(agentId),
const cancelled = await setRunStatus(run.id, "cancelled", {
finishedAt: new Date(),
error: "Cancelled by control plane",
errorCode: "cancelled",
});
await setWakeupStatus(run.wakeupRequestId, "cancelled", {
finishedAt: new Date(),
error: "Cancelled by control plane",
});
if (cancelled) {
await appendRunEvent(cancelled, 1, {
eventType: "lifecycle",
stream: "system",
level: "warn",
message: "run cancelled",
});
await releaseIssueExecutionAndPromote(cancelled);
}
runningProcesses.delete(run.id);
await finalizeAgentStatus(run.agentId, "cancelled");
await startNextQueuedRunForAgent(run.agentId);
return cancelled;
},
cancelActiveForAgent: async (agentId: string) => {
const runs = await db
.select()
.from(heartbeatRuns)
.where(and(eq(heartbeatRuns.agentId, agentId), inArray(heartbeatRuns.status, ["queued", "running"])));
for (const run of runs) {
await setRunStatus(run.id, "cancelled", {
finishedAt: new Date(),
error: "Cancelled due to agent pause",
errorCode: "cancelled",
});
await setWakeupStatus(run.wakeupRequestId, "cancelled", {
finishedAt: new Date(),
error: "Cancelled due to agent pause",
});
const running = runningProcesses.get(run.id);
if (running) {
running.child.kill("SIGTERM");
runningProcesses.delete(run.id);
}
await releaseIssueExecutionAndPromote(run);
}
return runs.length;
},
cancelBudgetScopeWork,
getActiveRunForAgent: async (agentId: string) => {
const [run] = await db