mcpctl/src/db/tests/agent-schema.test.ts

import { describe, it, expect, beforeAll, afterAll, beforeEach } from 'vitest';
import type { PrismaClient } from '@prisma/client';
import { setupTestDb, cleanupTestDb, clearAllTables } from './helpers.js';

describe('agent / chat-thread / chat-message schema', () => {
  let prisma: PrismaClient;

  beforeAll(async () => {
    prisma = await setupTestDb();
  }, 30_000);

  afterAll(async () => {
    await cleanupTestDb();
  });

  beforeEach(async () => {
    await clearAllTables(prisma);
  });

  async function makeUser(suffix = '') {
    return prisma.user.create({
      data: {
        email: `agent-test-${Date.now()}${suffix}@example.com`,
        name: 'Agent Tester',
        passwordHash: 'x',
      },
    });
  }

  async function makeLlm(name: string) {
    return prisma.llm.create({
      data: { name, type: 'openai', model: 'qwen3-thinking' },
    });
  }

  async function makeProject(ownerId: string, name: string) {
    return prisma.project.create({ data: { name, ownerId } });
  }

  async function makeAgent(opts: {
    name: string;
    llmId: string;
    ownerId: string;
    projectId?: string;
  }) {
    return prisma.agent.create({
      data: {
        name: opts.name,
        llmId: opts.llmId,
        ownerId: opts.ownerId,
        projectId: opts.projectId ?? null,
      },
    });
  }

  it('creates an agent with required fields and JSON defaults', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-default-fields');
    const agent = await makeAgent({ name: 'a1', llmId: llm.id, ownerId: user.id });

    expect(agent.id).toBeDefined();
    expect(agent.description).toBe('');
    expect(agent.systemPrompt).toBe('');
    expect(agent.defaultParams).toEqual({});
    expect(agent.extras).toEqual({});
    expect(agent.version).toBe(1);
  });

  it('enforces unique agent name', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-uniq');
    await makeAgent({ name: 'dup', llmId: llm.id, ownerId: user.id });
    await expect(
      makeAgent({ name: 'dup', llmId: llm.id, ownerId: user.id }),
    ).rejects.toThrow();
  });

  it('blocks deleting an Llm referenced by an agent (Restrict)', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-in-use');
    await makeAgent({ name: 'pinned', llmId: llm.id, ownerId: user.id });

    await expect(prisma.llm.delete({ where: { id: llm.id } })).rejects.toThrow();
  });

  it('sets agent.projectId NULL when its Project is deleted (SetNull)', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-setnull');
    const project = await makeProject(user.id, 'proj-detach');
    const agent = await makeAgent({
      name: 'detachable',
      llmId: llm.id,
      ownerId: user.id,
      projectId: project.id,
    });
    expect(agent.projectId).toBe(project.id);

    await prisma.project.delete({ where: { id: project.id } });
    const reloaded = await prisma.agent.findUnique({ where: { id: agent.id } });
    expect(reloaded?.projectId).toBeNull();
  });

  it('cascades thread + message delete when an Agent is deleted', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-cascade');
    const agent = await makeAgent({ name: 'doomed', llmId: llm.id, ownerId: user.id });
    const thread = await prisma.chatThread.create({
      data: { agentId: agent.id, ownerId: user.id, title: 't' },
    });
    await prisma.chatMessage.create({
      data: {
        threadId: thread.id,
        turnIndex: 0,
        role: 'user',
        content: 'hello',
      },
    });

    await prisma.agent.delete({ where: { id: agent.id } });

    expect(await prisma.chatThread.findUnique({ where: { id: thread.id } })).toBeNull();
    expect(await prisma.chatMessage.count({ where: { threadId: thread.id } })).toBe(0);
  });

  it('blocks duplicate (threadId, turnIndex)', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-turn-uniq');
    const agent = await makeAgent({ name: 'orderly', llmId: llm.id, ownerId: user.id });
    const thread = await prisma.chatThread.create({
      data: { agentId: agent.id, ownerId: user.id },
    });
    await prisma.chatMessage.create({
      data: { threadId: thread.id, turnIndex: 0, role: 'user', content: 'a' },
    });
    await expect(
      prisma.chatMessage.create({
        data: { threadId: thread.id, turnIndex: 0, role: 'assistant', content: 'b' },
      }),
    ).rejects.toThrow();
  });

  it('persists tool-call shape on assistant + tool turns', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-tools');
    const agent = await makeAgent({ name: 'toolish', llmId: llm.id, ownerId: user.id });
    const thread = await prisma.chatThread.create({
      data: { agentId: agent.id, ownerId: user.id },
    });

    await prisma.chatMessage.create({
      data: { threadId: thread.id, turnIndex: 0, role: 'user', content: 'do x' },
    });
    await prisma.chatMessage.create({
      data: {
        threadId: thread.id,
        turnIndex: 1,
        role: 'assistant',
        content: '',
        toolCalls: [
          { id: 'call_1', name: 'do_thing', arguments: { x: 1 } },
        ],
        status: 'pending',
      },
    });
    await prisma.chatMessage.create({
      data: {
        threadId: thread.id,
        turnIndex: 2,
        role: 'tool',
        content: 'ok',
        toolCallId: 'call_1',
      },
    });

    const messages = await prisma.chatMessage.findMany({
      where: { threadId: thread.id },
      orderBy: { turnIndex: 'asc' },
    });
    expect(messages).toHaveLength(3);
    expect(messages[1]?.toolCalls).toEqual([
      { id: 'call_1', name: 'do_thing', arguments: { x: 1 } },
    ]);
    expect(messages[2]?.toolCallId).toBe('call_1');
  });

  it('orders threads by lastTurnAt DESC for an agent', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-order');
    const agent = await makeAgent({ name: 'history', llmId: llm.id, ownerId: user.id });

    const t1 = await prisma.chatThread.create({
      data: { agentId: agent.id, ownerId: user.id, lastTurnAt: new Date(2000, 0, 1) },
    });
    const t2 = await prisma.chatThread.create({
      data: { agentId: agent.id, ownerId: user.id, lastTurnAt: new Date(2030, 0, 1) },
    });

    const ordered = await prisma.chatThread.findMany({
      where: { agentId: agent.id },
      orderBy: { lastTurnAt: 'desc' },
    });
    expect(ordered.map((t) => t.id)).toEqual([t2.id, t1.id]);
  });

  // ── Agent-direct prompts (Prompt.agentId) ──

  it('attaches a prompt directly to an agent (no project)', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-agent-prompt');
    const agent = await makeAgent({ name: 'directprompt', llmId: llm.id, ownerId: user.id });

    const prompt = await prisma.prompt.create({
      data: { name: 'agent-only', content: 'hi from agent', agentId: agent.id },
    });
    expect(prompt.agentId).toBe(agent.id);
    expect(prompt.projectId).toBeNull();
  });

  it('cascade-deletes agent-direct prompts when the agent is removed', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-agent-prompt-cascade');
    const agent = await makeAgent({ name: 'goingaway', llmId: llm.id, ownerId: user.id });
    await prisma.prompt.create({
      data: { name: 'p1', content: 'c', agentId: agent.id },
    });

    await prisma.agent.delete({ where: { id: agent.id } });

    expect(await prisma.prompt.count({ where: { agentId: agent.id } })).toBe(0);
  });

  it('enforces (name, agentId) uniqueness independently of (name, projectId)', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-uniq-prompt-agent');
    const agent = await makeAgent({ name: 'has-prompts', llmId: llm.id, ownerId: user.id });

    await prisma.prompt.create({
      data: { name: 'shared', content: 'a', agentId: agent.id },
    });
    await expect(
      prisma.prompt.create({
        data: { name: 'shared', content: 'b', agentId: agent.id },
      }),
    ).rejects.toThrow();

    // Same name on a different agent is fine.
    const agent2 = await makeAgent({ name: 'other-agent', llmId: llm.id, ownerId: user.id });
    const ok = await prisma.prompt.create({
      data: { name: 'shared', content: 'c', agentId: agent2.id },
    });
    expect(ok.id).toBeDefined();
  });

  // ── Personalities ──

  it('creates a personality bound to an agent', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-pers-1');
    const agent = await makeAgent({ name: 'with-pers', llmId: llm.id, ownerId: user.id });

    const p = await prisma.personality.create({
      data: { name: 'grumpy', description: 'curmudgeonly', agentId: agent.id },
    });
    expect(p.id).toBeDefined();
    expect(p.priority).toBe(5);
  });

  it('enforces (name, agentId) uniqueness on personalities', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-pers-uniq');
    const agent = await makeAgent({ name: 'pers-uniq', llmId: llm.id, ownerId: user.id });

    await prisma.personality.create({ data: { name: 'mode', agentId: agent.id } });
    await expect(
      prisma.personality.create({ data: { name: 'mode', agentId: agent.id } }),
    ).rejects.toThrow();
  });

  it('cascade-deletes personalities + bindings when the agent is removed', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-pers-cascade');
    const agent = await makeAgent({ name: 'kaboom', llmId: llm.id, ownerId: user.id });

    const personality = await prisma.personality.create({
      data: { name: 'snarky', agentId: agent.id },
    });
    const prompt = await prisma.prompt.create({
      data: { name: 'tone', content: 'be snarky', agentId: agent.id },
    });
    await prisma.personalityPrompt.create({
      data: { personalityId: personality.id, promptId: prompt.id },
    });

    await prisma.agent.delete({ where: { id: agent.id } });

    expect(await prisma.personality.count({ where: { id: personality.id } })).toBe(0);
    expect(await prisma.personalityPrompt.count({ where: { personalityId: personality.id } })).toBe(0);
    // Prompts directly attached to the agent also go via Prompt.agentId cascade.
    expect(await prisma.prompt.count({ where: { id: prompt.id } })).toBe(0);
  });

  it('SetNull on Agent.defaultPersonalityId when the personality is deleted', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-default-pers');
    const agent = await makeAgent({ name: 'has-default', llmId: llm.id, ownerId: user.id });
    const personality = await prisma.personality.create({
      data: { name: 'def', agentId: agent.id },
    });
    await prisma.agent.update({
      where: { id: agent.id },
      data: { defaultPersonalityId: personality.id },
    });

    await prisma.personality.delete({ where: { id: personality.id } });

    const reloaded = await prisma.agent.findUnique({ where: { id: agent.id } });
    expect(reloaded?.defaultPersonalityId).toBeNull();
  });

  // ── v3: Agent.kind virtual + lifecycle fields ──

  it('defaults a freshly inserted Agent to kind=public, status=active', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-default-kind');
    const agent = await makeAgent({ name: 'fresh', llmId: llm.id, ownerId: user.id });
    expect(agent.kind).toBe('public');
    expect(agent.status).toBe('active');
    expect(agent.providerSessionId).toBeNull();
    expect(agent.lastHeartbeatAt).toBeNull();
    expect(agent.inactiveSince).toBeNull();
  });

  it('persists kind=virtual + lifecycle fields together', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-pub-virtual');
    const now = new Date();
    const agent = await prisma.agent.create({
      data: {
        name: 'local-coder',
        llmId: llm.id,
        ownerId: user.id,
        kind: 'virtual',
        providerSessionId: 'sess-abc',
        lastHeartbeatAt: now,
        status: 'active',
      },
    });
    expect(agent.kind).toBe('virtual');
    expect(agent.providerSessionId).toBe('sess-abc');
    expect(agent.lastHeartbeatAt?.getTime()).toBe(now.getTime());
  });

  it('finds virtual agents by (kind, status) cheaply (GC sweep query)', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-gc-agent');
    await prisma.agent.create({ data: { name: 'pub-1', llmId: llm.id, ownerId: user.id } });
    await prisma.agent.create({
      data: { name: 'v-active', llmId: llm.id, ownerId: user.id, kind: 'virtual', providerSessionId: 's1' },
    });
    await prisma.agent.create({
      data: { name: 'v-inactive', llmId: llm.id, ownerId: user.id, kind: 'virtual', providerSessionId: 's2', status: 'inactive', inactiveSince: new Date() },
    });

    const stale = await prisma.agent.findMany({
      where: { kind: 'virtual', status: 'inactive' },
      select: { name: true },
    });
    expect(stale.map((a) => a.name)).toEqual(['v-inactive']);
  });

  it('finds agents by providerSessionId (used on mcplocal disconnect cascade)', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-sess-cascade');
    await prisma.agent.create({
      data: { name: 'a', llmId: llm.id, ownerId: user.id, kind: 'virtual', providerSessionId: 'shared' },
    });
    await prisma.agent.create({
      data: { name: 'b', llmId: llm.id, ownerId: user.id, kind: 'virtual', providerSessionId: 'shared' },
    });
    await prisma.agent.create({
      data: { name: 'c', llmId: llm.id, ownerId: user.id, kind: 'virtual', providerSessionId: 'other' },
    });

    const owned = await prisma.agent.findMany({
      where: { providerSessionId: 'shared' },
      select: { name: true },
      orderBy: { name: 'asc' },
    });
    expect(owned.map((a) => a.name)).toEqual(['a', 'b']);
  });

  it('binds the same prompt to multiple personalities of an agent', async () => {
    const user = await makeUser();
    const llm = await makeLlm('llm-shared-prompt');
    const agent = await makeAgent({ name: 'shared', llmId: llm.id, ownerId: user.id });
    const p1 = await prisma.personality.create({ data: { name: 'a', agentId: agent.id } });
    const p2 = await prisma.personality.create({ data: { name: 'b', agentId: agent.id } });
    const prompt = await prisma.prompt.create({
      data: { name: 'shared-text', content: 'reusable', agentId: agent.id },
    });

    await prisma.personalityPrompt.create({
      data: { personalityId: p1.id, promptId: prompt.id, priority: 9 },
    });
    await prisma.personalityPrompt.create({
      data: { personalityId: p2.id, promptId: prompt.id, priority: 1 },
    });

    expect(await prisma.personalityPrompt.count({ where: { promptId: prompt.id } })).toBe(2);

    // Same (personalityId, promptId) cannot collide.
    await expect(
      prisma.personalityPrompt.create({
        data: { personalityId: p1.id, promptId: prompt.id },
      }),
    ).rejects.toThrow();
  });
});