freezmo/app/api/generate/route.ts

import { z } from "zod";

import { extractNotes, generateActivities, notesToText, type PipelineStage } from "@/lib/ai";
import { AUTH_ERROR_MESSAGES, authErrorKind, getServerConfig, resolveAuth, resolveModelName } from "@/lib/config";
import { friendlyErrorMessage, OllamaError } from "@/lib/ollama";
import { rejectCrossOrigin, rejectOversizedBody } from "@/lib/http-guard";
import { clientKey, rateLimit, rateLimitResponse } from "@/lib/rate-limit";
import { 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({
  images: z.array(z.string().min(32).max(MAX_IMAGE_CHARS)).max(16).optional(),
  text: z.string().max(120_000).optional(),
  difficulty: difficultySchema.default("medium"),
  count: z.number().int().min(4).max(18).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 };

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 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);

  const auth = resolveAuth(config, request);
  if ("error" in auth) {
    return Response.json(
      { error: AUTH_ERROR_MESSAGES[auth.error], kind: authErrorKind(auth.error) },
      { status: 401 },
    );
  }

  let payload: unknown;
  try {
    payload = await request.json();
  } catch {
    return Response.json({ error: "Request body was not valid JSON." }, { status: 400 });
  }

  const parsed = requestSchema.safeParse(payload);
  if (!parsed.success) {
    return Response.json(
      {
        error: `Invalid request: ${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: "Add at least one photo of your notes or paste some text to build a deck." },
      { status: 400 },
    );
  }

  const requestedCount = count ?? config.defaultActivityCount;
  const encode = encoder();
  const abort = new AbortController();
  // A single hard deadline for the whole pipeline (two sequential model calls).
  const globalTimer = setTimeout(() => abort.abort(), config.timeoutMs * 2);

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

      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: config.visionModel,
            repairModel: config.textModel,
            auth: auth.auth,
            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: config.textModel,
          auth: auth.auth,
          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 ? resolveModelName(config.visionModel, auth.auth.directCloud) : null,
            textModel: resolveModelName(config.textModel, auth.auth.directCloud),
          },
        });
      } catch (error) {
        const message = friendlyErrorMessage(error, auth.auth.baseUrl);
        // 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 {
        clearTimeout(globalTimer);
        try {
          controller.close();
        } catch {
          /* already closed */
        }
      }
    },
    cancel() {
      clearTimeout(globalTimer);
      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;
}