app/src/__tests__/ai.test.ts

import { describe, it, expect, beforeEach, beforeAll, afterAll } from "vitest";
import Fastify from "fastify";
import cookie from "@fastify/cookie";
import rateLimit from "@fastify/rate-limit";
import jwt from "jsonwebtoken";
import prisma from "../config/database";
import { config as appConfig } from "../config/env";
import { securityHeaders } from "../middleware/security";
import aiRoutes from "../routes/admin/ai";
import {
  computeCostUsd,
  recordUsage,
  getMonthSpendUsd,
  getBudgetUsd,
  assertBudgetAvailable,
  isConfigured,
  listConversations,
  createConversation,
  getConversationMessages,
  appendConversationMessages,
  renameConversation,
  deleteConversation,
} from "../services/ai.service";

const KIND = "test_ai"; // marker so we only clean our own rows
const CONV_UID_A = 999999991; // synthetic, FK-free owner ids
const CONV_UID_B = 999999992;
const HIST_MARKER = "「test-conv」"; // marks HTTP-test conversations on real users

beforeEach(async () => {
  await prisma.ai_usage.deleteMany({ where: { kind: KIND } });
});
afterAll(async () => {
  await prisma.ai_usage.deleteMany({ where: { kind: KIND } });
  await prisma.ai_conversations.deleteMany({
    where: { user_id: { in: [CONV_UID_A, CONV_UID_B] } },
  });
  await prisma.ai_conversations.deleteMany({
    where: { title: { contains: HIST_MARKER } },
  });
});

describe("ai.service cost + budget", () => {
  it("computes Sonnet 4.6 cost from tokens", () => {
    // 1,000,000 input @ $3 + 1,000,000 output @ $15 = $18
    expect(
      computeCostUsd("claude-sonnet-4-6", 1_000_000, 1_000_000),
    ).toBeCloseTo(18, 6);
    expect(computeCostUsd("claude-sonnet-4-6", 4000, 500)).toBeCloseTo(
      0.0195,
      6,
    );
  });

  it("records usage with the computed cost", async () => {
    await recordUsage({
      userId: null,
      kind: KIND,
      model: "claude-sonnet-4-6",
      inputTokens: 4000,
      outputTokens: 500,
    });
    const rows = await prisma.ai_usage.findMany({ where: { kind: KIND } });
    expect(rows.length).toBe(1);
    expect(Number(rows[0].cost_usd)).toBeCloseTo(0.0195, 6);
  });

  it("sums only the current month's spend (excludes prior months)", async () => {
    // Measure the baseline so the assertion is robust to any OTHER current-month
    // rows that may already exist in the test DB (getMonthSpendUsd sums ALL
    // kinds, not just ours). We assert the DELTA we introduce, not an absolute
    // upper bound.
    const before = await getMonthSpendUsd();

    await recordUsage({
      userId: null,
      kind: KIND,
      model: "claude-sonnet-4-6",
      inputTokens: 1_000_000,
      outputTokens: 0,
    }); // +$3 this month
    // An old row (year 2000) must NOT count toward the current month.
    await prisma.ai_usage.create({
      data: {
        kind: KIND,
        model: "claude-sonnet-4-6",
        input_tokens: 1_000_000,
        output_tokens: 0,
        cost_usd: 3,
        created_at: new Date("2000-01-01T00:00:00Z"),
      },
    });

    const after = await getMonthSpendUsd();
    // Only the $3 current-month row moved the needle; the year-2000 $3 is
    // excluded. The delta is exactly $3 (not $6).
    expect(after - before).toBeCloseTo(3, 6);
  });

  it("assertBudgetAvailable blocks at/over budget, allows under", async () => {
    const budget = await getBudgetUsd();
    // Under budget → null (allowed)
    expect(await assertBudgetAvailable()).toBeNull();
    // Push spend over budget for this month, then it must block with 402.
    // Tag with a unique marker so we can delete EXACTLY this row at the end of
    // the test — that keeps the inflated spend from leaking into any other
    // current-month test regardless of execution order (the beforeEach kind=KIND
    // cleanup would catch it too, but cleaning in-test makes it order-proof).
    const overBudget = await prisma.ai_usage.create({
      data: {
        kind: KIND,
        model: "claude-sonnet-4-6",
        input_tokens: 0,
        output_tokens: 0,
        cost_usd: budget + 1,
        created_at: new Date(),
      },
    });
    try {
      const blocked = await assertBudgetAvailable();
      expect(blocked).not.toBeNull();
      expect(blocked?.status).toBe(402);
    } finally {
      // Remove the over-budget row immediately so it cannot inflate
      // getMonthSpendUsd for any subsequently-running test.
      await prisma.ai_usage
        .delete({ where: { id: overBudget.id } })
        .catch(() => {});
    }
  });

  it("isConfigured reflects the API key presence", () => {
    expect(typeof isConfigured()).toBe("boolean");
  });
});

