d3ro-voice/server/supabase/functions/_shared/embedding-quota.ts

// server/supabase/functions/_shared/embedding-quota.ts
// Cost guard for the paid OpenAI embedding calls made by embed-chunks and
// search-knowledge.
//
// Before this module both functions only checked the JWT: any signed-in account
// could insert an unbounded number of knowledge_chunks rows (RLS only checks
// document ownership, `content` is unbounded text) and have every one of them
// embedded, repeatedly, and search-knowledge embedded a 4,000-char query per
// call with no counter. stt/llm/realtime all go through quota.ts; embeddings
// did not.
//
// Policy (pure):   size caps + per-tier unit budget + batch planning.
// Port:            EmbeddingUsageStore (tier lookup + atomic counter).
// Adapter:         createSupabaseEmbeddingUsageStore (subscriptions + daily_usage
//                  via the existing service-role `increment_daily_usage` RPC).
//
// TODO(plan-catalog): the per-tier budget below belongs in
// packages/core/src/plan-catalog.ts `PLAN_QUOTA` (feature `knowledge_embedding`)
// and should be read from core-contract.generated.ts once it is added there.

import type { createClient } from '@supabase/supabase-js'
import type { PlanQuotaPeriod, PlanQuotaTier } from './core-contract.generated.ts'

/** daily_usage.feature key for embedding spend. */
export const EMBEDDING_QUOTA_FEATURE = 'knowledge_embedding'

export const EMBEDDING_LIMITS = Object.freeze({
  /** Longest chunk accepted for embedding (the OpenAI per-input cap is ~8k tokens). */
  maxChunkChars: 8_000,
  /** Most chunks one document may have. Desktop chunking stops at 2,000 chunks. */
  maxChunksPerDocument: 2_000,
  /** Longest search query accepted. */
  maxQueryChars: 4_000,
  /** Inputs per provider request. */
  maxBatchInputs: 100,
  /** Characters per provider request (keeps a batch under the per-request token cap). */
  maxBatchChars: 100_000,
  /** One quota unit = up to this many characters of embedded text. */
  charsPerUnit: 1_000,
})

export interface EmbeddingQuotaPolicy {
  /** Units per period. -1 = unlimited, 0 = not available. */
  limit: number
  period: PlanQuotaPeriod
}

/**
 * Units (1 unit = up to 1,000 chars) per tier. Cloud knowledge is offered on
 * every tier (web/mobile have no paywall for it), so free gets a small weekly
 * budget instead of none: 5,000 units ≈ 5M chars ≈ 20 maximum-size mobile
 * documents a week.
 */
export const EMBEDDING_QUOTA: Readonly<Record<PlanQuotaTier, EmbeddingQuotaPolicy>> = Object.freeze({
  free: { limit: 5_000, period: 'weekly' },
  pro: { limit: 50_000, period: 'daily' },
  pro_plus: { limit: 150_000, period: 'daily' },
  team: { limit: 300_000, period: 'daily' },
  enterprise: { limit: -1, period: 'daily' },
})

const QUOTA_TIERS: readonly PlanQuotaTier[] = ['free', 'pro', 'pro_plus', 'team', 'enterprise']

export function normalizeQuotaTier(value: unknown): PlanQuotaTier {
  return typeof value === 'string' && (QUOTA_TIERS as readonly string[]).includes(value)
    ? value as PlanQuotaTier
    : 'free'
}

export function embeddingPolicyFor(tier: PlanQuotaTier): EmbeddingQuotaPolicy {
  return EMBEDDING_QUOTA[tier] ?? { limit: 0, period: 'daily' }
}

/** Units one text costs: every started 1,000 characters, at least 1. */
export function embeddingUnitsFor(text: string): number {
  return Math.max(1, Math.ceil(text.length / EMBEDDING_LIMITS.charsPerUnit))
}

export function embeddingUnits(texts: readonly string[]): number {
  return texts.reduce((sum, text) => sum + embeddingUnitsFor(text), 0)
}

/** Index of the first chunk over `maxChunkChars`, or -1. */
export function firstOversizedChunk(chunks: readonly { content: string }[]): number {
  return chunks.findIndex((chunk) => chunk.content.length > EMBEDDING_LIMITS.maxChunkChars)
}

/** Split chunks into provider requests bounded by input count and characters. */
export function planEmbeddingBatches<T extends { content: string }>(chunks: readonly T[]): T[][] {
  const batches: T[][] = []
  let current: T[] = []
  let chars = 0
  for (const chunk of chunks) {
    const size = chunk.content.length
    if (
      current.length > 0
      && (current.length >= EMBEDDING_LIMITS.maxBatchInputs || chars + size > EMBEDDING_LIMITS.maxBatchChars)
    ) {
      batches.push(current)
      current = []
      chars = 0
    }
    current.push(chunk)
    chars += size
  }
  if (current.length > 0) batches.push(current)
  return batches
}

function isoDate(date: Date): string {
  return date.toISOString().slice(0, 10)
}

