liqa/src/mcp/mcp.service.ts

import { Injectable } from '@nestjs/common';
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
import { z } from 'zod';
import type { Request, Response } from 'express';
import { AuthService } from '../auth/auth.service';
import { SessionService } from '../session/session.service';
import { EnvironmentService } from '../environment/environment.service';
import type { EnvironmentUrls } from '../environment/environment.entity';
import { BrowserService } from '../browser/browser.service';
import { CodeExecutorService } from '../code-executor/code-executor.service';
import { ScenarioService } from '../scenario/scenario.service';

@Injectable()
export class McpService {
  constructor(
    private readonly authService: AuthService,
    private readonly sessionService: SessionService,
    private readonly environmentService: EnvironmentService,
    private readonly browserService: BrowserService,
    private readonly codeExecutor: CodeExecutorService,
    private readonly scenarioService: ScenarioService,
  ) {}

  private registerTools(server: McpServer): void {

    // ── Auth ──────────────────────────────────────────────────────────────────

    server.registerTool(
      'list_keys',
      { description: 'List available key identifiers from the keys directory' },
      async () => {
        const keys = this.authService.listKeys();
        return { content: [{ type: 'text' as const, text: JSON.stringify(keys) }] };
      },
    );

    server.registerTool(
      'login',
      {
        description: 'Log in using a file key against a named environment and store the session',
        inputSchema: {
          key: z.string().describe('Key identifier (filename without extension from keys/ dir)'),
          environmentName: z.string().describe('Environment name to resolve login/cabinet URLs'),
          sessionName: z.string().optional().describe('Session name to store credentials under. Auto-UUID if omitted.'),
        },
      },
      async ({ key, environmentName, sessionName }) => {
        const result = await this.authService.login(key, environmentName, sessionName);
        return { content: [{ type: 'text' as const, text: JSON.stringify(result) }] };
      },
    );

    // ── Sessions ──────────────────────────────────────────────────────────────

    server.registerTool(
      'list_sessions',
      { description: 'List all stored sessions (id, sessionName, createdAt, updatedAt)' },
      async () => {
        const sessions = await this.sessionService.findAll();
        return { content: [{ type: 'text' as const, text: JSON.stringify(sessions) }] };
      },
    );

    server.registerTool(
      'delete_session',
      {
        description: 'Delete a session by numeric ID',
        inputSchema: {
          id: z.number().int().describe('Session ID to delete'),
        },
      },
      async ({ id }) => {
        const sessions = await this.sessionService.findAll();
        if (!sessions.find(s => s.id === id)) {
          return { isError: true, content: [{ type: 'text' as const, text: `Session ${id} not found` }] };
        }
        await this.sessionService.remove(id);
        return { content: [{ type: 'text' as const, text: `Session ${id} deleted` }] };
      },
    );

    // ── Environments ──────────────────────────────────────────────────────────

    server.registerTool(
      'list_environments',
      { description: 'List all environments' },
      async () => {
        const envs = await this.environmentService.findAll();
        return { content: [{ type: 'text' as const, text: JSON.stringify(envs) }] };
      },
    );

    server.registerTool(
      'get_environment',
      {
        description: 'Get an environment record by ID',
        inputSchema: {
          id: z.number().int().describe('Environment ID'),
        },
      },
      async ({ id }) => {
        try {
          const env = await this.environmentService.findOne(id);
          return { content: [{ type: 'text' as const, text: JSON.stringify(env) }] };
        } catch {
          return { isError: true, content: [{ type: 'text' as const, text: `Environment ${id} not found` }] };
        }
      },
    );

    server.registerTool(
      'create_environment',
      {
        description: 'Create a new named environment with a set of URLs',
        inputSchema: {
          name: z.string().describe('Unique environment name, e.g. liquio-diia-stg'),
          urls: z.record(z.string(), z.string()).describe('Map of URL keys to URL strings (id_url, cabinet_url, admin_url, …)'),
        },
      },
      async ({ name, urls }) => {
        try {
          const env = await this.environmentService.create({ name, urls: urls as EnvironmentUrls });
          return { content: [{ type: 'text' as const, text: JSON.stringify(env) }] };
        } catch (err) {
          return { isError: true, content: [{ type: 'text' as const, text: (err as Error).message }] };
        }
      },
    );

    server.registerTool(
      'update_environment',
      {
        description: 'Update an existing environment (name and/or urls)',
        inputSchema: {
          id: z.number().int().describe('Environment ID to update'),
          name: z.string().optional().describe('New name'),
          urls: z.record(z.string(), z.string()).optional().describe('New URLs map'),
        },
      },
      async ({ id, name, urls }) => {
        try {
          const env = await this.environmentService.update(id, { name, urls: urls as EnvironmentUrls | undefined });
          return { content: [{ type: 'text' as const, text: JSON.stringify(env) }] };
        } catch (err) {
          return { isError: true, content: [{ type: 'text' as const, text: (err as Error).message }] };
        }
      },
    );

    server.registerTool(
      'delete_environment',
      {
        description: 'Delete an environment by ID',
        inputSchema: {
          id: z.number().int().describe('Environment ID to delete'),
        },
      },
      async ({ id }) => {
        try {
          await this.environmentService.remove(id);
          return { content: [{ type: 'text' as const, text: `Environment ${id} deleted` }] };
        } catch (err) {
          return { isError: true, content: [{ type: 'text' as const, text: (err as Error).message }] };
        }
      },
    );

    // ── Browser ───────────────────────────────────────────────────────────────

    server.registerTool(
      'open_url',
      {
        description: 'Open a URL using a stored session and return the page title and content',
        inputSchema: {
          sessionName: z.string().describe('Session name to restore cookies and localStorage from'),
          url: z.string().url().describe('URL to navigate to'),
          readerMode: z.boolean().optional().describe('Extract readable plain text instead of raw HTML'),
        },
      },
      async ({ sessionName, url, readerMode }) => {
        try {
          const result = await this.browserService.open(sessionName, url, readerMode ?? false);
          return { content: [{ type: 'text' as const, text: JSON.stringify(result) }] };
        } catch (err) {
          return { isError: true, content: [{ type: 'text' as const, text: (err as Error).message }] };
        }
      },
    );

    server.registerTool(
      'exec_code',
      {
        description: 'Execute arbitrary Playwright JavaScript with `page` and `context` in scope',
        inputSchema: {
          sessionName: z.string().describe('Session name to restore'),
          code: z.string().describe('JavaScript code body to execute (async-safe, may use `page` and `context`)'),
          url: z.string().url().optional().describe('Optional URL to navigate to before running code'),
        },
      },
      async ({ sessionName, code, url }) => {
        try {
          this.codeExecutor.validate(code);
          const result = await this.browserService.exec(sessionName, code, url);
          return { content: [{ type: 'text' as const, text: JSON.stringify(result) }] };
        } catch (err) {
          return { isError: true, content: [{ type: 'text' as const, text: (err as Error).message }] };
        }
      },
    );

    // ── Scenarios ─────────────────────────────────────────────────────────────

    server.registerTool(
      'list_scenarios',
      {
        description: 'List all scenarios (paginated)',
        inputSchema: {
          page: z.number().int().min(1).optional().describe('Page number (default 1)'),
          limit: z.number().int().min(1).optional().describe('Items per page (default 20)'),
        },
      },
      async ({ page, limit }) => {
        const result = await this.scenarioService.findAll({ page, limit });
        return { content: [{ type: 'text' as const, text: JSON.stringify(result) }] };
      },
    );

    server.registerTool(
      'get_scenario',
      {
        description: 'Get a scenario with its steps by ID',
        inputSchema: {
          id: z.number().int().describe('Scenario ID'),
        },
      },
      async ({ id }) => {
        try {
          const scenario = await this.scenarioService.findOne(id);
          return { content: [{ type: 'text' as const, text: JSON.stringify(scenario) }] };
        } catch (err) {
          return { isError: true, content: [{ type: 'text' as const, text: (err as Error).message }] };
        }
      },
    );

    server.registerTool(
      'create_scenario',
      {
        description: 'Create a new scenario',
        inputSchema: {
          name: z.string().describe('Scenario name'),
        },
      },
      async ({ name }) => {
        try {
          const scenario = await this.scenarioService.create({ name });
          return { content: [{ type: 'text' as const, text: JSON.stringify(scenario) }] };
        } catch (err) {
          return { isError: true, content: [{ type: 'text' as const, text: (err as Error).message }] };
        }
      },
    );

    server.registerTool(
      'update_scenario',
      {
        description: 'Update a scenario name',
        inputSchema: {
          id: z.number().int().describe('Scenario ID'),
          name: z.string().optional().describe('New name'),
        },
      },
      async ({ id, name }) => {
        try {
          const scenario = await this.scenarioService.update(id, { name });
          return { content: [{ type: 'text' as const, text: JSON.stringify(scenario) }] };
        } catch (err) {
          return { isError: true, content: [{ type: 'text' as const, text: (err as Error).message }] };
        }
      },
    );

    server.registerTool(
      'delete_scenario',
      {
        description: 'Delete a scenario by ID',
        inputSchema: {
          id: z.number().int().describe('Scenario ID'),
        },
      },
      async ({ id }) => {
        try {
          await this.scenarioService.remove(id);
          return { content: [{ type: 'text' as const, text: `Scenario ${id} deleted` }] };
        } catch (err) {
          return { isError: true, content: [{ type: 'text' as const, text: (err as Error).message }] };
        }
      },
    );

    server.registerTool(
      'create_scenario_step',
      {
        description: 'Add a step to a scenario',
        inputSchema: {
          scenarioId: z.number().int().describe('Parent scenario ID'),
          order: z.number().int().min(0).describe('Execution order (ascending)'),
          type: z.enum(['login', 'exec']).describe('Step type'),
          sessionName: z.string().describe('Session name used by this step'),
          execCode: z.string().optional().describe('Playwright JS code to execute (exec steps)'),
          validateCode: z.string().optional().describe('Validation JS code returning { success, description }'),
        },
      },
      async ({ scenarioId, ...dto }) => {
        try {
          const step = await this.scenarioService.createStep(scenarioId, dto);
          return { content: [{ type: 'text' as const, text: JSON.stringify(step) }] };
        } catch (err) {
          return { isError: true, content: [{ type: 'text' as const, text: (err as Error).message }] };
        }
      },
    );

    server.registerTool(
      'get_scenario_step',
      {
        description: 'Get a single step of a scenario',
        inputSchema: {
          scenarioId: z.number().int().describe('Scenario ID'),
          stepId: z.number().int().describe('Step ID'),
        },
      },
      async ({ scenarioId, stepId }) => {
        try {
          const step = await this.scenarioService.findStep(scenarioId, stepId);
          return { content: [{ type: 'text' as const, text: JSON.stringify(step) }] };
        } catch (err) {
          return { isError: true, content: [{ type: 'text' as const, text: (err as Error).message }] };
        }
      },
    );

    server.registerTool(
      'update_scenario_step',
      {
        description: 'Update a step within a scenario',
        inputSchema: {
          scenarioId: z.number().int().describe('Scenario ID'),
          stepId: z.number().int().describe('Step ID'),
          order: z.number().int().min(0).optional().describe('New execution order'),
          type: z.enum(['login', 'exec']).optional().describe('New step type'),
          sessionName: z.string().optional().describe('New session name'),
          execCode: z.string().optional().describe('New exec code'),
          validateCode: z.string().optional().describe('New validation code'),
        },
      },
      async ({ scenarioId, stepId, ...dto }) => {
        try {
          const step = await this.scenarioService.updateStep(scenarioId, stepId, dto);
          return { content: [{ type: 'text' as const, text: JSON.stringify(step) }] };
        } catch (err) {
          return { isError: true, content: [{ type: 'text' as const, text: (err as Error).message }] };
        }
      },
    );

    server.registerTool(
      'delete_scenario_step',
      {
        description: 'Delete a step from a scenario',
        inputSchema: {
          scenarioId: z.number().int().describe('Scenario ID'),
          stepId: z.number().int().describe('Step ID'),
        },
      },
      async ({ scenarioId, stepId }) => {
        try {
          await this.scenarioService.removeStep(scenarioId, stepId);
          return { content: [{ type: 'text' as const, text: `Step ${stepId} deleted from scenario ${scenarioId}` }] };
        } catch (err) {
          return { isError: true, content: [{ type: 'text' as const, text: (err as Error).message }] };
        }
      },
    );

    server.registerTool(
      'list_scenario_runs',
      {
        description: 'List runs for a scenario (paginated, optionally filtered by status)',
        inputSchema: {
          scenarioId: z.number().int().describe('Scenario ID'),
          status: z.enum(['pending', 'in_progress', 'pass', 'fail']).optional().describe('Filter by run status'),
          page: z.number().int().min(1).optional().describe('Page number (default 1)'),
          limit: z.number().int().min(1).optional().describe('Items per page (default 20)'),
        },
      },
      async ({ scenarioId, status, page, limit }) => {
        try {
          const result = await this.scenarioService.findRuns(scenarioId, { status, page, limit });
          return { content: [{ type: 'text' as const, text: JSON.stringify(result) }] };
        } catch (err) {
          return { isError: true, content: [{ type: 'text' as const, text: (err as Error).message }] };
        }
      },
    );

    server.registerTool(
      'run_scenario',
      {
        description: 'Trigger an immediate run of a scenario by ID',
        inputSchema: {
          id: z.number().int().describe('Scenario ID to run'),
        },
      },
      async ({ id }) => {
        try {
          const run = await this.scenarioService.createRun(id);
          return { content: [{ type: 'text' as const, text: JSON.stringify(run) }] };
        } catch (err) {
          return { isError: true, content: [{ type: 'text' as const, text: (err as Error).message }] };
        }
      },
    );

    // ── Transport ─────────────────────────────────────────────────────────────
  }

  async handle(req: Request, res: Response): Promise<void> {
    const server = new McpServer({ name: 'liquio-qa-bot', version: '1.0.0' });
    this.registerTools(server);

    const transport = new StreamableHTTPServerTransport({
      sessionIdGenerator: undefined, // stateless — no session management
    });

    await server.connect(transport);
    await transport.handleRequest(req, res, req.body);
  }
}