describe("ai.service conversations", () => {
  beforeEach(async () => {
    await prisma.ai_conversations.deleteMany({
      where: { user_id: { in: [CONV_UID_A, CONV_UID_B] } },
    });
  });

  // Narrow a service result to its error branch with a clear failure message.
  const expectError = (
    r: { error: string; status: number } | { data: unknown },
    status: number,
  ) => {
    if (!("error" in r))
      throw new Error(`expected an error result, got ${JSON.stringify(r)}`);
    expect(r.status).toBe(status);
  };

  it("creates, lists, appends (auto-title), and reads oldest→newest", async () => {
    const conv = await createConversation(CONV_UID_A);
    expect(conv.title).toBe("Nová konverzace");
    await appendConversationMessages(CONV_UID_A, conv.id, [
      { role: "user", content: "Kolik je hodin v Praze?" },
      { role: "assistant", content: "Nevím, ale…" },
    ]);
    const list = await listConversations(CONV_UID_A);
    expect(list).toHaveLength(1);
    expect(list[0].title).toBe("Kolik je hodin v Praze?"); // auto-titled
    const msgs = await getConversationMessages(CONV_UID_A, conv.id);
    if (!("data" in msgs)) throw new Error("expected messages, got an error");
    expect(msgs.data.map((m) => m.content)).toEqual([
      "Kolik je hodin v Praze?",
      "Nevím, ale…",
    ]);
  });

  it("does not overwrite a renamed title on later appends", async () => {
    const conv = await createConversation(CONV_UID_A);
    await renameConversation(CONV_UID_A, conv.id, "Moje téma");
    await appendConversationMessages(CONV_UID_A, conv.id, [
      { role: "user", content: "tohle by nemělo přepsat název" },
    ]);
    const list = await listConversations(CONV_UID_A);
    expect(list[0].title).toBe("Moje téma");
  });

  it("bumps updated_at on append", async () => {
    const conv = await createConversation(CONV_UID_A);
    // Force an old timestamp, then prove append refreshes it (TIMESTAMP(0) is
    // second-precision, so compare against a year rather than ms deltas).
    await prisma.ai_conversations.update({
      where: { id: conv.id },
      data: { updated_at: new Date("2000-01-01T00:00:00Z") },
    });
    await appendConversationMessages(CONV_UID_A, conv.id, [
      { role: "user", content: "ping" },
    ]);
    const after = (await listConversations(CONV_UID_A))[0].updated_at;
    expect(after.getFullYear()).toBeGreaterThan(2000);
  });

  it("isolates conversations per user (404 on every op across owners)", async () => {
    const a = await createConversation(CONV_UID_A);
    expect(await listConversations(CONV_UID_B)).toHaveLength(0);
    expectError(await getConversationMessages(CONV_UID_B, a.id), 404);
    expectError(
      await appendConversationMessages(CONV_UID_B, a.id, [
        { role: "user", content: "x" },
      ]),
      404,
    );
    expectError(await renameConversation(CONV_UID_B, a.id, "hack"), 404);
    expectError(await deleteConversation(CONV_UID_B, a.id), 404);
  });

  it("cascade-deletes messages with the conversation", async () => {
    const conv = await createConversation(CONV_UID_A);
    await appendConversationMessages(CONV_UID_A, conv.id, [
      { role: "user", content: "x" },
    ]);
    await deleteConversation(CONV_UID_A, conv.id);
    const remaining = await prisma.ai_chat_messages.count({
      where: { conversation_id: conv.id },
    });
    expect(remaining).toBe(0);
  });
});