/** First date (inclusive) of the usage window before today, or null for a daily period. */
export function usageWindowStart(period: PlanQuotaPeriod, now: Date): string | null {
  if (period !== 'weekly') return null
  const start = new Date(now.getTime())
  start.setUTCDate(start.getUTCDate() - 6)
  return isoDate(start)
}

// ── Port ────────────────────────────────────────────────────────────────

export interface EmbeddingUsageStore {
  readTier(userId: string): Promise<PlanQuotaTier>
  /** Atomically add `amount` (may be negative) to today's counter; returns the new count. */
  increment(userId: string, feature: string, amount: number): Promise<number>
  /** Sum of counters with fromDate <= date < beforeDate. */
  sumBetween(userId: string, feature: string, fromDate: string, beforeDate: string): Promise<number>
}

export interface EmbeddingReservation {
  allowed: boolean
  /** Units held by this reservation (0 when denied or nothing was requested). */
  units: number
  /** Usage in the current window including this reservation when allowed. */
  current: number
  limit: number
  period: PlanQuotaPeriod
  tier: PlanQuotaTier
  /** UTC date the units were counted on. */
  date: string
}

/**
 * Reserve `units` before spending. The counter is incremented first (atomic in
 * the database) and rolled back when the window total then exceeds the limit,
 * so concurrent requests can never admit more than the limit between them.
 * Throws when the store fails: callers must fail closed.
 */
export async function reserveEmbeddingUnits(
  store: EmbeddingUsageStore,
  userId: string,
  units: number,
  now: Date = new Date(),
): Promise<EmbeddingReservation> {
  const tier = await store.readTier(userId)
  const policy = embeddingPolicyFor(tier)
  const date = isoDate(now)
  const base = { limit: policy.limit, period: policy.period, tier, date }

  if (units <= 0) return { ...base, allowed: true, units: 0, current: 0 }
  if (policy.limit === 0) return { ...base, allowed: false, units: 0, current: 0 }

  const today = await store.increment(userId, EMBEDDING_QUOTA_FEATURE, units)
  if (policy.limit === -1) return { ...base, allowed: true, units, current: today }

  const windowStart = usageWindowStart(policy.period, now)
  const earlier = windowStart
    ? await store.sumBetween(userId, EMBEDDING_QUOTA_FEATURE, windowStart, date)
    : 0
  const used = earlier + today
  if (used > policy.limit) {
    await store.increment(userId, EMBEDDING_QUOTA_FEATURE, -units)
    return { ...base, allowed: false, units: 0, current: used - units }
  }
  return { ...base, allowed: true, units, current: used }
}

/**
 * Give back units that were reserved but not spent. Best effort: a failure
 * leaves the user slightly over-counted, never under-counted. A refund after
 * the UTC day rolled over is skipped so it cannot credit the new day.
 */
export async function refundEmbeddingUnits(
  store: EmbeddingUsageStore,
  userId: string,
  reservation: EmbeddingReservation,
  units: number,
  now: Date = new Date(),
): Promise<void> {
  const amount = Math.min(units, reservation.units)
  if (amount <= 0 || isoDate(now) !== reservation.date) return
  try {
    await store.increment(userId, EMBEDDING_QUOTA_FEATURE, -amount)
  } catch {
    // over-counting is the safe direction
  }
}

export function quotaExceededBody(reservation: EmbeddingReservation, requested: number): Record<string, unknown> {
  return {
    error: 'quota_exceeded',
    feature: EMBEDDING_QUOTA_FEATURE,
    requested,
    current: reservation.current,
    limit: reservation.limit,
    period: reservation.period,
    tier: reservation.tier,
  }
}

// ── Supabase adapter ────────────────────────────────────────────────────

type SupabaseClient = ReturnType<typeof createClient>

export function createSupabaseEmbeddingUsageStore(client: SupabaseClient): EmbeddingUsageStore {
  return {
    async readTier(userId) {
      const { data, error } = await client
        .from('subscriptions')
        .select('tier')
        .eq('user_id', userId)
        .maybeSingle()
      if (error) throw new Error('Failed to read subscription tier.')
      return normalizeQuotaTier((data as { tier?: unknown } | null)?.tier)
    },
    async increment(userId, feature, amount) {
      const { data, error } = await client.rpc('increment_daily_usage', {
        p_user_id: userId,
        p_feature: feature,
        p_amount: amount,
      })
      if (error || typeof data !== 'number') throw new Error('Failed to update embedding usage.')
      return data
    },
    async sumBetween(userId, feature, fromDate, beforeDate) {
      const { data, error } = await client
        .from('daily_usage')
        .select('count')
        .eq('user_id', userId)
        .eq('feature', feature)
        .gte('date', fromDate)
        .lt('date', beforeDate)
      if (error) throw new Error('Failed to read embedding usage.')
      const rows = (data ?? []) as Array<{ count: number | null }>
      return rows.reduce((sum, row) => sum + (row.count ?? 0), 0)
    },
  }
}