freezmo/app/api/homework/route.ts

import { z } from "zod";

import { buildHomeworkGuide, checkHomework, extractNotes, makeHomework, notesToText, type PipelineStage } from "@/lib/ai";
import { authErrorMessage, authErrorKind, getServerConfig, resolveAuth, resolveModelName } from "@/lib/config";
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, homeworkModeSchema, type HomeworkStep } from "@/lib/schemas";
import { uid } from "@/lib/utils";

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;
/** Generous but finite: ~16 compressed images plus JSON overhead. */
const MAX_BODY_BYTES = 33554432;

const requestSchema = z.object({
  mode: homeworkModeSchema,
  images: z.array(dataUrlSchema(MAX_IMAGE_CHARS)).max(16).optional(),
  text: z.string().max(120_000).optional(),
  /** Mode "check": the student's own attempt at the assignment. */
  solution: z.string().max(120_000).optional(),
  /** Mode "check": photos of the attempt (handwritten work, scanned pages). */
  solutionImages: z.array(dataUrlSchema(MAX_IMAGE_CHARS)).max(16).optional(),
  /** Language of the assignment (usually detected by the vision step). */
  language: z.string().max(60).optional(),
});

type StreamEvent =
  | { type: "stage"; stage: PipelineStage; message: string }
  | { type: "homework"; homework: 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`);
}

/** Same pipeline shape as /api/generate: vision (photos only) then one model call. */
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(`homework:${clientKey(request)}`, config.rateLimits.homework, config.rateLimits.windowMs);
  if (!limit.allowed) return rateLimitResponse(limit, locale);

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

  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 { mode, images, text, solution, solutionImages, language } = parsed.data;
  const imageList = (images ?? []).map(stripDataUrlPrefix).filter(Boolean);
  const solutionImageList = (solutionImages ?? []).map(stripDataUrlPrefix).filter(Boolean);

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

  // Concurrency + the account-wide model budget (lib/model-guard), taken
  // after every cheap rejection so a 400 can never leak a slot.
  const gate = acquireModelGate(request, config, locale);
  if ("response" in gate) return gate.response;

  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);
  // Heartbeats so the client's stall guard can tell "slow" from "dead".
  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 assignmentText: string;
        let detectedLanguage: string | undefined = language;
        let repaired = false;

        if (imageList.length > 0) {
          send({ type: "stage", stage: "vision", message: translate(locale, "homework.loading.reading") });
          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 }),
          });
          assignmentText = notesToText(extraction.notes);
          detectedLanguage = extraction.notes.language || language;
          repaired = extraction.repaired || extraction.languageRepaired;
        } else {
          assignmentText = text!.trim();
        }

        send({
          type: "stage",
          stage: "generation",
          message: translate(
            locale,
            mode === "make"
              ? "homework.loading.writing"
              : mode === "check"
                ? "homework.loading.checking"
                : "homework.loading.planning",
          ),
        });

        const meta = {
          repaired,
          language: detectedLanguage || null,
          visionModel: imageList.length ? resolveModelName(config.visionModel, auth.auth.directCloud) : null,
          textModel: resolveModelName(config.textModel, auth.auth.directCloud),
        };

        if (mode === "check") {
          /* Photos of the attempt (handwritten working, scanned pages) are
             transcribed first; whatever was typed stays as it is. */
          let solutionText = solution?.trim() ?? "";
          if (solutionImageList.length > 0) {
            send({
              type: "stage",
              stage: "vision",
              message: translate(locale, "homework.loading.readingSolution"),
            });
            const transcription = await extractNotes({
              model: config.visionModel,
              repairModel: config.textModel,
              auth: auth.auth,
              images: solutionImageList,
              signal: abort.signal,
              onProgress: (message) => send({ type: "stage", stage: "vision", message }),
            });
            solutionText = [solutionText, notesToText(transcription.notes)].filter(Boolean).join("\n\n");
          }

          const result = await checkHomework({
            model: config.textModel,
            auth: auth.auth,
            assignment: assignmentText,
            solution: solutionText,
            language: detectedLanguage,
            signal: abort.signal,
          });
          send({
            type: "homework",
            homework: {
              mode: "check",
              title: result.title,
              language: result.language || detectedLanguage,
              verdict: result.verdict,
              correct: result.correct,
              wrong: result.wrong,
              corrected: result.corrected,
              sourceText: assignmentText,
            },
            /* The full attempt (typed + transcribed) travels back so the client
               can store it: "Redo" then re-checks the same work without the
               photos — they are never stored. */
            meta: { ...meta, solutionText },
          });
        } else if (mode === "make") {
          const result = await makeHomework({
            model: config.textModel,
            auth: auth.auth,
            assignment: assignmentText,
            language: detectedLanguage,
            signal: abort.signal,
          });
          send({
            type: "homework",
            homework: {
              mode: "make",
              title: result.title,
              language: result.language || detectedLanguage,
              answer: result.answer,
              // The assignment as text (the vision transcription when photos
              // were used) so the client can store it for “Make it again”.
              sourceText: assignmentText,
            },
            meta,
          });
        } else {
          const guide = await buildHomeworkGuide({
            model: config.textModel,
            auth: auth.auth,
            assignment: assignmentText,
            language: detectedLanguage,
            signal: abort.signal,
          });
          const steps: HomeworkStep[] = guide.steps.map((step) => ({
            ...step,
            // xp: 0 keeps the shared activity components from rendering their
            // "+XP" badge — the guide deliberately has no XP or streak.
            checkpoint: step.checkpoint ? { ...step.checkpoint, id: uid("act"), xp: 0 } : undefined,
          }));
          send({
            type: "homework",
            homework: {
              mode: "guide",
              title: guide.title,
              language: guide.language || detectedLanguage,
              steps,
              sourceText: assignmentText,
            },
            meta,
          });
        }
      } catch (error) {
        const message = friendlyErrorMessage(error, auth.auth.baseUrl, locale);
        // Never log the request body: it may contain the student's assignment.
        console.error("[api/homework]", 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);
        gate.gate.release();
        try {
          controller.close();
        } catch {
          /* already closed */
        }
      }
    },
    cancel() {
      if (heartbeat) clearInterval(heartbeat);
      clearTimeout(globalTimer);
      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;
}