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 { SessionService } from "../session/session.service";
import { SessionContextService } from "../session/session-context.service";
import { EnvironmentService } from "../environment/environment.service";
import type { EnvironmentData } from "../environment/environment.entity";
import { CredentialService } from "../credential/credential.service";
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 sessionService: SessionService,
    private readonly sessionContextService: SessionContextService,
    private readonly environmentService: EnvironmentService,
    private readonly credentialService: CredentialService,
    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 {
    // ── 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.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.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 generic data",
        inputSchema: {
          name: z
            .string()
            .describe("Unique environment name, e.g. liquio-diia-stg"),
          data: z
            .record(z.string(), z.string())
            .describe("Map of string keys to string values"),
        },
      },
      async ({ name, data }) => {
        try {
          const env = await this.environmentService.create({
            name,
            data: data as EnvironmentData,
          });
          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 data)",
        inputSchema: {
          id: z.uuid().describe("Environment ID to update"),
          name: z.string().optional().describe("New name"),
          data: z
            .record(z.string(), z.string())
            .optional()
            .describe("New data map"),
        },
      },
      async ({ id, name, data }) => {
        try {
          const env = await this.environmentService.update(id, {
            name,
            data: data as EnvironmentData | 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.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 Playwright JS — script receives a single `context` argument (see scenario.md). Use `environmentId` (UUID) and `credentials` (alias → credential UUID) to resolve data from the DB and enable context.getEnv() / context.getCredential() / context.runSnippet().",
        inputSchema: {
          sessionName: z
            .string()
            .optional()
            .describe(
              "Session name to restore. Omit to run without a stored session.",
            ),
          code: z
            .string()
            .describe(
              "Async JavaScript body. Use `context.page` for Playwright, `context.getEnv(key)`, `context.getCredential(alias)`, `context.runSnippet(name, ...args)`.",
            ),
          url: z
            .string()
            .url()
            .optional()
            .describe("Optional URL to navigate to before running code"),
          environmentId: z
            .string()
            .uuid()
            .optional()
            .describe(
              "UUID of an existing environment entity. Its key/value data is available via context.env and context.getEnv().",
            ),
          credentials: z
            .record(z.string(), z.string())
            .optional()
            .describe(
              "Map of alias → credential UUID. Each credential is loaded from the DB and available via context.getCredential(alias).",
            ),
        },
      },
      async ({ sessionName, code, url, environmentId, credentials }) => {
        try {
          this.codeExecutor.validate(code);

          let environment: EnvironmentData | undefined;
          if (environmentId) {
            const env = await this.environmentService.findOne(environmentId);
            environment = env.data;
          }

          let resolvedCredentials: Record<string, unknown> | undefined;
          if (credentials && Object.keys(credentials).length > 0) {
            resolvedCredentials = {};
            for (const [alias, credId] of Object.entries(credentials)) {
              const cred = await this.credentialService.findOne(credId);
              resolvedCredentials[alias] = cred.data ? JSON.parse(cred.data) : {};
            }
          }

          const result = await this.browserService.exec(
            sessionName,
            code,
            url,
            environment,
            resolvedCredentials,
          );
          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.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.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.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.uuid().describe("Parent scenario ID"),
          order: z
            .number()
            .int()
            .min(0)
            .describe("Execution order (ascending)"),
          type: z.enum(["login", "exec", "sign"]).describe("Step type"),
          title: z.string().optional().describe("Optional step title"),
          execCode: z
            .string()
            .optional()
            .describe("Playwright JS code to execute (exec steps)"),
        },
      },
      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.uuid().describe("Scenario ID"),
          stepId: z.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.uuid().describe("Scenario ID"),
          stepId: z.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"),
          title: z.string().optional().describe("New step title"),
          execCode: z.string().optional().describe("New exec 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.uuid().describe("Scenario ID"),
          stepId: z.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.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.uuid().describe("Scenario ID to run"),
          environmentId: z.uuid().describe("Environment ID to run the scenario in"),
        },
      },
      async ({ id, environmentId }) => {
        try {
          const run = await this.scenarioService.createRun(id, environmentId);
          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.uuid().describe("Scenario ID"),
          runId: z.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.uuid().describe("Scenario ID"),
          runId: z.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(
      "list_scenario_credentials",
      {
        description: "List credentials assigned to a scenario",
        inputSchema: {
          scenarioId: z.uuid().describe("Scenario ID"),
        },
      },
      async ({ scenarioId }) => {
        try {
          const creds =
            await this.scenarioService.findScenarioCredentials(scenarioId);
          return {
            content: [{ type: "text" as const, text: JSON.stringify(creds) }],
          };
        } catch (err) {
          return {
            isError: true,
            content: [{ type: "text" as const, text: (err as Error).message }],
          };
        }
      },
    );

    server.registerTool(
      "add_scenario_credential",
      {
        description:
          "Add a credential to a scenario with an alias (used in snippets)",
        inputSchema: {
          scenarioId: z.uuid().describe("Scenario ID"),
          credentialId: z.uuid().describe("Credential ID to associate"),
          alias: z
            .string()
            .describe(
              "Alias used to reference this credential in step code, e.g. pkcs_key",
            ),
        },
      },
      async ({ scenarioId, credentialId, alias }) => {
        try {
          const result = await this.scenarioService.addScenarioCredential(
            scenarioId,
            { credentialId, alias },
          );
          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(
      "remove_scenario_credential",
      {
        description: "Remove a credential assignment from a scenario",
        inputSchema: {
          scenarioId: z.uuid().describe("Scenario ID"),
          scenarioCredentialId: z.uuid().describe("Scenario-credential assignment ID (not the credential ID)"),
        },
      },
      async ({ scenarioId, scenarioCredentialId }) => {
        try {
          await this.scenarioService.removeScenarioCredential(
            scenarioId,
            scenarioCredentialId,
          );
          return {
            content: [
              {
                type: "text" as const,
                text: `Credential assignment ${scenarioCredentialId} removed from scenario ${scenarioId}`,
              },
            ],
          };
        } catch (err) {
          return {
            isError: true,
            content: [{ type: "text" as const, text: (err as Error).message }],
          };
        }
      },
    );

    server.registerTool(
      "list_all_runs",
      {
        description:
          "List runs across all scenarios (paginated, filterable by status)",
        inputSchema: {
          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 ({ status, page, limit }) => {
        try {
          const result = await this.scenarioService.findAllRuns({
            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(
      "export_scenario",
      {
        description:
          "Export a scenario as a portable JSON payload (name + steps)",
        inputSchema: {
          id: z.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"),
              }),
            )
            .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();
    }
  }
}