file-agent-plus

import { tool, type ToolsProvider } from "@lmstudio/sdk";
import { z } from "zod";
import * as fs from "fs";
import * as path from "path";
import { listFiles, formatFilesList } from "./list-files";
import { regexSearchFiles } from "./ripgrep";
import { globalConfigSchematics } from "./config";

/**
 * Utility function to resolve paths relative to the plugin's base directory.
 */
function resolvePath(targetPath: string, baseDirectory: string): string {
  const resolved = path.resolve(baseDirectory || ".", targetPath);
  if (baseDirectory) {
    const base = path.resolve(baseDirectory);
    // Basic check to ensure the resolved path is within or equal to the base directory
    if (!resolved.startsWith(base + path.sep) && resolved !== base) {
      throw new Error(`Access denied: "${resolved}" is outside the allowed base directory "${base}".`);
    }
  }
  return resolved;
}

export const toolsProvider: ToolsProvider = async (ctl) => {
  const globalConfig = ctl.getGlobalPluginConfig(globalConfigSchematics);
  // Ensure baseDirectory exists and provide a fallback if necessary
  const baseDirectory: string = globalConfig.get("baseDirectory") ?? "";

  /* ----------------- Tool Definitions ----------------- */

  // FILE READ/WRITE TOOLS 
  const readFileTool = tool({
    name: "read_file",
    description: "Read the full contents of a file.",
    parameters: { path: z.string().describe("Path to the file to read.") }, // <-- Fixed comma placement
    implementation: async ({ path: filePath }, ctx) => {
      ctx.status("Reading file...");
      try { return fs.readFileSync(resolvePath(filePath, baseDirectory), "utf-8"); }
      catch (err: any) { return `Error: ${err.message}`; }
    }
  });

  const writeFileTool = tool({
    name: "write_file",
    description: "Write (or overwrite) a file with the given content. Creates parent directories if needed.",
    parameters: { // <-- Fixed comma placement
      path: z.string().describe("Path to the file to write."),
      content: z.string().describe("Content to write into the file.") 
    },
    implementation: async ({ path: filePath, content }, ctx) => {
      ctx.status("Writing file...");
      try {
        const resolved = resolvePath(filePath, baseDirectory);
        fs.mkdirSync(path.dirname(resolved), { recursive: true });
        fs.writeFileSync(resolved, content, "utf-8");
        return `File written: ${resolved}`;
      } catch (err: any) { return `Error: ${err.message}`; }
    }
  });

  // 🧪 IMPROVED: Global Replacement Tool
  const applyDiffTool = tool({ // <-- Critical comma fix here and below
    name: "apply_diff",
    description: "Apply a targeted search-and-replace edit to a file by replacing ALL instances of the old text with the new text.",
    parameters: { 
      path: z.string().describe("Path to the file to edit."),
      oldText: z.string().describe("Exact text to search for."), // <-- Comma added
      newText: z.string().describe("Text to replace oldText with.")  // <-- Last property, no comma needed
    },
    implementation: async ({ path: filePath, oldText, newText }, ctx) => {
      ctx.status("Applying diff...");
      try {
        const resolved = resolvePath(filePath, baseDirectory);
        const original = fs.readFileSync(resolved, "utf-8");

        // Escape special regex characters in the user-provided text
        const escapedOldText = oldText.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
        
        // Global replacement using RegExp constructor
        const regex = new RegExp(escapedOldText, 'g');
        const newContent = original.replace(regex, newText);

        if (original === newContent) {
            return `Success: The text "${oldText}" was not found in the file, or no change was necessary.`;
        }

        fs.writeFileSync(resolved, newContent, "utf-8");
        return `Diff applied successfully to ${resolved}. All instances of "${oldText}" were replaced with "${newText}".`;
      } catch (err: any) { return `Error: ${err.message}`; }
    }
  });

  const replaceInFileTool = tool({ // <-- Fixed comma placement and parameter definition
    name: "replace_in_file",
    description: "Make targeted edits using SEARCH/REPLACE blocks. Format: ------- SEARCH / [content] / ======= / [replacement] / +++++++ REPLACE",
    parameters: { 
      path: z.string().describe("Path to the file to modify."),
      diff: z.string().describe("One or more SEARCH/REPLACE blocks.") // <-- Fixed comma placement
    },
    implementation: async ({ path: filePath, diff }, ctx) => {
      ctx.status("Applying replace...");
      try {
        const resolved = resolvePath(filePath, baseDirectory);
        let fileContent = fs.readFileSync(resolved, "utf-8");

        // Normalise the content (CRLF and CR to LF) for consistent regex handling
        const normalize = (s: string) => s.replace(/\r\n/g, "\n").replace(/\r/g, "\n");
        let originalContent = fileContent;
        fileContent = normalize(fileContent);

        // Regex to capture blocks enclosed by the delimiters
        const blockRegex = /(?:-{3,}|<{3,})\s*SEARCH(?:\s*\/\s*|\s+)([\s\S]*?)(?:\s*\/\s*|\n)={3,}(?:\s*\/\s*|\n)([\s\S]*?)(?:\s*\/\s*|\n)\+{3,}\s*REPLACE/g;
        const matches = [...diff.matchAll(blockRegex)];

        if (matches.length === 0) {
             // Fallback for single-line inputs that might use slashes as separators
             const slashNormalized = diff.replace(/\s*\/\s*/g, "\n");
             const fallbackMatches = [...slashNormalized.matchAll(blockRegex)];
             if (fallbackMatches.length > 0) {
                 // If we found blocks via fallback, re-evaluate the matches list with these.
                  matches.push(...fallbackMatches);
             } else {
                  return "Error: No valid SEARCH/REPLACE blocks were found using the specified delimiters.";
             }
        }

        let tempContent = fileContent;
        for (const match of matches) {
            const searchText = (match[1] ?? "").trim();
            const replaceText = (match[2] ?? "").trim();
            if (!searchText || !replaceText) continue;

             // Use RegExp for replacement, ensuring global search ('g')
            const regex = new RegExp(escapeRegex(searchText), 'g');
            tempContent = tempContent.replace(regex, replaceText);
        }


        let finalContent: string;
        // Check if original content had Windows line endings to restore them for writing
        if (originalContent.includes("\r\n") && fileContent !== originalContent) {
             finalContent = tempContent.replace(/\n/g, "\r\n"); 
        } else {
            finalContent = tempContent;
        }

        fs.writeFileSync(resolved, finalContent, "utf-8");
        return "Replace applied successfully.";
      } catch (err: any) { return `Error during replacement: ${err.message}`; }
    }
  });

  // 🟢 IMPROVED: Robust appending with better newline detection.
  const appendToFileTool = tool({ // <-- Fixed comma placement and parameter definition
    name: "append_file",
    description: "Append content to the end of a file. Creates the file if it does not exist.",
    parameters: { 
      path: z.string().describe("Path to the file."),
      content: z.string().describe("Content to append."), // <-- Added comma
      add_newline: z.boolean().optional().default(true).describe("Ensure content starts on a new line.") // <-- Last property, no comma needed
    },
    implementation: async ({ path: filePath, content, add_newline }, ctx) => {
      ctx.status("Appending...");
      const resolved = resolvePath(filePath, baseDirectory);
      let toAppend = content;

      try {
        // 1. Check for file existence and read existing content safely
        let existingContent: string = "";
        if (fs.existsSync(resolved)) {
          existingContent = fs.readFileSync(resolved, "utf-8");
        }

        // 2. Determine if a leading newline is needed
        let needsNewline = add_newline;
        if (needsNewline && existingContent.length > 0) {
            const endsWithWhitespace = /[ \t\r\n]/.test(existingContent);
            // Prepend a newline only if the content isn't empty AND doesn't already end with whitespace/a newline character
            if (!endsWithWhitespace) {
                toAppend = "\n" + content; 
            }
        } else if (!needsNewline && existingContent.length > 0) {
             // If user explicitly disables newlines, but the file has content, we still append cleanly.
             toAppend = content;
        }

        fs.appendFileSync(resolved, toAppend, "utf-8");
        return `Content appended successfully to ${resolved}.`;
      } catch (err: any) { return `Error during file append operation: ${err.message}`; }
    }
  });


  // Standard/Unchanged Tools below this line...

  const copyFileTool = tool({
    name: "copy_file",
    description: "Copy a file from one location to another. Creates parent directories if needed.",
    parameters: { 
      source: z.string().describe("Path of the source file."), // <-- Fixed comma placement
      destination: z.string().describe("Path of the destination file."), // <-- Last property, no comma needed
      overwrite: z.boolean().optional().default(false).describe("Overwrite if exists. Default: false.") // <-- Corrected structure
    },
    implementation: async ({ source, destination, overwrite }, ctx) => {
      ctx.status("Copying file...");
      try {
        const src = resolvePath(source, baseDirectory);
        const dst = resolvePath(destination, baseDirectory);
        if (!fs.existsSync(src)) return `Error: source not found: ${src}`;
        // Check for existence and overwrite flag before copying
        if (!overwrite && fs.existsSync(dst)) return "Error: destination already exists. Use overwrite: true.";

        fs.mkdirSync(path.dirname(dst), { recursive: true });
        fs.copyFileSync(src, dst);
        return `File copied: ${src} -> ${dst}`;
      } catch (err: any) { return `Error: ${err.message}`; }
    }
  });

  const fileExistsTool = tool({
    name: "file_exists",
    description: "Check whether a file or directory exists. Returns existence status and type.",
    parameters: { path: z.string().describe("Path to check.") }, // <-- Fixed comma placement, added trailing comma for consistency
    implementation: async ({ path: filePath }, ctx) => {
      ctx.status("Checking...");
      try {
        const resolved = resolvePath(filePath, baseDirectory);
        if (!fs.existsSync(resolved)) return JSON.stringify({ exists: false, type: null, path: resolved });
        const s = fs.statSync(resolved);
        return JSON.stringify({ 
            exists: true, 
            type: s.isDirectory() ? "directory" : (s.isFile() ? "file" : "other"), 
            path: resolved 
        });
      } catch (err: any) { return `Error: ${err.message}`; }
    }
  });

  const createDirectoryTool = tool({
    name: "create_directory",
    description: "Create a directory (and all missing parent directories) at the given path.",
    parameters: { path: z.string().describe("Path of the directory to create.") }, // <-- Fixed comma placement
    implementation: async ({ path: dirPath }, ctx) => {
      ctx.status("Creating directory...");
      try {
        const resolved = resolvePath(dirPath, baseDirectory);
        fs.mkdirSync(resolved, { recursive: true });
        return `Directory created: ${resolved}`;
      } catch (err: any) { return `Error: ${err.message}`; }
    }
  });

  const listFilesTool = tool({
    name: "list_files",
    description: "List files and directories at a given path. Set recursive=true to include subdirectories.",
    parameters: { 
      path: z.string().describe("Directory to list."),
      recursive: z.boolean().optional().default(false).describe("Whether to list files recursively.") // <-- Fixed comma placement
    },
    implementation: async ({ path: dirPath, recursive }, ctx) => {
      ctx.status("Listing files...");
      try {
        const resolved = resolvePath(dirPath, baseDirectory);
        const [files, hitLimit] = await listFiles(resolved, recursive ?? false, 500);
        return formatFilesList(resolved, files, hitLimit);
      } catch (err: any) { return `Error: ${err.message}`; }
    }
  });

  const searchFilesTool = tool({
    name: "search_files",
    description: "Search for a regex pattern inside files using ripgrep.",
    parameters: { 
      path: z.string().describe("Directory to search in."),
      regex: z.string().describe("Regular expression to search for."), // <-- Fixed comma placement
      filePattern: z.string().optional().describe("Glob pattern to restrict which files are searched (e.g. '*.ts').") 
    },
    implementation: async ({ path: dirPath, regex, filePattern }, ctx) => {
      ctx.status("Searching files...");
      try {
        // Pass the current directory (__dirname) as cwd for searchFiles to correctly calculate paths
        const cwd = __dirname; 
        return await regexSearchFiles(cwd, resolvePath(dirPath, baseDirectory), regex, filePattern);
      } catch (err: any) { return `Error: ${err.message}`; }
    }
  });

  const deleteFileTool = tool({
    name: "delete_file",
    description: "Delete a file at the given path.",
    parameters: { path: z.string().describe("Path to the file to delete.") }, // <-- Fixed comma placement
    implementation: async ({ path: filePath }, ctx) => {
      ctx.status("Deleting file...");
      try {
        const resolved = resolvePath(filePath, baseDirectory);
        fs.rmSync(resolved, { force: true });
        return "File deleted.";
      } catch (err: any) { return `Error: ${err.message}`; }
    }
  });

  const moveFileTool = tool({
    name: "move_file",
    description: "Move or rename a file or directory.",
    parameters: { 
      source: z.string().describe("Current path of the file or directory."),
      destination: z.string().describe("New path (target location)") // <-- Fixed comma placement
    },
    implementation: async ({ source, destination }, ctx) => {
      ctx.status("Moving file...");
      try {
        const src = resolvePath(source, baseDirectory);
        const dst = resolvePath(destination, baseDirectory);

        // Ensure parent directory of destination exists before moving
        fs.mkdirSync(path.dirname(dst), { recursive: true }); 
        
        fs.renameSync(src, dst);
        return `Moved "${src}" -> "${dst}"`;
      } catch (err: any) { return `Error: ${err.message}`; }
    }
  });

  const fileInfoTool = tool({
    name: "file_info",
    description: "Get metadata about a file or directory (size, type, dates).",
    parameters: { path: z.string().describe("Path to inspect.") }, // <-- Fixed comma placement
    implementation: async ({ path: filePath }, ctx) => {
      ctx.status("Getting file info...");
      try {
        const resolved = resolvePath(filePath, baseDirectory);
        const stat = fs.statSync(resolved);
        return JSON.stringify({ 
            path: resolved, 
            type: stat.isDirectory() ? "directory" : (stat.isFile() ? "file" : "other"), 
            sizeBytes: stat.size, 
            created: stat.birthtime.toISOString(), 
            modified: stat.mtime.toISOString(), 
            accessed: stat.atime.toISOString() 
        }, null, 2);
      } catch (err: any) { return `Error: ${err.message}`; }
    }
  });

  return [
    readFileTool, writeFileTool, applyDiffTool, // <-- UPDATED/IMPROVED
    replaceInFileTool, appendToFileTool, copyFileTool, fileExistsTool, // <-- UPDATED/IMPROVED
    createDirectoryTool, listFilesTool, searchFilesTool, deleteFileTool, moveFileTool, fileInfoTool,
  ];
};

/** Helper function to escape special characters for use in RegExp constructor */
function escapeRegex(string: string): string {
    return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
}