proof-of-work/proof-of-work/app/api/ai/test/route.ts

import { NextRequest, NextResponse } from 'next/server';
import { z } from 'zod';
import { getCurrentUser } from '@/lib/auth';
import { prisma } from '@/lib/prisma';
import { getProvider, isCustomUrlProvider } from '@/lib/ai/providers';

/**
 * POST /api/ai/test
 *
 * Body (optional):
 *   {
 *     // If supplied: test this draft config without saving it.
 *     // Otherwise: test the actor's currently active config.
 *     provider?: string,
 *     model?: string,
 *     baseUrl?: string,
 *     apiKey?: string,
 *     // If supplied + apiKey is null: pull the saved key for that
 *     // profile (so the UI can test a saved profile by id without
 *     // forcing the user to re-type the key).
 *     useSavedKeyForId?: string,
 *   }
 *
 * Sends a tiny "say hi in 3 words" prompt. Reports latency, sample
 * reply (or finishReason if Gemini blocks it).
 *
 * Times out after 30s — long enough for cold Ollama starts.
 */

const TEST_TIMEOUT_MS = 30_000;

const bodySchema = z.object({
  provider: z.string().optional(),
  model: z.string().optional(),
  baseUrl: z.string().nullable().optional(),
  apiKey: z.string().nullable().optional(),
  useSavedKeyForId: z.string().optional(),
});

export async function POST(request: NextRequest) {
  const user = await getCurrentUser();
  if (!user) {
    return NextResponse.json({ ok: false, error: 'Unauthorized' }, { status: 401 });
  }

  const raw = await request.json().catch(() => ({}));
  const parsed = bodySchema.safeParse(raw);
  if (!parsed.success) {
    return NextResponse.json(
      { ok: false, error: 'Invalid body' },
      { status: 400 },
    );
  }
  const draft = parsed.data;

  // Resolve the config to test:
  //   1. If draft.provider is set → use the draft fields (testing
  //      a not-yet-saved config in the UI).
  //   2. Else if draft.useSavedKeyForId is set → load that profile.
  //   3. Else → use the active config (legacy single-config columns).
  let provider: string | null;
  let model: string | null;
  let baseUrl: string | null;
  let apiKey: string | null;

  if (draft.provider) {
    provider = draft.provider;
    model = draft.model ?? null;
    baseUrl = draft.baseUrl ?? null;
    apiKey = draft.apiKey ?? null;
    // Allow the UI to fill in just provider+model+baseUrl and have
    // us pull the saved key by profile id (so the user doesn't have
    // to retype it just to retest).
    if (draft.useSavedKeyForId && (apiKey == null || apiKey === '')) {
      const saved = await prisma.aIConfigProfile.findFirst({
        where: { id: draft.useSavedKeyForId, userId: user.id },
        select: { apiKey: true },
      });
      if (saved?.apiKey) apiKey = saved.apiKey;
    }
  } else if (draft.useSavedKeyForId) {
    const saved = await prisma.aIConfigProfile.findFirst({
      where: { id: draft.useSavedKeyForId, userId: user.id },
    });
    if (!saved) {
      return NextResponse.json(
        { ok: false, error: 'Config not found.' },
        { status: 404 },
      );
    }
    provider = saved.provider;
    model = saved.model;
    baseUrl = saved.baseUrl;
    apiKey = saved.apiKey;
  } else {
    const prefs = await prisma.userPreferences.findUnique({
      where: { userId: user.id },
      select: { aiProvider: true, aiModel: true, aiBaseUrl: true, aiApiKey: true },
    });
    provider = prefs?.aiProvider ?? null;
    model = prefs?.aiModel ?? null;
    baseUrl = prefs?.aiBaseUrl ?? null;
    apiKey = prefs?.aiApiKey ?? null;
  }

  if (!provider || !model) {
    return NextResponse.json(
      {
        ok: false,
        error: 'Pick a provider + model first.',
      },
      { status: 400 },
    );
  }

  // Testing an arbitrary base URL is the same SSRF surface as configuring
  // one — admin-only. Non-admins may only test fixed-URL cloud providers.
  if (!user.isAdmin && (baseUrl || isCustomUrlProvider(provider))) {
    return NextResponse.json(
      {
        ok: false,
        error:
          'Only an admin can test providers with a custom base URL (Ollama / OpenAI-compatible).',
      },
      { status: 403 },
    );
  }
  const providerImpl = getProvider(provider);
  if (!providerImpl) {
    return NextResponse.json(
      { ok: false, error: `Unknown provider: ${provider}` },
      { status: 400 },
    );
  }

  const controller = new AbortController();
  const timer = setTimeout(() => controller.abort(), TEST_TIMEOUT_MS);
  const t0 = Date.now();

  let sample = '';
  let tokensIn: number | undefined;
  let tokensOut: number | undefined;
  let providerError: string | null = null;

  try {
    for await (const chunk of providerImpl.generate({
      apiKey,
      baseUrl,
      model,
      systemPrompt:
        'You are a connectivity test. Reply with EXACTLY three words: "Hello there friend." Nothing else.',
      userPrompt: 'Say hi.',
      signal: controller.signal,
      // Generous output budget so thinking models (Gemini 2.5/3.x,
      // OpenAI o-series) actually have room to emit visible text after
      // their internal reasoning. Cheap because the prompt is tiny.
      maxOutputTokens: 4096,
    })) {
      if (chunk.type === 'text') sample += chunk.delta;
      else if (chunk.type === 'usage') {
        tokensIn = chunk.tokensIn;
        tokensOut = chunk.tokensOut;
      } else if (chunk.type === 'error') {
        providerError = chunk.message;
      }
    }
  } catch (e) {
    providerError =
      controller.signal.aborted
        ? `Timed out after ${Math.round(TEST_TIMEOUT_MS / 1000)}s`
        : (e as Error).message;
  } finally {
    clearTimeout(timer);
  }

  const ms = Date.now() - t0;

  if (providerError) {
    return NextResponse.json({ ok: false, error: providerError, ms }, { status: 200 });
  }
  if (!sample.trim()) {
    return NextResponse.json(
      {
        ok: false,
        error:
          'Empty reply. The provider returned a response with no text. ' +
          'For Gemini this often means a safety filter blocked the output ' +
          '(check the model name + try a flagship model). For thinking ' +
          'models the answer may have been spent on internal reasoning — ' +
          'try a non-thinking model.',
        ms,
      },
      { status: 200 },
    );
  }
  return NextResponse.json({
    ok: true,
    sample: sample.trim().slice(0, 200),
    tokensIn,
    tokensOut,
    ms,
  });
}