STCourses/packages/feature-ai/test/feature.test.ts

import { expect } from "bun:test";
import { cpSync, mkdtempSync, readFileSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join, resolve } from "node:path";
import type { AppConfig, AuthUser, JobRecord, ServerContext } from "@stc/contracts";
import { createApp, createJobQueue, resolveConfig } from "@stc/core";
import { createCourseRepository } from "@stc/course-format";
import type { Db } from "@stc/db";
import { forEachDb } from "../../db/test/helpers";
import { aiFeature, type AiFeatureOptions } from "../src/feature";
import { fakeProvider } from "./helpers";

const SAMPLE = resolve(import.meta.dir, "../../../fixtures/courses/sample");
const USERS: Record<string, AuthUser> = {
  owner: { id: "owner", email: "o@x", name: "Owner", roles: ["owner"], groups: [] },
  ann: { id: "ann", email: "a@x", name: "Ann", roles: ["learner"], groups: [] },
  bob: { id: "bob", email: "b@x", name: "Bob", roles: ["learner"], groups: [] },
  auth: { id: "auth", email: "w@x", name: "Writer", roles: ["author"], groups: [] },
};

async function setup(db: Db, turns: Parameters<typeof fakeProvider>[0] = [() => ({ text: "ok" })], extra: Partial<AiFeatureOptions> = {}, configOverrides: Partial<AppConfig> = {}) {
  const dir = mkdtempSync(join(tmpdir(), "stc-ai-feat-"));
  cpSync(SAMPLE, join(dir, "courses/sample"), { recursive: true });
  cpSync(SAMPLE, join(dir, "courses/secret"), { recursive: true });
  const m = readFileSync(join(SAMPLE, "stcourse.yaml"), "utf8");
  writeFileSync(join(dir, "courses/secret/stcourse.yaml"), m.replace("id: sample", "id: secret").replace("visibility: public", "visibility: restricted"));
  const courses = await createCourseRepository([join(dir, "courses")]);
  const fake = fakeProvider(turns);
  const keys: (string | null)[] = [];
  const ctx = {
    db, courses,
    config: resolveConfig({ port: 0, databaseUrl: "", dataDir: join(dir, "data"), courseDirs: [], secret: "test-secret", publicUrl: "http://x", ...configOverrides }),
    auth: { userFromRequest: async (req: Request) => USERS[req.headers.get("x-user") ?? ""] ?? null },
  } as unknown as ServerContext;
  const feature = aiFeature(ctx, { providerFactory: (_cfg, key) => { keys.push(key); return fake.provider; }, ...extra });
  const jobs = createJobQueue(db, feature.jobs!);
  ctx.jobs = jobs;
  const app = createApp(ctx, [feature]);
  const call = async (method: string, path: string, user?: string, body?: unknown) => {
    const headers: Record<string, string> = { origin: "http://x" };
    if (user) headers["x-user"] = user;
    if (body !== undefined) headers["content-type"] = "application/json";
    return app.fetch(new Request(`http://x${path}`, { method, headers, body: body === undefined ? undefined : JSON.stringify(body) }));
  };
  return { ctx, dir, call, fake, keys, feature };
}

const ndjson = async (r: Response) => (await r.text()).trim().split("\n").map((l) => JSON.parse(l));

forEachDb("AI is off by default and admin config is owner-only and hides keys", async (db) => {
  const { call } = await setup(db);
  expect(await (await call("GET", "/api/ai/status?courseId=sample", "ann")).json()).toMatchObject({ enabled: false, source: null });
  expect((await call("POST", "/api/ai/courses/sample/chat", "ann", { message: "hi" })).status).toBe(409);
  expect((await call("GET", "/api/ai/admin/config", "ann")).status).toBe(403);
  expect((await call("GET", "/api/ai/admin/config")).status).toBe(401);
  const put = await call("PUT", "/api/ai/admin/config", "owner", { provider: { kind: "anthropic", apiKey: "sk-ant-SERVER" }, mcpServers: [] });
  expect(put.status).toBe(200);
  const cfg = await (await call("GET", "/api/ai/admin/config", "owner")).text();
  expect(cfg).not.toContain("SERVER");
  expect(JSON.parse(cfg).provider).toEqual({ kind: "anthropic", hasApiKey: true });
  expect((await call("PUT", "/api/ai/admin/config", "owner", { provider: { kind: "nope" }, mcpServers: [] })).status).toBe(400);
  expect(await (await call("GET", "/api/ai/status?courseId=sample", "ann")).json()).toMatchObject({ enabled: true, source: "server", mentor: { available: true, on: true } });
});

