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 { byokAad, createSecretBox, legacyV1Encrypt, mcpAad, SERVER_KEY_AAD } from "../src/crypto";
import { migrationsDirs } from "../src/migrations";
import { aiSettingsRepo } from "../src/settings";
import { AI_SECRETS_UPGRADE, upgradeAiData } from "../src/upgrade";

const T = "2026-10-01T00:00:00.000Z";
const SECRET = "test-secret";
const box = createSecretBox(SECRET);
/** A secret as v1 stored it (one key, no additional data). */
const v1Token = (plain: string, secret = SECRET) => legacyV1Encrypt(secret, plain);
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), then an assistant with no row after it
  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 v1Token("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 v1Token("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.startsWith("v2.")).toBe(true); // re-sealed under the v2 scheme
  expect(await box.open(stored.apiKeyCipher, "server-api-key", SERVER_KEY_AAD)).toBe("sk-ant-SERVER");
  expect(stored.allowedBaseUrls).toEqual(["http://ollama:11434"]);
  type S = { id: string; env?: unknown; headers?: unknown; envCipher?: string; headersCipher?: string; secretCipher?: string };
  expect(stored.mcpServers.map((s: S) => [s.id, !!s.env, !!s.headers, !!s.envCipher, !!s.headersCipher, !!s.secretCipher]))
    .toEqual([["fs", false, false, true, false, false], ["web", false, false, false, true, false], ["bare", false, false, false, false, false]]);
  expect(JSON.parse(await box.open(stored.mcpServers[0].envCipher, "mcp", mcpAad("fs", "env")))).toEqual({ FS_TOKEN: "tok-PLAIN-123456" });
  expect((await db.sql`SELECT name FROM schema_migrations WHERE name = ${AI_SECRETS_UPGRADE}`)).toHaveLength(1);

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

forEachV1Db("v1.1 draft upgrade: secretCipher, v1 API keys and BYOK keys are re-sealed as v2, once", async (db) => {
  // The state the first v1.1 draft left: mentor step done, MCP secrets in one v1-token secretCipher.
  await migrate(db, aiMigrations);
  await db.sql`INSERT INTO schema_migrations (name, applied_at) VALUES ('ai_data_v11', ${T})`;
  await seedSettings(db, "server", "", {
    provider: { kind: "anthropic" }, apiKeyCipher: await v1Token("sk-ant-SERVER"),
    mcpServers: [
      { id: "fs", name: "Files", transport: "stdio", command: "x", secretCipher: await v1Token(JSON.stringify({ env: { T: "tok-DRAFT-123456" } })) },
      { id: "web", name: "Web", transport: "http", url: "https://w.example.com/mcp", secretCipher: await v1Token(JSON.stringify({ headers: { authorization: "Bearer DRAFT-abcdef" } })) },
      "not a server",
    ],
  });
  await db.sql`INSERT INTO user_api_keys (user_id, provider, ciphertext, created_at) VALUES ('u1', 'openai', ${await v1Token("sk-u1")}, ${T})`;
  await seedSettings(db, "user", "u1", { byok: { kind: "openai", model: "m" } });

  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("secretCipher");
  expect(raw).not.toContain("\"v1.");
  const ct = (await db.sql`SELECT ciphertext FROM user_api_keys WHERE user_id = 'u1'`)[0].ciphertext as string;
  expect(ct.startsWith("v2.")).toBe(true);
  expect(await box.open(ct, "byok", byokAad("u1", "openai"))).toBe("sk-u1");

  const repo = aiSettingsRepo(db, box);
  expect(await repo.serverApiKey()).toBe("sk-ant-SERVER");
  expect(await repo.userApiKey("u1")).toBe("sk-u1");
  expect(await repo.mcpServersWithSecrets()).toEqual([
    { id: "fs", name: "Files", transport: "stdio", command: "x", env: { T: "tok-DRAFT-123456" } },
    { id: "web", name: "Web", transport: "http", url: "https://w.example.com/mcp", headers: { authorization: "Bearer DRAFT-abcdef" } },
  ]);
  expect((await repo.getServer()).unreadableSecrets).toEqual([]);
  expect(JSON.parse(raw).mcpServers[2]).toBe("not a server"); // malformed entries are kept, not fatal

  const snapshot = raw;
  expect(await upgradeAiData(db, box)).toBe(false);
  expect((await db.sql`SELECT config_json FROM ai_settings WHERE scope = 'server'`)[0].config_json).toBe(snapshot);
});

forEachV1Db("v1 upgrade with a changed STC_SECRET keeps old tokens, stays pending, and finishes once the secret is restored", async (db) => {
  await seedSettings(db, "server", "", {
    provider: { kind: "anthropic" }, apiKeyCipher: await v1Token("sk-ant-SERVER", "old-secret"),
    mcpServers: [{ id: "fs", name: "Files", transport: "stdio", command: "x", env: { T: "tok-PLAIN-123456" } }],
  });
  await db.sql`INSERT INTO user_api_keys (user_id, provider, ciphertext, created_at) VALUES ('u1', 'anthropic', ${await v1Token("sk-u1", "old-secret")}, ${T})`;
  await seedSettings(db, "user", "u1", { byok: { kind: "anthropic" } });
  await migrate(db, aiMigrations);

  const repo = aiSettingsRepo(db, box); // the new secret
  expect(await repo.userApiKey("u1")).toBeNull();
  expect((await repo.getUser("u1")).byok).toEqual({ kind: "anthropic", hasApiKey: true, keyUnreadable: true });
  expect((await repo.getServer()).unreadableSecrets).toEqual([{ kind: "api-key" }]);
  // plaintext MCP values do not need the old secret: sealed already
  expect((await repo.mcpServersWithSecrets())[0]!.env).toEqual({ T: "tok-PLAIN-123456" });
  expect((await db.sql`SELECT name FROM schema_migrations WHERE name = ${AI_SECRETS_UPGRADE}`)).toHaveLength(0);
  expect(((await db.sql`SELECT ciphertext FROM user_api_keys`)[0].ciphertext as string).startsWith("v1.")).toBe(true);

  // The admin puts the previous secret back: the next start re-seals and completes.
  const restored = createSecretBox("old-secret");
  expect(await upgradeAiData(db, restored)).toBe(true);
  expect((await db.sql`SELECT name FROM schema_migrations WHERE name = ${AI_SECRETS_UPGRADE}`)).toHaveLength(1);
  const again = aiSettingsRepo(db, restored);
  expect(await again.userApiKey("u1")).toBe("sk-u1");
  expect(await again.serverApiKey()).toBe("sk-ant-SERVER");
});