Final updates and fixes.

Author: jorgebaralt

packages/api/src/api-router/model-pricing.ts

@@ -19,7 +19,7 @@ export const modelPricingRouter = {
           where: and(
             eq(ModelPricing.model_id, input.model_id),
             eq(ModelPricing.model_provider, input.model_provider),
-            lte(ModelPricing.effective_date, new Date()),
+            lte(ModelPricing.effective_date, new Date().toISOString()),
           ),
         });
 
@@ -51,7 +51,7 @@ export const modelPricingRouter = {
             eq(ModelPricing.model_id, input.model_id),
             eq(ModelPricing.model_provider, input.model_provider),
             eq(ModelPricing.unit_type, input.unit_type),
-            lte(ModelPricing.effective_date, new Date()),
+            lte(ModelPricing.effective_date, new Date().toISOString()),
           ),
         });
 
@@ -78,7 +78,6 @@ export const modelPricingRouter = {
           .onConflictDoUpdate({
             set: {
               price_per_unit: input.price_per_unit,
-              updated_at: new Date(),
             },
             target: [
               ModelPricing.model_id,

packages/api/src/api-router/usage.ts

@@ -7,13 +7,15 @@ import {
   UsagePeriod,
 } from "@acme/db/schema";
 import { TRPCError, type TRPCRouterRecord } from "@trpc/server";
+import { gte } from "drizzle-orm";
 import { z } from "zod/v4";
-import { getActiveSubscription, getFreePlan } from "../shared/subscription.js";
 import type { TRPCContext } from "../trpc.js";
 import { publicProcedure } from "../trpc.js";
 
 /**
- * Get or create the current usage period for a user
+ * Get the current usage period aggregate for a user
+ * Returns a complete aggregate with period, plan, and subscription data
+ * Periods are created by webhooks (pro) and cron jobs (free)
  */
 async function getCurrentUsagePeriod({
   ctx,
@@ -22,156 +24,63 @@ async function getCurrentUsagePeriod({
   ctx: TRPCContext;
   userId: string;
 }) {
-  return await ctx.db.transaction(async (tx) => {
-    // First, check if user has an active subscription
-    const activeSubscription = await getActiveSubscription({
-      ctx,
-      userId,
-    });
+  const now = new Date().toISOString();
 
-    if (activeSubscription) {
-      // Validate subscription has required fields
-      if (
-        !activeSubscription.periodStart ||
-        !activeSubscription.periodEnd ||
-        !activeSubscription.plan
-      ) {
-        throw new TRPCError({
-          code: "INTERNAL_SERVER_ERROR",
-          message: "Invalid subscription data",
-        });
-      }
-
-      // Pro user: try to find existing usage period for current subscription period
-      let usagePeriod = await tx.query.UsagePeriod.findFirst({
-        where: and(
-          eq(UsagePeriod.user_id, userId),
-          eq(UsagePeriod.subscription_id, activeSubscription.id),
-          eq(UsagePeriod.period_start, activeSubscription.periodStart),
-          eq(UsagePeriod.period_end, activeSubscription.periodEnd),
-        ),
-      });
-
-      if (!usagePeriod) {
-        // Create usage period for subscription period
-        const [newUsagePeriod] = await tx
-          .insert(UsagePeriod)
-          .values({
-            period_end: activeSubscription.periodEnd,
-            period_start: activeSubscription.periodStart,
-            plan_id: activeSubscription.plan.id,
-            subscription_id: activeSubscription.id,
-            user_id: userId,
-          })
-          .returning();
-
-        usagePeriod = newUsagePeriod;
-      }
-
-      if (!usagePeriod) {
-        throw new TRPCError({
-          code: "INTERNAL_SERVER_ERROR",
-          message: "Failed to create or retrieve usage period for subscription",
-        });
-      }
-
-      return usagePeriod;
-    } else {
-      // For free users, we will use a monthly usage period (1st of each month) for simplicity
-      const now = new Date();
-      const monthStart = new Date(now.getFullYear(), now.getMonth(), 1);
-      // The end of the month is the last day of the month at 23:59:59.999
-      const monthEnd = new Date(
-        now.getFullYear(),
-        now.getMonth() + 1,
-        0,
-        23,
-        59,
-        59,
-        999,
-      );
-
-      // Try to find existing usage period for current month
-      let usagePeriod = await tx.query.UsagePeriod.findFirst({
-        where: and(
-          eq(UsagePeriod.user_id, userId),
-          eq(UsagePeriod.period_start, monthStart),
-          eq(UsagePeriod.period_end, monthEnd),
-        ),
-      });
-
-      if (!usagePeriod) {
-        // Get the free plan for free users
-        const freePlan = await getFreePlan(ctx);
-
-        // Create usage period for free user
-        const [newUsagePeriod] = await tx
-          .insert(UsagePeriod)
-          .values({
-            period_end: monthEnd,
-            period_start: monthStart,
-            plan_id: freePlan?.id,
-            subscription_id: null, // Free users don't have a subscription_id
-            user_id: userId,
-          })
-          .returning();
-
-        usagePeriod = newUsagePeriod;
-      }
+  const usagePeriodAggregate = await ctx.db.query.UsagePeriod.findFirst({
+    orderBy: (fields, { desc }) => [
+      desc(fields.subscription_id),
+      desc(fields.created_at),
+    ],
+    where: and(
+      eq(UsagePeriod.user_id, userId),
+      lte(UsagePeriod.period_start, now),
+      gte(UsagePeriod.period_end, now),
+    ),
+    with: {
+      plan: true,
+      subscription: {
+        with: {
+          plan: true,
+        },
+      },
+      usageAggregate: true,
+    },
+  });
 
-      if (!usagePeriod) {
-        throw new TRPCError({
-          code: "INTERNAL_SERVER_ERROR",
-          message: "Failed to create or retrieve usage period for free user",
-        });
-      }
+  if (!usagePeriodAggregate) {
+    throw new TRPCError({
+      code: "PRECONDITION_FAILED",
+      message:
+        "No active usage period found. Usage periods are created automatically via webhooks and scheduled jobs.",
+    });
+  }
 
-      return usagePeriod;
-    }
-  });
+  return usagePeriodAggregate;
 }
 
 export const usageRouter = {
   checkUsage: publicProcedure
     .input(z.object({ user_id: z.string() }))
     .query(async ({ ctx, input }) => {
       try {
-        const activeSubscription = await getActiveSubscription({
+        const usagePeriodAggregate = await getCurrentUsagePeriod({
           ctx,
           userId: input.user_id,
         });
 
-        const usagePeriod = await getCurrentUsagePeriod({
-          ctx,
-          userId: input.user_id,
-        });
-
-        let quotaUsd: number;
-        if (activeSubscription?.plan) {
-          quotaUsd = activeSubscription.plan.quota / 100;
-        } else {
-          const freePlan = await getFreePlan(ctx);
-          if (!freePlan) {
-            throw new TRPCError({
-              code: "INTERNAL_SERVER_ERROR",
-              message: "No free plan found",
-            });
-          }
-          quotaUsd = freePlan.quota / 100; // Convert cents to dollars
+        const { plan } = usagePeriodAggregate;
+        if (!plan) {
+          throw new TRPCError({
+            code: "INTERNAL_SERVER_ERROR",
+            message: "Usage period missing plan data",
+          });
         }
 
-        // Get current usage for the period
-        let currentUsageUsd = 0;
-        const usageAggregate = await ctx.db.query.UsageAggregate.findFirst({
-          where: and(
-            eq(UsageAggregate.user_id, input.user_id),
-            eq(UsageAggregate.usage_period_id, usagePeriod.id),
-          ),
-        });
+        const quotaUsd = plan.quota / 100;
 
-        if (usageAggregate) {
-          currentUsageUsd = Number.parseFloat(usageAggregate.total_cost);
-        }
+        const currentUsageUsd = usagePeriodAggregate.usageAggregate
+          ? Number.parseFloat(usagePeriodAggregate.usageAggregate.total_cost)
+          : 0;
 
         const canUse = currentUsageUsd < quotaUsd;
         const remainingQuotaUsd = quotaUsd - currentUsageUsd;
@@ -181,8 +90,8 @@ export const usageRouter = {
           currentUsageUsd,
           quotaUsd,
           remainingQuotaUsd: Math.max(0, remainingQuotaUsd),
-          subscriptionType: activeSubscription ? "pro" : "free",
-          usagePeriodId: usagePeriod.id,
+          subscriptionType: usagePeriodAggregate.subscription ? "pro" : "free",
+          usagePeriodId: usagePeriodAggregate.id,
         };
       } catch (error) {
         console.error("Error checking usage:", error);
@@ -226,34 +135,40 @@ export const usageRouter = {
             return { message: "Already processed", success: true };
           }
 
-          // Get or create the current usage period
           const usagePeriod = await getCurrentUsagePeriod({
             ctx,
             userId: input.user_id,
           });
 
-          // Calculate total cost for this usage event
-          let totalCost = 0;
+          const eventDate = usageEvent.created_at
+            ? usageEvent.created_at
+            : new Date().toISOString();
+
+          // Fetch all pricing for this model upfront
+          const allPricing = await tx
+            .select()
+            .from(ModelPricing)
+            .where(
+              and(
+                eq(ModelPricing.model_id, usageEvent.model_id),
+                eq(ModelPricing.model_provider, usageEvent.model_provider),
+                lte(ModelPricing.effective_date, eventDate),
+              ),
+            )
+            .orderBy(desc(ModelPricing.effective_date));
+
+          // Create lookup map
+          const pricingMap = new Map<string, (typeof allPricing)[0]>();
+          for (const price of allPricing) {
+            if (!pricingMap.has(price.unit_type)) {
+              pricingMap.set(price.unit_type, price);
+            }
+          }
 
+          // Calculate total cost using in-memory pricing lookups
+          let totalCost = 0;
           for (const metric of usageEvent.metrics) {
-            // Get current pricing for this model and unit type
-            const eventDate = usageEvent.created_at
-              ? new Date(usageEvent.created_at)
-              : new Date();
-            const pricing = await tx
-              .select()
-              .from(ModelPricing)
-              .where(
-                and(
-                  eq(ModelPricing.model_id, usageEvent.model_id),
-                  eq(ModelPricing.model_provider, usageEvent.model_provider),
-                  eq(ModelPricing.unit_type, metric.unit),
-                  lte(ModelPricing.effective_date, eventDate),
-                ),
-              )
-              .orderBy(desc(ModelPricing.effective_date))
-              .limit(1)
-              .then((results) => results[0] || null);
+            const pricing = pricingMap.get(metric.unit);
 
             if (pricing) {
               const cost =
@@ -277,7 +192,6 @@ export const usageRouter = {
             .onConflictDoUpdate({
               set: {
                 total_cost: sql`${UsageAggregate.total_cost} + ${totalCost.toFixed(6)}`,
-                updated_at: new Date(),
               },
               target: [UsageAggregate.user_id, UsageAggregate.usage_period_id],
             });

packages/auth/src/stripe/plan-handler.ts

@@ -16,19 +16,19 @@ async function upsertPlan(product: Stripe.Product) {
     quota,
   };
 
-  const updateData = {
-    active: insertData.active,
-    description: insertData.description,
-    displayName: insertData.displayName,
-    metadata: insertData.metadata,
-    quota: insertData.quota,
-    updated_at: new Date(),
-  };
-
-  return db.insert(Plan).values(insertData).onConflictDoUpdate({
-    set: updateData,
-    target: Plan.id,
-  });
+  return db
+    .insert(Plan)
+    .values(insertData)
+    .onConflictDoUpdate({
+      set: {
+        active: insertData.active,
+        description: insertData.description,
+        displayName: insertData.displayName,
+        metadata: insertData.metadata,
+        quota: insertData.quota,
+      },
+      target: Plan.id,
+    });
 }
 
 async function deletePlan(productId: string) {

packages/auth/src/stripe/price-handler.ts

@@ -29,22 +29,22 @@ async function upsertPrice(price: Stripe.Price) {
     unitAmount: price.unit_amount ?? 0,
   };
 
-  const updateData = {
-    active: insertData.active,
-    currency: insertData.currency,
-    lookupKey: insertData.lookupKey,
-    metadata: insertData.metadata,
-    nickname: insertData.nickname,
-    recurring: insertData.recurring,
-    type: insertData.type,
-    unitAmount: insertData.unitAmount,
-    updatedAt: new Date(),
-  };
-
-  return db.insert(Price).values(insertData).onConflictDoUpdate({
-    set: updateData,
-    target: Price.id,
-  });
+  return db
+    .insert(Price)
+    .values(insertData)
+    .onConflictDoUpdate({
+      set: {
+        active: insertData.active,
+        currency: insertData.currency,
+        lookupKey: insertData.lookupKey,
+        metadata: insertData.metadata,
+        nickname: insertData.nickname,
+        recurring: insertData.recurring,
+        type: insertData.type,
+        unitAmount: insertData.unitAmount,
+      },
+      target: Price.id,
+    });
 }
 
 async function deletePrice(priceId: string) {