More efficient workspace tracker

This commit is contained in:
Matt Rubens
2024-12-27 17:05:08 -08:00
parent 5bf9a74cf7
commit db9827efae
3 changed files with 85 additions and 14 deletions

View File

@@ -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())
}
}

View File

@@ -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
})
})