liqa/server/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 { SessionContextService } from "../session/session-context.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";

import pkg from "../../package.json";

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

  private createServer(): McpServer {
    const server = new McpServer({ name: pkg.name, version: pkg.version });
    this.registerTools(server);
    return server;
  }

  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), 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)"),
          orderBy: z
            .enum([
              "id",
              "sessionName",
              "status",
              "lastUsedAt",
              "createdAt",
              "updatedAt",
            ])
            .optional()
            .describe("Field to order by (default id)"),
          orderDir: z
            .enum(["ASC", "DESC"])
            .optional()
            .describe("Sort direction (default ASC)"),
        },
      },
      async ({ page, limit, orderBy, orderDir }) => {
        const sessions = await this.sessionService.findAll({
          page,
          limit,
          orderBy,
          orderDir,
        });
        return {
          content: [{ type: "text" as const, text: JSON.stringify(sessions) }],
        };
      },
    );

    server.registerTool(
      "delete_session",
      {
        description: "Delete a session by numeric ID (closes it first if open)",
        inputSchema: {
          id: z.string().uuid().describe("Session ID to delete"),
        },
      },
      async ({ id }) => {
        try {
          await this.sessionContextService.delete(id);
        } catch (err) {
          return {
            isError: true,
            content: [{ type: "text" as const, text: (err as Error).message }],
          };
        }
        return {
          content: [{ type: "text" as const, text: `Session ${id} deleted` }],
        };
      },
    );

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

    server.registerTool(
      "list_environments",
      {
        description: "List all environments, 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)"),
          orderBy: z
            .enum(["id", "name", "createdAt", "updatedAt"])
            .optional()
            .describe("Field to order by (default id)"),
          orderDir: z
            .enum(["ASC", "DESC"])
            .optional()
            .describe("Sort direction (default ASC)"),
        },
      },
      async ({ page, limit, orderBy, orderDir }) => {
        const envs = await this.environmentService.findAll({
          page,
          limit,
          orderBy,
          orderDir,
        });
        return {
          content: [{ type: "text" as const, text: JSON.stringify(envs) }],
        };
      },
    );

    server.registerTool(
      "get_environment",
      {
        description: "Get an environment record by ID",
        inputSchema: {
          id: z.string().uuid().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.string().uuid().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.string().uuid().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()
            .optional()
            .describe(
              "Session name to restore cookies and localStorage from. Omit to open without a stored session.",
            ),
          url: z.string().url().describe("URL to navigate to"),
          readerMode: z
            .boolean()
            .optional()
            .describe("Extract readable plain text instead of raw HTML"),
          selector: z
            .string()
            .optional()
            .describe(
              "CSS selector whose matching element content is returned; applied before readerMode",
            ),
        },
      },
      async ({ sessionName, url, readerMode, selector }) => {
        try {
          const result = await this.browserService.open(
            sessionName,
            url,
            readerMode ?? false,
            selector,
          );
          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()
            .optional()
            .describe(
              "Session name to restore. Omit to run without a stored session.",
            ),
          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)"),
          orderBy: z
            .enum(["id", "name", "createdAt", "updatedAt"])
            .optional()
            .describe("Field to order by (default id)"),
          orderDir: z
            .enum(["ASC", "DESC"])
            .optional()
            .describe("Sort direction (default ASC)"),
        },
      },
      async ({ page, limit, orderBy, orderDir }) => {
        const result = await this.scenarioService.findAll({
          page,
          limit,
          orderBy,
          orderDir,
        });
        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.string().uuid().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.string().uuid().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.string().uuid().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.string().uuid().describe("Parent scenario ID"),
          order: z
            .number()
            .int()
            .min(0)
            .describe("Execution order (ascending)"),
          type: z.enum(["login", "exec", "sign"]).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.string().uuid().describe("Scenario ID"),
          stepId: z.string().uuid().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.string().uuid().describe("Scenario ID"),
          stepId: z.string().uuid().describe("Step ID"),
          order: z
            .number()
            .int()
            .min(0)
            .optional()
            .describe("New execution order"),
          type: z
            .enum(["login", "exec", "sign"])
            .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.string().uuid().describe("Scenario ID"),
          stepId: z.string().uuid().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.string().uuid().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.string().uuid().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 }],
          };
        }
      },
    );

    server.registerTool(
      "get_scenario_run",
      {
        description:
          "Get a specific scenario run with all step runs and their outputs",
        inputSchema: {
          scenarioId: z.string().uuid().describe("Scenario ID"),
          runId: z.string().uuid().describe("Run ID"),
        },
      },
      async ({ scenarioId, runId }) => {
        try {
          const run = await this.scenarioService.findRun(scenarioId, runId);
          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 }],
          };
        }
      },
    );

    server.registerTool(
      "wait_for_scenario_run",
      {
        description:
          "Block until a scenario run reaches pass or fail (max 5 min), then return the full run with step outputs",
        inputSchema: {
          scenarioId: z.string().uuid().describe("Scenario ID"),
          runId: z.string().uuid().describe("Run ID"),
        },
      },
      async ({ scenarioId, runId }) => {
        try {
          const run = await this.scenarioService.waitForRun(scenarioId, runId);
          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 }],
          };
        }
      },
    );

    server.registerTool(
      "export_scenario",
      {
        description:
          "Export a scenario as a portable JSON payload (name + steps)",
        inputSchema: {
          id: z.string().uuid().describe("Scenario ID to export"),
        },
      },
      async ({ id }) => {
        try {
          const exported = await this.scenarioService.exportScenario(id);
          return {
            content: [
              { type: "text" as const, text: JSON.stringify(exported) },
            ],
          };
        } catch (err) {
          return {
            isError: true,
            content: [{ type: "text" as const, text: (err as Error).message }],
          };
        }
      },
    );

    server.registerTool(
      "import_scenario",
      {
        description:
          "Import a scenario from an export payload, creating a new scenario with all its steps",
        inputSchema: {
          name: z.string().describe("Scenario name"),
          steps: z
            .array(
              z.object({
                order: z.number().int().min(0).describe("Execution order"),
                type: z.enum(["login", "exec", "sign"]).describe("Step type"),
                title: z.string().nullable().optional().describe("Step title"),
                execCode: z
                  .string()
                  .nullable()
                  .optional()
                  .describe("Exec/sign code"),
                validateCode: z
                  .string()
                  .nullable()
                  .optional()
                  .describe("Validation code"),
              }),
            )
            .describe("Ordered list of steps"),
        },
      },
      async ({ name, steps }) => {
        try {
          const scenario = await this.scenarioService.importScenario({
            name,
            steps: steps as Parameters<
              typeof this.scenarioService.importScenario
            >[0]["steps"],
          });
          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 }],
          };
        }
      },
    );

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

  async handle(req: Request, res: Response): Promise<void> {
    const server = this.createServer();
    const transport = new StreamableHTTPServerTransport({
      sessionIdGenerator: undefined,
    });
    await server.connect(transport);
    try {
      await transport.handleRequest(req, res, req.body);
    } finally {
      await transport.close();
      await server.close();
    }
  }
}