Fix budget auth and monthly spend rollups

server/src/services/companies.ts

@@ -1,4 +1,4 @@
-import { eq, count } from "drizzle-orm";
+import { and, count, eq, gte, inArray, lt, sql } from "drizzle-orm";
 import type { Db } from "@paperclipai/db";
 import {
   companies,
@@ -54,6 +54,49 @@ export function companyService(db: Db) {
     };
   }
 
+  function currentUtcMonthWindow(now = new Date()) {
+    const year = now.getUTCFullYear();
+    const month = now.getUTCMonth();
+    return {
+      start: new Date(Date.UTC(year, month, 1, 0, 0, 0, 0)),
+      end: new Date(Date.UTC(year, month + 1, 1, 0, 0, 0, 0)),
+    };
+  }
+
+  async function getMonthlySpendByCompanyIds(
+    companyIds: string[],
+    database: Pick<Db, "select"> = db,
+  ) {
+    if (companyIds.length === 0) return new Map<string, number>();
+    const { start, end } = currentUtcMonthWindow();
+    const rows = await database
+      .select({
+        companyId: costEvents.companyId,
+        spentMonthlyCents: sql<number>`coalesce(sum(${costEvents.costCents}), 0)::int`,
+      })
+      .from(costEvents)
+      .where(
+        and(
+          inArray(costEvents.companyId, companyIds),
+          gte(costEvents.occurredAt, start),
+          lt(costEvents.occurredAt, end),
+        ),
+      )
+      .groupBy(costEvents.companyId);
+    return new Map(rows.map((row) => [row.companyId, Number(row.spentMonthlyCents ?? 0)]));
+  }
+
+  async function hydrateCompanySpend<T extends { id: string; spentMonthlyCents: number }>(
+    rows: T[],
+    database: Pick<Db, "select"> = db,
+  ) {
+    const spendByCompanyId = await getMonthlySpendByCompanyIds(rows.map((row) => row.id), database);
+    return rows.map((row) => ({
+      ...row,
+      spentMonthlyCents: spendByCompanyId.get(row.id) ?? 0,
+    }));
+  }
+
   function getCompanyQuery(database: Pick<Db, "select">) {
     return database
       .select(companySelection)
@@ -104,13 +147,20 @@ export function companyService(db: Db) {
   }
 
   return {
-    list: () =>
-      getCompanyQuery(db).then((rows) => rows.map((row) => enrichCompany(row))),
+    list: async () => {
+      const rows = await getCompanyQuery(db);
+      const hydrated = await hydrateCompanySpend(rows);
+      return hydrated.map((row) => enrichCompany(row));
+    },
 
-    getById: (id: string) =>
-      getCompanyQuery(db)
+    getById: async (id: string) => {
+      const row = await getCompanyQuery(db)
         .where(eq(companies.id, id))
-        .then((rows) => (rows[0] ? enrichCompany(rows[0]) : null)),
+        .then((rows) => rows[0] ?? null);
+      if (!row) return null;
+      const [hydrated] = await hydrateCompanySpend([row], db);
+      return enrichCompany(hydrated);
+    },
 
     create: async (data: typeof companies.$inferInsert) => {
       const created = await createCompanyWithUniquePrefix(data);
@@ -118,7 +168,8 @@ export function companyService(db: Db) {
         .where(eq(companies.id, created.id))
         .then((rows) => rows[0] ?? null);
       if (!row) throw notFound("Company not found after creation");
-      return enrichCompany(row);
+      const [hydrated] = await hydrateCompanySpend([row], db);
+      return enrichCompany(hydrated);
     },
 
     update: (
@@ -175,10 +226,12 @@ export function companyService(db: Db) {
           await tx.delete(assets).where(eq(assets.id, existing.logoAssetId));
         }
 
-        return enrichCompany({
+        const [hydrated] = await hydrateCompanySpend([{
           ...updated,
           logoAssetId: logoAssetId === undefined ? existing.logoAssetId : logoAssetId,
-        });
+        }], tx);
+
+        return enrichCompany(hydrated);
       }),
 
     archive: (id: string) =>
@@ -193,7 +246,9 @@ export function companyService(db: Db) {
         const row = await getCompanyQuery(tx)
           .where(eq(companies.id, id))
           .then((rows) => rows[0] ?? null);
-        return row ? enrichCompany(row) : null;
+        if (!row) return null;
+        const [hydrated] = await hydrateCompanySpend([row], tx);
+        return enrichCompany(hydrated);
       }),
 
     remove: (id: string) =>