@tuturuuu/ai 0.0.10

package/src/credits/cap-output-tokens.ts

@@ -0,0 +1,90 @@
+import type { SupabaseClient } from '@tuturuuu/supabase';
+import { resolveGatewayModelId } from './model-mapping';
+
+/**
+ * Query the gateway model's output price per token from the database.
+ * Returns the flat output_price_per_token. For tiered pricing, returns
+ * the lowest tier cost as a conservative estimate.
+ */
+async function getOutputPricePerToken(
+  sbAdmin: SupabaseClient,
+  modelId: string
+): Promise<number | null> {
+  const gatewayId = resolveGatewayModelId(modelId);
+  const privateDb = sbAdmin.schema('private');
+
+  const { data, error } = await privateDb
+    .from('ai_gateway_models')
+    .select('output_price_per_token, output_tiers')
+    .or(`id.eq.${gatewayId},id.eq.google/${modelId}`)
+    .eq('is_enabled', true)
+    .limit(1)
+    .maybeSingle();
+
+  if (error || !data) return null;
+
+  // If tiered pricing exists, use the first (lowest) tier cost
+  if (
+    data.output_tiers &&
+    Array.isArray(data.output_tiers) &&
+    data.output_tiers.length > 0
+  ) {
+    const firstTier = data.output_tiers[0] as { cost?: string };
+    if (firstTier?.cost) {
+      return parseFloat(firstTier.cost);
+    }
+  }
+
+  return data.output_price_per_token ?? null;
+}
+
+/**
+ * Application-level safety cap: ensure maxOutputTokens won't exceed
+ * what the user can afford based on remaining credits and real gateway pricing.
+ *
+ * This is a defense-in-depth layer — the database function
+ * `check_ai_credit_allowance` applies the same cap, but this catches
+ * cases where the DB function is stale or fails open.
+ *
+ * @returns capped maxOutputTokens, or `null` if no output can be afforded
+ */
+export async function capMaxOutputTokensByCredits(
+  sbAdmin: SupabaseClient,
+  modelId: string,
+  maxOutputTokens: number | null,
+  remainingCredits: number,
+  markupMultiplier = 1.0
+): Promise<number | null> {
+  if (remainingCredits <= 0) return null;
+
+  const outputPricePerToken = await getOutputPricePerToken(sbAdmin, modelId);
+  if (!outputPricePerToken || outputPricePerToken <= 0) {
+    // Can't determine pricing — fall through to DB-layer cap
+    return maxOutputTokens;
+  }
+
+  // credits = (tokens * pricePerToken / 0.0001) * markup
+  // → tokens = (credits * 0.0001 / markup) / pricePerToken
+  const affordableTokens = Math.floor(
+    (remainingCredits * 0.0001) / markupMultiplier / outputPricePerToken
+  );
+
+  if (affordableTokens < 1) return null;
+  if (!maxOutputTokens) return affordableTokens;
+  return Math.min(maxOutputTokens, affordableTokens);
+}
+
+/**
+ * Pure computation version for unit testing (no DB dependency).
+ * Same formula as the async version, but takes price directly.
+ */
+export function computeAffordableTokens(
+  remainingCredits: number,
+  outputPricePerToken: number,
+  markupMultiplier = 1.0
+): number {
+  if (remainingCredits <= 0 || outputPricePerToken <= 0) return 0;
+  return Math.floor(
+    (remainingCredits * 0.0001) / markupMultiplier / outputPricePerToken
+  );
+}

package/src/credits/check-credits.ts

