Submatter/apps/web/server/routers/budgets.ts

import { TRPCError } from "@trpc/server";
import { and, eq, gte, sum } from "drizzle-orm";
import { z } from "zod";

import { budgetCaps, convoyEvents } from "@tasks/database/schema";

import { router, workspaceProcedure } from "@/server/trpc";

function startOfMonth(d = new Date()): Date {
  return new Date(d.getFullYear(), d.getMonth(), 1);
}

export const budgetsRouter = router({
  getSpend: workspaceProcedure.query(async ({ ctx }) => {
    const monthStart = startOfMonth();
    const rows = await ctx.db
      .select({
        modelAlias: convoyEvents.modelAlias,
        providerType: convoyEvents.providerType,
        modelTier: convoyEvents.modelTier,
        totalCost: sum(convoyEvents.estimatedCostUsd).as("total_cost"),
      })
      .from(convoyEvents)
      .where(
        and(
          eq(convoyEvents.workspaceId, ctx.workspace.id),
          gte(convoyEvents.occurredAt, monthStart),
        ),
      )
      .groupBy(
        convoyEvents.modelAlias,
        convoyEvents.providerType,
        convoyEvents.modelTier,
      );

    return rows.map((r) => ({
      modelAlias: r.modelAlias ?? "unknown",
      providerType: r.providerType ?? "unknown",
      modelTier: r.modelTier ?? "unknown",
      estimatedCostUsd: Number(r.totalCost ?? 0),
    }));
  }),

  getCap: workspaceProcedure.query(async ({ ctx }) => {
    const cap = await ctx.db.query.budgetCaps.findFirst({
      where: (c, { eq }) => eq(c.workspaceId, ctx.workspace.id),
    });
    return cap ?? { enabled: false, capUsd: null, period: "monthly" };
  }),

  upsertCap: workspaceProcedure
    .input(
      z.object({
        capUsd: z.number().positive().max(1_000_000),
        period: z.enum(["monthly"]).default("monthly"),
        enabled: z.boolean().default(true),
      }),
    )
    .mutation(async ({ ctx, input }) => {
      if (ctx.workspace.role !== "owner" && ctx.workspace.role !== "admin") {
        throw new TRPCError({
          code: "FORBIDDEN",
          message: "Only owners and admins can set budget caps.",
        });
      }

      const [upserted] = await ctx.db
        .insert(budgetCaps)
        .values({
          workspaceId: ctx.workspace.id,
          capUsd: input.capUsd.toString(),
          period: input.period,
          enabled: input.enabled,
        })
        .onConflictDoUpdate({
          target: [budgetCaps.workspaceId, budgetCaps.period],
          set: {
            capUsd: input.capUsd.toString(),
            enabled: input.enabled,
            updatedAt: new Date(),
          },
        })
        .returning();

      return upserted;
    }),

  deleteCap: workspaceProcedure.mutation(async ({ ctx }) => {
    if (ctx.workspace.role !== "owner" && ctx.workspace.role !== "admin") {
      throw new TRPCError({
        code: "FORBIDDEN",
        message: "Only owners and admins can delete budget caps.",
      });
    }

    await ctx.db
      .delete(budgetCaps)
      .where(eq(budgetCaps.workspaceId, ctx.workspace.id));

    return { ok: true };
  }),

  shouldPauseDispatch: workspaceProcedure
    .input(z.object({ workspaceId: z.string().uuid() }))
    .query(async ({ ctx, input }) => {
      if (input.workspaceId !== ctx.workspace.id) {
        throw new TRPCError({
          code: "FORBIDDEN",
          message: "Workspace mismatch.",
        });
      }

      const cap = await ctx.db.query.budgetCaps.findFirst({
        where: (c, { and, eq }) => and(eq(c.workspaceId, input.workspaceId), eq(c.enabled, true)),
      });

      if (!cap) return { shouldPause: false, reason: null, spent: 0, cap: null };

      const monthStart = startOfMonth();
      const spend = await ctx.db
        .select({ total: sum(convoyEvents.estimatedCostUsd) })
        .from(convoyEvents)
        .where(
          and(
            eq(convoyEvents.workspaceId, input.workspaceId),
            gte(convoyEvents.occurredAt, monthStart),
          ),
        );

      const spent = Number(spend[0]?.total ?? 0);
      const capUsd = Number(cap.capUsd);
      const shouldPause = spent >= capUsd;

      return {
        shouldPause,
        reason: shouldPause ? `Monthly budget cap $${capUsd} exceeded ($${spent} spent)` : null,
        spent,
        cap: capUsd,
      };
    }),
});

export type BudgetsRouter = typeof budgetsRouter;