describe("HTTP /api/admin/ai", () => {
  let app: Awaited<ReturnType<typeof buildAiApp>>;
  let adminToken: string;
  let noPermToken: string;
  let noPermRoleId: number;
  let noPermUserId: number;
  let savedKey: string;

  async function buildAiApp() {
    const a = Fastify({ logger: false });
    await a.register(cookie);
    await a.register(rateLimit, { max: 1000, timeWindow: "1 minute" });
    a.addHook("onRequest", securityHeaders);
    await a.register(aiRoutes, { prefix: "/api/admin/ai" });
    return a;
  }
  function token(user: {
    id: number;
    username: string;
    roleName: string | null;
  }) {
    return jwt.sign(
      { sub: user.id, username: user.username, role: user.roleName },
      appConfig.jwt.secret,
      { expiresIn: "15m" },
    );
  }

  beforeAll(async () => {
    savedKey = appConfig.anthropic.apiKey;
    app = await buildAiApp();
    const admin = await prisma.users.findFirst({
      where: { roles: { name: "admin" } },
      include: { roles: true },
    });
    if (!admin) throw new Error("admin not found");
    adminToken = token({
      id: admin.id,
      username: admin.username,
      roleName: admin.roles?.name ?? null,
    });
    const stamp = Date.now();
    const role = await prisma.roles.create({
      data: { name: `noperm_ai_${stamp}`, display_name: "No Perm AI" },
    });
    noPermRoleId = role.id;
    const u = await prisma.users.create({
      data: {
        username: `noperm_ai_${stamp}`,
        first_name: "No",
        last_name: "Perm",
        email: `noperm_ai_${stamp}@test.local`,
        password_hash:
          "$2a$10$invalidinvalidinvalidinvalidinvalidinvalidinvalidinvali",
        role_id: role.id,
        is_active: true,
      },
    });
    noPermUserId = u.id;
    noPermToken = token({
      id: u.id,
      username: u.username,
      roleName: role.name,
    });
  });
  afterAll(async () => {
    if (app) await app.close();
    (appConfig.anthropic as { apiKey: string }).apiKey = savedKey;
    if (noPermUserId)
      await prisma.users
        .deleteMany({ where: { id: noPermUserId } })
        .catch(() => {});
    if (noPermRoleId)
      await prisma.roles
        .deleteMany({ where: { id: noPermRoleId } })
        .catch(() => {});
  });

  it("GET /usage requires ai.use", async () => {
    const res = await app.inject({
      method: "GET",
      url: "/api/admin/ai/usage",
      headers: { Authorization: `Bearer ${noPermToken}` },
    });
    expect(res.statusCode).toBe(403);
  });

  it("GET /usage returns spend + budget for an admin", async () => {
    const res = await app.inject({
      method: "GET",
      url: "/api/admin/ai/usage",
      headers: { Authorization: `Bearer ${adminToken}` },
    });
    expect(res.statusCode).toBe(200);
    const body = res.json();
    expect(typeof body.data.budget_usd).toBe("number");
    expect(typeof body.data.month_spend_usd).toBe("number");
  });

  it("POST /chat returns 503 when AI is not configured", async () => {
    (appConfig.anthropic as { apiKey: string }).apiKey = "";
    const res = await app.inject({
      method: "POST",
      url: "/api/admin/ai/chat",
      headers: {
        Authorization: `Bearer ${adminToken}`,
        "Content-Type": "application/json",
      },
      payload: { messages: [{ role: "user", content: "ahoj" }] },
    });
    expect(res.statusCode).toBe(503);
    (appConfig.anthropic as { apiKey: string }).apiKey = savedKey;
  });

  it("GET /conversations requires ai.use", async () => {
    const res = await app.inject({
      method: "GET",
      url: "/api/admin/ai/conversations",
      headers: { Authorization: `Bearer ${noPermToken}` },
    });
    expect(res.statusCode).toBe(403);
  });

  it("conversation lifecycle over HTTP (create → append → list → rename → delete)", async () => {
    const created = await app.inject({
      method: "POST",
      url: "/api/admin/ai/conversations",
      headers: {
        Authorization: `Bearer ${adminToken}`,
        "Content-Type": "application/json",
      },
      payload: { title: `${HIST_MARKER} t` },
    });
    expect(created.statusCode).toBe(200);
    const id = created.json().data.conversation.id as number;

    const appended = await app.inject({
      method: "POST",
      url: `/api/admin/ai/conversations/${id}/messages`,
      headers: {
        Authorization: `Bearer ${adminToken}`,
        "Content-Type": "application/json",
      },
      payload: { messages: [{ role: "user", content: "ahoj" }] },
    });
    expect(appended.statusCode).toBe(200);

    const msgs = await app.inject({
      method: "GET",
      url: `/api/admin/ai/conversations/${id}/messages`,
      headers: { Authorization: `Bearer ${adminToken}` },
    });
    expect(
      msgs.json().data.messages.map((m: { content: string }) => m.content),
    ).toContain("ahoj");

    const renamed = await app.inject({
      method: "PATCH",
      url: `/api/admin/ai/conversations/${id}`,
      headers: {
        Authorization: `Bearer ${adminToken}`,
        "Content-Type": "application/json",
      },
      payload: { title: `${HIST_MARKER} renamed` },
    });
    expect(renamed.json().data.conversation.title).toBe(
      `${HIST_MARKER} renamed`,
    );

    const del = await app.inject({
      method: "DELETE",
      url: `/api/admin/ai/conversations/${id}`,
      headers: { Authorization: `Bearer ${adminToken}` },
    });
    expect(del.statusCode).toBe(200);
  });

  it("GET messages of a non-existent conversation → 404", async () => {
    const res = await app.inject({
      method: "GET",
      url: "/api/admin/ai/conversations/999999000/messages",
      headers: { Authorization: `Bearer ${adminToken}` },
    });
    expect(res.statusCode).toBe(404);
  });

  it("POST /conversations requires ai.use", async () => {
    const res = await app.inject({
      method: "POST",
      url: "/api/admin/ai/conversations",
      headers: {
        Authorization: `Bearer ${noPermToken}`,
        "Content-Type": "application/json",
      },
      payload: { title: "x" },
    });
    expect(res.statusCode).toBe(403);
  });

  it("GET /conversations lists the user's conversations", async () => {
    const created = await app.inject({
      method: "POST",
      url: "/api/admin/ai/conversations",
      headers: {
        Authorization: `Bearer ${adminToken}`,
        "Content-Type": "application/json",
      },
      payload: { title: `${HIST_MARKER} listed` },
    });
    const id = created.json().data.conversation.id as number;
    const list = await app.inject({
      method: "GET",
      url: "/api/admin/ai/conversations",
      headers: { Authorization: `Bearer ${adminToken}` },
    });
    expect(list.statusCode).toBe(200);
    expect(
      list.json().data.conversations.map((c: { id: number }) => c.id),
    ).toContain(id);
  });

  it("rejects an empty title (PATCH) and empty messages (POST) with 400", async () => {
    const patch = await app.inject({
      method: "PATCH",
      url: "/api/admin/ai/conversations/999999000",
      headers: {
        Authorization: `Bearer ${adminToken}`,
        "Content-Type": "application/json",
      },
      payload: { title: "" },
    });
    expect(patch.statusCode).toBe(400);

    const append = await app.inject({
      method: "POST",
      url: "/api/admin/ai/conversations/999999000/messages",
      headers: {
        Authorization: `Bearer ${adminToken}`,
        "Content-Type": "application/json",
      },
      payload: { messages: [] },
    });
    expect(append.statusCode).toBe(400);
  });
});