[feat] organization billing (#164)

Co-authored-by: Codebuff <noreply@codebuff.com>

Author: brandonkachen

.env.example

@@ -14,13 +14,11 @@ AWS_ACCESS_KEY_ID='your-aws-key-here'
 AWS_SECRET_ACCESS_KEY='your-aws-secret-access-key-here'
 
 # Stripe
-STRIPE_SECRET_KEY='your stripe secret key' 
-STRIPE_WEBHOOK_SECRET_KEY='your webhook secret key' 
-STRIPE_PRO_PRICE_ID='price_ for pro base subscription'
-STRIPE_PRO_OVERAGE_PRICE_ID='price_ for pro overage'
-STRIPE_MOAR_PRO_PRICE_ID='price_for moar pro base subscription'
-STRIPE_MOAR_PRO_OVERAGE_PRICE_ID='price_for moar pro overage'
-STRIPE_USAGE_PRICE_ID='price_for default subscription'
+STRIPE_SECRET_KEY='your stripe secret key'
+STRIPE_WEBHOOK_SECRET_KEY='your webhook secret key'
+STRIPE_TEAM_USAGE_PRICE_ID='price_...'
+STRIPE_TEAM_FEE_PRICE_ID='price_...'
+STRIPE_USAGE_PRICE_ID='price_...'
 NEXT_PUBLIC_STRIPE_PUBLISHABLE_KEY='your publishable key'
 NEXT_PUBLIC_STRIPE_CUSTOMER_PORTAL="https://billing.stripe.com/..."
 

backend/src/__tests__/auto-topup.test.ts

@@ -1,76 +1,65 @@
 import type { CreditBalance } from '@codebuff/billing'
 import { checkAndTriggerAutoTopup } from '@codebuff/billing'
 import * as billing from '@codebuff/billing'
-import { beforeEach, describe, expect, it, mock, afterEach, spyOn } from 'bun:test'
+import {
+  beforeEach,
+  describe,
+  expect,
+  it,
+  mock,
+  afterEach,
+  spyOn,
+} from 'bun:test'
 
 describe('Auto Top-up System', () => {
   describe('checkAndTriggerAutoTopup', () => {
     // Create fresh mocks for each test
     let dbMock: ReturnType<typeof mock>
     let balanceMock: ReturnType<typeof mock>
-    let paymentMethodsMock: ReturnType<typeof mock>
-    let paymentIntentMock: ReturnType<typeof mock>
+    let validateAutoTopupMock: ReturnType<typeof mock>
     let grantCreditsMock: ReturnType<typeof mock>
 
     beforeEach(() => {
-      // Clear previous mocks and set up logger again
-      mock.restore()
-
-      mock.module('../util/logger', () => ({
-        logger: {
-          debug: () => {},
-          error: () => {},
-          info: () => {},
-          warn: () => {},
-        },
-        withLoggerContext: async (context: any, fn: () => Promise<any>) => fn(),
-      }))
-
-      dbMock = mock(() => ({
-        id: 'test-user',
-        stripe_customer_id: 'cus_123',
-        auto_topup_enabled: true,
-        auto_topup_threshold: 100,
-        auto_topup_amount: 500,
-        next_quota_reset: new Date(),
-      }))
+      // Set up default mocks
+      dbMock = mock(() =>
+        Promise.resolve({
+          auto_topup_enabled: true,
+          auto_topup_threshold: 100,
+          auto_topup_amount: 500,
+          stripe_customer_id: 'cus_123',
+          next_quota_reset: new Date(),
+        })
+      )
 
       balanceMock = mock(() =>
         Promise.resolve({
           usageThisCycle: 0,
           balance: {
-            totalRemaining: 50, // Below threshold by default
+            totalRemaining: 50, // Below threshold
             totalDebt: 0,
             netBalance: 50,
             breakdown: {},
-          } as CreditBalance,
+          },
         })
       )
 
-      paymentMethodsMock = mock(() =>
+      validateAutoTopupMock = mock(() =>
         Promise.resolve({
-          data: [
-            {
-              id: 'pm_123',
-              card: {
-                exp_year: 2025,
-                exp_month: 12,
-              },
+          blockedReason: null,
+          validPaymentMethod: {
+            id: 'pm_123',
+            type: 'card',
+            card: {
+              exp_year: 2030,
+              exp_month: 12,
             },
-          ],
-        })
-      )
-
-      paymentIntentMock = mock(() =>
-        Promise.resolve({
-          status: 'succeeded',
-          id: 'pi_123',
+          },
         })
       )
 
       grantCreditsMock = mock(() => Promise.resolve())
 
-      // Set up module mocks with fresh mocks
+      // Mock the database
       mock.module('common/db', () => ({
         default: {
           query: {
@@ -87,15 +76,23 @@ describe('Auto Top-up System', () => {
       }))
 
       spyOn(billing, 'calculateUsageAndBalance').mockImplementation(balanceMock)
-      spyOn(billing, 'processAndGrantCredit').mockImplementation(grantCreditsMock)
+      spyOn(billing, 'validateAutoTopupStatus').mockImplementation(
+        validateAutoTopupMock
+      )
+      spyOn(billing, 'processAndGrantCredit').mockImplementation(
+        grantCreditsMock
+      )
 
+      // Mock Stripe payment intent creation
       mock.module('common/src/util/stripe', () => ({
         stripeServer: {
-          paymentMethods: {
-            list: paymentMethodsMock,
-          },
           paymentIntents: {
-            create: paymentIntentMock,
+            create: mock(() =>
+              Promise.resolve({
+                status: 'succeeded',
+                id: 'pi_123',
+              })
+            ),
           },
         },
       }))
@@ -110,8 +107,8 @@ describe('Auto Top-up System', () => {
       // Should check balance
       expect(balanceMock).toHaveBeenCalled()
 
-      // Should create payment intent
-      expect(paymentIntentMock).toHaveBeenCalled()
+      // Should validate auto top-up status
+      expect(validateAutoTopupMock).toHaveBeenCalled()
 
       // Should grant credits
       expect(grantCreditsMock).toHaveBeenCalled()
@@ -133,16 +130,23 @@ describe('Auto Top-up System', () => {
 
       // Update the spies with the new mock implementations
       spyOn(billing, 'calculateUsageAndBalance').mockImplementation(balanceMock)
-      spyOn(billing, 'processAndGrantCredit').mockImplementation(grantCreditsMock)
+      spyOn(billing, 'validateAutoTopupStatus').mockImplementation(
+        validateAutoTopupMock
+      )
+      spyOn(billing, 'processAndGrantCredit').mockImplementation(
+        grantCreditsMock
+      )
 
       await checkAndTriggerAutoTopup('test-user')
 
       // Should still check settings and balance
       expect(dbMock).toHaveBeenCalled()
       expect(balanceMock).toHaveBeenCalled()
 
-      // But should not create payment or grant credits
-      expect(paymentIntentMock.mock.calls.length).toBe(0)
+      // Should still validate auto top-up (this happens before balance check)
+      expect(validateAutoTopupMock.mock.calls.length).toBe(0)
+
+      // But should not grant credits
       expect(grantCreditsMock.mock.calls.length).toBe(0)
     })
 
@@ -162,7 +166,12 @@ describe('Auto Top-up System', () => {
 
       // Update the spies with the new mock implementations
       spyOn(billing, 'calculateUsageAndBalance').mockImplementation(balanceMock)
-      spyOn(billing, 'processAndGrantCredit').mockImplementation(grantCreditsMock)
+      spyOn(billing, 'validateAutoTopupStatus').mockImplementation(
+        validateAutoTopupMock
+      )
+      spyOn(billing, 'processAndGrantCredit').mockImplementation(
+        grantCreditsMock
+      )
 
       await checkAndTriggerAutoTopup('test-user')
 
@@ -172,27 +181,29 @@ describe('Auto Top-up System', () => {
       expect(grantCreditsMock.mock.calls[0]?.[1]).toBe(600)
     })
 
-    it('should disable auto-topup when payment fails', async () => {
-      // Set up payment failure mock
-      paymentIntentMock = mock(() =>
+    it('should disable auto-topup when validation fails', async () => {
+      // Set up validation failure mock
+      validateAutoTopupMock = mock(() =>
         Promise.resolve({
-          status: 'requires_payment_method',
+          blockedReason: 'No valid payment method found',
+          validPaymentMethod: null,
         })
       )
 
-      // Update the module mock
-      mock.module('common/src/util/stripe', () => ({
-        stripeServer: {
-          paymentMethods: {
-            list: paymentMethodsMock,
-          },
-          paymentIntents: {
-            create: paymentIntentMock,
-          },
-        },
-      }))
+      // Update the spy with the new mock implementation
+      spyOn(billing, 'validateAutoTopupStatus').mockImplementation(
+        validateAutoTopupMock
+      )
 
-      await expect(checkAndTriggerAutoTopup('test-user')).rejects.toThrow()
+      await expect(checkAndTriggerAutoTopup('test-user')).rejects.toThrow(
+        'No valid payment method found'
+      )
+
+      // Should have called validation
+      expect(validateAutoTopupMock).toHaveBeenCalled()
+
+      // Should not grant credits
+      expect(grantCreditsMock.mock.calls.length).toBe(0)
     })
 
     afterEach(() => {

backend/src/__tests__/usage-calculation.test.ts

@@ -118,6 +118,7 @@ describe('Usage Calculation System', () => {
       purchase: 0,
       referral: 0,
       admin: 0,
+      organization: 0,
     })
     expect(usageThisCycle).toBe(200) // 500 - 300 = 200 used
   })
@@ -163,6 +164,7 @@ describe('Usage Calculation System', () => {
       purchase: 0,
       referral: 0,
       admin: 0,
+      organization: 0,
     }) // No positive balances
   })
 
@@ -220,6 +222,7 @@ describe('Usage Calculation System', () => {
       purchase: 0, // No positive balance for purchase grant
       referral: 0,
       admin: 0,
+      organization: 0,
     })
 
     // Principals show original grant amounts
@@ -228,6 +231,7 @@ describe('Usage Calculation System', () => {
       purchase: 100,
       referral: 0,
       admin: 0,
+      organization: 0,
     })
 
     // Usage calculation: (200-100) + (100-(-50)) = 100 + 150 = 250

backend/src/api/org.ts

@@ -0,0 +1,60 @@
+import {
+  Request as ExpressRequest,
+  Response as ExpressResponse,
+  NextFunction,
+} from 'express'
+import { z } from 'zod'
+import { findOrganizationForRepository } from '@codebuff/billing'
+
+import { getUserIdFromAuthToken } from '../websockets/websocket-action'
+import { logger } from '@/util/logger'
+
+const isRepoCoveredRequestSchema = z.object({
+  owner: z.string(),
+  repo: z.string(),
+  remoteUrl: z.string(),
+})
+
+async function isRepoCoveredHandler(
+  req: ExpressRequest,
+  res: ExpressResponse,
+  next: NextFunction
+): Promise<void | ExpressResponse> {
+  try {
+    const { owner, repo, remoteUrl } = isRepoCoveredRequestSchema.parse(req.body)
+    
+    // Get user ID from Authorization header
+    const authHeader = req.headers.authorization
+    if (!authHeader || !authHeader.startsWith('Bearer ')) {
+      return res.status(401).json({ error: 'Missing or invalid authorization header' })
+    }
+    
+    const authToken = authHeader.substring(7) // Remove 'Bearer ' prefix
+    const userId = await getUserIdFromAuthToken(authToken)
+    
+    if (!userId) {
+      return res.status(401).json({ error: 'Invalid authentication token' })
+    }
+
+    // Check if repository is covered by an organization
+    const orgLookup = await findOrganizationForRepository(userId, remoteUrl)
+    
+    return res.status(200).json({
+      isCovered: orgLookup.found,
+      organizationName: orgLookup.organizationName,
+      organizationId: orgLookup.organizationId, // Keep organizationId for now, might be used elsewhere
+      organizationSlug: orgLookup.organizationSlug, // Add organizationSlug
+    })
+  } catch (error) {
+    logger.error({ error }, 'Error handling /api/orgs/is-repo-covered request')
+    if (error instanceof z.ZodError) {
+      return res
+        .status(400)
+        .json({ error: 'Invalid request body', issues: error.errors })
+    }
+    next(error)
+    return
+  }
+}
+
+export { isRepoCoveredHandler }

backend/src/api/usage.ts

@@ -10,11 +10,13 @@ import { eq } from 'drizzle-orm'
 
 import { checkAuth } from '../util/check-auth'
 import { genUsageResponse } from '../websockets/websocket-action'
+import { getOrganizationUsageResponse } from '@codebuff/billing'
 import { logger } from '@/util/logger'
 
 const usageRequestSchema = z.object({
   fingerprintId: z.string(),
   authToken: z.string().optional(),
+  orgId: z.string().optional(),
 })
 
 async function getUserIdFromAuthToken(
@@ -35,7 +37,7 @@ async function usageHandler(
   next: NextFunction
 ): Promise<void | ExpressResponse> {
   try {
-    const { fingerprintId, authToken } = usageRequestSchema.parse(req.body)
+    const { fingerprintId, authToken, orgId } = usageRequestSchema.parse(req.body)
     const clientSessionId = `api-${fingerprintId}-${Date.now()}`
 
     const authResult = await checkAuth({
@@ -59,6 +61,19 @@ async function usageHandler(
       return res.status(401).json({ message: 'Authentication failed' })
     }
 
+    // If orgId is provided, return organization usage data
+    if (orgId) {
+      try {
+        const orgUsageResponse = await getOrganizationUsageResponse(orgId, userId)
+        return res.status(200).json(orgUsageResponse)
+      } catch (error) {
+        logger.error({ error, orgId, userId }, 'Error fetching organization usage')
+        // If organization usage fails, fall back to personal usage
+        logger.info({ orgId, userId }, 'Falling back to personal usage due to organization error')
+      }
+    }
+
+    // Return personal usage data (default behavior)
     const usageResponse = await genUsageResponse(
       fingerprintId,
       userId,