add byok

Author: charleslien

common/src/constants/byok.ts

@@ -0,0 +1,2 @@
+export const BYOK_OPENROUTER_HEADER = 'x-openrouter-api-key'
+export const BYOK_OPENROUTER_ENV_VAR = 'CODEBUFF_BYOK_OPENROUTER'

packages/billing/src/balance-calculator.ts

@@ -381,6 +381,7 @@ export async function consumeCreditsAndAddAgentStep(params: {
 
   cost: number
   credits: number
+  byok: boolean
 
   inputTokens: number
   cacheCreationInputTokens: number | null
@@ -405,6 +406,7 @@ export async function consumeCreditsAndAddAgentStep(params: {
 
     cost,
     credits,
+    byok,
 
     inputTokens,
     cacheCreationInputTokens,
@@ -423,31 +425,37 @@ export async function consumeCreditsAndAddAgentStep(params: {
       await withSerializableTransaction({
         callback: async (tx) => {
           const now = new Date()
-          const activeGrants = await getOrderedActiveGrants({
-            ...params,
-            now,
-            conn: tx,
-          })
-
-          if (activeGrants.length === 0) {
-            logger.error(
-              { userId, credits },
-              'No active grants found to consume credits from',
-            )
-            throw new Error('No active grants found')
-          }
 
-          const result = await consumeFromOrderedGrants({
-            ...params,
-            creditsToConsume: credits,
-            grants: activeGrants,
-            tx,
-          })
+          let result: CreditConsumptionResult | null = null
+          consumeCredits: {
+            if (byok) {
+              break consumeCredits
+            }
+            const activeGrants = await getOrderedActiveGrants({
+              ...params,
+              now,
+              conn: tx,
+            })
 
-          if (userId === TEST_USER_ID) {
-            return { ...result, agentStepId: 'test-step-id' }
+            if (activeGrants.length === 0) {
+              logger.error(
+                { userId, credits },
+                'No active grants found to consume credits from',
+              )
+              throw new Error('No active grants found')
+            }
+
+            result = await consumeFromOrderedGrants({
+              ...params,
+              creditsToConsume: credits,
+              grants: activeGrants,
+              tx,
+            })
+
+            if (userId === TEST_USER_ID) {
+              return { ...result, agentStepId: 'test-step-id' }
+            }
           }
-          const stepId = crypto.randomUUID()
 
           try {
             await tx.insert(schema.message).values({
@@ -466,6 +474,7 @@ export async function consumeCreditsAndAddAgentStep(params: {
               output_tokens: outputTokens,
               cost: cost.toString(),
               credits,
+              byok,
               latency_ms: latencyMs,
               user_id: userId,
             })
@@ -474,7 +483,13 @@ export async function consumeCreditsAndAddAgentStep(params: {
             throw error
           }
 
-          return { ...result, agentStepId: stepId }
+          if (!result) {
+            result = {
+              consumed: 0,
+              fromPurchased: 0,
+            }
+          }
+          return { ...result, agentStepId: crypto.randomUUID() }
         },
         context: { userId, credits },
         logger,

packages/internal/src/db/migrations/0031_damp_hobgoblin.sql

@@ -0,0 +1 @@
+ALTER TABLE "message" ADD COLUMN "byok" boolean DEFAULT false NOT NULL;