bookwiz.io / lib / services / usage-tracker.ts
usage-tracker.ts
Raw
import { createServerSupabaseClient } from '@/lib/supabase'
import { getModelTier } from '@/lib/config/models'

export interface AIUsageData {
  user_id: string
  book_id?: string
  chat_id?: string
  model_name: string
  model_provider: string
  prompt_tokens?: number
  completion_tokens?: number
  total_tokens?: number
  cost_usd?: number
  request_type?: string
  success?: boolean
  error_message?: string
}

export interface UsageStats {
  booksCreated: number
  totalWords: number
  smartPrompts: number
  smartPromptsLimit: number
  fastPrompts: number
  fastPromptsLimit: number
}

export interface LimitCheckResult {
  canProceed: boolean
  currentUsage: number
  limit: number
  remainingUsage: number
  message?: string
}

/**
 * Usage Tracking Architecture:
 * 
 * 1. user_usage table: Fast counters for plan limits (smart_prompts_used, fast_prompts_used)
 *    - Used for real-time limit checking before requests
 *    - Updated by increment_usage() database function
 *    - Optimized for speed with single record per user per month
 * 
 * 2. ai_usage table: Detailed logs for analytics and billing
 *    - Records every AI request with full metadata (tokens, costs, models, etc.)
 *    - Used for detailed reporting, billing analysis, and audit trails
 *    - Optimized for analytics queries and historical data
 * 
 * This dual approach gives us both fast limit checking AND detailed analytics.
 */
class UsageTracker {
  /**
   * Check if user can make a request before processing it
   * This prevents requests that would exceed limits
   * Uses direct user_usage table query for real-time accuracy
   */
  async checkRequestLimits(
    userId: string, 
    modelName: string, 
    smartLimit: number, 
    fastLimit: number
  ): Promise<LimitCheckResult> {
    try {
      const tier = getModelTier(modelName)
      const limit = tier === 'smart' ? smartLimit : fastLimit
      
      // Query user_usage table directly for real-time accuracy
      const supabase = createServerSupabaseClient()
      const { data: usageRecords, error: usageError } = await supabase
        .from('user_usage')
        .select('*')
        .eq('user_id', userId)
        .order('updated_at', { ascending: false })
        .limit(1)

      if (usageError) {
        console.error('Failed to check request limits:', usageError)
        // In case of error, allow the request (fail open for user experience)
        return {
          canProceed: true,
          currentUsage: 0,
          limit,
          remainingUsage: limit,
          message: 'Unable to verify limits - proceeding with request'
        }
      }

      const currentUsage = usageRecords?.[0] || null
      
      // Get current usage for the specific tier
      let tierUsage = 0
      if (currentUsage) {
        tierUsage = tier === 'smart' 
          ? currentUsage.smart_prompts_used 
          : currentUsage.fast_prompts_used
      }
      
      const canProceed = tierUsage < limit || limit === 0
      const remainingUsage = Math.max(0, limit - tierUsage)
      
      return {
        canProceed,
        currentUsage: tierUsage,
        limit,
        remainingUsage,
        message: canProceed 
          ? undefined 
          : `You've reached your ${tier} AI limit for this month (${tierUsage}/${limit}). Upgrade your plan for more requests.`
      }
    } catch (error) {
      console.error('Error checking request limits:', error)
      // Fail open - allow request if we can't check limits
      const tier = getModelTier(modelName)
      const limit = tier === 'smart' ? smartLimit : fastLimit
      return {
        canProceed: true,
        currentUsage: 0,
        limit,
        remainingUsage: limit
      }
    }
  }

