freezmo/app/api/chat/route.ts

import { z } from "zod";

import { authErrorMessage, authErrorKind, getServerConfig, resolveAuth, resolveModelName } from "@/lib/config";
import { friendlyErrorMessage, ollamaChatStream, 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 { buildChatSystemPrompt } 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(),
});

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

  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 { messages, context } = parsed.data;

  // The model must come from the server's allowlist — never from the client.
  const allowed = config.chatModels.length > 0 ? config.chatModels : [config.textModel];
  const model = parsed.data.model ?? allowed[0];
  if (!allowed.some((entry) => entry.toLowerCase() === model.toLowerCase())) {
    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 system = buildChatSystemPrompt({
    context: keptContext.map((chunk) => ({
      deckTitle: chunk.deckTitle,
      section: chunk.section,
      text: chunk.text,
    })),
  });

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

      try {
        const result = await ollamaChatStream({
          model,
          auth: auth.auth,
          messages: [{ role: "system", content: system }, ...history],
          signal: abort.signal,
          onToken: (text) => send({ type: "token", text }),
          onThinking: (text) => send({ type: "thinking", text }),
        });
        send({ type: "done", model: resolveModelName(result.model, auth.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, auth.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);
        gate.gate.release();
        try {
          controller.close();
        } catch {
          /* already closed */
        }
      }
    },
    cancel() {
      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",
    },
  });
}