Momento/memento-note/app/api/user/api-keys/route.ts

import { NextRequest, NextResponse } from 'next/server';
import { z } from 'zod';
import { auth } from '@/auth';
import { getEffectiveTier } from '@/lib/entitlements';
import {
  getAllowedByokProviders,
  isByokProviderAllowed,
} from '@/lib/byok';
import {
  upsertUserApiKey,
  toPublicApiKey,
} from '@/lib/byok';
import { validateProviderApiKey } from '@/lib/byok/validate-key';
import { VALID_PROVIDERS, type AiGatewayProvider } from '@/lib/ai/router';
import { prisma } from '@/lib/prisma';

const createSchema = z.object({
  provider: z.string().min(1),
  apiKey: z.string().min(8),
  alias: z.string().max(120).optional(),
  model: z.string().max(120).optional(),
  baseUrl: z.string().url().optional(),
});

export async function GET() {
  const session = await auth();
  if (!session?.user?.id) {
    return NextResponse.json({ error: 'Unauthorized' }, { status: 401 });
  }

  const keys = await prisma.userAPIKey.findMany({
    where: { userId: session.user.id },
    orderBy: { provider: 'asc' },
  });

  return NextResponse.json({
    keys: keys.map(toPublicApiKey),
    allowedProviders: getAllowedByokProviders(await getEffectiveTier(session.user.id)),
  });
}

export async function POST(req: NextRequest) {
  const session = await auth();
  if (!session?.user?.id) {
    return NextResponse.json({ error: 'Unauthorized' }, { status: 401 });
  }

  const tier = await getEffectiveTier(session.user.id);
  if (tier === 'BASIC') {
    return NextResponse.json(
      { error: 'TIER_LIMITED', message: 'BYOK requires a Pro plan or higher' },
      { status: 403 },
    );
  }

  let body: z.infer<typeof createSchema>;
  try {
    body = createSchema.parse(await req.json());
  } catch {
    return NextResponse.json({ error: 'Invalid request body' }, { status: 400 });
  }

  const provider = body.provider as AiGatewayProvider;
  if (!VALID_PROVIDERS.has(provider)) {
    return NextResponse.json({ error: 'Unknown provider' }, { status: 400 });
  }

  if (!isByokProviderAllowed(tier, provider)) {
    return NextResponse.json(
      { error: 'TIER_LIMITED', message: `Provider "${provider}" is not available on your plan` },
      { status: 403 },
    );
  }

  const effectiveBaseUrl = provider === 'custom' ? body.baseUrl : undefined;
  if (provider !== 'custom' && body.baseUrl) {
    return NextResponse.json(
      { error: 'INVALID_REQUEST', message: 'baseUrl is only allowed for custom providers' },
      { status: 400 },
    );
  }

  try {
    await validateProviderApiKey(provider, body.apiKey, effectiveBaseUrl);
  } catch (err) {
    const message = err instanceof Error ? err.message : 'Invalid API key';
    return NextResponse.json({ error: 'INVALID_API_KEY', message }, { status: 400 });
  }

  const row = await upsertUserApiKey({
    userId: session.user.id,
    provider,
    plaintext: body.apiKey,
    alias: body.alias,
    model: body.model,
  });

  return NextResponse.json({ key: toPublicApiKey(row) });
}