fix(billing): address all code review findings from Phase 3

- Fix `any` type for logContext -> proper OrgAutoTopupLogContext interface
- Remove stale sync_failures comment that was misleading
- Preserve AutoTopupValidationError type when re-throwing (was losing error specificity)
- Add expiration validation to org payment method flow (was missing unlike user flow)
- Make auto_topup_threshold/amount nullable in orgs schema (consistent with users)
- Improve getOrSetDefaultPaymentMethod API to return { paymentMethodId, wasUpdated }
- Move message and adImpression tables from misc.ts to billing.ts (better domain organization)
- Remove trivial comments describing obvious code
- Add JSDoc explaining payment method type limitations
- Fix misleading index comment in orgRepo
- Extract isValidPaymentMethod() and filterValidPaymentMethods() helpers to eliminate duplication
- Remove unused error field from OrgAutoTopupLogContext interface
- Delete empty misc.ts file

All 81 billing tests pass, all 13 package typechecks pass

Author: brandonkachen

packages/billing/src/auto-topup-helpers.ts

@@ -5,6 +5,11 @@ import type Stripe from 'stripe'
 
 /**
  * Fetches both card and link payment methods for a Stripe customer.
+ *
+ * Note: Only 'card' and 'link' types are supported as these are the primary
+ * payment method types used for off-session automatic charges. Other types
+ * (e.g., 'us_bank_account', 'sepa_debit') may have different confirmation
+ * requirements that don't work well with auto-topup flows.
  */
 export async function fetchPaymentMethods(
   stripeCustomerId: string,
@@ -23,26 +28,41 @@ export async function fetchPaymentMethods(
   return [...cardPaymentMethods.data, ...linkPaymentMethods.data]
 }
 
+/**
+ * Checks if a payment method is valid for use.
+ * Cards are checked for expiration, link methods are always valid.
+ */
+export function isValidPaymentMethod(pm: Stripe.PaymentMethod): boolean {
+  if (pm.type === 'card') {
+    return (
+      pm.card?.exp_year !== undefined &&
+      pm.card.exp_month !== undefined &&
+      new Date(pm.card.exp_year, pm.card.exp_month - 1) > new Date()
+    )
+  }
+  if (pm.type === 'link') {
+    return true
+  }
+  return false
+}
+
+/**
+ * Filters payment methods to only include valid (non-expired) ones.
+ */
+export function filterValidPaymentMethods(
+  paymentMethods: Stripe.PaymentMethod[],
+): Stripe.PaymentMethod[] {
+  return paymentMethods.filter(isValidPaymentMethod)
+}
+
 /**
  * Finds the first valid (non-expired) payment method from a list.
  * Cards are checked for expiration, link methods are always valid.
  */
 export function findValidPaymentMethod(
   paymentMethods: Stripe.PaymentMethod[],
 ): Stripe.PaymentMethod | undefined {
-  return paymentMethods.find((pm) => {
-    if (pm.type === 'card') {
-      return (
-        pm.card?.exp_year &&
-        pm.card.exp_month &&
-        new Date(pm.card.exp_year, pm.card.exp_month - 1) > new Date()
-      )
-    }
-    if (pm.type === 'link') {
-      return true
-    }
-    return false
-  })
+  return paymentMethods.find(isValidPaymentMethod)
 }
 
 export interface PaymentIntentParams {
@@ -86,22 +106,25 @@ export async function createPaymentIntent(
   )
 }
 
+export interface GetOrSetDefaultPaymentMethodResult {
+  paymentMethodId: string
+  wasUpdated: boolean
+}
+
 /**
  * Gets the default payment method for a customer, or selects and sets the first available one.
- * Returns the payment method ID to use.
+ * Returns the payment method ID to use and whether it was newly set as default.
  */
 export async function getOrSetDefaultPaymentMethod(params: {
   stripeCustomerId: string
   paymentMethods: Stripe.PaymentMethod[]
   logger: Logger
   logContext: Record<string, unknown>
-}): Promise<string> {
+}): Promise<GetOrSetDefaultPaymentMethodResult> {
   const { stripeCustomerId, paymentMethods, logger, logContext } = params
 
-  // Get the customer to check for default payment method
   const customer = await stripeServer.customers.retrieve(stripeCustomerId)
 
-  // Check if there's already a default payment method
   if (
     customer &&
     !customer.deleted &&
@@ -112,7 +135,6 @@ export async function getOrSetDefaultPaymentMethod(params: {
         ? customer.invoice_settings.default_payment_method
         : customer.invoice_settings.default_payment_method.id
 
-    // Verify the default payment method is still valid and available
     const isDefaultValid = paymentMethods.some(
       (pm) => pm.id === defaultPaymentMethodId,
     )
@@ -122,20 +144,21 @@ export async function getOrSetDefaultPaymentMethod(params: {
         { ...logContext, paymentMethodId: defaultPaymentMethodId },
         'Using existing default payment method',
       )
-      return defaultPaymentMethodId
+      return { paymentMethodId: defaultPaymentMethodId, wasUpdated: false }
     }
   }
 
-  // Use the first available and set it as default
   const firstPaymentMethod = paymentMethods[0]
   const paymentMethodToUse = firstPaymentMethod.id
+  let wasUpdated = false
 
   try {
     await stripeServer.customers.update(stripeCustomerId, {
       invoice_settings: {
         default_payment_method: paymentMethodToUse,
       },
     })
+    wasUpdated = true
 
     logger.info(
       { ...logContext, paymentMethodId: paymentMethodToUse },
@@ -148,5 +171,5 @@ export async function getOrSetDefaultPaymentMethod(params: {
     )
   }
 
-  return paymentMethodToUse
+  return { paymentMethodId: paymentMethodToUse, wasUpdated }
 }

packages/billing/src/auto-topup.ts

@@ -10,6 +10,7 @@ import { eq } from 'drizzle-orm'
 import {
   createPaymentIntent,
   fetchPaymentMethods,
+  filterValidPaymentMethods,
   findValidPaymentMethod,
   getOrSetDefaultPaymentMethod,
 } from './auto-topup-helpers'
@@ -184,7 +185,6 @@ export async function checkAndTriggerAutoTopup(params: {
   const logContext = { userId }
 
   try {
-    // Get user info
     const user = await db.query.user.findFirst({
       where: eq(schema.user.id, userId),
       columns: {
@@ -206,7 +206,6 @@ export async function checkAndTriggerAutoTopup(params: {
       return undefined
     }
 
-    // Calculate balance
     const { balance } = await calculateUsageAndBalance({
       ...params,
       quotaResetDate: user.next_quota_reset ?? new Date(0),
@@ -252,12 +251,13 @@ export async function checkAndTriggerAutoTopup(params: {
       `Auto-top-up needed for user ${userId}. Will attempt to purchase ${amountToTopUp} credits.`,
     )
 
-    // Validate payment method
     const { blockedReason, validPaymentMethod } =
       await validateAutoTopupStatus(params)
 
     if (blockedReason || !validPaymentMethod) {
-      throw new Error(blockedReason || 'Auto top-up is not available.')
+      throw new AutoTopupValidationError(
+        blockedReason || 'Auto top-up is not available.',
+      )
     }
 
     try {
@@ -341,12 +341,22 @@ async function getOrganizationPaymentMethod(params: {
     )
   }
 
-  return getOrSetDefaultPaymentMethod({
+  const validPaymentMethods = filterValidPaymentMethods(allPaymentMethods)
+
+  if (validPaymentMethods.length === 0) {
+    throw new AutoTopupPaymentError(
+      'No valid (non-expired) payment methods available for organization',
+    )
+  }
+
+  const { paymentMethodId } = await getOrSetDefaultPaymentMethod({
     stripeCustomerId,
-    paymentMethods: allPaymentMethods,
+    paymentMethods: validPaymentMethods,
     logger,
     logContext,
   })
+
+  return paymentMethodId
 }
 
 async function processOrgAutoTopupPayment(params: {
@@ -412,13 +422,21 @@ async function processOrgAutoTopupPayment(params: {
   )
 }
 
+interface OrgAutoTopupLogContext {
+  organizationId: string
+  userId: string
+  currentBalance?: number
+  threshold?: number | null
+  amountToTopUp?: number
+}
+
 export async function checkAndTriggerOrgAutoTopup(params: {
   organizationId: string
   userId: string
   logger: Logger
 }): Promise<void> {
   const { organizationId, userId, logger } = params
-  const logContext: any = { organizationId, userId }
+  const logContext: OrgAutoTopupLogContext = { organizationId, userId }
 
   try {
     const org = await getOrganizationSettings(organizationId)
@@ -471,8 +489,6 @@ export async function checkAndTriggerOrgAutoTopup(params: {
         stripeCustomerId: org.stripe_customer_id,
       })
     } catch (error) {
-      // Auto-topup failures are automatically logged to sync_failures table
-      // by the existing error handling in processOrgAutoTopupPayment
       logger.error(
         { ...logContext, error },
         'Organization auto top-up payment failed',

packages/internal/src/db/schema/billing.ts

@@ -1,12 +1,17 @@
 import { sql } from 'drizzle-orm'
 import {
+  boolean,
   index,
   integer,
+  jsonb,
+  numeric,
   pgTable,
   text,
   timestamp,
 } from 'drizzle-orm/pg-core'
 
+import type { SQL } from 'drizzle-orm'
+
 import { grantTypeEnum } from './enums'
 import { org } from './organizations'
 import { user } from './users'
@@ -69,3 +74,80 @@ export const syncFailure = pgTable(
       .where(sql`${table.retry_count} < 5`),
   ],
 )
+
+// Usage tracking table - stores LLM message costs and token usage
+export const message = pgTable(
+  'message',
+  {
+    id: text('id').primaryKey(),
+    finished_at: timestamp('finished_at', { mode: 'date' }).notNull(),
+    client_id: text('client_id'),
+    client_request_id: text('client_request_id'),
+    model: text('model').notNull(),
+    agent_id: text('agent_id'),
+    request: jsonb('request'),
+    lastMessage: jsonb('last_message').generatedAlwaysAs(
+      (): SQL => sql`${message.request} -> -1`,
+    ),
+    reasoning_text: text('reasoning_text'),
+    response: jsonb('response').notNull(),
+    input_tokens: integer('input_tokens').notNull().default(0),
+    cache_creation_input_tokens: integer('cache_creation_input_tokens'),
+    cache_read_input_tokens: integer('cache_read_input_tokens')
+      .notNull()
+      .default(0),
+    reasoning_tokens: integer('reasoning_tokens'),
+    output_tokens: integer('output_tokens').notNull(),
+    cost: numeric('cost', { precision: 100, scale: 20 }).notNull(),
+    credits: integer('credits').notNull(),
+    byok: boolean('byok').notNull().default(false),
+    latency_ms: integer('latency_ms'),
+    user_id: text('user_id').references(() => user.id, { onDelete: 'cascade' }),
+    org_id: text('org_id').references(() => org.id, { onDelete: 'cascade' }),
+    repo_url: text('repo_url'),
+  },
+  (table) => [
+    index('message_user_id_idx').on(table.user_id),
+    index('message_finished_at_user_id_idx').on(
+      table.finished_at,
+      table.user_id,
+    ),
+    index('message_org_id_idx').on(table.org_id),
+    index('message_org_id_finished_at_idx').on(table.org_id, table.finished_at),
+  ],
+)
+
+// Ad impression tracking - grants credits to users for viewing ads
+export const adImpression = pgTable(
+  'ad_impression',
+  {
+    id: text('id')
+      .primaryKey()
+      .$defaultFn(() => crypto.randomUUID()),
+    user_id: text('user_id')
+      .notNull()
+      .references(() => user.id, { onDelete: 'cascade' }),
+    ad_text: text('ad_text').notNull(),
+    title: text('title').notNull(),
+    cta: text('cta').notNull().default(''),
+    url: text('url').notNull(),
+    favicon: text('favicon').notNull(),
+    click_url: text('click_url').notNull(),
+    imp_url: text('imp_url').notNull().unique(),
+    payout: numeric('payout', { precision: 10, scale: 6 }).notNull(),
+    credits_granted: integer('credits_granted').notNull(),
+    grant_operation_id: text('grant_operation_id'),
+    served_at: timestamp('served_at', { mode: 'date', withTimezone: true })
+      .notNull()
+      .defaultNow(),
+    impression_fired_at: timestamp('impression_fired_at', {
+      mode: 'date',
+      withTimezone: true,
+    }),
+    clicked_at: timestamp('clicked_at', { mode: 'date', withTimezone: true }),
+  },
+  (table) => [
+    index('idx_ad_impression_user').on(table.user_id, table.served_at),
+    index('idx_ad_impression_imp_url').on(table.imp_url),
+  ],
+)

packages/internal/src/db/schema/index.ts

@@ -4,4 +4,4 @@ export * from './users'
 export * from './organizations'
 export * from './billing'
 export * from './agents'
-export * from './misc'
+// Note: misc.ts is now empty - message and adImpression moved to billing.ts