shockbot/mcp.ts

import { existsSync, readFileSync, readdirSync } from "node:fs";
import { execFileSync } from "node:child_process";
import { join, resolve, relative } from "node:path";
import type { Tool } from "ollama";
import {
  getFileContents,
  listDir as giteaListDir,
  searchCode,
} from "./gitea.ts";
import type { PRContext } from "./types.ts";

const MAX_FILE_LINES = 200;
const MAX_GREP_RESULTS = 50;

export const TOOLS: Tool[] = [
  {
    type: "function",
    function: {
      name: "read_file",
      description: "Read the contents of a file in the repository.",
      parameters: {
        type: "object",
        required: ["path"],
        properties: {
          path: {
            type: "string",
            description: "Repo-relative path to the file",
          },
        },
      },
    },
  },
  {
    type: "function",
    function: {
      name: "list_dir",
      description: "List files and directories at a path in the repository.",
      parameters: {
        type: "object",
        required: ["path"],
        properties: {
          path: { type: "string", description: "Repo-relative directory path" },
        },
      },
    },
  },
  {
    type: "function",
    function: {
      name: "find_symbol",
      description:
        "Search for a symbol, function, or pattern across the repository.",
      parameters: {
        type: "object",
        required: ["symbol"],
        properties: {
          symbol: {
            type: "string",
            description: "Symbol name or regex pattern to search for",
          },
        },
      },
    },
  },
];

export class ToolServer {
  private _callCount = 0;
  private readonly pr: PRContext;
  private readonly workspace: string | null;
  private readonly maxToolCalls: number;

  constructor(pr: PRContext, workspace: string | null, maxToolCalls: number) {
    this.pr = pr;
    this.workspace = workspace;
    this.maxToolCalls = maxToolCalls;
  }

  get tools(): Tool[] {
    return TOOLS;
  }

  get callCount(): number {
    return this._callCount;
  }

  get atLimit(): boolean {
    return this._callCount >= this.maxToolCalls;
  }

  resetCallCount(): void {
    this._callCount = 0;
  }

  async execute(name: string, args: Record<string, unknown>): Promise<string> {
    if (this._callCount >= this.maxToolCalls) {
      return "[tool call limit reached — conclude your review now]";
    }
    this._callCount++;

    try {
      switch (name) {
        case "read_file":
          return await this.readFile(String(args["path"] ?? ""));
        case "list_dir":
          return await this.listDir(String(args["path"] ?? ""));
        case "find_symbol":
          return await this.findSymbol(String(args["symbol"] ?? ""));
        default:
          return `[unknown tool: ${name}]`;
      }
    } catch (err) {
      return `[error: ${err instanceof Error ? err.message : String(err)}]`;
    }
  }

  // Validates path is repo-relative with no traversal or absolute prefix.
  private safePath(path: string): string | null {
    if (!path) return null;
    if (path.startsWith("/") || path.includes("..")) return null;
    return path.replace(/\/+/g, "/").replace(/\/$/, "");
  }

  private async readFile(path: string): Promise<string> {
    const safe = this.safePath(path);
    if (!safe) return "[invalid path]";

    if (this.workspace) {
      const full = join(this.workspace, safe);
      // Verify resolved path stays inside workspace
      if (!resolve(full).startsWith(resolve(this.workspace)))
        return "[invalid path]";
      if (!existsSync(full)) return "[file not found]";
      const content = readFileSync(full, "utf8");
      const lines = content.split("\n");
      const truncated = lines.slice(0, MAX_FILE_LINES).join("\n");
      return lines.length > MAX_FILE_LINES
        ? `${truncated}\n... (truncated, ${lines.length - MAX_FILE_LINES} lines omitted)`
        : truncated;
    }

    return await getFileContents(this.pr, safe);
  }

  private async listDir(path: string): Promise<string> {
    const safe = this.safePath(path || ".");
    if (!safe) return "[invalid path]";

    if (this.workspace) {
      const full = join(this.workspace, safe);
      if (!resolve(full).startsWith(resolve(this.workspace)))
        return "[invalid path]";
      if (!existsSync(full)) return "[directory not found]";
      const entries = readdirSync(full, { withFileTypes: true });
      return entries
        .map((e) => `${e.isDirectory() ? "d" : "f"} ${e.name}`)
        .join("\n");
    }

    const entries = await giteaListDir(this.pr, safe);
    return entries.join("\n");
  }

  private async findSymbol(symbol: string): Promise<string> {
    if (!symbol) return "[no symbol provided]";

    if (this.workspace) {
      try {
        // execFileSync avoids shell injection — symbol is passed as an argument, not interpolated
        const out = execFileSync(
          "grep",
          [
            "-r",
            "-n",
            "--include=*.ts",
            "--include=*.tsx",
            "--include=*.js",
            "--include=*.jsx",
            "--include=*.py",
            "--include=*.go",
            "--include=*.rs",
            "--include=*.rb",
            "--include=*.java",
            "--include=*.kt",
            symbol,
            this.workspace,
          ],
          { encoding: "utf8", timeout: 5000, maxBuffer: 256 * 1024 },
        );
        const lines = out
          .split("\n")
          .filter(Boolean)
          .slice(0, MAX_GREP_RESULTS);
        // Strip workspace prefix so paths are repo-relative
        return (
          lines.map((l) => l.replace(this.workspace! + "/", "")).join("\n") ||
          "[not found]"
        );
      } catch {
        return "[not found]";
      }
    }

    const results = await searchCode(this.pr, symbol);
    return results.join("\n") || "[not found]";
  }
}

export function createMcpServer(
  pr: PRContext,
  maxToolCalls: number,
): ToolServer {
  console.log("Initializing tool server...");
  const ws =
    process.env.GITEA_WORKSPACE ?? process.env.GITHUB_WORKSPACE ?? null;
  const workspace = ws && existsSync(ws) ? ws : null;
  return new ToolServer(pr, workspace, maxToolCalls);
}