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

import { beforeAll, expect, test } from "bun:test";
import { join } from "node:path";
import type { ChatMessage } from "@stc/contracts";
import { createCourseRepository, type CourseRepository, type LoadedCourse } from "@stc/course-format";
import { buildSystemPrompt, citationsFor, readMentorText } from "../src/assistant/prompt";
import { courseTools, combineTools } from "../src/assistant/tools";
import { runAssistant, type AssistantEvent } from "../src/assistant/loop";
import { buildCourseIndex, type CourseIndex } from "../src/retrieval/course-index";
import { fakeProvider } from "./helpers";

const FIXTURES = join(import.meta.dir, "../../../fixtures/courses");
let repo: CourseRepository;
let course: LoadedCourse;
let index: CourseIndex;
beforeAll(async () => {
  repo = await createCourseRepository([FIXTURES]);
  course = repo.get("sample")!;
  index = await buildCourseIndex(course, "/nonexistent");
});

test("mentor rules are in the system prompt when mentor mode is on", async () => {
  const mentorText = await readMentorText(repo, course);
  expect(mentorText).toContain("never full solutions");
  const on = buildSystemPrompt({ course, mentorText, mentorOn: true, sources: index.search("add", 2) });
  expect(on).toContain("## Course mentor rules");
  expect(on).toContain("Give hints, never full solutions.");
  expect(on).toContain("[1] lessons/04-exercise.md");
  const off = buildSystemPrompt({ course, mentorText, mentorOn: false, sources: [] });
  expect(off).not.toContain("never full solutions");
  expect(off).toContain("Mentor mode is off");
});

test("citations are deduplicated by path and lesson", () => {
  const c = citationsFor([
    { path: "a.md", kind: "lesson", text: "x", lessonId: "l1", lessonTitle: "L1" },
    { path: "a.md", kind: "lesson", text: "y", lessonId: "l1", lessonTitle: "L1" },
    { path: "notes", kind: "note", text: "z", lessonId: "l1" },
    { path: "b.js", kind: "code", text: "z" },
  ]);
  expect(c).toEqual([
    { path: "a.md", kind: "lesson", lessonId: "l1", lessonTitle: "L1" },
    { path: "b.js", kind: "code" },
  ]);
});

test("course tools search, read lessons and files, list notes, and refuse solutions", async () => {
  const t = courseTools({ repo, course, index, listNotes: async () => [{ id: "n", lessonId: "m1-intro", body: "my note", anchorKind: "none", anchorValue: null }] });
  expect(t.specs.map((s) => s.name).sort()).toEqual(["list_notes", "read_file", "read_lesson", "search_course"]);
  expect(await t.run("search_course", { query: "sum of its arguments" })).toContain("lessons/04-exercise.md");
  expect(t.sources.some((s) => s.path === "lessons/04-exercise.md")).toBe(true);
  expect(await t.run("read_lesson", { lessonId: "m2-ex-add" })).toContain("Implement `add(a, b)`");
  expect(await t.run("read_lesson", { lessonId: "nope" })).toMatch(/No lesson/);
  expect(await t.run("read_file", { path: "exercises/add/starter.js" })).toContain("add");
  expect(await t.run("read_file", { path: "exercises/add/solution.js" })).toMatch(/authors only/);
  expect(await t.run("read_file", { path: "../../../etc/passwd" })).toMatch(/No such file/);
  expect(await t.run("list_notes", {})).toContain("my note");
});

test("the loop runs tools and feeds results back until the model answers", async () => {
  const t = courseTools({ repo, course, index, listNotes: async () => [] });
  const { provider, seen } = fakeProvider([
    () => ({ text: "Searching.", toolCalls: [{ id: "c1", name: "search_course", input: { query: "add arguments" } }] }),
    (msgs) => ({ text: `Found it in lessons/04-exercise.md. (${msgs.filter((m) => m.role === "tool").length} tool result)` }),
  ]);
  const messages: ChatMessage[] = [{ role: "system", content: "sys" }, { role: "user", content: "How do I add?" }];
  const events: AssistantEvent[] = [];
  for await (const e of runAssistant({ provider, messages, tools: combineTools([t]) })) events.push(e);
  expect(events.filter((e) => e.type === "tool")).toEqual([{ type: "tool", name: "search_course" }]);
  const done = events.find((e) => e.type === "done") as { type: "done"; text: string };
  expect(done.text).toBe("Searching.Found it in lessons/04-exercise.md. (1 tool result)");
  const second = seen[1]!.messages;
  expect(second.at(-2)).toMatchObject({ role: "assistant", toolCallId: "c1", name: "search_course" });
  expect(second.at(-1)!.role).toBe("tool");
  expect(second.at(-1)!.content).toContain("lessons/04-exercise.md");
  expect(seen[0]!.tools!.map((x) => x.name)).toContain("search_course");
});

