Fix budget incident resolution edge cases

packages/db/src/migrations/0034_fat_dormammu.sql

@@ -0,0 +1,2 @@
+DROP INDEX "budget_incidents_policy_window_threshold_idx";--> statement-breakpoint
+CREATE UNIQUE INDEX "budget_incidents_policy_window_threshold_idx" ON "budget_incidents" USING btree ("policy_id","window_start","threshold_type") WHERE "budget_incidents"."status" <> 'dismissed';

packages/db/src/migrations/meta/0033_snapshot.json

@@ -8918,6 +8918,110 @@
       "policies": {},
       "checkConstraints": {},
       "isRLSEnabled": false
+    },
+    "public.company_logos": {
+      "name": "company_logos",
+      "schema": "",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": true,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "company_id": {
+          "name": "company_id",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "asset_id": {
+          "name": "asset_id",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "created_at": {
+          "name": "created_at",
+          "type": "timestamp with time zone",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "now()"
+        },
+        "updated_at": {
+          "name": "updated_at",
+          "type": "timestamp with time zone",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "now()"
+        }
+      },
+      "indexes": {
+        "company_logos_company_uq": {
+          "name": "company_logos_company_uq",
+          "columns": [
+            {
+              "expression": "company_id",
+              "isExpression": false,
+              "asc": true,
+              "nulls": "last"
+            }
+          ],
+          "isUnique": true,
+          "concurrently": false,
+          "method": "btree",
+          "with": {}
+        },
+        "company_logos_asset_uq": {
+          "name": "company_logos_asset_uq",
+          "columns": [
+            {
+              "expression": "asset_id",
+              "isExpression": false,
+              "asc": true,
+              "nulls": "last"
+            }
+          ],
+          "isUnique": true,
+          "concurrently": false,
+          "method": "btree",
+          "with": {}
+        }
+      },
+      "foreignKeys": {
+        "company_logos_company_id_companies_id_fk": {
+          "name": "company_logos_company_id_companies_id_fk",
+          "tableFrom": "company_logos",
+          "tableTo": "companies",
+          "columnsFrom": [
+            "company_id"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "cascade",
+          "onUpdate": "no action"
+        },
+        "company_logos_asset_id_assets_id_fk": {
+          "name": "company_logos_asset_id_assets_id_fk",
+          "tableFrom": "company_logos",
+          "tableTo": "assets",
+          "columnsFrom": [
+            "asset_id"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "cascade",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {},
+      "policies": {},
+      "checkConstraints": {},
+      "isRLSEnabled": false
     }
   },
   "enums": {},

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

@@ -239,6 +239,13 @@
       "when": 1773664961967,
       "tag": "0033_shiny_black_tarantula",
       "breakpoints": true
+    },
+    {
+      "idx": 34,
+      "version": "7",
+      "when": 1773697572188,
+      "tag": "0034_fat_dormammu",
+      "breakpoints": true
     }
   ]
-}
+}

packages/db/src/schema/budget_incidents.ts

@@ -1,3 +1,4 @@
+import { sql } from "drizzle-orm";
 import { index, integer, pgTable, text, timestamp, uuid, uniqueIndex } from "drizzle-orm/pg-core";
 import { approvals } from "./approvals.js";
 import { budgetPolicies } from "./budget_policies.js";
@@ -36,6 +37,6 @@ export const budgetIncidents = pgTable(
       table.policyId,
       table.windowStart,
       table.thresholdType,
-    ),
+    ).where(sql`${table.status} <> 'dismissed'`),
   }),
 );

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

@@ -218,4 +218,94 @@ describe("budgetService", () => {
       reason: "Company is paused because its budget hard-stop was reached.",
     });
   });