  /**
   * Record AI usage for billing and analytics
   */
  async recordUsage(data: AIUsageData): Promise<boolean> {
    try {
      const tier = getModelTier(data.model_name)
      
      // Ensure token values are numbers and calculate total tokens correctly for constraint
      const promptTokens = Number(data.prompt_tokens) || 0
      const completionTokens = Number(data.completion_tokens) || 0
      // Always calculate total tokens to ensure constraint is satisfied
      const totalTokens = promptTokens + completionTokens
      
      const supabase = createServerSupabaseClient()
      
      // First attempt with chat_id
      let { error } = await supabase
        .from('ai_usage')
        .insert({
          user_id: data.user_id,
          book_id: data.book_id || null,
          chat_id: data.chat_id || null,
          model_name: data.model_name,
          model_tier: tier,
          model_provider: data.model_provider,
          prompt_tokens: promptTokens,
          completion_tokens: completionTokens,
          total_tokens: totalTokens,
          cost_usd: data.cost_usd || 0,
          request_type: data.request_type || 'chat',
          success: data.success !== false, // Default to true
          error_message: data.error_message || null,
        })

      // If we got a foreign key constraint error on chat_id, retry without it
      if (error && error.code === '23503' && error.message.includes('chat_id')) {
        console.warn('Chat ID not found in database, recording usage without chat_id:', data.chat_id)
        
        const retryResult = await supabase
          .from('ai_usage')
          .insert({
            user_id: data.user_id,
            book_id: data.book_id || null,
            chat_id: null, // Set to null to avoid constraint
            model_name: data.model_name,
            model_tier: tier,
            model_provider: data.model_provider,
            prompt_tokens: promptTokens,
            completion_tokens: completionTokens,
            total_tokens: totalTokens,
            cost_usd: data.cost_usd || 0,
            request_type: data.request_type || 'chat',
            success: data.success !== false,
            error_message: data.error_message || null,
          })
        
        error = retryResult.error
      }

      if (error) {
        console.error('Failed to record AI usage:', error)
        console.error('Data that failed to insert:', {
          user_id: data.user_id,
          model_name: data.model_name,
          model_tier: tier,
          model_provider: data.model_provider,
        })
        // Return false but don't throw - usage tracking shouldn't break the main flow
        return false
      }
      
      // ALSO increment the user_usage table for plan limit tracking (only on successful requests)
      if (data.success !== false) {
        const featureType = tier === 'smart' ? 'smart_prompts' : 'fast_prompts'
        
        try {
          const { error: incrementError } = await supabase
            .rpc('increment_usage', {
              user_uuid: data.user_id,
              feature_type_param: featureType,
              amount_param: 1,
              metadata_param: {
                model_name: data.model_name,
                model_provider: data.model_provider,
                tokens: totalTokens,
                request_type: data.request_type || 'chat'
              },
              skip_limit_check: true // Skip limit check since we already validated before the request
            })

          if (incrementError) {
            console.error('Failed to increment user usage counter:', incrementError)
            // Don't fail the whole operation - the ai_usage record is more important
          }
        } catch (incrementError) {
          console.error('Error incrementing user usage counter:', incrementError)
          // Don't fail the whole operation
        }
      }
      
      return true
    } catch (error) {
      console.error('Error recording AI usage:', error)
      // Silent fail - usage tracking is non-critical
      return false
    }
  }