test("unknown tools and provider errors are reported, and the step limit holds", async () => {
  const t = courseTools({ repo, course, index, listNotes: async () => [] });
  const looping = fakeProvider([() => ({ toolCalls: [{ id: crypto.randomUUID(), name: "nope", input: {} }] })]);
  const events: AssistantEvent[] = [];
  for await (const e of runAssistant({ provider: looping.provider, messages: [{ role: "user", content: "x" }], tools: combineTools([t]), maxSteps: 3 })) events.push(e);
  expect(looping.seen).toHaveLength(3);
  expect(looping.seen[1]!.messages.at(-1)!.content).toMatch(/Unknown tool/);
  expect(events.at(-1)!.type).toBe("done");

  const broken = { id: "b", async *chat() { throw new Error("boom"); } };
  const ev2: AssistantEvent[] = [];
  for await (const e of runAssistant({ provider: broken, messages: [{ role: "user", content: "x" }], tools: combineTools([t]) })) ev2.push(e);
  expect(ev2).toEqual([{ type: "error", message: "boom" }]);
});

test("my_progress reports the learner's progress per lesson", async () => {
  const t = courseTools({
    repo, course, index, listNotes: async () => [],
    getProgress: async () => ({ "m1-intro": { state: "done", positionSeconds: null, updatedAt: "2026-10-30" }, "m1-video": { state: "in-progress", positionSeconds: 75, updatedAt: "2026-10-30" } }),
  });
  expect(t.specs.map((s) => s.name)).toContain("my_progress");
  const out = await t.run("my_progress", {});
  expect(out).toContain("Introduction (m1-intro): done");
  expect(out).toContain("Watch the clip (m1-video): in progress, at 1:15");
  expect(out).toContain("Adding numbers (m2-text): not started");
  expect(out).toMatch(/1 of \d+ lessons done/);
});

// The solution guard must hold for every spelling of a path: it checks the canonical real path.
test("read_file refuses solutions however the path is written, including through symlinks", async () => {
  const { cpSync, mkdtempSync, mkdirSync, readFileSync, rmSync, symlinkSync, writeFileSync } = await import("node:fs");
  const { tmpdir } = await import("node:os");
  const dir = mkdtempSync(join(tmpdir(), "stc-ai-guard-"));
  try {
    const root = join(dir, "courses/sample");
    cpSync(join(FIXTURES, "sample"), root, { recursive: true });
    const yaml = join(root, "stcourse.yaml");
    writeFileSync(yaml, readFileSync(yaml, "utf8").replace("solution: exercises/add/solution.js", "solution: exercises/add/answer.js"));
    writeFileSync(join(root, "exercises/add/answer.js"), "export const add = (a, b) => a + b; // ANSWER\n");
    mkdirSync(join(root, "solutions"));
    writeFileSync(join(root, "solutions/ref.js"), "// REFERENCE SOLUTION\n");
    writeFileSync(join(root, ".env"), "TOKEN=hidden\n");
    symlinkSync("exercises/add/answer.js", join(root, "peek.js"));
    symlinkSync("solutions", join(root, "material"));
    symlinkSync("../solutions/ref.js", join(root, "lessons/extra.js"));
    const r = await createCourseRepository([join(dir, "courses")]);
    const c = r.get("sample")!;
    expect(c.status).toBe("ok");
    const t = courseTools({ repo: r, course: c, index: await buildCourseIndex(c, "/nonexistent"), listNotes: async () => [] });
    const read = (path: string) => t.run("read_file", { path });

    const refused = [
      "exercises/add/answer.js",
      "exercises/add/../add/answer.js",
      "./exercises/add/answer.js",
      "exercises//add/answer.js",
      "./exercises/./add/answer.js",
      "/exercises/add/answer.js",
      "exercises\\add\\answer.js",
      "exercises%2Fadd%2Fanswer.js",
      "exercises/add/%2E%2E/add/answer.js",
      "exercises%252Fadd%252Fanswer.js",
      "%2E/exercises/add/answer.js",
      "peek.js",
      "material/ref.js",
      "lessons/extra.js",
      "lessons/../solutions/ref.js",
      "exercises/add/solution.js",
      ".env",
      "lessons/../.env",
    ];
    for (const p of refused) {
      const out = await read(p);
      expect({ p, out }).toEqual({ p, out: expect.stringMatching(/authors only|not available|No such file/) });
      expect(out).not.toContain("ANSWER");
      expect(out).not.toContain("REFERENCE");
      expect(out).not.toContain("hidden");
    }
    expect(await read("exercises/add/answer.js")).toMatch(/authors only/);
    expect(await read("peek.js")).toMatch(/authors only/);
    expect(await read("../../etc/passwd")).toMatch(/No such file/);
    // Ordinary files still read, in any spelling.
    expect(await read("exercises/add/starter.js")).toContain("add");
    expect(await read("./exercises/./add/../add/starter.js")).toContain("add");
    expect(await read("exercises%2Fadd%2Fstarter.js")).toContain("add");
    expect(t.sources.at(-1)?.path).toBe("exercises/add/starter.js");
    // The index never holds them either.
    const idx = await buildCourseIndex(c, "/nonexistent");
    const hits = idx.search("ANSWER REFERENCE SOLUTION add", 50);
    expect(hits.filter((x) => /answer|ref\.js|peek|material|extra/.test(x.path) || x.text.includes("ANSWER") || x.text.includes("REFERENCE"))).toEqual([]);
  } finally {
    rmSync(dir, { recursive: true, force: true });
  }
});