forEachDb("BYOK enables AI for that user only and is used before the server key", async (db) => {
  const { call, keys } = await setup(db);
  expect((await call("PUT", "/api/ai/account/key", "ann", { kind: "anthropic", apiKey: "sk-ann" })).status).toBe(200);
  const acct = await (await call("GET", "/api/ai/account", "ann")).text();
  expect(acct).not.toContain("sk-ann");
  expect(JSON.parse(acct).byok).toEqual({ kind: "anthropic", hasApiKey: true });
  expect(await (await call("GET", "/api/ai/status?courseId=sample", "ann")).json()).toMatchObject({ enabled: true, source: "byok" });
  expect(await (await call("GET", "/api/ai/status?courseId=sample", "bob")).json()).toMatchObject({ enabled: false });
  await call("POST", "/api/ai/courses/sample/chat", "ann", { message: "hello" }).then((r) => r.text());
  expect(keys).toEqual(["sk-ann"]);
  expect((await call("DELETE", "/api/ai/account/key", "ann")).status).toBe(200);
  expect(await (await call("GET", "/api/ai/status?courseId=sample", "ann")).json()).toMatchObject({ enabled: false });
  expect((await call("PUT", "/api/ai/account/key", "ann", { kind: "anthropic" })).status).toBe(400);
});

forEachDb("chat cites course files, applies mentor rules, and stores the conversation", async (db) => {
  const { call, fake } = await setup(db, [
    () => ({ toolCalls: [{ id: "t1", name: "read_lesson", input: { lessonId: "m2-ex-add" } }] }),
    () => ({ text: "Which operator combines two numbers? See [lessons/04-exercise.md · Adding numbers]." }),
  ]);
  await call("PUT", "/api/ai/admin/config", "owner", { provider: { kind: "anthropic", apiKey: "k" }, mcpServers: [] });
  const r = await call("POST", "/api/ai/courses/sample/chat", "ann", { message: "How do I return the sum of its arguments?" });
  expect(r.status).toBe(200);
  expect(r.headers.get("content-type")).toContain("application/x-ndjson");
  const events = await ndjson(r);
  expect(events[0]).toMatchObject({ type: "meta", mentorOn: true });
  expect(events[0].citations.map((c: { path: string }) => c.path)).toContain("lessons/04-exercise.md");
  expect(events.some((e) => e.type === "tool" && e.name === "read_lesson")).toBe(true);
  const done = events.find((e) => e.type === "done");
  expect(done.text).toContain("Which operator");
  expect(done.citations.some((c: { path: string; lessonId?: string }) => c.path === "lessons/04-exercise.md")).toBe(true);
  const system = fake.seen[0]!.messages[0]!;
  expect(system.role).toBe("system");
  expect(system.content).toContain("Give hints, never full solutions.");
  expect(fake.seen[0]!.tools!.map((t) => t.name)).toEqual(expect.arrayContaining(["search_course", "read_lesson", "list_notes", "my_progress"]));

  const convs = await (await call("GET", "/api/ai/courses/sample/conversations", "ann")).json();
  expect(convs).toHaveLength(1);
  const conv = await (await call("GET", `/api/ai/conversations/${convs[0].id}`, "ann")).json();
  expect(conv.messages.map((m: { role: string }) => m.role)).toEqual(["user", "assistant"]);
  expect(conv.messages[1].citations.some((c: { path: string }) => c.path === "lessons/04-exercise.md")).toBe(true);
  expect((await call("GET", `/api/ai/conversations/${convs[0].id}`, "bob")).status).toBe(404);

  // Learner switches mentor mode off for this course: rules leave the prompt.
  expect((await call("PUT", "/api/ai/courses/sample/mentor", "ann", { off: true })).status).toBe(200);
  expect(await (await call("GET", "/api/ai/status?courseId=sample", "ann")).json()).toMatchObject({ mentor: { available: true, on: false } });
  const r2 = await ndjson(await call("POST", "/api/ai/courses/sample/chat", "ann", { message: "Show me the full solution", conversationId: convs[0].id }));
  expect(r2[0]).toMatchObject({ type: "meta", mentorOn: false, conversationId: convs[0].id });
  const sys2 = fake.seen.at(-1)!.messages[0]!.content;
  expect(sys2).not.toContain("never full solutions");
  expect(sys2).toContain("Mentor mode is off");
  // history is replayed
  expect(fake.seen.at(-1)!.messages.some((m) => m.role === "assistant" && m.content.includes("Which operator"))).toBe(true);
  // Bob still has mentor mode on
  expect(await (await call("GET", "/api/ai/status?courseId=sample", "bob")).json()).toMatchObject({ mentor: { on: true } });

  expect((await call("DELETE", `/api/ai/conversations/${convs[0].id}`, "ann")).status).toBe(200);
  expect(await (await call("GET", "/api/ai/courses/sample/conversations", "ann")).json()).toEqual([]);
});

