mcpctl/src/mcpd/tests/chat-service-membership.test.ts

/**
 * `mcpctl chat <agent> --project <p>` — the same person, viewing one of the
 * teams they joined (docs/agentic-teams.md §4).
 *
 * The regression these lock down: the AgentProject junction existed for a whole
 * release without the chat path reading it, so `--project` on an agent could
 * only ever mean the one project pinned on the Agent row.
 */
import type { ChatMessage, ChatThread, Prompt } from '@prisma/client';
import { describe, it, expect, vi } from 'vitest';

import type { IChatRepository } from '../src/repositories/chat.repository.js';
import type { IPromptRepository } from '../src/repositories/prompt.repository.js';
import type { AgentService, ResolvedMembership } from '../src/services/agent.service.js';
import { ChatService, type ChatToolDispatcher } from '../src/services/chat.service.js';
import type { LlmAdapterRegistry } from '../src/services/llm/dispatcher.js';
import type { LlmAdapter, NonStreamingResult } from '../src/services/llm/types.js';
import type { LlmService } from '../src/services/llm.service.js';
import { NotFoundError } from '../src/services/mcp-server.service.js';

const NOW = new Date('2026-01-01T00:00:00Z');

function mockChatRepo(): IChatRepository {
  const msgs: ChatMessage[] = [];
  let n = 1;
  return {
    createThread: vi.fn(async ({ agentId, ownerId }) => ({
      id: `thread-${String(n++)}`, agentId, projectId: null, ownerId,
      title: '', lastTurnAt: NOW, createdAt: NOW, updatedAt: NOW,
    } as ChatThread)),
    findThread: vi.fn(async () => null),
    listThreadsByAgent: vi.fn(async () => []),
    listMessages: vi.fn(async () => []),
    appendMessage: vi.fn(async (input: { threadId: string; role: string; content: string }) => {
      const m = {
        id: `msg-${String(n++)}`, threadId: input.threadId, turnIndex: msgs.length,
        role: input.role, content: input.content, toolCalls: null, toolCallId: null,
        status: 'complete', createdAt: NOW,
      } as unknown as ChatMessage;
      msgs.push(m);
      return m;
    }),
    updateStatus: vi.fn(async () => ({}) as ChatMessage),
    markPendingAsError: vi.fn(async () => 0),
    touchThread: vi.fn(async () => undefined),
    nextTurnIndex: vi.fn(async () => msgs.length),
  } as unknown as IChatRepository;
}

function mockPromptRepo(rows: Prompt[] = []): IPromptRepository {
  return {
    findAll: vi.fn(async () => rows),
    findByAgent: vi.fn(async () => []),
  } as unknown as IPromptRepository;
}

/**
 * The agent: pinned to `home-brain`, legacy-pinned to project `legacy`, and a
 * member of `sre` (with its own brain) and `payments` (without one).
 */
const MEMBERSHIPS: Record<string, ResolvedMembership> = {
  sre: { projectId: 'proj-sre', projectName: 'sre', role: 'sre', llmName: 'sre-brain' },
  payments: { projectId: 'proj-payments', projectName: 'payments', role: 'security', llmName: null },
};

function mockAgents(): AgentService {
  return {
    getByName: vi.fn(async (name: string) => ({
      id: `agent-${name}`, name, description: '', systemPrompt: 'You are alice.',
      llm: { id: 'llm-home-brain', name: 'home-brain' },
      project: { id: 'proj-legacy', name: 'legacy' },
      defaultPersonality: null, proxyModelName: null, toolset: 'project',
      defaultParams: {}, extras: {}, ownerId: 'owner-1', version: 1,
      createdAt: NOW, updatedAt: NOW,
    })),
    resolveMembership: vi.fn(async (_name: string, projectName?: string) => {
      if (projectName === undefined) {
        return { projectId: 'proj-legacy', projectName: 'legacy', role: '', llmName: null };
      }
      const found = MEMBERSHIPS[projectName];
      if (found === undefined) {
        throw new NotFoundError(`Agent 'alice' is not a member of project '${projectName}'.`);
      }
      return found;
    }),
  } as unknown as AgentService;
}