@@ -0,0 +1,232 @@
+import type {
+  AiFeature,
+  CreditErrorCode,
+} from '@tuturuuu/ai/credits/constants';
+import {
+  matchesAllowedModel,
+  resolveGatewayModelId,
+} from '@tuturuuu/ai/credits/model-mapping';
+import type {
+  CreditCheckResult,
+  CreditDeductionResult,
+  DeductCreditsParams,
+} from '@tuturuuu/ai/credits/types';
+import { createAdminClient } from '@tuturuuu/supabase/next/server';
+import {
+  decrementAiCreditChargeInFlight,
+  hasAiCreditChargeInFlight,
+  incrementAiCreditChargeInFlight,
+  invalidateAiCreditSnapshot,
+  isAiCreditSnapshotUsable,
+  readAiCreditSnapshot,
+} from '@tuturuuu/utils/ai-temp-auth';
+
+type DeductAiCreditsRpcRow = {
+  success?: boolean;
+  credits_deducted?: number | string;
+  remaining_credits?: number | string;
+  error_code?: string | null;
+};
+
+type RpcError = { message?: string } | null;
+
+/**
+ * Pre-flight check: can this workspace use AI credits for the given model/feature?
+ * Returns allowance info including remaining credits and effective maxOutputTokens.
+ */
+export async function checkAiCredits(
+  wsId: string | undefined,
+  modelId: string,
+  feature: AiFeature,
+  opts?: { userId?: string; estimatedInputTokens?: number }
+): Promise<CreditCheckResult> {
+  if (!wsId) {
+    return {
+      allowed: false,
+      remainingCredits: 0,
+      tier: 'FREE',
+      maxOutputTokens: null,
+      errorCode: 'CREDIT_CHECK_FAILED',
+      errorMessage: 'Workspace ID is missing.',
+    };
+  }
+
+  const gatewayModelId = resolveGatewayModelId(modelId);
+
+  if (opts?.userId) {
+    const snapshot = await readAiCreditSnapshot({
+      wsId,
+      userId: opts.userId,
+    });
+    const inFlight = await hasAiCreditChargeInFlight({
+      wsId,
+      userId: opts.userId,
+    });
+    if (
+      isAiCreditSnapshotUsable(snapshot, { inFlight }) &&
+      (snapshot.allowedFeatures.length === 0 ||
+        snapshot.allowedFeatures.includes(feature)) &&
+      matchesAllowedModel(gatewayModelId, snapshot.allowedModels)
+    ) {
+      return {
+        allowed: snapshot.remainingCredits > 0,
+        remainingCredits: snapshot.remainingCredits,
+        tier: snapshot.tier,
+        maxOutputTokens: snapshot.maxOutputTokens,
+        errorCode: snapshot.remainingCredits > 0 ? null : 'CREDITS_EXHAUSTED',
+        errorMessage:
+          snapshot.remainingCredits > 0
+            ? null
+            : 'AI credits exhausted. Please upgrade your plan or purchase more credits.',
+      };
+    }
+  }
+
+  const sbAdmin = await createAdminClient();
+  const { data, error } = await sbAdmin.rpc('check_ai_credit_allowance', {
+    p_ws_id: wsId,
+    p_model_id: gatewayModelId,
+    p_feature: feature,
+    ...(opts?.estimatedInputTokens != null
+      ? { p_estimated_input_tokens: opts.estimatedInputTokens }
+      : {}),
+    ...(opts?.userId ? { p_user_id: opts.userId } : {}),
+  });
+
+  if (error) {
+    console.error('Error checking AI credits:', error);
+    return {
+      allowed: false,
+      remainingCredits: 0,
+      tier: 'FREE',
+      maxOutputTokens: null,
+      errorCode: 'CREDIT_CHECK_FAILED',
+      errorMessage: 'AI credit check failed. Please try again.',
+    };
+  }
+
+  const row = Array.isArray(data) ? data[0] : data;
+  if (!row) {
+    return {
+      allowed: false,
+      remainingCredits: 0,
+      tier: 'FREE',
+      maxOutputTokens: null,
+      errorCode: 'CREDIT_CHECK_FAILED',
+      errorMessage: 'AI credit check returned no result. Please try again.',
+    };
+  }
+
+  return {
+    allowed: row.allowed ?? true,
+    remainingCredits: Number(row.remaining_credits ?? 0),
+    tier: row.tier ?? 'FREE',
+    maxOutputTokens: row.max_output_tokens ?? null,
+    errorCode: (row.error_code as CreditErrorCode) ?? null,
+    errorMessage: row.error_message ?? null,
+  };
+}
+
+/**
+ * Deduct credits after a successful AI execution.
+ * Should be called in the `onFinish` callback or after `generateObject` completes.
+ */
+export async function deductAiCredits(
+  params: DeductCreditsParams
+): Promise<CreditDeductionResult> {
+  if (!params.wsId) {
+    return {
+      success: false,
+      creditsDeducted: 0,
+      remainingCredits: 0,
+      errorCode: 'DEDUCTION_FAILED',
+    };
+  }
+
+  const wsId = params.wsId;
+  const sbAdmin = await createAdminClient();
+  const gatewayModelId = resolveGatewayModelId(params.modelId);
+  let inFlightMarked = false;
+
+  if (params.userId) {
+    inFlightMarked = await incrementAiCreditChargeInFlight({
+      wsId,
+      userId: params.userId,
+    });
+  }
+
+  let data: DeductAiCreditsRpcRow[] | DeductAiCreditsRpcRow | null = null;
+  let error: RpcError = null;
+
+  try {
+    const response = await sbAdmin.rpc('deduct_ai_credits', {
+      p_ws_id: wsId,
+      p_model_id: gatewayModelId,
+      p_input_tokens: params.inputTokens,
+      p_output_tokens: params.outputTokens,
+      p_reasoning_tokens: params.reasoningTokens ?? 0,
+      p_feature: params.feature,
+      ...(params.executionId ? { p_execution_id: params.executionId } : {}),
+      ...(params.chatMessageId
+        ? { p_chat_message_id: params.chatMessageId }
+        : {}),
+      ...(params.metadata
+        ? {
+            p_metadata:
+              params.metadata as unknown as import('@tuturuuu/types').Json,
+          }
+        : {}),
+      ...(params.userId ? { p_user_id: params.userId } : {}),
+      ...(params.imageCount ? { p_image_count: params.imageCount } : {}),
+      ...(params.searchCount ? { p_search_count: params.searchCount } : {}),
+    });
+    data = response.data as
+      | DeductAiCreditsRpcRow[]
+      | DeductAiCreditsRpcRow
+      | null;
+    error = response.error as RpcError;
+  } finally {
+    if (inFlightMarked && params.userId) {
+      await decrementAiCreditChargeInFlight({
+        wsId,
+        userId: params.userId,
+      });
+    }
+  }
+
+  if (error) {
+    console.error('Error deducting AI credits:', error);
+    return {
+      success: false,
+      creditsDeducted: 0,
+      remainingCredits: 0,
+      errorCode: 'DEDUCTION_FAILED',
+    };
+  }
+
+  const row = Array.isArray(data) ? data[0] : data;
+  if (!row) {
+    return {
+      success: false,
+      creditsDeducted: 0,
+      remainingCredits: 0,
+      errorCode: 'NO_RESULT',
+    };
+  }
+
+  const result = {
+    success: row.success ?? false,
+    creditsDeducted: Number(row.credits_deducted ?? 0),
+    remainingCredits: Number(row.remaining_credits ?? 0),
+    errorCode: row.error_code ?? null,
+  };
+
+  if (result.success && params.wsId && params.userId) {
+    await invalidateAiCreditSnapshot({
+      wsId,
+      userId: params.userId,
+    });
+  }
+
+  return result;
+}