forEachDb("course visibility and sign-in are enforced", async (db) => {
  const { call } = await setup(db);
  await call("PUT", "/api/ai/admin/config", "owner", { provider: { kind: "anthropic", apiKey: "k" }, mcpServers: [] });
  expect((await call("POST", "/api/ai/courses/secret/chat", "ann", { message: "hi" })).status).toBe(404);
  expect((await call("POST", "/api/ai/courses/sample/chat", undefined, { message: "hi" })).status).toBe(401);
  expect((await call("POST", "/api/ai/courses/nope/chat", "ann", { message: "hi" })).status).toBe(404);
  expect((await call("POST", "/api/ai/courses/secret/chat", "owner", { message: "hi" })).status).toBe(200);
});

forEachDb("generation is for authors and writes drafts, not course files", async (db) => {
  const doc = "=== FILE: questions/adding.md ===\n## Check your understanding\n1. Predict: add(2, 3)?\n=== END FILE ===";
  const { call, ctx, dir } = await setup(db, [() => ({ text: doc })]);
  await call("PUT", "/api/ai/admin/config", "owner", { provider: { kind: "anthropic", apiKey: "k" }, mcpServers: [] });
  const body = { courseId: "sample", kind: "questions", style: "pure-snake", topic: "Adding", sources: [{ kind: "lesson", lessonId: "m2-text" }] };
  expect((await call("POST", "/api/ai/studio/generate", "ann", body)).status).toBe(403);
  expect((await call("GET", "/api/ai/studio/sources?courseId=sample", "ann")).status).toBe(403);
  const src = await (await call("GET", "/api/ai/studio/sources?courseId=sample", "auth")).json();
  expect(src.lessons.map((l: { id: string }) => l.id)).toContain("m2-text");
  expect(src.files).toContain("code/S01 - First steps/hello.sh");
  expect((await call("POST", "/api/ai/studio/generate", "auth", { ...body, kind: "poem" })).status).toBe(400);
  const g = await call("POST", "/api/ai/studio/generate", "auth", body);
  expect(g.status).toBe(202);
  const { jobId } = await g.json();
  // run the queued job by hand (the worker is not started in tests)
  const job = (await ctx.jobs.get(jobId)) as JobRecord;
  expect(job.kind).toBe("ai.generate");
  const feature = aiFeature(ctx, { providerFactory: () => fakeProvider([() => ({ text: doc })]).provider });
  const result = (await feature.jobs!["ai.generate"]!(job.payload, { reportProgress: async () => {}, signal: new AbortController().signal })) as { location: string; paths: string[] };
  expect(result.paths).toEqual(["questions/adding.md"]);
  expect(result.location.startsWith(join(dir, "data", "drafts", "sample"))).toBe(true);
  expect(await Bun.file(join(dir, "courses/sample/questions/adding.md")).exists()).toBe(false);
  expect((await call("GET", `/api/ai/studio/jobs/${jobId}`, "auth")).status).toBe(200);
  expect((await call("GET", `/api/ai/studio/jobs/${jobId}`, "bob")).status).toBe(404);
});

