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