freezmo/app/api/chat/route.ts

import { z } from "zod";

import { canonicalModelName, getServerConfig, resolveModelName } from "@/lib/config";
import { friendlyErrorMessage, ollamaChatStream, ollamaListModels, OllamaError } from "@/lib/ollama";
import { providerChatStream } from "@/lib/openai-compat";
import { resolveModelTarget } from "@/lib/providers-server";
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 { buildChatSystemPrompt, buildTeachSummarySystemPrompt, buildTeachSystemPrompt } from "@/lib/prompts";
import { estimateTokens, trimHistory } from "@/lib/chat-trim";

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

/** Context chunks (~7k tokens) plus a trimmed window plus JSON overhead. */
const MAX_BODY_BYTES = 1_000_000;

/** Hard cap on the attached context, whatever the client sends. */
const CONTEXT_TOKEN_BUDGET = 9000;

const requestSchema = z.object({
  model: z.string().min(1).max(120).optional(),
  messages: z
    .array(
      z.object({
        role: z.enum(["user", "assistant"]),
        content: z.string().min(1).max(20_000),
      }),
    )
    .min(1)
    .max(60),
  context: z
    .array(
      z.object({
        deckId: z.string().max(120),
        deckTitle: z.string().max(200),
        section: z.string().max(200),
        text: z.string().min(1).max(12_000),
      }),
    )
    .max(40)
    .optional(),
  /** Teach it back: "teach" is one student turn, "teach_summary" the recap. */
  mode: z.enum(["teach", "teach_summary"]).optional(),
  teach: z
    .object({
      topic: z.string().min(1).max(200),
      difficulty: z.enum(["friendly", "skeptical"]),
      exchange: z.number().int().min(0).max(60),
      limit: z.number().int().min(1).max(60),
      cards: z.array(z.string().max(1000)).max(80),
    })
    .optional(),
});

type StreamEvent =
  | { type: "token"; text: string }
  | { type: "thinking"; text: string }
  | { type: "done"; model: string }
  | { type: "error"; error: string; kind?: string };

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

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

  const limit = rateLimit(`chat:${clientKey(request)}`, config.rateLimits.chat, 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;

  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 { messages, context, mode, teach } = parsed.data;
  if (mode && !teach) {
    return Response.json(
      { error: translate(locale, "server.invalidRequest", { detail: "teach mode needs the teach payload" }) },
      { status: 400 },
    );
  }

  // The model must be one the server actually offers — the configured
  // favourites, or anything the user's Ollama lists. The picker shows exactly
  // that list, so nothing selectable is ever rejected; an unknown model still
  // gets the 400.
  const allowed = config.chatModels.length > 0 ? config.chatModels : [config.textModel];
  const model = parsed.data.model ?? allowed[0];
  const wanted = canonicalModelName(model);
  const inFavourites = allowed.some((entry) => canonicalModelName(entry) === wanted);
  if (target.kind === "provider") {
    // The provider's own catalogue is the authority (the picker lists exactly
    // what /models returned); an unknown name comes back as its 404.
  } else if (!inFavourites) {
    let listed = false;
    try {
      const models = await ollamaListModels(target.auth, config.healthTimeoutMs);
      listed = models.some((entry) => canonicalModelName(entry.name) === wanted);
    } catch {
      listed = false;
    }
    if (!listed) {
      return Response.json(
        { error: translate(locale, "server.chatModelNotAllowed"), kind: "bad_model" },
        { status: 400 },
      );
    }
  }

  const last = messages[messages.length - 1];
  if (last.role !== "user") {
    return Response.json({ error: translate(locale, "server.chatNeedsQuestion") }, { status: 400 });
  }

  // Cap the context, then trim the conversation window.
  const keptContext: typeof context = [];
  let contextTokens = 0;
  for (const chunk of context ?? []) {
    const cost = estimateTokens(chunk.text);
    if (contextTokens + cost > CONTEXT_TOKEN_BUDGET) break;
    keptContext.push(chunk);
    contextTokens += cost;
  }
  const history = trimHistory(messages);

  const notesContext = keptContext.map((chunk) => ({
    deckTitle: chunk.deckTitle,
    section: chunk.section,
    text: chunk.text,
  }));
  const system =
    mode === "teach" && teach
      ? buildTeachSystemPrompt({
          topic: teach.topic,
          difficulty: teach.difficulty,
          cards: teach.cards,
          context: notesContext,
          exchange: teach.exchange,
          limit: teach.limit,
        })
      : mode === "teach_summary" && teach
        ? buildTeachSummarySystemPrompt({
            topic: teach.topic,
            difficulty: teach.difficulty,
            cards: teach.cards,
            context: notesContext,
          })
        : buildChatSystemPrompt({ context: notesContext });

  // 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();
  const globalTimer = setTimeout(() => abort.abort(), config.timeoutMs);
  const onRequestAbort = () => abort.abort();
  if (request.signal.aborted) abort.abort();
  else request.signal.addEventListener("abort", onRequestAbort, { once: true });

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

      try {
        const messages = [{ role: "system" as const, content: system }, ...history];
        const result =
          target.kind === "provider"
            ? await providerChatStream({
                target: target.target,
                model,
                messages,
                signal: abort.signal,
                onToken: (text) => send({ type: "token", text }),
                onThinking: (text) => send({ type: "thinking", text }),
              })
            : await ollamaChatStream({
                model,
                auth: target.auth,
                messages,
                signal: abort.signal,
                onToken: (text) => send({ type: "token", text }),
                onThinking: (text) => send({ type: "thinking", text }),
              });
        send({
          type: "done",
          model: target.kind === "provider" ? result.model : resolveModelName(result.model, target.auth.directCloud),
        });
      } catch (error) {
        // A client-side stop aborts the upstream call: nothing to report.
        if (abort.signal.aborted && request.signal.aborted) {
          /* stopped by the user */
        } else {
          const message = friendlyErrorMessage(error, target.kind === "provider" ? target.target.baseUrl : target.auth.baseUrl, locale);
          console.error("[api/chat]", error instanceof OllamaError ? `${error.kind}: ${error.message}` : error);
          send({
            type: "error",
            error: message,
            kind: error instanceof OllamaError ? error.kind : "unknown",
          });
        }
      } finally {
        clearTimeout(globalTimer);
        request.signal.removeEventListener("abort", onRequestAbort);
        gate.gate.release();
        try {
          controller.close();
        } catch {
          /* already closed */
        }
      }
    },
    cancel() {
      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",
    },
  });
}