Merge pull request #228 from RooVetGit/more_efficient_filetracker

More efficient workspace tracker

Author: mrubens

.changeset/early-icons-roll.md

@@ -0,0 +1,5 @@
+---
+"roo-cline": patch
+---
+
+More efficient workspace tracker

src/integrations/workspace/WorkspaceTracker.ts

@@ -4,12 +4,14 @@ import { listFiles } from "../../services/glob/list-files"
 import { ClineProvider } from "../../core/webview/ClineProvider"
 
 const cwd = vscode.workspace.workspaceFolders?.map((folder) => folder.uri.fsPath).at(0)
+const MAX_INITIAL_FILES = 1_000
 
 // Note: this is not a drop-in replacement for listFiles at the start of tasks, since that will be done for Desktops when there is no workspace selected
 class WorkspaceTracker {
 	private providerRef: WeakRef<ClineProvider>
 	private disposables: vscode.Disposable[] = []
 	private filePaths: Set<string> = new Set()
+	private updateTimer: NodeJS.Timeout | null = null
 
 	constructor(provider: ClineProvider) {
 		this.providerRef = new WeakRef(provider)
@@ -21,8 +23,8 @@ class WorkspaceTracker {
 		if (!cwd) {
 			return
 		}
-		const [files, _] = await listFiles(cwd, true, 1_000)
-		files.forEach((file) => this.filePaths.add(this.normalizeFilePath(file)))
+		const [files, _] = await listFiles(cwd, true, MAX_INITIAL_FILES)
+		files.slice(0, MAX_INITIAL_FILES).forEach((file) => this.filePaths.add(this.normalizeFilePath(file)))
 		this.workspaceDidUpdate()
 	}
 
@@ -49,16 +51,23 @@ class WorkspaceTracker {
 	}
 
 	private workspaceDidUpdate() {
-		if (!cwd) {
-			return
+		if (this.updateTimer) {
+			clearTimeout(this.updateTimer)
 		}
-		this.providerRef.deref()?.postMessageToWebview({
-			type: "workspaceUpdated",
-			filePaths: Array.from(this.filePaths).map((file) => {
-				const relativePath = path.relative(cwd, file).toPosix()
-				return file.endsWith("/") ? relativePath + "/" : relativePath
+
+		this.updateTimer = setTimeout(() => {
+			if (!cwd) {
+				return
+			}
+			this.providerRef.deref()?.postMessageToWebview({
+				type: "workspaceUpdated",
+				filePaths: Array.from(this.filePaths).map((file) => {
+					const relativePath = path.relative(cwd, file).toPosix()
+					return file.endsWith("/") ? relativePath + "/" : relativePath
+				})
 			})
-		})
+			this.updateTimer = null
+		}, 300) // Debounce for 300ms
 	}
 
 	private normalizeFilePath(filePath: string): string {
@@ -67,6 +76,11 @@ class WorkspaceTracker {
 	}
 
 	private async addFilePath(filePath: string): Promise<string> {
+		// Allow for some buffer to account for files being created/deleted during a task
+		if (this.filePaths.size >= MAX_INITIAL_FILES * 2) {
+			return filePath
+		}
+
 		const normalizedPath = this.normalizeFilePath(filePath)
 		try {
 			const stat = await vscode.workspace.fs.stat(vscode.Uri.file(normalizedPath))
@@ -87,6 +101,10 @@ class WorkspaceTracker {
 	}
 
 	public dispose() {
+		if (this.updateTimer) {
+			clearTimeout(this.updateTimer)
+			this.updateTimer = null
+		}
 		this.disposables.forEach((d) => d.dispose())
 	}
 }

src/integrations/workspace/__tests__/WorkspaceTracker.test.ts

@@ -38,9 +38,12 @@ describe("WorkspaceTracker", () => {
 
     beforeEach(() => {
         jest.clearAllMocks()
+        jest.useFakeTimers()
 
         // Create provider mock
-        mockProvider = { postMessageToWebview: jest.fn() } as any
+        mockProvider = {
+            postMessageToWebview: jest.fn().mockResolvedValue(undefined)
+        } as unknown as ClineProvider & { postMessageToWebview: jest.Mock }
 
         // Create tracker instance
         workspaceTracker = new WorkspaceTracker(mockProvider)
@@ -51,17 +54,20 @@ describe("WorkspaceTracker", () => {
         ;(listFiles as jest.Mock).mockResolvedValue(mockFiles)
 
         await workspaceTracker.initializeFilePaths()
+        jest.runAllTimers()
 
         expect(mockProvider.postMessageToWebview).toHaveBeenCalledWith({
             type: "workspaceUpdated",
-            filePaths: ["file1.ts", "file2.ts"]
+            filePaths: expect.arrayContaining(["file1.ts", "file2.ts"])
         })
+        expect((mockProvider.postMessageToWebview as jest.Mock).mock.calls[0][0].filePaths).toHaveLength(2)
     })
 
     it("should handle file creation events", async () => {
         // Get the creation callback and call it
         const [[callback]] = mockOnDidCreate.mock.calls
         await callback({ fsPath: "/test/workspace/newfile.ts" })
+        jest.runAllTimers()
 
         expect(mockProvider.postMessageToWebview).toHaveBeenCalledWith({
             type: "workspaceUpdated",
@@ -73,10 +79,12 @@ describe("WorkspaceTracker", () => {
         // First add a file
         const [[createCallback]] = mockOnDidCreate.mock.calls
         await createCallback({ fsPath: "/test/workspace/file.ts" })
+        jest.runAllTimers()
 
         // Then delete it
         const [[deleteCallback]] = mockOnDidDelete.mock.calls
         await deleteCallback({ fsPath: "/test/workspace/file.ts" })
+        jest.runAllTimers()
 
         // The last call should have empty filePaths
         expect(mockProvider.postMessageToWebview).toHaveBeenLastCalledWith({
@@ -91,15 +99,55 @@ describe("WorkspaceTracker", () => {
 
         const [[callback]] = mockOnDidCreate.mock.calls
         await callback({ fsPath: "/test/workspace/newdir" })
+        jest.runAllTimers()
 
         expect(mockProvider.postMessageToWebview).toHaveBeenCalledWith({
             type: "workspaceUpdated",
-            filePaths: ["newdir"]
+            filePaths: expect.arrayContaining(["newdir"])
         })
+        const lastCall = (mockProvider.postMessageToWebview as jest.Mock).mock.calls.slice(-1)[0]
+        expect(lastCall[0].filePaths).toHaveLength(1)
     })
 
-    it("should clean up watchers on dispose", () => {
+    it("should respect file limits", async () => {
+        // Create array of unique file paths for initial load
+        const files = Array.from({ length: 1001 }, (_, i) => `/test/workspace/file${i}.ts`)
+        ;(listFiles as jest.Mock).mockResolvedValue([files, false])
+        
+        await workspaceTracker.initializeFilePaths()
+        jest.runAllTimers()
+
+        // Should only have 1000 files initially
+        const expectedFiles = Array.from({ length: 1000 }, (_, i) => `file${i}.ts`).sort()
+        const calls = (mockProvider.postMessageToWebview as jest.Mock).mock.calls
+        
+        expect(mockProvider.postMessageToWebview).toHaveBeenCalledWith({
+            type: "workspaceUpdated",
+            filePaths: expect.arrayContaining(expectedFiles)
+        })
+        expect(calls[0][0].filePaths).toHaveLength(1000)
+
+        // Should allow adding up to 2000 total files
+        const [[callback]] = mockOnDidCreate.mock.calls
+        for (let i = 0; i < 1000; i++) {
+            await callback({ fsPath: `/test/workspace/extra${i}.ts` })
+        }
+        jest.runAllTimers()
+
+        const lastCall = (mockProvider.postMessageToWebview as jest.Mock).mock.calls.slice(-1)[0]
+        expect(lastCall[0].filePaths).toHaveLength(2000)
+
+        // Adding one more file beyond 2000 should not increase the count
+        await callback({ fsPath: "/test/workspace/toomany.ts" })
+        jest.runAllTimers()
+
+        const finalCall = (mockProvider.postMessageToWebview as jest.Mock).mock.calls.slice(-1)[0]
+        expect(finalCall[0].filePaths).toHaveLength(2000)
+    })
+
+    it("should clean up watchers and timers on dispose", () => {
         workspaceTracker.dispose()
         expect(mockDispose).toHaveBeenCalled()
+        jest.runAllTimers() // Ensure any pending timers are cleared
     })
 })