fix: enforce fallback_to_a_la_carte preference and move block grant after validation

Author: brandonkachen

web/src/app/api/v1/chat/completions/_post.ts

@@ -18,7 +18,18 @@ import type {
   LoggerWithContextFn,
 } from '@codebuff/common/types/contracts/logger'
 
-import type { BlockGrantResult } from '@codebuff/billing/subscription'
+import type {
+  BlockGrantResult,
+} from '@codebuff/billing/subscription'
+import {
+  isWeeklyLimitError,
+  isBlockExhaustedError,
+} from '@codebuff/billing/subscription'
+
+export type GetUserPreferencesFn = (params: {
+  userId: string
+  logger: Logger
+}) => Promise<{ fallbackToALaCarte: boolean }>
 import type { NextRequest } from 'next/server'
 
 import type { ChatCompletionRequestBody } from '@/llm-api/types'
@@ -81,6 +92,7 @@ export async function postChatCompletions(params: {
   fetch: typeof globalThis.fetch
   insertMessageBigquery: InsertMessageBigqueryFn
   ensureSubscriberBlockGrant?: (params: { userId: string; logger: Logger }) => Promise<BlockGrantResult | null>
+  getUserPreferences?: GetUserPreferencesFn
 }) {
   const {
     req,
@@ -92,6 +104,7 @@ export async function postChatCompletions(params: {
     fetch,
     insertMessageBigquery,
     ensureSubscriberBlockGrant,
+    getUserPreferences,
   } = params
   let { logger } = params
 
@@ -186,19 +199,6 @@ export async function postChatCompletions(params: {
       logger,
     })
 
-    // For subscribers, ensure a block grant exists before checking balance.
-    // This is done here because block grants should only start when the user begins working.
-    if (ensureSubscriberBlockGrant) {
-      try {
-        await ensureSubscriberBlockGrant({ userId, logger })
-      } catch (error) {
-        logger.error(
-          { error: getErrorObject(error), userId },
-          'Error ensuring subscription block grant',
-        )
-      }
-    }
-
     // Check user credits
     const {
       balance: { totalRemaining },
@@ -281,6 +281,59 @@ export async function postChatCompletions(params: {
       )
     }
 
+    // For subscribers, ensure a block grant exists before processing the request.
+    // This is done AFTER validation so malformed requests don't start a new 5-hour block.
+    if (ensureSubscriberBlockGrant) {
+      try {
+        const blockGrantResult = await ensureSubscriberBlockGrant({ userId, logger })
+        
+        // Check if user hit subscription limit and should be rate-limited
+        if (blockGrantResult && (isWeeklyLimitError(blockGrantResult) || isBlockExhaustedError(blockGrantResult))) {
+          // Fetch user's preference for falling back to a-la-carte credits
+          const preferences = getUserPreferences
+            ? await getUserPreferences({ userId, logger })
+            : { fallbackToALaCarte: true } // Default to allowing a-la-carte if no preference function
+          
+          if (!preferences.fallbackToALaCarte) {
+            const resetTime = blockGrantResult.resetsAt
+            const resetCountdown = formatQuotaResetCountdown(resetTime.toISOString())
+            const limitType = isWeeklyLimitError(blockGrantResult) ? 'weekly' : '5-hour session'
+            
+            trackEvent({
+              event: AnalyticsEvent.CHAT_COMPLETIONS_INSUFFICIENT_CREDITS,
+              userId,
+              properties: {
+                reason: 'subscription_limit_no_fallback',
+                limitType,
+                fallbackToALaCarte: false,
+              },
+              logger,
+            })
+            
+            return NextResponse.json(
+              {
+                error: 'rate_limit_exceeded',
+                message: `Subscription ${limitType} limit reached. Your limit resets ${resetCountdown}. Enable "Continue with credits" in the CLI to use a-la-carte credits.`,
+              },
+              { status: 429 },
+            )
+          }
+          // If fallbackToALaCarte is true, continue to use a-la-carte credits
+          logger.info(
+            { userId, limitType: isWeeklyLimitError(blockGrantResult) ? 'weekly' : 'session' },
+            'Subscriber hit limit, falling back to a-la-carte credits',
+          )
+        }
+      } catch (error) {
+        logger.error(
+          { error: getErrorObject(error), userId },
+          'Error ensuring subscription block grant',
+        )
+        // Fail open: if we can't check the subscription status, allow the request to proceed
+        // This is intentional - we prefer to allow requests rather than block legitimate users
+      }
+    }
+
     const openrouterApiKey = req.headers.get(BYOK_OPENROUTER_HEADER)
 
     // Handle streaming vs non-streaming

web/src/app/api/v1/chat/completions/route.ts

@@ -2,15 +2,29 @@ import { insertMessageBigquery } from '@codebuff/bigquery'
 import { ensureSubscriberBlockGrant } from '@codebuff/billing/subscription'
 import { getUserUsageData } from '@codebuff/billing/usage-service'
 import { trackEvent } from '@codebuff/common/analytics'
+import db from '@codebuff/internal/db'
+import * as schema from '@codebuff/internal/db/schema'
+import { eq } from 'drizzle-orm'
 
 import { postChatCompletions } from './_post'
 
+import type { GetUserPreferencesFn } from './_post'
 import type { NextRequest } from 'next/server'
 
 import { getAgentRunFromId } from '@/db/agent-run'
 import { getUserInfoFromApiKey } from '@/db/user'
 import { logger, loggerWithContext } from '@/util/logger'
 
+const getUserPreferences: GetUserPreferencesFn = async ({ userId }) => {
+  const userPrefs = await db.query.user.findFirst({
+    where: eq(schema.user.id, userId),
+    columns: { fallback_to_a_la_carte: true },
+  })
+  return {
+    fallbackToALaCarte: userPrefs?.fallback_to_a_la_carte ?? false,
+  }
+}
+
 export async function POST(req: NextRequest) {
   return postChatCompletions({
     req,
@@ -23,5 +37,6 @@ export async function POST(req: NextRequest) {
     fetch,
     insertMessageBigquery,
     ensureSubscriberBlockGrant,
+    getUserPreferences,
   })
 }