fix: revert unrelated refactors from referral migration

- Revert org-billing.ts to main branch version (expiration bucket
  tracking for orgs was not part of the migration plan)
- Make breakdownByExpiration and principalsByExpiration optional in
  CreditBalance interface (only needed for user balance, not orgs)
- Add CreditBalanceWithExpiration type for code requiring full balance
- Add guard in usage-display.tsx for optional expiration fields

Author: brandonkachen

packages/billing/src/balance-calculator.ts

@@ -28,7 +28,16 @@ export interface CreditBalance {
   netBalance: number
   breakdown: Record<GrantType, number>
   principals: Record<GrantType, number>
-  // Split by expiration for UI display
+  // Split by expiration for UI display (optional - only populated for user balance, not org balance)
+  breakdownByExpiration?: Record<CreditExpirationBucket, Record<GrantType, number>>
+  principalsByExpiration?: Record<CreditExpirationBucket, Record<GrantType, number>>
+}
+
+/**
+ * CreditBalance with required expiration bucket fields.
+ * Returned by createInitialCreditBalance() and calculateUsageAndBalance().
+ */
+export interface CreditBalanceWithExpiration extends CreditBalance {
   breakdownByExpiration: Record<CreditExpirationBucket, Record<GrantType, number>>
   principalsByExpiration: Record<CreditExpirationBucket, Record<GrantType, number>>
 }
@@ -37,7 +46,7 @@ export interface CreditBalance {
  * Creates an initialized CreditBalance object with all values set to 0.
  * Used by both user and organization balance calculations to avoid code duplication.
  */
-export function createInitialCreditBalance(): CreditBalance {
+export function createInitialCreditBalance(): CreditBalanceWithExpiration {
   const breakdown: Record<GrantType, number> = {} as Record<GrantType, number>
   const principals: Record<GrantType, number> = {} as Record<GrantType, number>
 

packages/billing/src/org-billing.ts

@@ -1,19 +1,16 @@
 import { GRANT_PRIORITIES } from '@codebuff/common/constants/grant-priorities'
+import { GrantTypeValues } from '@codebuff/common/types/grant'
 import db from '@codebuff/internal/db'
 import * as schema from '@codebuff/internal/db/schema'
 import { withAdvisoryLockTransaction } from '@codebuff/internal/db/transaction'
 import { env } from '@codebuff/internal/env'
 import { stripeServer } from '@codebuff/internal/util/stripe'
 import { and, asc, gt, isNull, or, eq } from 'drizzle-orm'
 
-import {
-  consumeFromOrderedGrants,
-  createInitialCreditBalance,
-} from './balance-calculator'
+import { consumeFromOrderedGrants } from './balance-calculator'
 
 import type {
   CreditBalance,
-  CreditExpirationBucket,
   CreditUsageAndBalance,
   CreditConsumptionResult,
 } from './balance-calculator'
@@ -187,8 +184,29 @@ export async function calculateOrganizationUsageAndBalance(
   // Get all relevant grants for the organization
   const grants = await getOrderedActiveOrganizationGrants(withDefaults)
 
-  // Initialize balance structure with all values set to 0
-  const balance = createInitialCreditBalance()
+  // Initialize breakdown and principals with all grant types set to 0
+  const initialBreakdown: Record<GrantType, number> = {} as Record<
+    GrantType,
+    number
+  >
+  const initialPrincipals: Record<GrantType, number> = {} as Record<
+    GrantType,
+    number
+  >
+
+  for (const type of GrantTypeValues) {
+    initialBreakdown[type] = 0
+    initialPrincipals[type] = 0
+  }
+
+  // Initialize balance structure
+  const balance: CreditBalance = {
+    totalRemaining: 0,
+    totalDebt: 0,
+    netBalance: 0,
+    breakdown: initialBreakdown,
+    principals: initialPrincipals,
+  }
 
   // Calculate both metrics in one pass
   let usageThisCycle = 0
@@ -211,17 +229,9 @@ export async function calculateOrganizationUsageAndBalance(
     // Add to balance if grant is currently active
     if (!grant.expires_at || grant.expires_at > now) {
       balance.principals[grantType] += grant.principal
-
-      // Classify by expiration: renewable if expires_at is set, nonRenewable if null
-      const bucket: CreditExpirationBucket = grant.expires_at
-        ? 'renewable'
-        : 'nonRenewable'
-      balance.principalsByExpiration[bucket][grantType] += grant.principal
-
       if (grant.balance > 0) {
         totalPositiveBalance += grant.balance
         balance.breakdown[grantType] += grant.balance
-        balance.breakdownByExpiration[bucket][grantType] += grant.balance
       } else if (grant.balance < 0) {
         totalDebt += Math.abs(grant.balance)
       }

web/src/app/profile/components/usage-display.tsx

@@ -246,10 +246,13 @@ export const UsageDisplay = ({
   if (isLoading) {
     return <UsageDisplaySkeleton />
   }
-  const { totalDebt, netBalance } = balance
+  const { totalDebt, netBalance, breakdownByExpiration, principalsByExpiration } = balance
 
-  // Get expiration buckets from balance (data-driven, not hardcoded)
-  const { breakdownByExpiration, principalsByExpiration } = balance
+  // Handle case where expiration buckets might not be provided (e.g., org balance)
+  // This should not happen for user balance, but TypeScript requires the check
+  if (!breakdownByExpiration || !principalsByExpiration) {
+    return <UsageDisplaySkeleton />
+  }
 
   const renewableTypes = getTypesInBucket('renewable', principalsByExpiration, breakdownByExpiration)
   const nonRenewableTypes = getTypesInBucket('nonRenewable', principalsByExpiration, breakdownByExpiration)