diff --git a/src/core/tools/__tests__/readFileTool.spec.ts b/src/core/tools/__tests__/readFileTool.spec.ts index 44be1d3b92..58e3ead445 100644 --- a/src/core/tools/__tests__/readFileTool.spec.ts +++ b/src/core/tools/__tests__/readFileTool.spec.ts @@ -10,6 +10,7 @@ import { isBinaryFile } from "isbinaryfile" import { ReadFileToolUse, ToolParamName, ToolResponse } from "../../../shared/tools" import { readFileTool } from "../readFileTool" import { formatResponse } from "../../prompts/responses" +import { tiktoken } from "../../../utils/tiktoken" vi.mock("path", async () => { const originalPath = await vi.importActual("path") @@ -35,18 +36,27 @@ vi.mock("../../../integrations/misc/read-lines") let mockInputContent = "" // First create all the mocks -vi.mock("../../../integrations/misc/extract-text") +vi.mock("../../../integrations/misc/extract-text", () => ({ + extractTextFromFile: vi.fn(), + addLineNumbers: vi.fn(), + getSupportedBinaryFormats: vi.fn(() => [".pdf", ".docx", ".ipynb"]), +})) vi.mock("../../../services/tree-sitter") +vi.mock("../../../utils/tiktoken") + +// Import the mocked functions +import { addLineNumbers, getSupportedBinaryFormats } from "../../../integrations/misc/extract-text" // Then create the mock functions -const addLineNumbersMock = vi.fn().mockImplementation((text, startLine = 1) => { +const addLineNumbersMock = vi.mocked(addLineNumbers) +addLineNumbersMock.mockImplementation((text: string, startLine = 1) => { if (!text) return "" const lines = typeof text === "string" ? text.split("\n") : [text] - return lines.map((line, i) => `${startLine + i} | ${line}`).join("\n") + return lines.map((line: string, i: number) => `${startLine + i} | ${line}`).join("\n") }) -const extractTextFromFileMock = vi.fn() -const getSupportedBinaryFormatsMock = vi.fn(() => [".pdf", ".docx", ".ipynb"]) +const extractTextFromFileMock = vi.mocked(extractTextFromFile) +const getSupportedBinaryFormatsMock = vi.mocked(getSupportedBinaryFormats) vi.mock("../../ignore/RooIgnoreController", () => ({ RooIgnoreController: class { @@ -520,3 +530,317 @@ describe("read_file tool XML output structure", () => { }) }) }) + +describe("read_file tool with large file safeguard", () => { + // Test data + const testFilePath = "test/largefile.txt" + const absoluteFilePath = "/test/largefile.txt" + + // Mocked functions + const mockedCountFileLines = vi.mocked(countFileLines) + const mockedReadLines = vi.mocked(readLines) + const mockedExtractTextFromFile = vi.mocked(extractTextFromFile) + const mockedIsBinaryFile = vi.mocked(isBinaryFile) + const mockedPathResolve = vi.mocked(path.resolve) + const mockedTiktoken = vi.mocked(tiktoken) + + const mockCline: any = {} + let mockProvider: any + let toolResult: ToolResponse | undefined + + beforeEach(() => { + vi.clearAllMocks() + + mockedPathResolve.mockReturnValue(absoluteFilePath) + mockedIsBinaryFile.mockResolvedValue(false) + + mockProvider = { + getState: vi.fn(), + deref: vi.fn().mockReturnThis(), + } + + mockCline.cwd = "/" + mockCline.task = "Test" + mockCline.providerRef = mockProvider + mockCline.rooIgnoreController = { + validateAccess: vi.fn().mockReturnValue(true), + } + mockCline.say = vi.fn().mockResolvedValue(undefined) + mockCline.ask = vi.fn().mockResolvedValue({ response: "yesButtonClicked" }) + mockCline.fileContextTracker = { + trackFileContext: vi.fn().mockResolvedValue(undefined), + } + mockCline.recordToolUsage = vi.fn().mockReturnValue(undefined) + mockCline.recordToolError = vi.fn().mockReturnValue(undefined) + + toolResult = undefined + }) + + async function executeReadFileTool( + params: Partial = {}, + options: { + maxReadFileLine?: number + totalLines?: number + tokenCount?: number + } = {}, + ): Promise { + const maxReadFileLine = options.maxReadFileLine ?? -1 + const totalLines = options.totalLines ?? 5 + const tokenCount = options.tokenCount ?? 100 + + mockProvider.getState.mockResolvedValue({ maxReadFileLine }) + mockedCountFileLines.mockResolvedValue(totalLines) + mockedTiktoken.mockResolvedValue(tokenCount) + + const argsContent = `${testFilePath}` + + const toolUse: ReadFileToolUse = { + type: "tool_use", + name: "read_file", + params: { args: argsContent, ...params }, + partial: false, + } + + await readFileTool( + mockCline, + toolUse, + mockCline.ask, + vi.fn(), + (result: ToolResponse) => { + toolResult = result + }, + (_: ToolParamName, content?: string) => content ?? "", + ) + + return toolResult + } + + describe("when file has many lines and high token count", () => { + it("should apply safeguard and read only first 2000 lines", async () => { + // Setup - large file with high token count + const largeFileContent = Array(1500).fill("This is a line of text").join("\n") + const partialContent = Array(2000).fill("This is a line of text").join("\n") + + mockedExtractTextFromFile.mockResolvedValue(largeFileContent) + mockedReadLines.mockResolvedValue(partialContent) + + // Setup addLineNumbers mock for this test + addLineNumbersMock.mockImplementation((text: string) => { + const lines = text.split("\n") + return lines.map((line: string, i: number) => `${i + 1} | ${line}`).join("\n") + }) + + // Execute with high line count and token count + const result = await executeReadFileTool( + {}, + { + maxReadFileLine: -1, + totalLines: 1500, + tokenCount: 60000, // Above threshold + }, + ) + + // Verify safeguard was applied + expect(mockedTiktoken).toHaveBeenCalled() + expect(mockedReadLines).toHaveBeenCalledWith(absoluteFilePath, 1999, 0) + + // Verify the result contains the safeguard notice + expect(result).toContain("This file contains 1500 lines and approximately 60,000 tokens") + expect(result).toContain("Showing only the first 2000 lines to preserve context space") + expect(result).toContain(``) + }) + + it("should not apply safeguard when token count is below threshold", async () => { + // Setup - large file but with low token count + const fileContent = Array(1500).fill("Short").join("\n") + const numberedContent = fileContent + .split("\n") + .map((line, i) => `${i + 1} | ${line}`) + .join("\n") + + mockedExtractTextFromFile.mockImplementation(() => Promise.resolve(numberedContent)) + + // Execute with high line count but low token count + const result = await executeReadFileTool( + {}, + { + maxReadFileLine: -1, + totalLines: 1500, + tokenCount: 30000, // Below threshold + }, + ) + + // Verify safeguard was NOT applied + expect(mockedTiktoken).toHaveBeenCalled() + expect(mockedReadLines).not.toHaveBeenCalled() + expect(mockedExtractTextFromFile).toHaveBeenCalled() + + // Verify no safeguard notice + expect(result).not.toContain("preserve context space") + expect(result).toContain(``) + }) + + it("should not apply safeguard for files under 1000 lines", async () => { + // Setup - file with less than 1000 lines + const fileContent = Array(999).fill("This is a line of text").join("\n") + const numberedContent = fileContent + .split("\n") + .map((line, i) => `${i + 1} | ${line}`) + .join("\n") + + mockedExtractTextFromFile.mockImplementation(() => Promise.resolve(numberedContent)) + + // Execute + const result = await executeReadFileTool( + {}, + { + maxReadFileLine: -1, + totalLines: 999, + tokenCount: 100000, // Even with high token count + }, + ) + + // Verify tiktoken was NOT called (optimization) + expect(mockedTiktoken).not.toHaveBeenCalled() + expect(mockedReadLines).not.toHaveBeenCalled() + expect(mockedExtractTextFromFile).toHaveBeenCalled() + + // Verify no safeguard notice + expect(result).not.toContain("preserve context space") + expect(result).toContain(``) + }) + + it("should apply safeguard for very large files even if token counting fails", async () => { + // Setup - very large file and token counting fails + const partialContent = Array(2000).fill("This is a line of text").join("\n") + + mockedExtractTextFromFile.mockResolvedValue("Large content") + mockedReadLines.mockResolvedValue(partialContent) + + // Setup addLineNumbers mock for partial content + addLineNumbersMock.mockImplementation((text: string) => { + const lines = text.split("\n") + return lines.map((line: string, i: number) => `${i + 1} | ${line}`).join("\n") + }) + + // Set up the provider state + mockProvider.getState.mockResolvedValue({ maxReadFileLine: -1 }) + mockedCountFileLines.mockResolvedValue(6000) + + // IMPORTANT: Set up tiktoken to reject AFTER other mocks are set + mockedTiktoken.mockRejectedValue(new Error("Token counting failed")) + + const argsContent = `${testFilePath}` + + const toolUse: ReadFileToolUse = { + type: "tool_use", + name: "read_file", + params: { args: argsContent }, + partial: false, + } + + await readFileTool( + mockCline, + toolUse, + mockCline.ask, + vi.fn(), + (result: ToolResponse) => { + toolResult = result + }, + (_: ToolParamName, content?: string) => content ?? "", + ) + + // Verify safeguard was applied despite token counting failure + expect(mockedTiktoken).toHaveBeenCalled() + expect(mockedReadLines).toHaveBeenCalledWith(absoluteFilePath, 1999, 0) + + // Verify the result contains the safeguard notice (without token count) + expect(toolResult).toContain("This file contains 6000 lines") + expect(toolResult).toContain("Showing only the first 2000 lines to preserve context space") + expect(toolResult).toContain(``) + }) + + it("should not apply safeguard when maxReadFileLine is not -1", async () => { + // Setup + const fileContent = Array(2000).fill("This is a line of text").join("\n") + mockedExtractTextFromFile.mockResolvedValue(fileContent) + + // Execute with maxReadFileLine = 500 (not -1) + const result = await executeReadFileTool( + {}, + { + maxReadFileLine: 500, + totalLines: 2000, + tokenCount: 100000, + }, + ) + + // Verify tiktoken was NOT called + expect(mockedTiktoken).not.toHaveBeenCalled() + + // The normal maxReadFileLine logic should apply + expect(mockedReadLines).toHaveBeenCalled() + }) + + it("should handle line ranges correctly with safeguard", async () => { + // When line ranges are specified, safeguard should not apply + const rangeContent = "Line 100\nLine 101\nLine 102" + mockedReadLines.mockResolvedValue(rangeContent) + + const argsContent = `${testFilePath}100-102` + + const toolUse: ReadFileToolUse = { + type: "tool_use", + name: "read_file", + params: { args: argsContent }, + partial: false, + } + + mockProvider.getState.mockResolvedValue({ maxReadFileLine: -1 }) + mockedCountFileLines.mockResolvedValue(10000) + + await readFileTool( + mockCline, + toolUse, + mockCline.ask, + vi.fn(), + (result: ToolResponse) => { + toolResult = result + }, + (_: ToolParamName, content?: string) => content ?? "", + ) + + // Verify tiktoken was NOT called for range reads + expect(mockedTiktoken).not.toHaveBeenCalled() + expect(toolResult).toContain(``) + expect(toolResult).not.toContain("preserve context space") + }) + }) + + describe("safeguard thresholds", () => { + it("should use correct thresholds for line count and token count", async () => { + // Test boundary conditions + + // Just below line threshold - no token check + await executeReadFileTool({}, { totalLines: 1000, maxReadFileLine: -1 }) + expect(mockedTiktoken).not.toHaveBeenCalled() + + // Just above line threshold - token check performed + vi.clearAllMocks() + mockedExtractTextFromFile.mockResolvedValue("content") + await executeReadFileTool({}, { totalLines: 1001, maxReadFileLine: -1, tokenCount: 40000 }) + expect(mockedTiktoken).toHaveBeenCalled() + + // Token count just below threshold - no safeguard + expect(toolResult).not.toContain("preserve context space") + + // Token count just above threshold - safeguard applied + vi.clearAllMocks() + mockedExtractTextFromFile.mockResolvedValue("content") + mockedReadLines.mockResolvedValue("partial content") + await executeReadFileTool({}, { totalLines: 1001, maxReadFileLine: -1, tokenCount: 50001 }) + expect(mockedReadLines).toHaveBeenCalled() + expect(toolResult).toContain("preserve context space") + }) + }) +}) diff --git a/src/core/tools/readFileTool.ts b/src/core/tools/readFileTool.ts index 6de8dd5642..b9036351e0 100644 --- a/src/core/tools/readFileTool.ts +++ b/src/core/tools/readFileTool.ts @@ -14,6 +14,7 @@ import { readLines } from "../../integrations/misc/read-lines" import { extractTextFromFile, addLineNumbers, getSupportedBinaryFormats } from "../../integrations/misc/extract-text" import { parseSourceCodeDefinitionsForFile } from "../../services/tree-sitter" import { parseXml } from "../../utils/xml" +import { tiktoken } from "../../utils/tiktoken" export function getReadFileToolDescription(blockName: string, blockParams: any): string { // Handle both single path and multiple files via args @@ -516,13 +517,60 @@ export async function readFileTool( continue } - // Handle normal file read - const content = await extractTextFromFile(fullPath) - const lineRangeAttr = ` lines="1-${totalLines}"` + // Handle normal file read with safeguard for large files + // Define thresholds for the safeguard + const LARGE_FILE_LINE_THRESHOLD = 1000 // Consider files with more than 1000 lines as "large" + const MAX_TOKEN_THRESHOLD = 50000 // ~50% of a typical 100k context window + const FALLBACK_MAX_LINES = 2000 // Default number of lines to read when applying safeguard + + // Check if we should apply the safeguard + let shouldApplySafeguard = false + let safeguardNotice = "" + let linesToRead = totalLines + + if (maxReadFileLine === -1 && totalLines > LARGE_FILE_LINE_THRESHOLD) { + // File has many lines and we're trying to read the full file + // Perform token count check + try { + const fullContent = await extractTextFromFile(fullPath) + const tokenCount = await tiktoken([{ type: "text", text: fullContent }]) + + if (tokenCount > MAX_TOKEN_THRESHOLD) { + shouldApplySafeguard = true + linesToRead = FALLBACK_MAX_LINES + safeguardNotice = `This file contains ${totalLines} lines and approximately ${tokenCount.toLocaleString()} tokens, which could consume a significant portion of the context window. Showing only the first ${FALLBACK_MAX_LINES} lines to preserve context space. Use line_range if you need to read specific sections.\n` + } + } catch (error) { + // If token counting fails, apply safeguard based on line count alone + console.warn(`Failed to count tokens for large file ${relPath}:`, error) + if (totalLines > LARGE_FILE_LINE_THRESHOLD * 5) { + // For very large files (>5000 lines), apply safeguard anyway + shouldApplySafeguard = true + linesToRead = FALLBACK_MAX_LINES + safeguardNotice = `This file contains ${totalLines} lines, which could consume a significant portion of the context window. Showing only the first ${FALLBACK_MAX_LINES} lines to preserve context space. Use line_range if you need to read specific sections.\n` + } + } + } + + let content: string + let lineRangeAttr: string + + if (shouldApplySafeguard) { + // Read partial file with safeguard + content = addLineNumbers(await readLines(fullPath, linesToRead - 1, 0)) + lineRangeAttr = ` lines="1-${linesToRead}"` + } else { + // Read full file as normal + content = await extractTextFromFile(fullPath) + lineRangeAttr = ` lines="1-${totalLines}"` + } + let xmlInfo = totalLines > 0 ? `\n${content}\n` : `` if (totalLines === 0) { xmlInfo += `File is empty\n` + } else if (safeguardNotice) { + xmlInfo += safeguardNotice } // Track file read