package/src/credits/constants.ts

@@ -0,0 +1,30 @@
+/** 1 credit = this many USD */
+export const CREDIT_UNIT_USD = 0.0001;
+
+/** AI feature slugs used throughout the credit system */
+export const AI_FEATURES = [
+  'chat',
+  'generate',
+  'task_journal',
+  'email_draft',
+  'image_generation',
+  'embeddings',
+] as const;
+
+export type AiFeature = (typeof AI_FEATURES)[number];
+
+/** Error codes returned by credit checks */
+export const CREDIT_ERROR_CODES = {
+  CREDITS_EXHAUSTED: 'CREDITS_EXHAUSTED',
+  MODEL_NOT_ALLOWED: 'MODEL_NOT_ALLOWED',
+  MODEL_DISABLED: 'MODEL_DISABLED',
+  MODEL_PRICING_UNAVAILABLE: 'MODEL_PRICING_UNAVAILABLE',
+  FEATURE_NOT_ALLOWED: 'FEATURE_NOT_ALLOWED',
+  DAILY_LIMIT_REACHED: 'DAILY_LIMIT_REACHED',
+  NO_ALLOCATION: 'NO_ALLOCATION',
+  NO_BALANCE: 'NO_BALANCE',
+  CREDIT_CHECK_FAILED: 'CREDIT_CHECK_FAILED',
+} as const;
+
+export type CreditErrorCode =
+  (typeof CREDIT_ERROR_CODES)[keyof typeof CREDIT_ERROR_CODES];

package/src/credits/index.ts

@@ -0,0 +1,46 @@
+export {
+  capMaxOutputTokensByCredits,
+  computeAffordableTokens,
+} from './cap-output-tokens';
+export { checkAiCredits, deductAiCredits } from './check-credits';
+export {
+  AI_FEATURES,
+  type AiFeature,
+  CREDIT_ERROR_CODES,
+  CREDIT_UNIT_USD,
+  type CreditErrorCode,
+} from './constants';
+export {
+  matchesAllowedModel,
+  resolveGatewayModelId,
+  toBareModelName,
+  toGatewayModelId,
+} from './model-mapping';
+export {
+  commitFixedAiCreditReservation,
+  commitMeteredEmbeddingCredits,
+  releaseFixedAiCreditReservation,
+  releaseMeteredEmbeddingCredits,
+  reserveFixedAiCredits,
+  reserveMeteredEmbeddingCredits,
+} from './reservations';
+export {
+  type EffectivePlanModel,
+  type PlanModelCapability,
+  PlanModelResolutionError,
+  resolvePlanModel,
+  selectEffectivePlanModel,
+} from './resolve-plan-model';
+export { syncGatewayModels } from './sync-gateway-models';
+export type {
+  AiCreditStatus,
+  CreditAllocation,
+  CreditCheckResult,
+  CreditDeductionResult,
+  CreditReservationCommitResult,
+  CreditReservationReleaseResult,
+  CreditReservationResult,
+  DeductCreditsParams,
+  FeatureAccess,
+  MeteredEmbeddingReservationResult,
+} from './types';

