Harden budget enforcement and migration startup

cli/src/__tests__/company-delete.test.ts

@@ -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,

packages/db/src/migrations/0033_shiny_black_tarantula.sql

@@ -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;

packages/db/src/migrations/meta/_journal.json

@@ -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
     }
   ]
 }

packages/db/src/schema/companies.ts

@@ -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),

packages/shared/src/types/company.ts

@@ -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;

scripts/dev-runner.mjs

@@ -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,

server/package.json

@@ -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",

server/src/__tests__/budgets-service.test.ts

@@ -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.",
+    });
+  });
+});

server/src/__tests__/costs-service.test.ts

@@ -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,
 }));
 

server/src/index.ts

@@ -83,8 +83,7 @@ export async function startServer(): Promise<StartedServer> {
     | "skipped"
     | "already applied"
     | "applied (empty database)"
-    | "applied (pending migrations)"
+    | "applied (pending migrations)";
-    | "pending migrations skipped";
   
   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(
+        throw new Error(
-          { pendingMigrations: state.pendingMigrations },
+          `${label} has pending migrations (${formatPendingMigrationSummary(state.pendingMigrations)}). ` +
-          `${label} has pending migrations; continuing without applying. Run pnpm db:migrate to apply before startup.`,
+            "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(
+      throw new Error(
-        { pendingMigrations: state.pendingMigrations },
+        `${label} has pending migrations (${formatPendingMigrationSummary(state.pendingMigrations)}). ` +
-        `${label} has pending migrations; continuing without applying. Run pnpm db:migrate to apply before startup.`,
+          "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}`);

server/src/routes/costs.ts

@@ -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);
 

server/src/services/budgets.ts

@@ -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",
+      paused: row.status === "paused" || Boolean(row.pausedAt),
-      pauseReason: row.status === "paused" ? "budget" : null,
+      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.",
         };
       }
 

server/src/services/costs.ts

@@ -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) {
+export function costService(db: Db, budgetHooks: BudgetServiceHooks = {}) {
-  const budgets = budgetService(db);
+  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;

server/src/services/heartbeat.ts

@@ -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) => {
+    cancelRun: (runId: string) => cancelRunInternal(runId),
-      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);
+    cancelActiveForAgent: (agentId: string) => cancelActiveForAgentInternal(agentId),
-      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", {
+    cancelBudgetScopeWork,
-        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;
-    },
 
     getActiveRunForAgent: async (agentId: string) => {
       const [run] = await db