Momento/memento-note/lib/byok.ts

import { prisma } from '@/lib/prisma';
import { decryptApiKey, encryptApiKey, hashApiKey } from '@/lib/crypto';
import {
  VALID_PROVIDERS,
  type AiGatewayProvider,
} from '@/lib/ai/router';
import { getProviderConfigKeys } from '@/lib/ai/factory';
import type { SubscriptionTier } from '@/lib/entitlements';

const PRO_BYOK_PROVIDERS: readonly AiGatewayProvider[] = [
  'openai',
  'anthropic',
  'deepseek',
  'openrouter',
  'minimax',
  'zai',
];

const BUSINESS_BYOK_PROVIDERS: readonly AiGatewayProvider[] = [
  ...VALID_PROVIDERS,
].filter((p) => p !== 'ollama' && p !== 'lmstudio') as AiGatewayProvider[];

export function getAllowedByokProviders(
  tier: SubscriptionTier,
): readonly AiGatewayProvider[] {
  if (tier === 'BASIC') return [];
  if (tier === 'PRO') return PRO_BYOK_PROVIDERS;
  return BUSINESS_BYOK_PROVIDERS;
}

export function isByokProviderAllowed(
  tier: SubscriptionTier,
  provider: string,
): boolean {
  return getAllowedByokProviders(tier).includes(provider as AiGatewayProvider);
}

export async function hasAnyActiveByok(userId: string): Promise<boolean> {
  const count = await prisma.userAPIKey.count({
    where: { userId, isActive: true },
  });
  return count > 0;
}

export async function getActiveByokKey(userId: string, provider: string) {
  return prisma.userAPIKey.findFirst({
    where: { userId, provider, isActive: true },
  });
}

export async function resolveByokApiKey(
  userId: string,
  providerType: string,
): Promise<{ plaintext: string; provider: string } | null> {
  const row = await getActiveByokKey(userId, providerType);
  if (!row) return null;
  try {
    const plaintext = await decryptApiKey(row.encryptedKey);
    return { plaintext, provider: row.provider };
  } catch (err) {
    console.error('[byok] Failed to decrypt key for provider', providerType, err);
    return null;
  }
}

export async function applyByokToConfig(
  billingUserId: string,
  providerType: string,
  config: Record<string, string>,
): Promise<{ config: Record<string, string>; usedByok: boolean }> {
  const byok = await resolveByokApiKey(billingUserId, providerType);
  if (!byok) return { config, usedByok: false };

  const { apiKeyConfigKey } = getProviderConfigKeys(providerType);
  if (!apiKeyConfigKey) return { config, usedByok: false };

  return {
    config: { ...config, [apiKeyConfigKey]: byok.plaintext },
    usedByok: true,
  };
}

export async function upsertUserApiKey(params: {
  userId: string;
  provider: AiGatewayProvider;
  plaintext: string;
  alias?: string;
  model?: string;
}) {
  const encryptedKey = await encryptApiKey(params.plaintext);
  const keyHash = hashApiKey(params.plaintext);

  return prisma.userAPIKey.upsert({
    where: {
      userId_provider: {
        userId: params.userId,
        provider: params.provider,
      },
    },
    create: {
      userId: params.userId,
      provider: params.provider,
      alias: params.alias ?? '',
      encryptedKey,
      keyHash,
      model: params.model ?? null,
      isActive: true,
    },
    update: {
      alias: params.alias ?? '',
      encryptedKey,
      keyHash,
      model: params.model ?? null,
      isActive: true,
    },
  });
}

export function toPublicApiKey(row: {
  provider: string;
  alias: string;
  model: string | null;
  isActive: boolean;
  lastUsedAt: Date | null;
  createdAt: Date;
  updatedAt: Date;
}) {
  return {
    provider: row.provider,
    alias: row.alias,
    model: row.model,
    isActive: row.isActive,
    lastUsedAt: row.lastUsedAt,
    createdAt: row.createdAt,
    updatedAt: row.updatedAt,
  };
}