  /**
   * Get current month usage statistics for a user
   */
  async getUserUsageStats(userId: string): Promise<UsageStats | null> {
    try {
      const supabase = createServerSupabaseClient()
      
      // CACHE FIX: Use direct table query instead of RPC function to avoid caching
      // First get the current period using RPC (this is less likely to be cached)
      const { data: periodData, error: periodError } = await supabase
        .rpc('get_current_usage_period', { user_uuid: userId })

      if (periodError) {
        console.error('Failed to get current period:', periodError)
        return null
      }

      const period = periodData?.[0]
      if (!period) {
        console.error('No period data returned')
        return null
      }

      // CACHE FIX: Query user_usage table directly with timestamp filter to get fresh data
      const { data: usageRecords, error: usageError } = await supabase
        .from('user_usage')
        .select('*')
        .eq('user_id', userId)
        .eq('period_start', period.period_start)
        .eq('period_end', period.period_end)
        .order('updated_at', { ascending: false })
        .limit(1)

      if (usageError) {
        console.error('Failed to get current usage:', usageError)
        return null
      }

      // Get the current usage record, create if doesn't exist
      let currentUsage = usageRecords?.[0] || null
      
      // If no record exists, create one using RPC (this is rare)
      if (!currentUsage) {
        const { data: createdUsage, error: createError } = await supabase
          .rpc('get_or_create_current_usage', { user_uuid: userId })
          
        if (createError) {
          console.error('Failed to create usage record:', createError)
          return null
        }
        currentUsage = createdUsage
      }
      
      // Get total books created (all time)
      const { data: booksData, error: booksError } = await supabase
        .from('books')
        .select('id')
        .eq('user_id', userId)

      const booksCreated = booksData?.length || 0

      if (booksError) {
        console.error('Error getting books count:', booksError)
      }

      // Get total words across all books
      const { data: wordsData, error: wordsError } = await supabase
        .from('books')
        .select('word_count')
        .eq('user_id', userId)

      const totalWords = wordsData?.reduce((sum, book) => sum + (book.word_count || 0), 0) || 0

      if (wordsError) {
        console.error('Error getting total words:', wordsError)
      }

      // Get user's subscription to determine limits (handle free tier users)
      const { data: subscriptionData, error: subscriptionError } = await supabase
        .from('subscriptions')
        .select('price_id, status, current_period_start, current_period_end')
        .eq('user_id', userId)
        .eq('status', 'active')
        .order('created_at', { ascending: false })
        .limit(1)

      if (subscriptionError) {
        console.error('Error getting subscription:', subscriptionError)
        // Continue with free tier limits if subscription query fails
      }

      // Get the first subscription or null if none exists
      const subscription = subscriptionData && subscriptionData.length > 0 ? subscriptionData[0] : null

      // Default to free tier limits
      const { PRICING_TIERS } = await import('@/lib/stripe')
      let smartPromptsLimit = PRICING_TIERS.FREE.maxSmartPrompts
      let fastPromptsLimit = PRICING_TIERS.FREE.hasUnlimitedFastPrompts ? 999999 : PRICING_TIERS.FREE.maxFastPrompts

      // Set limits based on subscription
      let planName = 'FREE'
      if (subscription?.price_id) {
        // Import here to avoid circular dependencies
        const { getPlanByPriceId } = await import('@/lib/stripe')
        const plan = getPlanByPriceId(subscription.price_id)
        
        if (plan) {
          planName = plan.name || subscription.price_id
          smartPromptsLimit = plan.maxSmartPrompts
          fastPromptsLimit = plan.hasUnlimitedFastPrompts ? 999999 : plan.maxFastPrompts
        } else {
          console.warn('⚠️ Plan not found for price_id:', subscription.price_id)
        }
      }
      
      const result = {
        booksCreated,
        totalWords,
        smartPrompts: currentUsage?.smart_prompts_used || 0,
        smartPromptsLimit,
        fastPrompts: currentUsage?.fast_prompts_used || 0,
        fastPromptsLimit
      }
      
      return result
    } catch (error) {
      console.error('Error getting usage stats:', error)
      return null
    }
  }

  /**
   * Get detailed usage breakdown for a specific month
   */
  async getMonthlyUsage(userId: string, month?: Date): Promise<any> {
    try {
      const targetMonth = month || new Date()
      const startDate = new Date(targetMonth.getFullYear(), targetMonth.getMonth(), 1)
      const endDate = new Date(targetMonth.getFullYear(), targetMonth.getMonth() + 1, 1)
      
      const supabase = createServerSupabaseClient()
      
      // Use direct query without .group() since Supabase JS client doesn't support it
      // Get raw data and group manually
      const { data: rawData, error } = await supabase
        .from('ai_usage')
        .select(`
          model_name,
          model_tier,
          model_provider,
          request_type,
          total_tokens,
          cost_usd
        `)
        .eq('user_id', userId)
        .gte('created_at', startDate.toISOString())
        .lt('created_at', endDate.toISOString())
        .eq('success', true)

      if (error) {
        console.error('Failed to get monthly usage:', error)
        return null
      }

      // Group the data manually in JavaScript
      const groupedData = new Map()
      
      rawData?.forEach(record => {
        const key = `${record.model_name}|${record.model_tier}|${record.model_provider}|${record.request_type}`
        
        if (!groupedData.has(key)) {
          groupedData.set(key, {
            model_name: record.model_name,
            model_tier: record.model_tier,
            model_provider: record.model_provider,
            request_type: record.request_type,
            usage_count: 0,
            total_tokens: 0,
            total_cost: 0
          })
        }
        
        const group = groupedData.get(key)
        group.usage_count += 1
        group.total_tokens += record.total_tokens || 0
        group.total_cost += record.cost_usd || 0
      })

      return Array.from(groupedData.values())
    } catch (error) {
      console.error('Error getting monthly usage:', error)
      return null
    }
  }


}

export const usageTracker = new UsageTracker()