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

import { expect } from "bun:test";
import { migrate, type Db } from "@stc/db";
// Core migrations only: the ai_* tables exactly as v1 shipped them.
import { forEachDb as forEachV1Db } from "../../db/test/helpers";
import { createSecretBox } from "../src/crypto";
import { migrationsDirs } from "../src/migrations";
import { aiSettingsRepo } from "../src/settings";
import { upgradeAiData } from "../src/upgrade";

const T = "2026-10-01T00:00:00.000Z";
const box = createSecretBox("test-secret");
const aiMigrations = { "feature-ai": migrationsDirs };

async function seedMessage(db: Db, id: string, conv: string, role: string, content: string) {
  await db.sql`INSERT INTO ai_messages (id, conversation_id, role, content, created_at) VALUES (${id}, ${conv}, ${role}, ${content}, ${T})`;
}

forEachV1Db("v1 upgrade: sources rows fold into citations_json on the assistant row before them", async (db) => {
  const cites = JSON.stringify([{ path: "lessons/04-exercise.md", title: "Adding numbers", lessonId: "m2-ex-add" }]);
  await db.sql`INSERT INTO ai_conversations (id, user_id, course_id, title, created_at) VALUES ('c1', 'u1', 'sample', 'Hi', ${T})`;
  await db.sql`INSERT INTO ai_conversations (id, user_id, course_id, title, created_at) VALUES ('c2', 'u1', 'sample', 'Two', ${T})`;
  // ids sort in the order v1 wrote them (UUIDv7); c1: user, assistant, sources, user, assistant (no sources)
  await seedMessage(db, "0001-a", "c1", "user", "How do I add?");
  await seedMessage(db, "0001-b", "c1", "assistant", "Which operator?");
  await seedMessage(db, "0001-c", "c1", "sources", cites);
  await seedMessage(db, "0001-d", "c1", "user", "And then?");
  await seedMessage(db, "0001-e", "c1", "assistant", "Try it.");
  // c2: an orphan sources row (no assistant before it) and an assistant whose next row belongs to c1's ids range
  await seedMessage(db, "0002-a", "c2", "sources", "[]");
  await seedMessage(db, "0002-b", "c2", "assistant", "Alone.");

  const applied = await migrate(db, aiMigrations);
  expect(applied).toContain("0001_ai_message_citations.sql");

  const rows = await db.sql`SELECT id, role, content, citations_json FROM ai_messages ORDER BY id`;
  expect(rows.map((r: { id: string; role: string }) => `${r.id}:${r.role}`)).toEqual([
    "0001-a:user", "0001-b:assistant", "0001-d:user", "0001-e:assistant", "0002-b:assistant",
  ]);
  const byId = Object.fromEntries(rows.map((r: { id: string; citations_json: string | null }) => [r.id, r.citations_json]));
  expect(JSON.parse(byId["0001-b"]!)).toEqual(JSON.parse(cites));
  expect(byId["0001-e"]).toBeNull();
  expect(byId["0002-b"]).toBeNull();
  expect(byId["0001-a"]).toBeNull();
  expect(await migrate(db, aiMigrations)).toEqual([]);
});

async function seedSettings(db: Db, scope: "server" | "user", owner: string, config: unknown) {
  const json = typeof config === "string" ? config : JSON.stringify(config);
  await db.sql`INSERT INTO ai_settings (scope, owner_id, mentor_off, config_json, updated_at) VALUES (${scope}, ${owner}, 0, ${json}, ${T})`;
}

