Jorge/usage aggregate #141
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
create new table schemas new procedures
- update usage-event schema and add total cost
rmolinamir
reviewed
Oct 8, 2025
packages/api/src/api-router/usage.ts
Outdated
Comment on lines
27
to
53
| const activeSubscription = await getActiveSubscription({ | ||
| ctx, | ||
| userId, | ||
| }); | ||
|
|
||
| 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), | ||
| ), | ||
| }); |
Contributor
There was a problem hiding this comment.
I think we should get both the activeSubscription and the usagePeriod in 1 query; otherwise, we are doing two round-trips to the DB, which is a lot of lag. We'll notice things are slower if we use this everywhere. I think we can get both at the same time with Drizzle relationships - worst case scenario we can write a SQL query lol.
packages/api/src/api-router/usage.ts
Outdated
Comment on lines
94
to
127
| // 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; | ||
| } | ||
|
|
||
| if (!usagePeriod) { | ||
| throw new TRPCError({ | ||
| code: "INTERNAL_SERVER_ERROR", | ||
| message: "Failed to create or retrieve usage period for free user", | ||
| }); | ||
| } |
Contributor
There was a problem hiding this comment.
I don't like doing a write here since this is supposed to be a query; if we combine reads with writes here, it will be very hard to optimize later. This is meant to be a middleware being used everywhere 🤔
Contributor
There was a problem hiding this comment.
Maybe:
- When users create accounts, create a usage period too
- When users subscribe, update the period?
- When user subscriptions end, create a new period?
- Then we could reset usage in a cron job for free users 🤔
This way we separate writes/reads.
| recurring: insertData.recurring, | ||
| type: insertData.type, | ||
| unitAmount: insertData.unitAmount, | ||
| updatedAt: new Date(), |
Comment on lines
26
to
29
| created_at: timestamp().defaultNow(), | ||
| updated_at: timestamp() | ||
| .defaultNow() | ||
| .$onUpdateFn(() => new Date()), |
Contributor
There was a problem hiding this comment.
Should look like this:
Suggested change
| created_at: timestamp().defaultNow(), | |
| updated_at: timestamp() | |
| .defaultNow() | |
| .$onUpdateFn(() => new Date()), | |
| created_at: t | |
| .timestamp({ mode: "string", withTimezone: true }) | |
| .defaultNow() | |
| .notNull(), | |
| updated_at: t | |
| .timestamp({ mode: "string", withTimezone: true }) | |
| .defaultNow() | |
| .notNull() | |
| .$onUpdateFn(() => sql`now()`), |
Comment on lines
30
to
33
| created_at: timestamp().defaultNow(), | ||
| updated_at: timestamp() | ||
| .defaultNow() | ||
| .$onUpdateFn(() => new Date()), |
Contributor
There was a problem hiding this comment.
Same date-related changes here.
| usage_period_id: t | ||
| .uuid() | ||
| .notNull() | ||
| .references(() => UsagePeriod.id, { onDelete: "cascade" }), |
Contributor
There was a problem hiding this comment.
I don't think we should ever delete usage when periods are deleted for tracking purposes 🤔
We should probably never delete periods either.
Comment on lines
18
to
21
| created_at: timestamp().defaultNow(), | ||
| updated_at: timestamp() | ||
| .defaultNow() | ||
| .$onUpdateFn(() => new Date()), |
Contributor
There was a problem hiding this comment.
Same date related changes could be done here.
- free user: created on create account, new periods created via cronjob - paid user: created once user subscribes
rmolinamir
approved these changes
Oct 24, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Add Usage Tracking & Aggregation System
Overview
Implements comprehensive usage tracking for AI model consumption with cost calculation, quota management, and webhook-based processing.
Changes
Database Schema
usage_event,usage_period,usage_aggregate,model_pricingAPI Endpoints
api.usage.checkUsage- Check user quota against current usageapi.usage.trackModelUsage- Record AI model usage eventsapi.usage.processUsageEvent- Process events and calculate costsapi.modelPricing.*- Manage model pricing dataAI Integration
Key Features
Infrastructure
/api/supabase/usageNext Steps
checkUsagein the correct places to prevent/allow usage if needed