Sybil-2/server/tests/message-generation.test.ts

import assert from "node:assert/strict";
import { execFileSync } from "node:child_process";
import { mkdtempSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { dirname, join, resolve } from "node:path";
import test from "node:test";
import { fileURLToPath } from "node:url";
import Fastify from "fastify";
import sensible from "@fastify/sensible";
import { withMessageGenerations } from "../src/message-generation.js";

const serverRoot = resolve(dirname(fileURLToPath(import.meta.url)), "..");
const databaseDir = mkdtempSync(join(tmpdir(), "sybil-message-generation-"));
process.env.DATABASE_URL = `file:${join(databaseDir, "test.db")}`;
process.env.ADMIN_TOKEN = "message-generation-test-token";
process.env.OPENAI_API_KEY = "synthetic-openai-key";
process.env.XAI_API_KEY = "synthetic-xai-key";
writeFileSync(join(databaseDir, "test.db"), "");
execFileSync(process.execPath, [join(serverRoot, "node_modules/prisma/build/index.js"), "migrate", "deploy"], {
  cwd: serverRoot, env: { ...process.env, PRISMA_HIDE_UPDATE_MESSAGE: "1" }, stdio: "pipe",
});
const [{ registerRoutes }, { prisma }, { runMultiplex }, { runMultiplexStream }] = await Promise.all([
  import("../src/routes.js"), import("../src/db.js"), import("../src/llm/multiplexer.js"), import("../src/llm/streaming.js"),
]);
const app = Fastify({ logger: false });
await app.register(sensible);
await registerRoutes(app);
await app.ready();
const auth = { authorization: "Bearer message-generation-test-token" };
test.after(async () => { await app.close(); await prisma.$disconnect(); rmSync(databaseDir, { recursive: true, force: true }); });

test("streaming and non-streaming replies retain their own model across selection changes and forks", async t => {
  t.mock.method(globalThis, "fetch", async (input: RequestInfo | URL, init?: RequestInit) => {
    const url = input instanceof Request ? input.url : String(input);
    assert.ok(["https://api.openai.com/v1/chat/completions", "https://api.x.ai/v1/chat/completions"].includes(url), url);
    const body = JSON.parse(input instanceof Request ? await input.text() : String(init?.body));
    const content = `Reply from ${body.model}`;
    return body.stream
      ? new Response(`data: ${JSON.stringify({ choices: [{ index: 0, delta: { content }, finish_reason: null }] })}\n\ndata: [DONE]\n\n`, { headers: { "content-type": "text/event-stream" } })
      : Response.json({ choices: [{ message: { role: "assistant", content } }] });
  });
  const chat = await prisma.chat.create({ data: { title: "Mixed models" } });
  const messages = [{ role: "user" as const, content: "Answer" }];
  const events = [];
  for await (const event of runMultiplexStream({ chatId: chat.id, provider: "openai", model: "model-one", clientRequestId: "first-turn", messages, enabledTools: [] })) events.push(event);
  assert.equal(events.at(-1)?.type, "done");
  await runMultiplex({ chatId: chat.id, provider: "xai", model: "model-two", messages, enabledTools: [] });
  await prisma.chat.update({ where: { id: chat.id }, data: { lastUsedProvider: "gemini", lastUsedModel: "model-next" } });
  const detail = await app.inject({ method: "GET", url: `/v1/chats/${chat.id}`, headers: auth });
  assert.equal(detail.statusCode, 200);
  const saved = detail.json().chat.messages;
  assert.deepEqual(saved.map((m: any) => m.metadata.generation), [{ provider: "openai", model: "model-one" }, { provider: "xai", model: "model-two" }]);
  assert.equal(saved[0].metadata.clientRequestId, "first-turn");
  const fork = await app.inject({ method: "POST", url: `/v1/chats/${chat.id}/fork`, headers: auth, payload: { messageId: saved[1].id } });
  assert.equal(fork.statusCode, 200);
  const forkDetail = await app.inject({ method: "GET", url: `/v1/chats/${fork.json().chat.id}`, headers: auth });
  assert.deepEqual(forkDetail.json().chat.messages.map((m: any) => m.metadata.generation), saved.map((m: any) => m.metadata.generation));
  assert.equal(forkDetail.json().chat.messages[0].metadata.clientRequestId, undefined);
});

test("legacy call matches are projected without rewriting history and survive a new fork", async () => {
  const chat = await prisma.chat.create({ data: {
    title: "Older reply", lastUsedProvider: "gemini", lastUsedModel: "different-current-model",
    messages: { create: { role: "assistant", content: "Older reply", metadata: { clientRequestId: "legacy-request", citations: ["https://example.test"] } } },
    calls: { create: { provider: "anthropic", model: "original-model", request: { clientRequestId: "legacy-request" }, latencyMs: 10 } },
  }, include: { messages: true } });
  const detail = await app.inject({ method: "GET", url: `/v1/chats/${chat.id}`, headers: auth });
  assert.deepEqual(detail.json().chat.messages[0].metadata.generation, { provider: "anthropic", model: "original-model" });
  assert.deepEqual((await prisma.message.findUniqueOrThrow({ where: { id: chat.messages[0]!.id } })).metadata, chat.messages[0]!.metadata);
  const fork = await app.inject({ method: "POST", url: `/v1/chats/${chat.id}/fork`, headers: auth, payload: {} });
  const forkDetail = await app.inject({ method: "GET", url: `/v1/chats/${fork.json().chat.id}`, headers: auth });
  assert.deepEqual(forkDetail.json().chat.messages[0].metadata, { citations: ["https://example.test"], generation: { provider: "anthropic", model: "original-model" } });
});

test("missing, ambiguous, failed and unfinished calls never guess a reply's model", () => {
  const call = { provider: "openai", model: "model-one", request: { clientRequestId: "request" }, error: null, latencyMs: 10 };
  const message = { role: "assistant", metadata: { clientRequestId: "request" } };
  for (const calls of [[], [call, { ...call, provider: "xai" }], [{ ...call, error: "failed" }], [{ ...call, latencyMs: null }]]) {
    assert.deepEqual(withMessageGenerations([message], calls), [message]);
  }
  const existing = { ...message, metadata: { ...message.metadata, generation: { provider: "gemini", model: "original" } } };
  const aside = { ...message, metadata: { ...message.metadata, agentAside: true } };
  const user = { ...message, role: "user" };
  assert.deepEqual(withMessageGenerations([existing, aside, user], [call]), [existing, aside, user]);
});