Harden budget enforcement and migration startup

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