+
+  it("uses live observed spend when raising a budget incident", async () => {
+    const dbStub = createDbStub([
+      [{
+        id: "incident-1",
+        companyId: "company-1",
+        policyId: "policy-1",
+        amountObserved: 120,
+        approvalId: "approval-1",
+      }],
+      [{
+        id: "policy-1",
+        companyId: "company-1",
+        scopeType: "company",
+        scopeId: "company-1",
+        metric: "billed_cents",
+        windowKind: "calendar_month_utc",
+      }],
+      [{ total: 150 }],
+    ]);
+
+    const service = budgetService(dbStub.db as any);
+
+    await expect(
+      service.resolveIncident(
+        "company-1",
+        "incident-1",
+        { action: "raise_budget_and_resume", amount: 140 },
+        "board-user",
+      ),
+    ).rejects.toThrow("New budget must exceed current observed spend");
+  });
+
+  it("syncs company monthly budget when raising and resuming a company incident", async () => {
+    const now = new Date();
+    const dbStub = createDbStub([
+      [{
+        id: "incident-1",
+        companyId: "company-1",
+        policyId: "policy-1",
+        scopeType: "company",
+        scopeId: "company-1",
+        metric: "billed_cents",
+        windowKind: "calendar_month_utc",
+        windowStart: now,
+        windowEnd: now,
+        thresholdType: "hard",
+        amountLimit: 100,
+        amountObserved: 120,
+        status: "open",
+        approvalId: "approval-1",
+        resolvedAt: null,
+        createdAt: now,
+        updatedAt: now,
+      }],
+      [{
+        id: "policy-1",
+        companyId: "company-1",
+        scopeType: "company",
+        scopeId: "company-1",
+        metric: "billed_cents",
+        windowKind: "calendar_month_utc",
+        amount: 100,
+      }],
+      [{ total: 120 }],
+      [{ id: "approval-1", status: "approved" }],
+      [{
+        companyId: "company-1",
+        name: "Paperclip",
+        status: "paused",
+        pauseReason: "budget",
+        pausedAt: now,
+      }],
+    ]);
+
+    const service = budgetService(dbStub.db as any);
+    await service.resolveIncident(
+      "company-1",
+      "incident-1",
+      { action: "raise_budget_and_resume", amount: 175 },
+      "board-user",
+    );
+
+    expect(dbStub.updateSet).toHaveBeenCalledWith(
+      expect.objectContaining({
+        budgetMonthlyCents: 175,
+        updatedAt: expect.any(Date),
+      }),
+    );
+  });
 });

server/src/services/budgets.ts

@@ -878,24 +878,33 @@ export function budgetService(db: Db, hooks: BudgetServiceHooks = {}) {
       const policy = await getPolicyRow(incident.policyId);
       if (input.action === "raise_budget_and_resume") {
         const nextAmount = Math.max(0, Math.floor(input.amount ?? 0));
-        if (nextAmount <= incident.amountObserved) {
+        const currentObserved = await computeObservedAmount(db, policy);
+        if (nextAmount <= currentObserved) {
           throw unprocessable("New budget must exceed current observed spend");
         }
 
+        const now = new Date();
         await db
           .update(budgetPolicies)
           .set({
             amount: nextAmount,
             isActive: true,
             updatedByUserId: actorUserId,
-            updatedAt: new Date(),
+            updatedAt: now,
           })
           .where(eq(budgetPolicies.id, policy.id));
 
+        if (policy.scopeType === "company" && policy.windowKind === "calendar_month_utc") {
+          await db
+            .update(companies)
+            .set({ budgetMonthlyCents: nextAmount, updatedAt: now })
+            .where(eq(companies.id, policy.scopeId));
+        }
+
         if (policy.scopeType === "agent" && policy.windowKind === "calendar_month_utc") {
           await db
             .update(agents)
-            .set({ budgetMonthlyCents: nextAmount, updatedAt: new Date() })
+            .set({ budgetMonthlyCents: nextAmount, updatedAt: now })
             .where(eq(agents.id, policy.scopeId));
         }
 
@@ -904,8 +913,8 @@ export function budgetService(db: Db, hooks: BudgetServiceHooks = {}) {
           .update(budgetIncidents)
           .set({
             status: "resolved",
-            resolvedAt: new Date(),
-            updatedAt: new Date(),
+            resolvedAt: now,
+            updatedAt: now,
           })
           .where(and(eq(budgetIncidents.policyId, policy.id), eq(budgetIncidents.status, "open")));