// SSRF guard: a learner's own key (BYOK) may only go to the official endpoints or admin-allowed origins.
forEachDb("BYOK base URLs are limited to official endpoints and allowed origins", async (db) => {
  const { call } = await setup(db, undefined, {}, { aiAllowedBaseUrls: ["http://ollama.internal:11434", "https://llm.example.org"] });
  const put = (body: unknown) => call("PUT", "/api/ai/account/key", "ann", body);
  // Internal or arbitrary hosts are refused for either kind.
  for (const baseUrl of ["http://127.0.0.1:8080/v1", "http://169.254.169.254/latest", "http://localhost:11434/v1", "https://api.anthropic.com.evil.test", "http://ollama.internal:11435/v1", "https://ollama.internal:11434/v1", "file:///etc/passwd"]) {
    expect((await put({ kind: "openai", model: "m", apiKey: "k", baseUrl })).status).toBe(400);
  }
  expect((await put({ kind: "anthropic", apiKey: "k", baseUrl: "http://127.0.0.1:9/" })).status).toBe(400);
  expect((await call("GET", "/api/ai/account", "ann")).status).toBe(200);
  expect((await (await call("GET", "/api/ai/account", "ann")).json()).byok).toBeNull();
  // Official endpoints and allowlisted origins (any path on them) are accepted.
  expect((await put({ kind: "anthropic", apiKey: "k" })).status).toBe(200);
  expect((await put({ kind: "anthropic", apiKey: "k", baseUrl: "https://api.anthropic.com/" })).status).toBe(200);
  expect((await put({ kind: "openai", model: "m", apiKey: "k" })).status).toBe(200);
  expect(await (await call("GET", "/api/ai/account", "ann")).json()).toMatchObject({ byok: { kind: "openai", model: "m", hasApiKey: true } }); // official endpoint: stored as "no base URL"
  expect((await put({ kind: "openai", model: "m", apiKey: "k", baseUrl: "http://ollama.internal:11434/v1" })).status).toBe(200);
  expect((await put({ kind: "openai", model: "m", apiKey: "k", baseUrl: "https://llm.example.org/v1/" })).status).toBe(200);
  // The admin's server-scope provider may use any endpoint (Ollama on localhost).
  expect((await call("PUT", "/api/ai/admin/config", "owner", { provider: { kind: "openai", model: "llama", baseUrl: "http://127.0.0.1:11434/v1" }, mcpServers: [] })).status).toBe(200);
});

forEachDb("the admin allowlist in the server config enables an origin for BYOK, and removing it disables the key", async (db) => {
  const { call } = await setup(db);
  const put = () => call("PUT", "/api/ai/account/key", "ann", { kind: "openai", model: "m", apiKey: "k", baseUrl: "http://gpu-box:8000/v1" });
  expect((await put()).status).toBe(400);
  expect((await call("PUT", "/api/ai/admin/config", "owner", { provider: null, mcpServers: [], allowedBaseUrls: ["http://gpu-box:8000/v1"] })).status).toBe(400); // origins only
  expect((await call("PUT", "/api/ai/admin/config", "ann", { provider: null, mcpServers: [], allowedBaseUrls: ["http://gpu-box:8000"] })).status).toBe(403);
  expect((await call("PUT", "/api/ai/admin/config", "owner", { provider: null, mcpServers: [], allowedBaseUrls: ["http://gpu-box:8000"] })).status).toBe(200);
  expect((await (await call("GET", "/api/ai/admin/config", "owner")).json()).allowedBaseUrls).toEqual(["http://gpu-box:8000"]);
  expect((await put()).status).toBe(200);
  expect(await (await call("GET", "/api/ai/status?courseId=sample", "ann")).json()).toMatchObject({ enabled: true, source: "byok" });
  // Saving the config without the field keeps the list; an empty list removes it, and the saved key stops being used.
  await call("PUT", "/api/ai/admin/config", "owner", { provider: null, mcpServers: [] });
  expect((await (await call("GET", "/api/ai/admin/config", "owner")).json()).allowedBaseUrls).toEqual(["http://gpu-box:8000"]);
  await call("PUT", "/api/ai/admin/config", "owner", { provider: null, mcpServers: [], allowedBaseUrls: [] });
  expect(await (await call("GET", "/api/ai/status?courseId=sample", "ann")).json()).toMatchObject({ enabled: false });
});