package/src/credits/model-mapping.ts

@@ -0,0 +1,92 @@
+/**
+ * Maps bare model names used in existing code to gateway model IDs.
+ * The gateway uses `provider/model-name` format (e.g., `google/gemini-2.5-flash`),
+ * while existing code uses bare names (e.g., `gemini-2.5-flash`).
+ */
+
+export const GEMINI_31_FLASH_LITE_MODEL = 'gemini-3.1-flash-lite';
+export const GEMINI_31_FLASH_LITE_GATEWAY_MODEL =
+  'google/gemini-3.1-flash-lite';
+export const GEMINI_31_FLASH_LITE_PREVIEW_MODEL =
+  'gemini-3.1-flash-lite-preview';
+export const GEMINI_31_FLASH_LITE_PREVIEW_GATEWAY_MODEL =
+  'google/gemini-3.1-flash-lite-preview';
+export const GEMINI_3_FLASH_MODEL = 'gemini-3-flash';
+export const GEMINI_3_FLASH_GATEWAY_MODEL = 'google/gemini-3-flash';
+
+/** Normalize retired model aliases to the current stable model ids. */
+export function normalizeStableModelId(modelId: string): string {
+  const slashIndex = modelId.indexOf('/');
+  const providerPrefix =
+    slashIndex === -1 ? '' : `${modelId.slice(0, slashIndex + 1)}`;
+  const bareModelId =
+    slashIndex === -1 ? modelId : modelId.slice(slashIndex + 1);
+
+  if (
+    bareModelId === GEMINI_31_FLASH_LITE_PREVIEW_MODEL ||
+    bareModelId === GEMINI_3_FLASH_MODEL
+  ) {
+    return `${providerPrefix}${GEMINI_31_FLASH_LITE_MODEL}`;
+  }
+
+  return modelId;
+}
+
+/** Convert a bare model name + provider to a gateway model ID */
+export function toGatewayModelId(provider: string, modelName: string): string {
+  return `${provider}/${normalizeStableModelId(modelName)}`;
+}
+
+/**
+ * Resolve a model name to a gateway model ID.
+ * If the name already contains '/', it's assumed to be in gateway format.
+ * Otherwise, prepend the provider (defaults to 'google').
+ */
+export function resolveGatewayModelId(
+  modelName: string,
+  provider?: string
+): string {
+  const stableModelName = normalizeStableModelId(modelName);
+  if (stableModelName.includes('/')) return stableModelName;
+  return `${provider || 'google'}/${stableModelName}`;
+}
+
+/**
+ * Checks whether a model is allowed by an allocation list.
+ * Empty lists mean "all models".
+ */
+export function matchesAllowedModel(
+  modelName: string,
+  allowedModels: string[]
+): boolean {
+  if (allowedModels.length === 0) return true;
+
+  const gatewayModelId = resolveGatewayModelId(modelName);
+  const bareModelName = toBareModelName(gatewayModelId);
+
+  return allowedModels.some((allowedModel) => {
+    const normalizedAllowedModel = resolveGatewayModelId(allowedModel);
+    return (
+      normalizedAllowedModel === gatewayModelId ||
+      toBareModelName(normalizedAllowedModel) === bareModelName
+    );
+  });
+}
+
+export function isGoogleModelId(modelId: string): boolean {
+  const slashIndex = modelId.indexOf('/');
+  if (slashIndex === -1) return true;
+
+  const provider = modelId.slice(0, slashIndex);
+  return provider === 'google' || provider === 'google-vertex';
+}
+
+/**
+ * Extract the bare model name from a gateway model ID.
+ * e.g., 'google/gemini-2.5-flash' → 'gemini-2.5-flash'
+ */
+export function toBareModelName(gatewayModelId: string): string {
+  const slashIndex = gatewayModelId.indexOf('/');
+  if (slashIndex === -1) return gatewayModelId;
+  return gatewayModelId.slice(slashIndex + 1);
+}