forEachV1Db("v1 upgrade: per-course mentor toggles move from config_json to ai_mentor_prefs", async (db) => {
  await seedSettings(db, "user", "u1", { byok: { kind: "anthropic" }, mentorOff: { sample: true, other: false } });
  await seedSettings(db, "user", "u2", { byok: null, mentorOff: {} });
  await seedSettings(db, "user", "u3", "{not json");
  await db.sql`INSERT INTO user_api_keys (user_id, provider, ciphertext, created_at) VALUES ('u1', 'anthropic', ${await box.encrypt("sk-u1")}, ${T})`;

  expect(await migrate(db, aiMigrations)).toContain("0002_ai_mentor_prefs.sql");
  expect(await upgradeAiData(db, box)).toBe(true);

  const prefs = await db.sql`SELECT user_id, course_id, off FROM ai_mentor_prefs ORDER BY user_id, course_id`;
  expect(prefs.map((r: { user_id: string; course_id: string; off: number }) => [r.user_id, r.course_id, Number(r.off)])).toEqual([["u1", "sample", 1]]);
  const cfg = async (u: string) => (await db.sql`SELECT config_json FROM ai_settings WHERE scope = 'user' AND owner_id = ${u}`)[0].config_json as string;
  expect(JSON.parse(await cfg("u1"))).toEqual({ byok: { kind: "anthropic" } });
  expect(JSON.parse(await cfg("u2"))).toEqual({ byok: null });
  expect(await cfg("u3")).toBe("{not json");

  const repo = aiSettingsRepo(db, box);
  expect(await repo.isMentorOff("u1", "sample")).toBe(true);
  expect(await repo.isMentorOff("u1", "other")).toBe(false);
  expect(await repo.getUser("u1")).toEqual({ byok: { kind: "anthropic", hasApiKey: true }, mentorOff: { sample: true } });

  // idempotent: the marker stops a second run
  expect(await upgradeAiData(db, box)).toBe(false);
  expect((await db.sql`SELECT COUNT(*) AS n FROM ai_mentor_prefs`)[0].n.toString()).toBe("1");
  expect((await db.sql`SELECT name FROM schema_migrations WHERE name = 'ai_data_v11'`)).toHaveLength(1);
});

forEachV1Db("v1 upgrade: the settings repo upgrades on first use", async (db) => {
  await seedSettings(db, "user", "u1", { byok: null, mentorOff: { sample: true } });
  await migrate(db, aiMigrations);
  expect(await aiSettingsRepo(db, box).isMentorOff("u1", "sample")).toBe(true);
  expect(JSON.parse((await db.sql`SELECT config_json FROM ai_settings WHERE owner_id = 'u1'`)[0].config_json)).toEqual({ byok: null });
});

forEachV1Db("v1 upgrade: plaintext MCP env and headers are encrypted at rest", async (db) => {
  const apiKeyCipher = await box.encrypt("sk-ant-SERVER");
  await seedSettings(db, "server", "", {
    provider: { kind: "anthropic" }, apiKeyCipher, allowedBaseUrls: ["http://ollama:11434"],
    mcpServers: [
      { id: "fs", name: "Files", transport: "stdio", command: "mcp-fs", args: ["--root", "/x"], env: { FS_TOKEN: "tok-PLAIN-123456" } },
      { id: "web", name: "Web", transport: "http", url: "https://mcp.example.com/mcp", headers: { authorization: "Bearer hdr-PLAIN-abcdef" }, courseIds: ["sample"] },
      { id: "bare", name: "Bare", transport: "http", url: "https://bare.example.com/mcp" },
    ],
  });
  await migrate(db, aiMigrations);
  expect(await upgradeAiData(db, box)).toBe(true);

  const raw = (await db.sql`SELECT config_json FROM ai_settings WHERE scope = 'server'`)[0].config_json as string;
  expect(raw).not.toContain("PLAIN");
  const stored = JSON.parse(raw);
  expect(stored.apiKeyCipher).toBe(apiKeyCipher);
  expect(stored.allowedBaseUrls).toEqual(["http://ollama:11434"]);
  expect(stored.mcpServers.map((s: { id: string; env?: unknown; headers?: unknown; secretCipher?: string }) => [s.id, !!s.env, !!s.headers, !!s.secretCipher]))
    .toEqual([["fs", false, false, true], ["web", false, false, true], ["bare", false, false, false]]);

  const repo = aiSettingsRepo(db, box);
  expect(await repo.serverApiKey()).toBe("sk-ant-SERVER");
  expect(await repo.mcpServersWithSecrets()).toEqual([
    { id: "fs", name: "Files", transport: "stdio", command: "mcp-fs", args: ["--root", "/x"], env: { FS_TOKEN: "tok-PLAIN-123456" } },
    { id: "web", name: "Web", transport: "http", url: "https://mcp.example.com/mcp", headers: { authorization: "Bearer hdr-PLAIN-abcdef" }, courseIds: ["sample"] },
    { id: "bare", name: "Bare", transport: "http", url: "https://bare.example.com/mcp" },
  ]);
  const shown = JSON.stringify(await repo.getServer());
  expect(shown).not.toContain("PLAIN");
});