freezmo/app/api/generate/route.ts

import { z } from "zod";

import { activityBatches, extractNotes, generateActivities, notesToText, type PipelineStage } from "@/lib/ai";
import { getServerConfig, resolveModelName } from "@/lib/config";
import { INERT_OLLAMA_AUTH, resolveModelTarget } from "@/lib/providers-server";
import { friendlyErrorMessage, OllamaError } from "@/lib/ollama";
import { rejectCrossOrigin, rejectOversizedBody, requestLocale } from "@/lib/http-guard";
import { clientKey, rateLimit, rateLimitResponse } from "@/lib/rate-limit";
import { acquireModelGate } from "@/lib/model-guard";
import { translate } from "@/lib/i18n-dict";
import { dataUrlSchema, difficultySchema } from "@/lib/schemas";
import { uid } from "@/lib/utils";
import { xpForActivity } from "@/lib/xp";

export const runtime = "nodejs";
export const dynamic = "force-dynamic";
export const maxDuration = 300;

/** ~8 MB per image after client-side resizing; up to 16 images (photos + PDF pages). */
const MAX_IMAGE_CHARS = 8_000_000;

const requestSchema = z.object({

  textModel: z.string().min(1).max(200).optional(),
  visionModel: z.string().min(1).max(200).optional(),
  images: z.array(dataUrlSchema(MAX_IMAGE_CHARS)).max(16).optional(),
  text: z.string().max(120_000).optional(),
  difficulty: difficultySchema.default("medium"),
  /** One deck; a big count is built in several model calls (lib/ai.ts). */
  count: z.number().int().min(4).max(30).optional(),
  focus: z.string().max(300).optional(),
  /** Language the deck must be written in (usually detected by the vision step). */
  language: z.string().max(60).optional(),
  mode: z.enum(["create", "more"]).default("create"),
  existing: z
    .object({
      title: z.string().max(120),
      types: z.array(z.string()).max(30),
      count: z.number().int().min(0).max(100),
    })
    .optional(),
});

type StreamEvent =
  | { type: "stage"; stage: PipelineStage; message: string }
  | { type: "deck"; deck: unknown; meta: Record<string, unknown> }
  | { type: "error"; error: string; kind?: string }
  | { type: "ping" };

function encoder() {
  const enc = new TextEncoder();
  return (event: StreamEvent) => enc.encode(`${JSON.stringify(event)}\n`);
}

/** Generous but finite: ~16 compressed images plus JSON overhead. */
const MAX_BODY_BYTES = 33554432;