forEachDb("provider errors on a user's own key show the status and a fixed message, never the upstream body", async (db) => {
  const logs: string[] = [];
  const seen: { url: string; redirect?: RequestRedirect }[] = [];
  const fakeFetch = (async (input: RequestInfo | URL, init?: RequestInit) => {
    seen.push({ url: String(input), redirect: init?.redirect });
    return new Response(JSON.stringify({ error: { message: "INTERNAL-SECRET-BODY from upstream" } }), { status: 502, headers: { "content-type": "application/json" } });
  }) as typeof fetch;
  const { call } = await setup(db, undefined, { providerFactory: undefined, fetch: fakeFetch, log: (m) => logs.push(m) }, { aiAllowedBaseUrls: ["http://ollama.internal:11434"] });
  expect((await call("PUT", "/api/ai/account/key", "ann", { kind: "openai", model: "m", apiKey: "k", baseUrl: "http://ollama.internal:11434/v1" })).status).toBe(200);
  const text = await (await call("POST", "/api/ai/courses/sample/chat", "ann", { message: "hello" })).text();
  expect(text).not.toContain("INTERNAL-SECRET-BODY");
  const events = text.trim().split("\n").map((l) => JSON.parse(l));
  expect(events.find((e) => e.type === "error").message).toBe("502: The AI provider returned an error. Check your key, model and endpoint under Account, AI.");
  expect(seen).toEqual([{ url: "http://ollama.internal:11434/v1/chat/completions", redirect: "manual" }]);
  expect(logs.some((l) => l.includes("INTERNAL-SECRET-BODY") && l.includes("ann"))).toBe(true);

  // The admin's server provider keeps the full message for the owner to debug.
  await call("DELETE", "/api/ai/account/key", "ann");
  await call("PUT", "/api/ai/admin/config", "owner", { provider: { kind: "openai", model: "m", baseUrl: "http://127.0.0.1:11434/v1" }, mcpServers: [] });
  const t2 = await (await call("POST", "/api/ai/courses/sample/chat", "owner", { message: "hello" })).text();
  expect(t2).toContain("502: INTERNAL-SECRET-BODY");
});

forEachDb("authoring uses core canView: the global author role reaches restricted courses (R14)", async (db) => {
  const { call } = await setup(db);
  await call("PUT", "/api/ai/admin/config", "owner", { provider: { kind: "anthropic", apiKey: "k" }, mcpServers: [] });
  const body = { courseId: "secret", kind: "questions", style: "pure-snake", topic: "Adding", sources: [{ kind: "lesson", lessonId: "m2-text" }] };
  const courses = async (u: string) => ((await (await call("GET", "/api/ai/studio/courses", u)).json()) as { id: string }[]).map((c) => c.id).sort();
  // The global author role views every course under core canView, so it may author the restricted one too.
  expect(await courses("auth")).toEqual(["sample", "secret"]);
  expect((await call("GET", "/api/ai/studio/sources?courseId=secret", "auth")).status).toBe(200);
  expect((await call("POST", "/api/ai/studio/generate", "auth", body)).status).toBe(202);
  // The owner sees everything.
  expect(await courses("owner")).toEqual(["sample", "secret"]);
  expect((await call("POST", "/api/ai/studio/generate", "owner", body)).status).toBe(202);
  // A course_authors row makes a learner an author of that course only.
  expect((await call("POST", "/api/ai/studio/generate", "bob", body)).status).toBe(403);
  await db.sql`INSERT INTO course_authors (course_id, user_id) VALUES ('secret', 'bob')`;
  expect(await courses("bob")).toEqual(["secret"]);
  expect((await call("POST", "/api/ai/studio/generate", "bob", body)).status).toBe(202);
  // A grant to view alone does not make a learner an author.
  await db.sql`INSERT INTO course_acl (course_id, principal_type, principal_id) VALUES ('secret', 'user', 'ann')`;
  expect((await call("POST", "/api/ai/studio/generate", "ann", body)).status).toBe(403);
});