gemma-4-tools

import { tool, text, type Tool, type ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import { existsSync, statSync } from "fs";
import { readFile, writeFile, readdir, mkdir } from "fs/promises";
import { join, resolve, relative } from "path";
import { exec } from "child_process";
import { promisify } from "util";
import { isPathSafe, MAX_FILE_SIZE, MAX_WRITE_SIZE } from "./safety";

const execAsync = promisify(exec);

export async function toolsProvider(ctl: ToolsProviderController): Promise<Array<Tool>> {
  const tools = [];

  const readFileTool = tool({
    name: "read_file",
    description: text`
      Read the contents of a file at the specified path.
      Returns the file content as a string. Max 100KB.
    `,
    parameters: {
      path: z.string().describe("Absolute or relative path to the file"),
    },
    implementation: async ({ path }: { path: string }) => {
      try {
        const wd = ctl.getWorkingDirectory();
        const resolved = resolve(wd, path);
        const check = isPathSafe(resolved, wd);
        if (!check.safe) {
          return `Error: ${check.reason}`;
        }
        if (!existsSync(resolved)) {
          return `Error: File not found at ${resolved}`;
        }
        if (statSync(resolved).size > MAX_FILE_SIZE) {
          return `File is too large (>100KB). Reading first 100KB:\n\n${(await readFile(resolved, "utf-8")).slice(0, MAX_FILE_SIZE)}`;
        }
        return await readFile(resolved, "utf-8");
      } catch (err: any) {
        return `Error reading file: ${err.message}`;
      }
    },
  });

  const writeFileTool = tool({
    name: "write_file",
    description: text`
      Write or overwrite a file with the given content.
      Creates parent directories if they don't exist.
      Use this to update system prompts, configs, or code files.
    `,
    parameters: {
      path: z.string().describe("Path to the file (absolute or relative to working dir)"),
      content: z.string().describe("Full file content to write"),
    },
    implementation: async ({ path, content }: { path: string; content: string }) => {
      try {
        const wd = ctl.getWorkingDirectory();
        const resolved = resolve(wd, path);
        const check = isPathSafe(resolved, wd);
        if (!check.safe) {
          return `Error: ${check.reason}`;
        }
        if (content.length > MAX_WRITE_SIZE) {
          return `Error: Content too large (${content.length} bytes). Max 1MB.`;
        }
        const dir = resolved.substring(0, resolved.lastIndexOf("\\"));
        await mkdir(dir, { recursive: true });
        await writeFile(resolved, content, "utf-8");
        return `Successfully wrote ${content.length} bytes to ${resolved}`;
      } catch (err: any) {
        return `Error writing file: ${err.message}`;
      }
    },
  });

  const listDirTool = tool({
    name: "list_directory",
    description: text`
      List files and directories at the specified path.
      Returns names and types (file/directory).
    `,
    parameters: {
      path: z.string().describe("Path to list").default("."),
    },
    implementation: async ({ path }: { path: string }) => {
      try {
        const wd = ctl.getWorkingDirectory();
        const resolved = resolve(wd, path);
        const check = isPathSafe(resolved, wd);
        if (!check.safe) {
          return `Error: ${check.reason}`;
        }
        if (!existsSync(resolved)) {
          return `Error: Path not found at ${resolved}`;
        }
        const entries = await readdir(resolved, { withFileTypes: true });
        const lines = entries.map((e) => {
          const type = e.isDirectory() ? "dir" : "file";
          return `  [${type}] ${e.name}`;
        });
        return `Contents of ${resolved}:\n${lines.join("\n")}`;
      } catch (err: any) {
        return `Error listing directory: ${err.message}`;
      }
    },
  });

  const searchFilesTool = tool({
    name: "search_files",
    description: text`
      Find files matching an extension or name pattern recursively.
      Example extensions: .md, .json, .js, .ts, .py, .txt
    `,
    parameters: {
      extension: z.string().describe("File extension to search for (e.g. .md, .json)"),
      rootDir: z.string().describe("Root directory to start search from").default("."),
      maxResults: z.number().int().min(1).max(200).default(50),
    },
    implementation: async ({ extension, rootDir, maxResults }: { extension: string; rootDir: string; maxResults: number }) => {
      try {
        const wd = ctl.getWorkingDirectory();
        const resolved = resolve(wd, rootDir);
        const check = isPathSafe(resolved, wd);
        if (!check.safe) {
          return `Error: ${check.reason}`;
        }
        const results: string[] = [];
        const walk = async (dir: string) => {
          if (results.length >= maxResults) return;
          const entries = await readdir(dir, { withFileTypes: true });
          for (const entry of entries) {
            if (results.length >= maxResults) break;
            const fullPath = join(dir, entry.name);
            if (entry.isDirectory() && !entry.name.startsWith(".") && entry.name !== "node_modules") {
              await walk(fullPath);
            } else if (entry.isFile() && entry.name.endsWith(extension)) {
              results.push(fullPath);
            }
          }
        };
        await walk(resolved);
        return results.length === 0
          ? `No files with extension '${extension}' found under ${resolved}`
          : `Found ${results.length} files:\n${results.join("\n")}`;
      } catch (err: any) {
        return `Error searching files: ${err.message}`;
      }
    },
  });

  const grepFilesTool = tool({
    name: "grep_files",
    description: text`
      Search file contents using a regex pattern.
      Returns matching file paths, line numbers, and line content.
      Skips node_modules, .git, and binary files.
      Max 100 results.
    `,
    parameters: {
      pattern: z.string().describe("Regex pattern to search for (case-insensitive)"),
      rootDir: z.string().describe("Root directory to search in").default("."),
      maxResults: z.number().int().min(1).max(100).default(50),
      include: z.string().describe("Only search files matching this glob (e.g. *.ts, *.md)").optional(),
    },
    implementation: async ({ pattern, rootDir, maxResults, include }: { pattern: string; rootDir: string; maxResults: number; include?: string }) => {
      try {
        const wd = ctl.getWorkingDirectory();
        const resolved = resolve(wd, rootDir);
        const check = isPathSafe(resolved, wd);
        if (!check.safe) {
          return `Error: ${check.reason}`;
        }

        let cmd = `rg --no-heading --line-number --max-count ${maxResults} --smart-case`;
        if (include) {
          cmd += ` --glob "${include}"`;
        }
        cmd += ` --glob '!node_modules' --glob '!.git' --glob '!*.bin'`;
        cmd += ` -e "${pattern.replace(/"/g, '\\"')}" "${resolved}"`;

        const { stdout, stderr } = await execAsync(cmd, { timeout: 15000 });
        if (stderr && !stdout) {
          return `Error searching: ${stderr}`;
        }
        const lines = stdout.trim().split("\n").filter(Boolean);
        if (lines.length === 0) {
          return `No matches found for pattern '${pattern}' under ${resolved}`;
        }
        const truncated = lines.slice(0, maxResults);
        let result = `Found ${lines.length} matches for '${pattern}':\n`;
        result += truncated.join("\n");
        if (lines.length > maxResults) {
          result += `\n... and ${lines.length - maxResults} more matches`;
        }
        return result;
      } catch (err: any) {
        if (err.message?.includes("stdout maxBuffer")) {
          return `Too many results for pattern '${pattern}'. Try a more specific pattern.`;
        }
        if (err.message?.includes("not found") || err.message?.includes("not recognized")) {
          return "Error: ripgrep (rg) is not installed. Install it from https://github.com/BurntSushi/ripgrep";
        }
        return `Error searching: ${err.message}`;
      }
    },
  });

  const fetchUrlTool = tool({
    name: "fetch_url",
    description: text`
      Fetch content from a URL (HTTP/HTTPS only).
      Max 500KB response size. 15 second timeout.
      Blocks internal/private IPs for security.
    `,
    parameters: {
      url: z.string().url().describe("Full URL to fetch (must start with http:// or https://)"),
    },
    implementation: async ({ url }: { url: string }) => {
      try {
        const parsed = new URL(url);
        if (parsed.protocol !== "http:" && parsed.protocol !== "https:") {
          return "Error: Only http and https URLs are supported";
        }

        const hostname = parsed.hostname.toLowerCase();
        const isPrivate =
          hostname === "localhost" ||
          hostname === "127.0.0.1" ||
          hostname === "0.0.0.0" ||
          hostname === "[::1]" ||
          hostname.startsWith("10.") ||
          hostname.startsWith("192.168.") ||
          hostname.startsWith("172.16.") ||
          hostname.endsWith(".local") ||
          hostname.endsWith(".internal");

        if (isPrivate) {
          return "Error: Fetching from private/internal IPs is not allowed for security";
        }

        const controller = new AbortController();
        const timeout = setTimeout(() => controller.abort(), 15000);

        const response = await fetch(url, {
          signal: controller.signal,
          headers: {
            "User-Agent": "Gemma4ToolsPlugin/1.0",
            "Accept": "text/html,text/plain,application/json,*/*",
          },
        });
        clearTimeout(timeout);

        if (!response.ok) {
          return `Error: HTTP ${response.status} ${response.statusText}`;
        }

        const text = await response.text();
        const maxLen = 512000;
        if (text.length > maxLen) {
          return `Response too large (${text.length} bytes). Returning first 500KB:\n\n${text.slice(0, maxLen)}`;
        }

        const contentType = response.headers.get("content-type") ?? "unknown";
        return `Content-Type: ${contentType}\nURL: ${url}\n\n${text}`;
      } catch (err: any) {
        if (err.name === "AbortError") {
          return "Error: Request timed out after 15 seconds";
        }
        return `Error fetching URL: ${err.message}`;
      }
    },
  });

  const diffFilesTool = tool({
    name: "diff_files",
    description: text`
      Show line-by-line diff between two files.
      Returns unified diff format with added/removed lines.
      Useful for reviewing code changes before writing.
    `,
    parameters: {
      pathA: z.string().describe("First file path (absolute or relative)"),
      pathB: z.string().describe("Second file path (absolute or relative)"),
    },
    implementation: async ({ pathA, pathB }: { pathA: string; pathB: string }) => {
      try {
        const wd = ctl.getWorkingDirectory();
        const resolvedA = resolve(wd, pathA);
        const resolvedB = resolve(wd, pathB);

        const checkA = isPathSafe(resolvedA, wd);
        const checkB = isPathSafe(resolvedB, wd);
        if (!checkA.safe) return `Error: ${checkA.reason}`;
        if (!checkB.safe) return `Error: ${checkB.reason}`;

        if (!existsSync(resolvedA)) return `Error: File not found: ${resolvedA}`;
        if (!existsSync(resolvedB)) return `Error: File not found: ${resolvedB}`;

        const contentA = await readFile(resolvedA, "utf-8");
        const contentB = await readFile(resolvedB, "utf-8");

        if (contentA === contentB) {
          return "Files are identical";
        }

        const linesA = contentA.split("\n");
        const linesB = contentB.split("\n");

        const relA = relative(wd, resolvedA);
        const relB = relative(wd, resolvedB);

        let diff = `--- ${relA}\n+++ ${relB}\n`;
        let i = 0, j = 0;
        let contextLines = 3;

        while (i < linesA.length && j < linesB.length) {
          if (linesA[i] === linesB[j]) {
            i++;
            j++;
          } else {
            const startI = i;
            const startJ = j;

            while (i < linesA.length && j < linesB.length && linesA[i] !== linesB[j]) {
              i++;
              j++;
            }

            const lenA = i - startI;
            const lenB = j - startJ;

            diff += `@@ -${startI + 1},${lenA} +${startJ + 1},${lenB} @@\n`;
            for (let k = startI; k < i; k++) {
              diff += `-${linesA[k]}\n`;
            }
            for (let k = startJ; k < j; k++) {
              diff += `+${linesB[k]}\n`;
            }

            if (i < linesA.length) {
              const ctxEnd = Math.min(i + contextLines, linesA.length);
              for (let k = i; k < ctxEnd; k++) {
                diff += ` ${linesA[k]}\n`;
              }
              i = ctxEnd;
              j = i;
            }
          }
        }

        if (i < linesA.length) {
          diff += `@@ -${i + 1},${linesA.length - i} +${j + 1},0 @@\n`;
          for (let k = i; k < linesA.length; k++) {
            diff += `-${linesA[k]}\n`;
          }
        }
        if (j < linesB.length) {
          diff += `@@ -${i + 1},0 +${j + 1},${linesB.length - j} @@\n`;
          for (let k = j; k < linesB.length; k++) {
            diff += `+${linesB[k]}\n`;
          }
        }

        return diff;
      } catch (err: any) {
        return `Error diffing files: ${err.message}`;
      }
    },
  });

  tools.push(readFileTool, writeFileTool, listDirTool, searchFilesTool, grepFilesTool, fetchUrlTool, diffFilesTool);
  return tools;
}