Create a credit block when you send a message

Author: jahooma

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

@@ -17,6 +17,11 @@ import type {
   Logger,
   LoggerWithContextFn,
 } from '@codebuff/common/types/contracts/logger'
+
+import type {
+  BlockGrantResult,
+  SubscriptionRow,
+} from '@codebuff/billing/subscription'
 import type { NextRequest } from 'next/server'
 
 import type { ChatCompletionRequestBody } from '@/llm-api/types'
@@ -78,6 +83,8 @@ export async function postChatCompletions(params: {
   getAgentRunFromId: GetAgentRunFromIdFn
   fetch: typeof globalThis.fetch
   insertMessageBigquery: InsertMessageBigqueryFn
+  getActiveSubscription?: (params: { userId: string; logger: Logger }) => Promise<SubscriptionRow | null>
+  ensureActiveBlockGrant?: (params: { userId: string; subscription: SubscriptionRow; logger: Logger }) => Promise<BlockGrantResult>
 }) {
   const {
     req,
@@ -88,6 +95,8 @@ export async function postChatCompletions(params: {
     getAgentRunFromId,
     fetch,
     insertMessageBigquery,
+    getActiveSubscription,
+    ensureActiveBlockGrant,
   } = params
   let { logger } = params
 
@@ -182,6 +191,22 @@ export async function postChatCompletions(params: {
       logger,
     })
 
+    // For subscribers, ensure a block grant exists before checking balance.
+    // This is done here block grants should only start when the user begins working.
+    if (getActiveSubscription && ensureActiveBlockGrant) {
+      try {
+        const activeSub = await getActiveSubscription({ userId, logger })
+        if (activeSub) {
+          await ensureActiveBlockGrant({ userId, subscription: activeSub, logger })
+        }
+      } catch (error) {
+        logger.error(
+          { error: getErrorObject(error), userId },
+          'Error ensuring subscription block grant',
+        )
+      }
+    }
+
     // Check user credits
     const {
       balance: { totalRemaining },

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

@@ -1,4 +1,8 @@
 import { insertMessageBigquery } from '@codebuff/bigquery'
+import {
+  ensureActiveBlockGrant,
+  getActiveSubscription,
+} from '@codebuff/billing/subscription'
 import { getUserUsageData } from '@codebuff/billing/usage-service'
 import { trackEvent } from '@codebuff/common/analytics'
 
@@ -21,5 +25,7 @@ export async function POST(req: NextRequest) {
     getAgentRunFromId,
     fetch,
     insertMessageBigquery,
+    getActiveSubscription,
+    ensureActiveBlockGrant,
   })
 }