function mockLlms(): LlmService {
  const row = (name: string): Record<string, unknown> => ({
    id: `llm-${name}`, name, type: 'openai', model: `${name}-model`,
    url: '', tier: 'fast', description: '',
    apiKeySecretId: null, apiKeySecretKey: null, extraConfig: {}, poolName: null,
    kind: 'public', providerSessionId: null, status: 'active',
    lastHeartbeatAt: null, inactiveSince: null,
    version: 1, createdAt: NOW, updatedAt: NOW,
  });
  return {
    getByName: vi.fn(async (name: string) => ({ ...row(name), apiKeyRef: null })),
    findByPoolName: vi.fn(async (poolName: string) => [row(poolName)]),
    resolveApiKey: vi.fn(async () => 'fake-key'),
  } as unknown as LlmService;
}

function reply(content: string): NonStreamingResult {
  return {
    status: 200,
    body: { id: 'c', object: 'chat.completion', choices: [{ index: 0, message: { role: 'assistant', content }, finish_reason: 'stop' }] },
  };
}

function build() {
  const listTools = vi.fn(async () => []);
  const tools = { listTools, callTool: vi.fn(async () => ({})) } as unknown as ChatToolDispatcher;
  const llms = mockLlms();
  const promptRepo = mockPromptRepo();
  const adapter: LlmAdapter = {
    kind: 'scripted',
    infer: vi.fn(async () => reply('ok')),
    stream: async function*() { yield { data: '[DONE]', done: true }; },
  } as unknown as LlmAdapter;
  const service = new ChatService(
    mockAgents(), llms, { get: () => adapter } as unknown as LlmAdapterRegistry,
    mockChatRepo(), promptRepo, tools,
  );
  return { service, llms, listTools, promptRepo, adapter };
}

describe('agent chat scoped to a joined project', () => {
  it('uses the MEMBERSHIP project for tools, not the legacy pin', async () => {
    const { service, listTools } = build();
    await service.chat({ agentName: 'alice', ownerId: 'u1', userMessage: 'hi', projectName: 'payments' });
    expect(listTools).toHaveBeenCalledWith('proj-payments');
  });

  it('falls back to the legacy pin when no project is named', async () => {
    const { service, listTools } = build();
    await service.chat({ agentName: 'alice', ownerId: 'u1', userMessage: 'hi' });
    expect(listTools).toHaveBeenCalledWith('proj-legacy');
  });

  it('the per-membership brain wins over the agent\'s pinned Llm', async () => {
    const { service, llms } = build();
    const res = await service.chat({ agentName: 'alice', ownerId: 'u1', userMessage: 'hi', projectName: 'sre' });
    expect(llms.getByName).toHaveBeenCalledWith('sre-brain');
    expect(llms.getByName).not.toHaveBeenCalledWith('home-brain');
    // Reported back to the caller, so `mcpctl chat` can show which brain answered.
    expect(res.llm).toBe('sre-brain');
  });

  it('keeps the agent\'s own brain on a membership with no override', async () => {
    const { service, llms } = build();
    const res = await service.chat({ agentName: 'alice', ownerId: 'u1', userMessage: 'hi', projectName: 'payments' });
    expect(llms.getByName).toHaveBeenCalledWith('home-brain');
    expect(res.llm).toBe('home-brain');
  });

  it('and on bare chat, with no project at all', async () => {
    const { service, llms } = build();
    await service.chat({ agentName: 'alice', ownerId: 'u1', userMessage: 'hi' });
    expect(llms.getByName).toHaveBeenCalledWith('home-brain');
  });

  it('refuses a project the agent has not joined — before spending a token', async () => {
    const { service, adapter } = build();
    await expect(
      service.chat({ agentName: 'alice', ownerId: 'u1', userMessage: 'hi', projectName: 'finance' }),
    ).rejects.toThrow(/not a member of project 'finance'/);
    expect(adapter.infer).not.toHaveBeenCalled();
  });

  it('scopes the streaming path the same way', async () => {
    const { service, listTools } = build();
    const chunks = [];
    for await (const c of service.chatStream({ agentName: 'alice', ownerId: 'u1', userMessage: 'hi', projectName: 'sre' })) {
      chunks.push(c);
    }
    expect(listTools).toHaveBeenCalledWith('proj-sre');
  });
});