export async function POST(request: Request) {
  const locale = requestLocale(request);
  const crossOrigin = rejectCrossOrigin(request);
  if (crossOrigin) return crossOrigin;
  const oversized = rejectOversizedBody(request, MAX_BODY_BYTES);
  if (oversized) return oversized;

  const config = getServerConfig();

  // Cheap requests are rejected before any work happens.
  const limit = rateLimit(`generate:${clientKey(request)}`, config.rateLimits.generate, config.rateLimits.windowMs);
  if (!limit.allowed) return rateLimitResponse(limit, locale);

  // Ollama (server-configured) or a cloud provider the browser sent a key for.
  const resolved = resolveModelTarget(config, request, locale);
  if ("error" in resolved) {
    return Response.json({ error: resolved.error.message, kind: resolved.error.kind }, { status: resolved.error.status });
  }
  const target = resolved.target;
  const provider = target.kind === "provider" ? target.target : undefined;
  const baseAuth = target.kind === "provider" ? INERT_OLLAMA_AUTH : target.auth;

  let payload: unknown;
  try {
    payload = await request.json();
  } catch {
    return Response.json({ error: translate(locale, "server.bodyNotJson") }, { status: 400 });
  }

  const parsed = requestSchema.safeParse(payload);
  if (!parsed.success) {
    return Response.json(
      {
        error: translate(locale, "server.invalidRequest", {
          detail: parsed.error.issues
            .slice(0, 3)
            .map((issue) => `${issue.path.join(".") || "body"} ${issue.message}`)
            .join("; "),
        }),
      },
      { status: 400 },
    );
  }

  const { images, text, difficulty, count, focus, mode, existing, language } = parsed.data;
  const imageList = (images ?? []).map(stripDataUrlPrefix).filter(Boolean);

  if (imageList.length === 0 && !text?.trim()) {
    return Response.json(
      { error: translate(locale, "server.needContent") },
      { status: 400 },
    );
  }

  // Concurrency + the account-wide model budget (lib/model-guard). Taken
  // only now, after every cheap rejection, so a 400 can never leak a slot.
  // Provider runs carry the model the user picked in Settings → Providers;
  // the server's Ollama models stay the default.
  const providerTextModel = provider ? parsed.data.textModel ?? parsed.data.visionModel : undefined;
  const providerVisionModel = provider ? parsed.data.visionModel ?? parsed.data.textModel : undefined;
  if (provider && !providerTextModel) {
    return Response.json(
      { error: translate(locale, "server.providerModelNeeded"), kind: "bad_provider" },
      { status: 400 },
    );
  }
  const textModel = provider ? providerTextModel! : config.textModel;
  const visionModel = provider ? providerVisionModel ?? providerTextModel! : config.visionModel;

  const gate = acquireModelGate(request, config, locale);
  if ("response" in gate) return gate.response;

  const requestedCount = count ?? config.defaultActivityCount;
  const encode = encoder();
  const abort = new AbortController();
  // A single hard deadline for the whole pipeline: the vision step (when there
  // are photos) plus one model call per activity batch.
  const modelCalls = activityBatches(requestedCount).length + 1;
  const globalTimer = setTimeout(() => abort.abort(), config.timeoutMs * modelCalls);
  const onRequestAbort = () => abort.abort();
  if (request.signal.aborted) abort.abort();
  else request.signal.addEventListener("abort", onRequestAbort, { once: true });
  // Heartbeats: one quiet model call can run for minutes, and the client's
  // stall guard needs to tell "slow" from "dead" — a silent stream is how the
  // wizard ended up loading forever. Cleared with the stream either way.
  let heartbeat: ReturnType<typeof setInterval> | null = null;

  const stream = new ReadableStream<Uint8Array>({
    async start(controller) {
      const send = (event: StreamEvent) => {
        try {
          controller.enqueue(encode(event));
        } catch {
          /* client went away */
        }
      };

      heartbeat = setInterval(() => send({ type: "ping" }), 15_000);
      try {
        let sourceText: string;
        let deckTitle: string | undefined;
        let subject: string | undefined;
        let summary: string | undefined;
        let detectedLanguage: string | undefined = language;
        let repaired = false;

        if (imageList.length > 0) {
          send({ type: "stage", stage: "vision", message: "Reading your notes…" });
          const extraction = await extractNotes({
            model: visionModel,
            repairModel: textModel,
            auth: baseAuth,
            provider,
            images: imageList,
            text,
            signal: abort.signal,
            onProgress: (message) => send({ type: "stage", stage: "vision", message }),
          });
          sourceText = notesToText(extraction.notes);
          deckTitle = extraction.notes.title;
          subject = extraction.notes.subject;
          summary = extraction.notes.summary;
          detectedLanguage = extraction.notes.language || language;
          repaired = repaired || extraction.repaired || extraction.languageRepaired;
          send({
            type: "stage",
            stage: "generation",
            message: "Finding the key ideas…",
          });
        } else {
          sourceText = text!.trim();
          send({ type: "stage", stage: "generation", message: "Finding the key ideas…" });
        }

        const generated = await generateActivities({
          model: textModel,
          auth: baseAuth,
          provider,
          notes: sourceText,
          difficulty,
          count: requestedCount,
          language: detectedLanguage,
          existing: mode === "more" ? existing : undefined,
          focus,
          signal: abort.signal,
        });

        const activities = generated.deck.activities.map((activity) => ({
          ...activity,
          id: uid("act"),
          xp: xpForActivity(activity.type, difficulty),
        }));

        send({
          type: "deck",
          deck: {
            title: generated.deck.title || deckTitle || "Your study deck",
            subject: generated.deck.subject ?? subject,
            summary: generated.deck.summary ?? summary,
            language: generated.deck.language || detectedLanguage,
            sourceText,
            difficulty,
            activities,
          },
          meta: {
            requested: generated.requested,
            generated: activities.length,
            repaired,
            language: generated.deck.language || detectedLanguage || null,
            visionModel: imageList.length ? (provider ? visionModel : resolveModelName(config.visionModel, baseAuth.directCloud)) : null,
            textModel: provider ? textModel : resolveModelName(config.textModel, baseAuth.directCloud),
          },
        });
      } catch (error) {
        const message = friendlyErrorMessage(error, baseAuth.baseUrl || provider?.baseUrl || '', locale);
        // Never log the request body: it may contain the user's notes.
        console.error("[api/generate]", error instanceof OllamaError ? `${error.kind}: ${error.message}` : error);
        send({
          type: "error",
          error: message,
          kind: error instanceof OllamaError ? error.kind : "unknown",
        });
      } finally {
        if (heartbeat) clearInterval(heartbeat);
        clearTimeout(globalTimer);
        request.signal.removeEventListener("abort", onRequestAbort);
        gate.gate.release();
        try {
          controller.close();
        } catch {
          /* already closed */
        }
      }
    },
    cancel() {
      if (heartbeat) clearInterval(heartbeat);
      clearTimeout(globalTimer);
      request.signal.removeEventListener("abort", onRequestAbort);
      gate.gate.release();
      abort.abort();
    },
  });

  return new Response(stream, {
    headers: {
      "Content-Type": "application/x-ndjson; charset=utf-8",
      "Cache-Control": "no-store, no-transform",
      Connection: "keep-alive",
    },
  });
}

function stripDataUrlPrefix(value: string): string {
  const commaIndex = value.indexOf(",");
  if (value.startsWith("data:") && commaIndex !== -1) return value.slice(commaIndex + 1);
  return value;
}