Momento/memento-note/lib/ai-quota.ts

import { willUseByokForLane } from '@/lib/ai/provider-for-user'
import type { AiFeatureLane } from '@/lib/ai/router'
import { getSystemConfig } from '@/lib/config'
import {
  reserveUsageOrThrow,
  releaseUsage,
  checkSessionEntitlementOrThrow,
  QuotaExceededError,
  QuotaServiceUnavailableError,
} from '@/lib/entitlements'
import type { FeatureName } from '@/lib/quota-utils'
import { emitAiUsageChanged } from '@/lib/ai-usage-sync'

export type ReserveAiUsageOptions = {
  lane?: AiFeatureLane
  /** Defaults to userId — host-pays brainstorm bills this account. */
  billingUserId?: string
}

/** Reserve one AI unit; skips when BYOK will be used for the lane. */
export async function reserveAiUsageOrThrow(
  userId: string,
  feature: FeatureName,
  options?: ReserveAiUsageOptions,
): Promise<{ usedByok: boolean; billedUserId: string }> {
  const billedUserId = options?.billingUserId ?? userId
  const lane = options?.lane ?? 'chat'
  const config = await getSystemConfig()
  const { usedByok } = await willUseByokForLane(lane, config, billedUserId)
  if (!usedByok) {
    await reserveUsageOrThrow(billedUserId, feature)
  }
  return { usedByok, billedUserId }
}

/** Host-pays: reserve on session owner with guest metadata on 402. */
export async function reserveSessionAiUsageOrThrow(
  billingOwnerId: string,
  triggeredByUserId: string,
  isGuestActor: boolean,
  feature: FeatureName,
  options?: { lane?: AiFeatureLane },
): Promise<{ usedByok: boolean }> {
  const lane = options?.lane ?? 'chat'
  const config = await getSystemConfig()
  const { usedByok } = await willUseByokForLane(lane, config, billingOwnerId)
  if (!usedByok) {
    await checkSessionEntitlementOrThrow(
      billingOwnerId,
      triggeredByUserId,
      isGuestActor,
      feature,
    )
  }
  return { usedByok }
}

/** Run fn under quota; releases the unit if fn throws (unless BYOK). */
export async function withAiQuota<T>(
  userId: string,
  feature: FeatureName,
  fn: () => Promise<T>,
  options?: ReserveAiUsageOptions,
): Promise<T> {
  const { usedByok, billedUserId } = await reserveAiUsageOrThrow(userId, feature, options)
  try {
    const result = await fn()
    if (!usedByok) {
      emitAiUsageChanged()
    }
    return result
  } catch (err) {
    if (!usedByok) {
      await releaseUsage(billedUserId, feature)
    }
    throw err
  }
}

/** Host-pays collaborative flows (brainstorm) with rollback on failure. */
export async function withSessionAiQuota<T>(
  billingOwnerId: string,
  triggeredByUserId: string,
  isGuestActor: boolean,
  feature: FeatureName,
  fn: () => Promise<T>,
  options?: { lane?: AiFeatureLane },
): Promise<T> {
  const { usedByok } = await reserveSessionAiUsageOrThrow(
    billingOwnerId,
    triggeredByUserId,
    isGuestActor,
    feature,
    options,
  )
  try {
    const result = await fn()
    if (!usedByok) {
      emitAiUsageChanged()
    }
    return result
  } catch (err) {
    if (!usedByok) {
      await releaseUsage(billingOwnerId, feature)
    }
    throw err
  }
}

export function isQuotaError(err: unknown): err is QuotaExceededError {
  return err instanceof QuotaExceededError
}

export function quotaExceededJson(err: QuotaExceededError) {
  return err.toJSON()
}

export function quotaErrorStatus() {
  return 402
}

/** Standard Next.js / Response handlers for quota errors. */
export function handleQuotaHttpError(err: unknown): Response | null {
  if (err instanceof QuotaExceededError) {
    return Response.json(err.toJSON(), { status: 402 })
  }
  if (err instanceof QuotaServiceUnavailableError) {
    return Response.json({ error: err.code }, { status: 503 })
  }
  return null
}

export function notifyAiUsageChanged(): void {
  emitAiUsageChanged()
}