src / toolsProvider.ts
import { tool, Tool, ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import { configSchematics } from "./config";
import { existsSync, statSync, createReadStream } from "fs";
import { writeFile, readFile, mkdir, stat } from "fs/promises";
import { regexSearchFiles } from "./ripgrep";
import { listFiles, formatFilesList } from "./list-files";
import * as readline from "readline";
import * as path from "path";
// --- Constants ---
const MAX_EDITABLE_SIZE_BYTES = 10 * 1024 * 1024; // 10MB
const MAX_READ_SIZE_BYTES = 50 * 1024 * 1024; // 50MB
const TRUNCATION_MESSAGE = "\n...[Content truncated due to size. Use start_line/end_line to read specific sections]...";
// --- Helper: Path Resolver ---
function resolveTargetFile(userPath: string, configRoot: string): string {
const root = configRoot && configRoot.trim() !== "" ? path.resolve(configRoot) : process.cwd();
return path.resolve(root, userPath);
}
// --- Search/Replace Helpers ---
const SEARCH_BLOCK_START_REGEX = /^([-]{3,}|[<]{3,})\s?SEARCH\s*>?$/;
const SEARCH_BLOCK_END_REGEX = /^[=]{3,}$/;
const REPLACE_BLOCK_END_REGEX = /^([+]{3,}|[>]{3,})\s?REPLACE\s*>?$/;
const SEARCH_BLOCK_START = "------- SEARCH";
const SEARCH_BLOCK_END = "=======";
const REPLACE_BLOCK_END = "+++++++ REPLACE";
function isSearchBlockStart(line: string): boolean {
return SEARCH_BLOCK_START_REGEX.test(line.trim());
}
function isSearchBlockEnd(line: string): boolean {
return SEARCH_BLOCK_END_REGEX.test(line.trim());
}
function isReplaceBlockEnd(line: string): boolean {
return REPLACE_BLOCK_END_REGEX.test(line.trim());
}
export function fixModelHtmlEscaping(text: string): string {
if (!text) return "";
const map: { [key: string]: string } = { ">": ">", "<": "<", """: '"', "&": "&", "'": "'" };
return text.replace(/>|<|"|&|'/g, (m) => map[m]);
}
export function removeInvalidChars(text: string): string {
return text.replace(/\uFFFD/g, "");
}
// Cleans up the input diff (removes markdown fences, fixes HTML entities)
function sanitizeDiffInput(diff: string): string {
let clean = diff;
// 1. Remove wrapping markdown code blocks
if (clean.trimStart().startsWith("```")) {
const firstLineEnd = clean.indexOf("\n");
if (firstLineEnd !== -1) {
clean = clean.substring(firstLineEnd + 1);
}
}
clean = clean.trimEnd();
if (clean.endsWith("```")) {
clean = clean.substring(0, clean.lastIndexOf("```"));
}
// 2. Fix HTML escaping
clean = fixModelHtmlEscaping(clean);
return clean;
}
async function readLines(filePath: string, startLine: number, endLine: number): Promise<string> {
const fileStream = createReadStream(filePath);
const rl = readline.createInterface({ input: fileStream, crlfDelay: Infinity });
const lines: string[] = [];
let currentLine = 1;
for await (const line of rl) {
if (currentLine >= startLine && currentLine <= endLine) lines.push(line);
if (currentLine > endLine) { rl.close(); break; }
currentLine++;
}
return lines.join("\n");
}
// --- Matching Logic ---
// Helper: Collapse whitespace for fuzzy matching (e.g. "fn test ( )" == "fn test ()")
function normalizeLine(line: string): string {
return line.trim().replace(/\s+/g, ' ');
}
// Strategy 1: Fuzzy Line Match (ignores indentation and multiple spaces)
function lineFuzzyFallbackMatch(originalContent: string, searchContent: string, startIndex: number): [number, number] | false {
const originalLines = originalContent.split("\n");
const searchLines = searchContent.split("\n");
// Remove trailing empty line if present
if (searchLines.length > 0 && searchLines[searchLines.length - 1] === "") searchLines.pop();
if (searchLines.length === 0) return false;
let startLineNum = 0;
let currentIndex = 0;
// Fast forward to startIndex
while (currentIndex < startIndex && startLineNum < originalLines.length) {
currentIndex += originalLines[startLineNum].length + 1; // +1 for newline
startLineNum++;
}
// Iterate through original file lines looking for a block match
for (let i = startLineNum; i <= originalLines.length - searchLines.length; i++) {
let matches = true;
for (let j = 0; j < searchLines.length; j++) {
if (normalizeLine(originalLines[i + j]) !== normalizeLine(searchLines[j])) {
matches = false;
break;
}
}
if (matches) {
// Calculate start index
let matchStartIndex = 0;
for (let k = 0; k < i; k++) matchStartIndex += originalLines[k].length + 1;
// Calculate end index
let matchEndIndex = matchStartIndex;
for (let k = 0; k < searchLines.length; k++) matchEndIndex += originalLines[i + k].length + 1;
// Safety: Subtract 1 from end index to exclude the last newline if it wasn't in searchContent
// (This is a simplification, but generally safe for block replacement)
return [matchStartIndex, matchEndIndex];
}
}
return false;
}
// Strategy 2: Anchored Match (Match first 3 lines and last 3 lines, ignore middle)
// Useful for large blocks where the AI hallucinates a line in the middle.
function blockAnchorExpandedMatch(originalContent: string, searchContent: string, startIndex: number): [number, number] | false {
const originalLines = originalContent.split("\n");
const searchLines = searchContent.split("\n");
if (searchLines.length > 0 && searchLines[searchLines.length - 1] === "") searchLines.pop();
if (searchLines.length < 5) return false; // Only use for blocks > 4 lines
const ANCHOR_SIZE = 2; // Match first 2 and last 2 lines
let startLineNum = 0;
let currentIndex = 0;
while (currentIndex < startIndex && startLineNum < originalLines.length) {
currentIndex += originalLines[startLineNum].length + 1;
startLineNum++;
}
for (let i = startLineNum; i <= originalLines.length - searchLines.length; i++) {
// Check Start Anchors
let startMatches = true;
for (let j = 0; j < ANCHOR_SIZE; j++) {
if (normalizeLine(originalLines[i + j]) !== normalizeLine(searchLines[j])) { startMatches = false; break; }
}
if (!startMatches) continue;
// Check End Anchors
let endMatches = true;
for (let j = 0; j < ANCHOR_SIZE; j++) {
const oIdx = i + searchLines.length - 1 - j;
const sIdx = searchLines.length - 1 - j;
if (normalizeLine(originalLines[oIdx]) !== normalizeLine(searchLines[sIdx])) { endMatches = false; break; }
}
if (endMatches) {
let matchStartIndex = 0;
for (let k = 0; k < i; k++) matchStartIndex += originalLines[k].length + 1;
let matchEndIndex = matchStartIndex;
for (let k = 0; k < searchLines.length; k++) matchEndIndex += originalLines[i + k].length + 1;
return [matchStartIndex, matchEndIndex];
}
}
return false;
}
enum ProcessingState { Idle = 0, StateSearch = 1 << 0, StateReplace = 1 << 1 }
class NewFileContentConstructor {
private originalContent: string;
private isFinal: boolean;
private state: number;
private pendingNonStandardLines: string[] = [];
private result: string = "";
private lastProcessedIndex: number = 0;
private currentSearchContent: string = "";
private searchMatchIndex: number = -1;
private searchEndIndex: number = -1;
constructor(originalContent: string, isFinal: boolean) {
this.originalContent = originalContent;
this.isFinal = isFinal;
this.state = ProcessingState.Idle;
}
private resetForNextBlock() {
this.state = ProcessingState.Idle;
this.currentSearchContent = "";
this.searchMatchIndex = -1;
this.searchEndIndex = -1;
}
private findLastMatchingLineIndex(regx: RegExp, lineLimit: number) {
for (let i = lineLimit; i > 0; ) {
i--;
if (this.pendingNonStandardLines[i].match(regx)) return i;
}
return -1;
}
private updateProcessingState(newState: ProcessingState) {
this.state |= newState;
}
private isSearchingActive(): boolean { return (this.state & ProcessingState.StateSearch) === ProcessingState.StateSearch; }
private isReplacingActive(): boolean { return (this.state & ProcessingState.StateReplace) === ProcessingState.StateReplace; }
private hasPendingNonStandardLines(limit: number): boolean { return this.pendingNonStandardLines.length - limit < this.pendingNonStandardLines.length; }
public processLine(line: string) { this.internalProcessLine(line, true, this.pendingNonStandardLines.length); }
public getResult() {
if (this.isFinal && this.lastProcessedIndex < this.originalContent.length) {
this.result += this.originalContent.slice(this.lastProcessedIndex);
}
if (this.isFinal && this.state !== ProcessingState.Idle) {
throw new Error("File processing incomplete - SEARCH/REPLACE operations still active. Did you forget a closing marker?");
}
return this.result;
}
private internalProcessLine(line: string, canWrite: boolean, limit: number): number {
let removeLineCount = 0;
if (isSearchBlockStart(line)) {
removeLineCount = this.trimPendingNonStandardTrailingEmptyLines(limit);
if (removeLineCount > 0) limit -= removeLineCount;
if (this.hasPendingNonStandardLines(limit)) {
this.tryFixSearchReplaceBlock(limit);
canWrite && (this.pendingNonStandardLines.length = 0);
}
this.updateProcessingState(ProcessingState.StateSearch);
this.currentSearchContent = "";
} else if (isSearchBlockEnd(line)) {
if (!this.isSearchingActive()) {
this.tryFixSearchBlock(limit);
canWrite && (this.pendingNonStandardLines.length = 0);
}
this.updateProcessingState(ProcessingState.StateReplace);
this.beforeReplace();
} else if (isReplaceBlockEnd(line)) {
if (!this.isReplacingActive()) {
this.tryFixReplaceBlock(limit);
canWrite && (this.pendingNonStandardLines.length = 0);
}
this.lastProcessedIndex = this.searchEndIndex;
this.resetForNextBlock();
} else {
if (this.isReplacingActive()) {
if (this.searchMatchIndex !== -1) this.result += line + "\n";
} else if (this.isSearchingActive()) {
this.currentSearchContent += line + "\n";
} else {
if (canWrite) this.pendingNonStandardLines.push(line);
}
}
return removeLineCount;
}
private beforeReplace() {
if (!this.currentSearchContent) {
if (this.originalContent.length === 0) { this.searchMatchIndex = 0; this.searchEndIndex = 0; }
else { this.searchMatchIndex = 0; this.searchEndIndex = this.originalContent.length; }
} else {
// 1. Try Exact Match
const exactIndex = this.originalContent.indexOf(this.currentSearchContent, this.lastProcessedIndex);
if (exactIndex !== -1) {
this.searchMatchIndex = exactIndex;
this.searchEndIndex = exactIndex + this.currentSearchContent.length;
} else {
// 2. Try Fuzzy Line Match (Whitespace tolerant)
const fuzzyMatch = lineFuzzyFallbackMatch(this.originalContent, this.currentSearchContent, this.lastProcessedIndex);
if (fuzzyMatch) {
[this.searchMatchIndex, this.searchEndIndex] = fuzzyMatch;
} else {
// 3. Try Expanded Anchor Match (Match ends, ignore middle errors)
const anchorMatch = blockAnchorExpandedMatch(this.originalContent, this.currentSearchContent, this.lastProcessedIndex);
if (anchorMatch) {
[this.searchMatchIndex, this.searchEndIndex] = anchorMatch;
} else {
// 4. Construct a helpful error message
const searchLines = this.currentSearchContent.split('\n');
const firstLine = searchLines[0] || "";
throw new Error(`The SEARCH block could not be found.
Matching failed for block starting with:
"${firstLine.trim().substring(0, 50)}..."
Potential causes:
1. The AI hallucinated incorrect whitespace/indentation (try copying the file content exactly).
2. The file content changed since it was read.
3. The block contains comments/formatting that don't match exactly.
`);
}
}
}
}
// Safety check
if (this.searchMatchIndex < this.lastProcessedIndex) {
throw new Error(`The SEARCH block matched content that was already processed (lines ${this.lastProcessedIndex} to ${this.searchMatchIndex}). Duplicate block?`);
}
this.result += this.originalContent.slice(this.lastProcessedIndex, this.searchMatchIndex);
}
private tryFixSearchBlock(limit: number): number {
let removeLineCount = 0;
const idx = this.findLastMatchingLineIndex(SEARCH_BLOCK_START_REGEX, limit);
if (idx !== -1) {
const fixLines = this.pendingNonStandardLines.slice(idx, limit);
fixLines[0] = SEARCH_BLOCK_START;
for (const line of fixLines) removeLineCount += this.internalProcessLine(line, false, idx);
} else throw new Error(`Invalid REPLACE marker detected without a matching SEARCH start.`);
return removeLineCount;
}
private tryFixReplaceBlock(limit: number): number {
let removeLineCount = 0;
const idx = this.findLastMatchingLineIndex(SEARCH_BLOCK_END_REGEX, limit);
if (idx !== -1) {
const fixLines = this.pendingNonStandardLines.slice(idx - removeLineCount, limit - removeLineCount);
fixLines[0] = SEARCH_BLOCK_END;
for (const line of fixLines) removeLineCount += this.internalProcessLine(line, false, idx - removeLineCount);
} else throw new Error(`Malformed REPLACE block. Missing '======='.`);
return removeLineCount;
}
private tryFixSearchReplaceBlock(limit: number): number {
let removeLineCount = 0;
const idx = this.findLastMatchingLineIndex(REPLACE_BLOCK_END_REGEX, limit);
if (idx === limit - 1) {
const fixLines = this.pendingNonStandardLines.slice(idx - removeLineCount, limit - removeLineCount);
fixLines[fixLines.length - 1] = REPLACE_BLOCK_END;
for (const line of fixLines) removeLineCount += this.internalProcessLine(line, false, idx - removeLineCount);
} else throw new Error("Malformed SEARCH/REPLACE block structure");
return removeLineCount;
}
private trimPendingNonStandardTrailingEmptyLines(limit: number): number {
let removed = 0;
let i = Math.min(limit, this.pendingNonStandardLines.length) - 1;
while (i >= 0 && this.pendingNonStandardLines[i].trim() === "") { this.pendingNonStandardLines.pop(); removed++; i--; }
return removed;
}
}
async function constructNewFileContentV2(diff: string, orig: string, isFinal: boolean): Promise<string> {
const constructor = new NewFileContentConstructor(orig, isFinal);
const lines = diff.split("\n");
// Cleanup closing fences that might look like markers
const lastLine = lines[lines.length - 1];
if (lines.length > 0 &&
(lastLine.startsWith("-") || lastLine.startsWith("<") || lastLine.startsWith("=") || lastLine.startsWith("+") || lastLine.startsWith(">")) &&
!isSearchBlockStart(lastLine) && !isSearchBlockEnd(lastLine) && !isReplaceBlockEnd(lastLine)) {
lines.pop();
}
for (const line of lines) constructor.processLine(line);
return constructor.getResult();
}
// --- TOOLS PROVIDER ---
export async function toolsProvider(ctl: ToolsProviderController) {
const config = ctl.getPluginConfig(configSchematics);
const tools: Tool[] = [];
const workspaceRoot = config.get("workspaceRoot");
// 1. REPLACE IN FILE
const replaceInFileTool = tool({
name: "replace_in_file",
description: `Request to replace sections of content in an existing file using SEARCH/REPLACE blocks.
IMPORTANT: You must use the following format:
<<<<<<< SEARCH
[exact content to find]
=======
[new content to replace with]
>>>>>>> REPLACE
Files are resolved relative to: ${workspaceRoot || "Current Process Directory"}`,
parameters: {
path: z.string().describe("The file path."),
diff: z.string().describe("The SEARCH/REPLACE blocks."),
},
implementation: async ({ path, diff }) => {
try {
const filePath = resolveTargetFile(path, workspaceRoot);
if (!existsSync(filePath)) return `Error: could not find file at ${filePath}`;
const stats = await stat(filePath);
if (stats.size > MAX_EDITABLE_SIZE_BYTES) return `Error: File too large to edit.`;
const originalContent = await readFile(filePath, "utf-8");
// 1. Sanitize (remove fences, fix HTML)
const sanitizedDiff = sanitizeDiffInput(diff);
// 2. Validate Format
const hasSearchMarker = /([-]{3,}|[<]{3,})\s?SEARCH/.test(sanitizedDiff);
const hasReplaceMarker = /([=]{3,})/.test(sanitizedDiff);
if (!hasSearchMarker || !hasReplaceMarker) {
return `Error: Invalid diff format. The tool refused your input because you used a standard Git Diff or missed the markers.
You MUST use the SEARCH/REPLACE block format shown below:
<<<<<<< SEARCH
[exact lines to remove]
=======
[new lines to insert]
>>>>>>> REPLACE`;
}
const newContent = await constructNewFileContentV2(sanitizedDiff, originalContent, true);
await writeFile(filePath, newContent, "utf-8");
return "Diff applied successfully.";
} catch (error) { return `Error applying diff: ${(error as Error).message}`; }
},
});
tools.push(replaceInFileTool);
// 2. WRITE FILE
const writeFileTool = tool({
name: "write_file",
description: `Request to write content to a file. Overwrites existing files.
Files are resolved relative to: ${workspaceRoot || "Current Process Directory"}`,
parameters: {
path: z.string().describe("The file path."),
content: z.string().describe("The complete content."),
},
implementation: async ({ path, content }) => {
try {
const filePath = resolveTargetFile(path, workspaceRoot);
const dirPath = filePath.substring(0, filePath.lastIndexOf(process.platform === 'win32' ? "\\" : "/"));
if (dirPath) await mkdir(dirPath, { recursive: true });
// Clean up content
let newContent = content;
if (newContent.trimStart().startsWith("```")) {
newContent = newContent.replace(/^```[^\n]*\n/, "");
}
if (newContent.trimEnd().endsWith("```")) {
newContent = newContent.substring(0, newContent.lastIndexOf("```"));
}
newContent = fixModelHtmlEscaping(newContent);
newContent = removeInvalidChars(newContent);
await writeFile(filePath, newContent, "utf-8");
return "File written successfully.";
} catch (error) { return `Error: ${(error as Error).message}`; }
},
});
tools.push(writeFileTool);
// 3. READ FILE
const readFileTool = tool({
name: "read_file",
description: `Request to read a file.
OPTIMIZATION: For files > ${(MAX_READ_SIZE_BYTES / 1024 / 1024).toFixed(0)}MB, use start_line/end_line.`,
parameters: {
path: z.string().describe("The file path."),
start_line: z.number().optional().describe("Start line (1-based)."),
end_line: z.number().optional().describe("End line (inclusive)."),
},
implementation: async ({ path, start_line, end_line }) => {
try {
const filePath = resolveTargetFile(path, workspaceRoot);
if (!existsSync(filePath)) return `Error: could not find file at ${filePath}`;
const stats = await stat(filePath);
if (start_line !== undefined && end_line !== undefined) {
if (start_line < 1 || end_line < start_line) return `Error: Invalid line range.`;
const MAX_LINES = 5000;
if (end_line - start_line > MAX_LINES) return `Error: Range too large. Limit ${MAX_LINES}.`;
const text = await readLines(filePath, start_line, end_line);
return { text: text, meta: `Showing lines ${start_line} to ${end_line}` };
}
if (stats.size > MAX_READ_SIZE_BYTES) return `Error: File too large. Use start_line/end_line.`;
const text = await readFile(filePath, "utf-8");
if (text.length > 500000) return { text: text.slice(0, 500000) + TRUNCATION_MESSAGE, warning: "Truncated." };
return { text };
} catch (error) { return `Error: ${(error as Error).message}`; }
},
});
tools.push(readFileTool);
// 4. SEARCH FILES
const searchFilesTool = tool({
name: "search_files",
description: `Request to perform a regex search across files.`,
parameters: {
path: z.string().describe("The directory path."),
regex: z.string().describe("Regex pattern (JavaScript syntax)."),
file_pattern: z.string().optional().describe("Glob pattern (e.g. '*.ts')."),
},
implementation: async ({ path: inputPath, regex, file_pattern }) => {
try {
const filePath = resolveTargetFile(inputPath, workspaceRoot);
if (!existsSync(filePath) || !statSync(filePath).isDirectory()) return `Error: '${filePath}' is not a valid directory`;
const results = await regexSearchFiles(workspaceRoot || process.cwd(), filePath, regex, file_pattern);
return results;
} catch (error) { return `Error: ${(error as Error).message}`; }
},
});
tools.push(searchFilesTool);
// 5. LIST FILES IN DIRECTORY (Replaces 'list_files')
const listFilesInDirectoryTool = tool({
name: "list_files_in_directory",
description: `Request to list files and directories in the specified path (not recursive).`,
parameters: {
path: z.string().describe("The directory path."),
},
implementation: async ({ path: inputPath }) => {
try {
const filePath = resolveTargetFile(inputPath, workspaceRoot);
if (!existsSync(filePath) || !statSync(filePath).isDirectory()) return `Error: '${filePath}' is not a valid directory`;
// Call listFiles with recursive=false
const [filePaths, wasLimited] = await listFiles(filePath, false, 2000);
return formatFilesList(filePath, filePaths, wasLimited);
} catch (error) { return `Error: ${(error as Error).message}`; }
},
});
tools.push(listFilesInDirectoryTool);
// 6. LIST FILES RECURSIVELY (New Tool)
const listFilesRecursivelyTool = tool({
name: "list_files_recursively",
description: `Request to list files and directories recursively starting from the specified path.`,
parameters: {
path: z.string().describe("The directory path."),
},
implementation: async ({ path: inputPath }) => {
try {
const filePath = resolveTargetFile(inputPath, workspaceRoot);
if (!existsSync(filePath) || !statSync(filePath).isDirectory()) return `Error: '${filePath}' is not a valid directory`;
// Call listFiles with recursive=true
const [filePaths, wasLimited] = await listFiles(filePath, true, 2000);
return formatFilesList(filePath, filePaths, wasLimited);
} catch (error) { return `Error: ${(error as Error).message}`; }
},
});
tools.push(listFilesRecursivelyTool);
return tools;
}