From f533948817a179dbf290265e288c8ea56c1c10c6 Mon Sep 17 00:00:00 2001 From: Thanh Doan Date: Sun, 31 May 2026 10:25:24 +0700 Subject: [PATCH] fix: stop multi-root/huge-workspace hang and crash (#7) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Opening the panel previously scanned every workspace root recursively, shipped the whole tree in one postMessage, and rendered every node regardless of collapse state — OOMing the host/webview on large or multi-root workspaces. Backend (extension): - Lazy APIs: getWorkspaceRoots (shallow), getDirectoryChildren, listFilesUnderUri; getWorkspaceFileTree kept (capped) for legacy/tests - Default heavy-dir excludes (node_modules, dist, build, .next, ...) - Async gitignore I/O via loadIgnoreRulesForRoot - Scan caps (depth/nodes/list) + AbortSignal + coalesced getFileTree - copyContext uses generateFileMapFromSelections (no full-tree cache) Webview: - Collapse-gated render; only expanded folders render children - Lazy expand fetches children via getDirectoryChildren and merges them - Select-all uses extension listFilesUnderUri; truncation banner - Iterative getAllDescendantPaths Fixes found while validating: - buildTreeIndex crashed on shallow (unloaded) folders whose subItems is undefined — guard with `subItems ?? []` - Seeded-expanded roots never fetched children on first render/refresh — auto-load expanded-but-unloaded roots (guarded by loadingFolderUris) Tests: - Backend file-system.test.ts: excludes/gitignore/abort/shallow roots (workspace-scoped cases self-skip if workspaceFolders can't be stubbed) - Webview: shallow-root auto-load, no double-request, no auto-load of already-loaded roots; buildTreeIndex shallow-folder coverage Co-Authored-By: Claude Opus 4.8 (1M context) --- src/prompts/index.ts | 81 +++ src/providers/file-explorer/index.ts | 248 ++++--- src/providers/file-explorer/types.ts | 13 + src/test/suite/utils/file-system.test.ts | 210 ++++++ src/utils/file-system.ts | 640 +++++++++++++----- src/webview-ui/src/App.tsx | 85 ++- .../__tests__/context-tab.layout.test.tsx | 2 + .../__tests__/utils-descendants.test.ts | 25 + .../file-explorer/__tests__/index.test.tsx | 105 ++- .../__tests__/tree-index.test.ts | 33 +- .../__tests__/tree-merge.test.ts | 39 ++ .../__tests__/tree-node.test.tsx | 26 + .../context-tab/file-explorer/index.tsx | 170 ++++- .../file-explorer/list-files-under-uri.ts | 41 ++ .../context-tab/file-explorer/tree-index.ts | 6 +- .../context-tab/file-explorer/tree-merge.ts | 28 + .../context-tab/file-explorer/tree-node.tsx | 18 +- .../src/components/context-tab/index.tsx | 9 + .../src/components/context-tab/utils.ts | 15 +- src/webview-ui/src/utils/mock.ts | 98 ++- 20 files changed, 1559 insertions(+), 333 deletions(-) create mode 100644 src/test/suite/utils/file-system.test.ts create mode 100644 src/webview-ui/src/components/context-tab/__tests__/utils-descendants.test.ts create mode 100644 src/webview-ui/src/components/context-tab/file-explorer/__tests__/tree-merge.test.ts create mode 100644 src/webview-ui/src/components/context-tab/file-explorer/list-files-under-uri.ts create mode 100644 src/webview-ui/src/components/context-tab/file-explorer/tree-merge.ts diff --git a/src/prompts/index.ts b/src/prompts/index.ts index 1763c83..34f76b4 100644 --- a/src/prompts/index.ts +++ b/src/prompts/index.ts @@ -41,6 +41,87 @@ function filterSelectedTree( return filterItems(items) } +interface PathEntry { + fsPath: string + segments: string[] +} + +/** + * Builds a file map from selected URIs without requiring a full workspace tree. + */ +export function generateFileMapFromSelections( + selectedUris: Set, +): string { + const folders = vscode.workspace.workspaceFolders ?? [] + if (folders.length === 0 || selectedUris.size === 0) { + return '' + } + + const entries: PathEntry[] = [] + + for (const uriString of selectedUris) { + const uri = vscode.Uri.parse(uriString) + let matchedRoot: vscode.WorkspaceFolder | undefined + let relative = '' + + for (const folder of folders) { + const rootPath = folder.uri.fsPath + const filePath = uri.fsPath + if (filePath === rootPath) { + matchedRoot = folder + relative = '' + break + } + const prefix = `${rootPath}${path.sep}` + if (filePath.startsWith(prefix)) { + matchedRoot = folder + relative = filePath.slice(prefix.length) + break + } + } + + if (!matchedRoot) continue + + const segments = relative + ? relative.split(path.sep).filter((s) => s.length > 0) + : [] + entries.push({ + fsPath: matchedRoot.uri.fsPath, + segments, + }) + } + + entries.sort((a, b) => { + const rootCmp = a.fsPath.localeCompare(b.fsPath) + if (rootCmp !== 0) return rootCmp + return a.segments.join('/').localeCompare(b.segments.join('/')) + }) + + const lines: string[] = [] + let currentRoot = '' + + for (const entry of entries) { + if (entry.fsPath !== currentRoot) { + if (currentRoot !== '') lines.push('') + lines.push(entry.fsPath) + currentRoot = entry.fsPath + } + + let prefix = '' + for (let i = 0; i < entry.segments.length; i++) { + const isLast = i === entry.segments.length - 1 + const connector = isLast ? '└── ' : '├── ' + lines.push(prefix + connector + entry.segments[i]!) + prefix += isLast ? ' ' : '│ ' + } + } + + if (lines.length > 0 && lines[lines.length - 1] === '') { + lines.pop() + } + return lines.join('\n') +} + // --- Exported Functions --- /** diff --git a/src/providers/file-explorer/index.ts b/src/providers/file-explorer/index.ts index 8a4cf55..18b3b98 100644 --- a/src/providers/file-explorer/index.ts +++ b/src/providers/file-explorer/index.ts @@ -3,21 +3,26 @@ import * as vscode from 'vscode' import path from 'node:path' import { generateFileContents, - generateFileMap, + generateFileMapFromSelections, generatePrompt, } from '../../prompts' import { telemetry } from '../../services/telemetry' import { countManyWithInfo, encodeText } from '../../services/token-counter' -import type { VscodeTreeItem } from '../../types' -import { getWorkspaceFileTree } from '../../utils/file-system' +import { + getDirectoryChildren, + getWorkspaceRoots, + listFilesUnderUri, +} from '../../utils/file-system' import { getLanguageIdFromPath } from '../../utils/language-detection' import { parseXmlResponse } from '../../utils/xml-parser' import { applyFileActions } from './file-action-handler' import { getHtmlForWebview } from './html-generator' import type { CopyContextPayload, + GetDirectoryChildrenPayload, GetFileTreePayload, GetTokenCountsPayload, + ListFilesUnderUriPayload, OpenFilePayload, SaveSettingsPayload, UpdateSettingsPayload, @@ -28,11 +33,25 @@ export class FileExplorerWebviewProvider implements vscode.WebviewViewProvider { private _view?: vscode.WebviewView private _context: vscode.ExtensionContext - private _fullTreeCache: VscodeTreeItem[] = [] // Cache for the full file tree - private _isBuildingTree = false // Prevent overlapping tree builds - private _treeBuildTimeout: NodeJS.Timeout | null = null // Timeout for tree building - // Always-excluded patterns that never show in the UI; keep minimal (.git only) - private readonly _excludedDirs = ['.git', '.hg', '.svn'] + private _isBuildingTree = false + private _treeBuildAbort: AbortController | null = null + private _treeBuildTimeout: NodeJS.Timeout | null = null + private _pendingFileTreePayload: GetFileTreePayload | undefined + private readonly _excludedDirs = [ + '.git', + '.hg', + '.svn', + 'node_modules', + 'dist', + 'build', + '.next', + 'out', + 'vendor', + 'target', + '.turbo', + '.cache', + 'coverage', + ] private static readonly EXCLUDED_FOLDERS_KEY = 'overwrite.excludedFolders' private static readonly READ_GITIGNORE_KEY = 'overwrite.readGitignore' @@ -130,9 +149,18 @@ export class FileExplorerWebviewProvider implements vscode.WebviewViewProvider { } break case 'getFileTree': - // Payload may contain excluded folders await this._handleGetFileTree(message.payload as GetFileTreePayload) break + case 'getDirectoryChildren': + await this._handleGetDirectoryChildren( + message.payload as GetDirectoryChildrenPayload, + ) + break + case 'listFilesUnderUri': + await this._handleListFilesUnderUri( + message.payload as ListFilesUnderUriPayload, + ) + break case 'getSettings': await this._handleGetSettings() break @@ -400,91 +428,161 @@ export class FileExplorerWebviewProvider implements vscode.WebviewViewProvider { } } + private _resolveExcludedDirs(payload?: { + excludedFolders?: string + }): string[] { + const excludedFoldersArray = payload?.excludedFolders + ? payload.excludedFolders + .split(/\r?\n/) + .map((line) => line.trim()) + .filter((line) => line.length > 0 && !line.startsWith('#')) + : this._loadSettings() + .excludedFolders.split(/\r?\n/) + .map((line) => line.trim()) + .filter((line) => line.length > 0 && !line.startsWith('#')) + return [...this._excludedDirs, ...excludedFoldersArray] + } + + private _readGitignoreFromPayload(payload?: { + readGitignore?: boolean + }): boolean { + return payload?.readGitignore ?? this._loadSettings().readGitignore + } + /** - * Fetches the file tree and sends it to the webview. + * Fetches shallow workspace roots and sends them to the webview. */ private async _handleGetFileTree( payload?: GetFileTreePayload, ): Promise { if (!this._view) return - try { - // Clear any existing timeout - if (this._treeBuildTimeout) { - clearTimeout(this._treeBuildTimeout) - this._treeBuildTimeout = null - } + if (this._isBuildingTree) { + this._pendingFileTreePayload = payload + return + } - if (this._isBuildingTree) { - // Set a timeout to prevent infinite blocking - this._treeBuildTimeout = setTimeout(() => { - this._isBuildingTree = false - this._treeBuildTimeout = null - console.warn('Tree building operation timed out after 30 seconds') - this._view?.webview.postMessage({ - command: 'showError', - message: 'Tree building operation timed out. Please try again.', - }) - vscode.window.showErrorMessage( - 'Tree building operation timed out. Please try again.', - ) - }, 30000) // 30 second timeout + await this._runGetFileTree(payload) + } - return - } + private async _runGetFileTree(payload?: GetFileTreePayload): Promise { + if (!this._view) return + + this._isBuildingTree = true + this._treeBuildAbort?.abort() + this._treeBuildAbort = new AbortController() + const signal = this._treeBuildAbort.signal - this._isBuildingTree = true + if (this._treeBuildTimeout) { + clearTimeout(this._treeBuildTimeout) + } + this._treeBuildTimeout = setTimeout(() => { + this._treeBuildAbort?.abort() + }, 30000) - // Set a timeout for this operation - const timeoutPromise = new Promise((_, reject) => { - this._treeBuildTimeout = setTimeout(() => { - reject( - new Error('Tree building operation timed out after 30 seconds'), - ) - }, 30000) + try { + const allExcludedDirs = this._resolveExcludedDirs(payload) + const result = await getWorkspaceRoots(allExcludedDirs, { + useGitignore: this._readGitignoreFromPayload(payload), + signal, }) - // Race the tree building against the timeout - const excludedFoldersArray = payload?.excludedFolders - ? payload.excludedFolders - .split(/\r?\n/) - .map((line) => line.trim()) - .filter((line) => line.length > 0 && !line.startsWith('#')) - : this._loadSettings() - .excludedFolders.split(/\r?\n/) - .map((line) => line.trim()) - .filter((line) => line.length > 0 && !line.startsWith('#')) - - // Combine with default excluded dirs - const allExcludedDirs = [...this._excludedDirs, ...excludedFoldersArray] - - const workspaceFiles = await Promise.race([ - getWorkspaceFileTree(allExcludedDirs, { - useGitignore: - payload?.readGitignore ?? this._loadSettings().readGitignore, - }), - timeoutPromise, - ]) - - this._fullTreeCache = workspaceFiles // Cache the full tree + if (signal.aborted) { + return + } + this._view.webview.postMessage({ command: 'updateFileTree', - payload: workspaceFiles, + payload: { + tree: result.roots, + truncated: result.truncated, + }, }) } catch (error) { + if (signal.aborted) { + this._view.webview.postMessage({ + command: 'showError', + message: 'Tree building operation timed out. Please try again.', + }) + return + } this._handleError(error, 'Error getting workspace files') - // Optionally send specific error state to webview this._view.webview.postMessage({ command: 'showError', message: `Error getting workspace files: ${error instanceof Error ? error.message : String(error)}`, }) } finally { - // Clear the flag and timeout this._isBuildingTree = false if (this._treeBuildTimeout) { clearTimeout(this._treeBuildTimeout) this._treeBuildTimeout = null } + this._treeBuildAbort = null + + const pending = this._pendingFileTreePayload + this._pendingFileTreePayload = undefined + if (pending !== undefined) { + await this._runGetFileTree(pending) + } + } + } + + private async _handleGetDirectoryChildren( + payload: GetDirectoryChildrenPayload, + ): Promise { + if (!this._view || !payload?.parentUri) return + + try { + const allExcludedDirs = this._resolveExcludedDirs(payload) + const result = await getDirectoryChildren( + payload.parentUri, + allExcludedDirs, + { + useGitignore: this._readGitignoreFromPayload(payload), + }, + ) + + this._view.webview.postMessage({ + command: 'updateDirectoryChildren', + parentUri: payload.parentUri, + children: result.roots, + truncated: result.truncated, + }) + } catch (error) { + this._handleError(error, 'Error loading directory children') + } + } + + private async _handleListFilesUnderUri( + payload: ListFilesUnderUriPayload, + ): Promise { + if (!this._view || !payload?.parentUri || !payload.requestId) return + + try { + const allExcludedDirs = this._resolveExcludedDirs(payload) + const result = await listFilesUnderUri( + payload.parentUri, + allExcludedDirs, + { + useGitignore: this._readGitignoreFromPayload(payload), + }, + ) + + this._view.webview.postMessage({ + command: 'listFilesUnderUriResponse', + requestId: payload.requestId, + uris: result.uris, + truncated: result.truncated, + }) + } catch (error) { + this._handleError(error, 'Error listing files under folder') + this._view.webview.postMessage({ + command: 'listFilesUnderUriResponse', + requestId: payload.requestId, + uris: [], + truncated: false, + error: error instanceof Error ? error.message : String(error), + }) } } @@ -550,23 +648,7 @@ export class FileExplorerWebviewProvider implements vscode.WebviewViewProvider { includeXml, }) - // Ensure fullTreeCache is populated - if (this._fullTreeCache.length === 0) { - console.log('Full tree cache empty, fetching...') - // Refetch if cache is empty (should ideally not happen if getFileTree was called) - await this._handleGetFileTree() - if (this._fullTreeCache.length === 0) { - throw new Error('Failed to populate file tree cache.') - } - } - - // Generate components using imported functions - // TODO: Update generateFileMap and generateFileContents to handle selectedUriStrings and multi-root logic - // For now, this will likely break or produce incorrect results until those functions are updated. - const fileMap = generateFileMap( - this._fullTreeCache, // Use cached tree, which is now multi-root - selectedUriStrings, // Pass Set of URI strings - ) + const fileMap = generateFileMapFromSelections(selectedUriStrings) const fileContents = await generateFileContents( selectedUriStrings, // Pass Set of URI strings ) diff --git a/src/providers/file-explorer/types.ts b/src/providers/file-explorer/types.ts index 52f6851..a320a8e 100644 --- a/src/providers/file-explorer/types.ts +++ b/src/providers/file-explorer/types.ts @@ -20,6 +20,19 @@ export interface GetFileTreePayload { readGitignore?: boolean // Whether to respect .gitignore when building the tree } +export interface GetDirectoryChildrenPayload { + parentUri: string + excludedFolders?: string + readGitignore?: boolean +} + +export interface ListFilesUnderUriPayload { + parentUri: string + requestId: string + excludedFolders?: string + readGitignore?: boolean +} + export interface SaveSettingsPayload { excludedFolders: string readGitignore: boolean diff --git a/src/test/suite/utils/file-system.test.ts b/src/test/suite/utils/file-system.test.ts new file mode 100644 index 0000000..4a65746 --- /dev/null +++ b/src/test/suite/utils/file-system.test.ts @@ -0,0 +1,210 @@ +import * as assert from 'node:assert' +import * as fsp from 'node:fs/promises' +import * as os from 'node:os' +import * as path from 'node:path' +import { after, before, describe, it } from 'mocha' +import * as vscode from 'vscode' +import { + MAX_LIST_FILES, + MAX_TREE_DEPTH, + MAX_TREE_NODES, + getDirectoryChildren, + getWorkspaceRoots, + listFilesUnderUri, + loadIgnoreRulesForRoot, +} from '../../../utils/file-system' + +// Mirrors the heavy-dir defaults the provider passes (see FileExplorerProvider). +const DEFAULT_EXCLUDES = [ + 'node_modules', + 'dist', + 'build', + '.next', + 'out', + 'vendor', + 'target', + '.turbo', + '.cache', + 'coverage', +] + +const baseDir = path.join(os.tmpdir(), 'overwrite-fs-tests') + +async function writeFixture(rel: string, content = ''): Promise { + const abs = path.join(baseDir, rel) + await fsp.mkdir(path.dirname(abs), { recursive: true }) + await fsp.writeFile(abs, content, 'utf8') +} + +/** + * Attempts to point `vscode.workspace.workspaceFolders` at a temp dir. Returns a + * restore fn, or null if the property cannot be redefined in this VS Code build + * (in which case the workspace-scoped tests self-skip rather than fail). + */ +function tryStubWorkspaceFolders(dir: string): (() => void) | null { + try { + const folders = [{ uri: vscode.Uri.file(dir), name: 'fixture', index: 0 }] + const original = Object.getOwnPropertyDescriptor( + vscode.workspace, + 'workspaceFolders', + ) + Object.defineProperty(vscode.workspace, 'workspaceFolders', { + configurable: true, + get: () => folders, + }) + if ( + vscode.workspace.workspaceFolders?.[0]?.uri.fsPath !== + vscode.Uri.file(dir).fsPath + ) { + throw new Error('workspaceFolders stub did not take effect') + } + return () => { + if (original) { + Object.defineProperty(vscode.workspace, 'workspaceFolders', original) + } + } + } catch { + return null + } +} + +describe('file-system: constants', () => { + it('exports sane scan limits', () => { + assert.ok(MAX_TREE_DEPTH >= 1) + assert.ok(MAX_TREE_NODES >= 1000) + assert.ok(MAX_LIST_FILES >= 1000) + assert.ok(MAX_LIST_FILES <= MAX_TREE_NODES) + }) +}) + +describe('file-system: loadIgnoreRulesForRoot', () => { + before(async () => { + await fsp.rm(baseDir, { recursive: true, force: true }) + await fsp.mkdir(baseDir, { recursive: true }) + await writeFixture('.gitignore', 'ignored\n*.log\n') + }) + after(async () => { + await fsp.rm(baseDir, { recursive: true, force: true }) + }) + + it('reads .gitignore lines when gitignore is enabled', async () => { + const lines = await loadIgnoreRulesForRoot(baseDir, true) + assert.ok(lines.includes('ignored')) + assert.ok(lines.includes('*.log')) + }) + + it('returns no rules when gitignore is disabled', async () => { + const lines = await loadIgnoreRulesForRoot(baseDir, false) + assert.deepStrictEqual(lines, []) + }) +}) + +describe('file-system: abort handling (no workspace required)', () => { + it('getDirectoryChildren rejects when the signal is already aborted', async () => { + const controller = new AbortController() + controller.abort() + await assert.rejects( + () => + getDirectoryChildren('file:///anything', DEFAULT_EXCLUDES, { + signal: controller.signal, + }), + /Aborted/, + ) + }) + + it('listFilesUnderUri rejects when the signal is already aborted', async () => { + const controller = new AbortController() + controller.abort() + await assert.rejects( + () => + listFilesUnderUri('file:///anything', DEFAULT_EXCLUDES, { + signal: controller.signal, + }), + /Aborted/, + ) + }) +}) + +describe('file-system: lazy tree APIs (workspace-scoped)', () => { + let restore: (() => void) | null = null + + before(async function () { + await fsp.rm(baseDir, { recursive: true, force: true }) + await fsp.mkdir(baseDir, { recursive: true }) + await writeFixture('src/app.ts', 'export const a = 1') + await writeFixture('src/util.ts', 'export const u = 2') + await writeFixture('keep.ts', 'keep') + await writeFixture('node_modules/dep/index.js', 'dep') + await writeFixture('dist/out.js', 'out') + await writeFixture('.gitignore', 'ignored\n') + await writeFixture('ignored/secret.ts', 'secret') + + restore = tryStubWorkspaceFolders(baseDir) + if (!restore) { + this.skip() + } + }) + + after(async () => { + restore?.() + await fsp.rm(baseDir, { recursive: true, force: true }) + }) + + it('getWorkspaceRoots returns one shallow root per workspace folder', async () => { + const { roots, truncated } = await getWorkspaceRoots(DEFAULT_EXCLUDES) + assert.strictEqual(truncated, false) + assert.strictEqual(roots.length, 1) + assert.strictEqual(roots[0]?.value, vscode.Uri.file(baseDir).toString()) + // Shallow: roots carry no children until expanded. + assert.strictEqual(roots[0]?.subItems, undefined) + }) + + it('getDirectoryChildren applies default excludes and .gitignore', async () => { + const parent = vscode.Uri.file(baseDir).toString() + const { roots: children } = await getDirectoryChildren( + parent, + DEFAULT_EXCLUDES, + ) + const labels = children.map((c) => c.label) + + assert.ok(labels.includes('src'), 'src present') + assert.ok(labels.includes('keep.ts'), 'keep.ts present') + assert.ok(!labels.includes('node_modules'), 'node_modules excluded') + assert.ok(!labels.includes('dist'), 'dist excluded') + assert.ok(!labels.includes('ignored'), 'gitignored dir excluded') + + // Folder children are returned shallow (no nested subItems). + const src = children.find((c) => c.label === 'src') + assert.ok(src) + assert.strictEqual(src?.subItems, undefined) + }) + + it('listFilesUnderUri returns files excluding heavy and gitignored dirs', async () => { + const target = vscode.Uri.file(baseDir).toString() + const { uris, truncated } = await listFilesUnderUri( + target, + DEFAULT_EXCLUDES, + ) + assert.strictEqual(truncated, false) + + const rel = uris.map((u) => + path.relative(baseDir, vscode.Uri.parse(u).fsPath), + ) + + assert.ok(rel.includes('keep.ts')) + assert.ok(rel.includes(path.join('src', 'app.ts'))) + assert.ok(rel.includes(path.join('src', 'util.ts'))) + assert.ok( + !rel.some((r) => r.startsWith(`node_modules${path.sep}`)), + 'no node_modules files', + ) + assert.ok( + !rel.some((r) => r.startsWith(`dist${path.sep}`)), + 'no dist files', + ) + assert.ok( + !rel.some((r) => r.startsWith(`ignored${path.sep}`)), + 'no gitignored files', + ) + }) +}) diff --git a/src/utils/file-system.ts b/src/utils/file-system.ts index f4df93a..03b1b27 100644 --- a/src/utils/file-system.ts +++ b/src/utils/file-system.ts @@ -1,14 +1,20 @@ -import { execFileSync } from 'node:child_process' +import { execFile } from 'node:child_process' import * as fs from 'node:fs' import * as os from 'node:os' -import * as path from 'node:path' // Still needed for path.relative and path.join for ignore patterns +import * as path from 'node:path' +import { promisify } from 'node:util' import ignore from 'ignore' import * as vscode from 'vscode' import type { VscodeTreeItem } from '../types' +const execFileAsync = promisify(execFile) + +export const MAX_TREE_DEPTH = 12 +export const MAX_TREE_NODES = 25_000 +export const MAX_LIST_FILES = 10_000 + // Comprehensive list of binary file extensions const BINARY_EXTENSIONS = new Set([ - // Images '.jpg', '.jpeg', '.png', @@ -21,7 +27,6 @@ const BINARY_EXTENSIONS = new Set([ '.ico', '.heic', '.avif', - // Videos '.mp4', '.avi', '.mov', @@ -32,7 +37,6 @@ const BINARY_EXTENSIONS = new Set([ '.m4v', '.3gp', '.ogv', - // Audio '.mp3', '.wav', '.flac', @@ -42,7 +46,6 @@ const BINARY_EXTENSIONS = new Set([ '.m4a', '.opus', '.oga', - // Archives '.zip', '.rar', '.7z', @@ -54,7 +57,6 @@ const BINARY_EXTENSIONS = new Set([ '.cab', '.dmg', '.iso', - // Executables '.exe', '.dll', '.so', @@ -64,7 +66,6 @@ const BINARY_EXTENSIONS = new Set([ '.rpm', '.msi', '.pkg', - // Documents '.pdf', '.doc', '.docx', @@ -75,13 +76,11 @@ const BINARY_EXTENSIONS = new Set([ '.odt', '.ods', '.odp', - // Fonts '.ttf', '.otf', '.woff', '.woff2', '.eot', - // Other binary formats '.bin', '.dat', '.db', @@ -93,7 +92,6 @@ const BINARY_EXTENSIONS = new Set([ '.obj', ]) -// Magic number signatures for common binary formats const MAGIC_NUMBERS = [ { signature: [0xff, 0xd8, 0xff], format: 'JPEG' }, { signature: [0x89, 0x50, 0x4e, 0x47], format: 'PNG' }, @@ -106,6 +104,74 @@ const MAGIC_NUMBERS = [ { signature: [0xca, 0xfe, 0xba, 0xbe], format: 'Mach-O' }, ] +const FOLDER_ICONS = { + branch: 'folder', + leaf: 'file', + open: 'folder-opened', +} + +const FILE_ICONS = { + branch: 'file', + leaf: 'file', + open: 'file', +} + +export interface FileTreeOptions { + useGitignore?: boolean + signal?: AbortSignal +} + +export interface FileTreeResult { + roots: VscodeTreeItem[] + truncated: boolean +} + +interface IgnoreContext { + ign: ignore.Ignore + userIgnore: ignore.Ignore + rootUri: vscode.Uri +} + +interface ScanBudget { + nodesVisited: number + truncated: boolean +} + +function throwIfAborted(signal?: AbortSignal): void { + if (signal?.aborted) { + const err = new Error('Aborted') + err.name = 'AbortError' + throw err + } +} + +function checkBudget( + budget: ScanBudget, + signal?: AbortSignal, + depth?: number, +): boolean { + throwIfAborted(signal) + if (budget.truncated) return false + if (depth !== undefined && depth > MAX_TREE_DEPTH) { + budget.truncated = true + return false + } + if (budget.nodesVisited >= MAX_TREE_NODES) { + budget.truncated = true + return false + } + return true +} + +function bumpBudget(budget: ScanBudget): boolean { + budget.nodesVisited++ + if (budget.nodesVisited >= MAX_TREE_NODES) { + budget.truncated = true + return false + } + return true +} + function checkMagicNumbers(chunk: Uint8Array): boolean { for (const { signature } of MAGIC_NUMBERS) { if (chunk.length >= signature.length) { @@ -129,26 +195,19 @@ function analyzeByteContent(chunk: Uint8Array): boolean { let nullByteCount = 0 for (const byte of chunk) { - // Count null bytes if (byte === 0) { nullByteCount++ - } - // Count non-printable characters (excluding common whitespace) - else if (byte < 32 && byte !== 9 && byte !== 10 && byte !== 13) { + } else if (byte < 32 && byte !== 9 && byte !== 10 && byte !== 13) { nonPrintableCount++ - } - // Very high bytes that are uncommon in text - else if (byte > 126) { + } else if (byte > 126) { nonPrintableCount++ } } - // If more than 1% null bytes, likely binary if (nullByteCount > chunk.length * 0.01) { return true } - // If more than 30% non-printable characters, likely binary if (nonPrintableCount > chunk.length * 0.3) { return true } @@ -157,18 +216,12 @@ function analyzeByteContent(chunk: Uint8Array): boolean { } export function looksBinary(chunk: Uint8Array): boolean { - // First check magic numbers (most reliable) if (checkMagicNumbers(chunk)) { return true } - - // Then analyze byte content return analyzeByteContent(chunk) } -/** - * Checks if a file is binary using extension-based detection first, then content analysis - */ export async function isBinaryFile(uri: vscode.Uri): Promise { try { const stats = await vscode.workspace.fs.stat(uri) @@ -176,57 +229,36 @@ export async function isBinaryFile(uri: vscode.Uri): Promise { return false } - // First, check if the file extension indicates a binary file (fastest method) const fileName = uri.path.toLowerCase() const extension = fileName.substring(fileName.lastIndexOf('.')) if (BINARY_EXTENSIONS.has(extension)) { return true } - // If extension is unknown, read first 8KB to check for binary content const content = await vscode.workspace.fs.readFile(uri) const chunk = content.slice(0, 8000) return looksBinary(chunk) } catch { - // If we can't read the file, assume it's not binary return false } } -// Define icons for files and folders (can be customized further) -const FOLDER_ICONS = { - branch: 'folder', // Codicon for closed folder - leaf: 'file', // Codicon for file - open: 'folder-opened', // Codicon for opened folder -} - -const FILE_ICONS = { - branch: 'file', // Placeholder, not typically used for files by tree components - leaf: 'file', // Codicon for file - open: 'file', // Placeholder, not typically used for files by tree components -} - -// Expand common shell-style tokens in paths from git config (e.g., '~/.config/git/ignore'). function resolveExcludesPath(input: string): string { let p = input.trim() if (p.startsWith('~')) { p = path.join(os.homedir(), p.slice(1)) } - // Expand $HOME or $USERPROFILE p = p.replace(/\$(HOME|USERPROFILE)/g, (_m, name) => { const env = process.env[String(name)] return env ? env : _m }) - // Expand Windows-style %VAR% p = p.replace(/%([^%]+)%/g, (_m, name) => { const env = process.env[name] return env ? env : _m }) - // Remove surrounding quotes if present without tricky escaping if (p.length >= 2) { const first = p.charCodeAt(0) const last = p.charCodeAt(p.length - 1) - // 34 = '"', 39 = '\'' if ((first === 34 && last === 34) || (first === 39 && last === 39)) { p = p.slice(1, -1) } @@ -234,81 +266,313 @@ function resolveExcludesPath(input: string): string { return path.resolve(p) } +/** Loads gitignore-related rules for a workspace root (async). */ +export async function loadIgnoreRulesForRoot( + rootFsPath: string, + useGitignore: boolean, +): Promise { + const allIgnoreLines: string[] = [] + + if (!useGitignore) { + return allIgnoreLines + } + + const gitignorePath = path.join(rootFsPath, '.gitignore') + try { + const bytes = await vscode.workspace.fs.readFile( + vscode.Uri.file(gitignorePath), + ) + const content = Buffer.from(bytes).toString('utf8') + allIgnoreLines.push(...content.split(/\r?\n/)) + } catch { + // No .gitignore at repo root + } + + const repoExcludePath = path.join(rootFsPath, '.git', 'info', 'exclude') + try { + const content = await fs.promises.readFile(repoExcludePath, 'utf8') + allIgnoreLines.push(...content.split(/\r?\n/)) + } catch { + // Not present + } + + try { + const { stdout } = await execFileAsync( + 'git', + ['config', '--get', 'core.excludesFile'], + { + cwd: rootFsPath, + encoding: 'utf8', + maxBuffer: 1024 * 1024, + }, + ) + const excludesPath = stdout.trim() + if (excludesPath) { + const resolved = resolveExcludesPath(excludesPath) + try { + await fs.promises.access(resolved) + const content = await fs.promises.readFile(resolved, 'utf8') + allIgnoreLines.push(...content.split(/\r?\n/)) + } catch { + // unreadable + } + } + } catch { + const candidates = [ + path.join(os.homedir(), '.config', 'git', 'ignore'), + path.join(os.homedir(), '.gitignore_global'), + path.join(os.homedir(), '.gitignore'), + ] + for (const p of candidates) { + try { + await fs.promises.access(p) + const content = await fs.promises.readFile(p, 'utf8') + allIgnoreLines.push(...content.split(/\r?\n/)) + break + } catch { + // try next + } + } + } + + return allIgnoreLines +} + +function buildIgnoreContext( + rootUri: vscode.Uri, + excludedDirs: string[], + allIgnoreLines: string[], +): IgnoreContext { + let ign = ignore() + if (allIgnoreLines.length > 0) { + ign = ignore().add(allIgnoreLines) + } + ign.add(['.git', '.hg', '.svn']) + + const userIgnore = ignore().add(excludedDirs) + return { ign, userIgnore, rootUri } +} + +function findWorkspaceFolderForUri( + uri: vscode.Uri, +): vscode.WorkspaceFolder | undefined { + const folders = vscode.workspace.workspaceFolders + if (!folders) return undefined + for (const folder of folders) { + const root = folder.uri.fsPath + const target = uri.fsPath + if (target === root || target.startsWith(`${root}${path.sep}`)) { + return folder + } + } + return folders[0] +} + +async function createIgnoreContextForUri( + uri: vscode.Uri, + excludedDirs: string[], + options?: FileTreeOptions, +): Promise { + const folder = findWorkspaceFolderForUri(uri) + if (!folder) { + throw new Error('No workspace folder for URI') + } + const useGitignore = options?.useGitignore !== false + const lines = await loadIgnoreRulesForRoot(folder.uri.fsPath, useGitignore) + return buildIgnoreContext(folder.uri, excludedDirs, lines) +} + +function shouldIgnorePath( + relativePathForIgnore: string, + ctx: IgnoreContext, +): boolean { + return ( + ctx.ign.ignores(relativePathForIgnore) || + ctx.userIgnore.ignores(relativePathForIgnore) + ) +} + +function sortEntries( + entries: [string, vscode.FileType][], +): [string, vscode.FileType][] { + return entries.sort((a, b) => { + if ( + a[1] === vscode.FileType.Directory && + b[1] !== vscode.FileType.Directory + ) + return -1 + if ( + a[1] !== vscode.FileType.Directory && + b[1] === vscode.FileType.Directory + ) + return 1 + return a[0].localeCompare(b[0]) + }) +} + /** - * Recursively reads a directory using vscode.workspace.fs and builds a tree structure. - * @param currentUri The URI of the directory to read. - * @param rootUri The URI of the workspace root for this directory (for .gitignore path relativity). - * @param excludedDirs An array of additional directory patterns to exclude (like .gitignore patterns). - * @param ign The ignore object from the 'ignore' package for .gitignore rules. - * @param userIgnore The ignore object from the 'ignore' package for user-defined excluded patterns. - * @returns A promise that resolves to an array of VscodeTreeItem objects. + * Shallow workspace roots (one node per workspace folder, no children loaded). */ +export async function getWorkspaceRoots( + _excludedDirs: string[], + _options?: FileTreeOptions, +): Promise { + const workspaceFolders = vscode.workspace.workspaceFolders + if (!workspaceFolders || workspaceFolders.length === 0) { + vscode.window.showInformationMessage('No workspace folder open.') + return { roots: [], truncated: false } + } + + const roots: VscodeTreeItem[] = [] + for (const folder of workspaceFolders) { + roots.push({ + label: folder.name, + value: folder.uri.toString(), + icons: FOLDER_ICONS, + }) + } + return { roots, truncated: false } +} + +/** + * Reads one directory level for lazy tree expansion. + */ +export async function getDirectoryChildren( + parentUriString: string, + excludedDirs: string[], + options?: FileTreeOptions, +): Promise { + const signal = options?.signal + throwIfAborted(signal) + + const parentUri = vscode.Uri.parse(parentUriString) + const ctx = await createIgnoreContextForUri(parentUri, excludedDirs, options) + const budget: ScanBudget = { nodesVisited: 0, truncated: false } + + const children = await readDirectoryOneLevel( + parentUri, + ctx, + budget, + signal, + true, + ) + return { roots: children, truncated: budget.truncated } +} + +async function readDirectoryOneLevel( + currentUri: vscode.Uri, + ctx: IgnoreContext, + budget: ScanBudget, + signal?: AbortSignal, + lazyFolders = false, +): Promise { + const items: VscodeTreeItem[] = [] + + if (!checkBudget(budget, signal)) { + return items + } + + try { + const entries = sortEntries( + await vscode.workspace.fs.readDirectory(currentUri), + ) + + for (const [name, type] of entries) { + if (!checkBudget(budget, signal)) break + + const entryUri = vscode.Uri.joinPath(currentUri, name) + const relativePathForIgnore = path.relative( + ctx.rootUri.fsPath, + entryUri.fsPath, + ) + + if (shouldIgnorePath(relativePathForIgnore, ctx)) { + continue + } + + if (!bumpBudget(budget)) break + + const item: VscodeTreeItem = { + label: name, + value: entryUri.toString(), + } + + if (type === vscode.FileType.Directory) { + item.icons = FOLDER_ICONS + if (!lazyFolders) { + // subItems filled by recursive reader + } + } else if (type === vscode.FileType.File) { + item.icons = FILE_ICONS + } + + items.push(item) + } + } catch (error: unknown) { + const errorMessage = error instanceof Error ? error.message : String(error) + console.error( + `Error reading directory ${currentUri.fsPath}: ${errorMessage}`, + ) + } + + return items +} + async function readDirectoryRecursiveForRoot( currentUri: vscode.Uri, rootUri: vscode.Uri, - excludedDirs: string[], - ign: ignore.Ignore, - userIgnore: ignore.Ignore, + ctx: IgnoreContext, + budget: ScanBudget, + depth: number, + signal?: AbortSignal, ): Promise { const items: VscodeTreeItem[] = [] + if (!checkBudget(budget, signal, depth)) { + return items + } + try { - const entries = await vscode.workspace.fs.readDirectory(currentUri) + const entries = sortEntries( + await vscode.workspace.fs.readDirectory(currentUri), + ) - // Sort entries: directories first, then alphabetically by name - const sortedEntries = entries.sort((a, b) => { - if ( - a[1] === vscode.FileType.Directory && - b[1] !== vscode.FileType.Directory - ) - return -1 - if ( - a[1] !== vscode.FileType.Directory && - b[1] === vscode.FileType.Directory - ) - return 1 - return a[0].localeCompare(b[0]) - }) + for (const [name, type] of entries) { + if (!checkBudget(budget, signal, depth)) break - for (const [name, type] of sortedEntries) { const entryUri = vscode.Uri.joinPath(currentUri, name) const relativePathForIgnore = path.relative( rootUri.fsPath, entryUri.fsPath, ) - // Check .gitignore - if (ign.ignores(relativePathForIgnore)) { + if (shouldIgnorePath(relativePathForIgnore, ctx)) { continue } - // Check user-defined excluded patterns - if (userIgnore.ignores(relativePathForIgnore)) { - continue - } + if (!bumpBudget(budget)) break const item: VscodeTreeItem = { label: name, - value: entryUri.toString(), // Use full URI string as the value - // icons: type === vscode.FileType.Directory ? FOLDER_ICONS : { leaf: 'file' }, // Simplified, vscode-tree might handle this + value: entryUri.toString(), } if (type === vscode.FileType.Directory) { - item.icons = FOLDER_ICONS // Apply folder icons + item.icons = FOLDER_ICONS const subItems = await readDirectoryRecursiveForRoot( entryUri, rootUri, - excludedDirs, - ign, - userIgnore, + ctx, + budget, + depth + 1, + signal, ) if (subItems.length > 0) { item.subItems = subItems } } else if (type === vscode.FileType.File) { - item.icons = FILE_ICONS // Apply file icon + item.icons = FILE_ICONS } - // Symlinks and Unknown types are currently ignored but could be handled items.push(item) } @@ -317,135 +581,135 @@ async function readDirectoryRecursiveForRoot( console.error( `Error reading directory ${currentUri.fsPath}: ${errorMessage}`, ) - // Optionally, inform the user if a specific directory is unreadable - // vscode.window.showWarningMessage(`Could not read directory: ${currentUri.fsPath}`); } return items } /** - * Gets the file tree structure for all workspace folders, respecting .gitignore for each. - * @param excludedDirs An array of additional directory names to exclude globally. - * @returns A promise that resolves to an array of VscodeTreeItem objects, where each top-level item is a workspace root. + * Full recursive tree (capped). Used by tests and legacy paths. */ export async function getWorkspaceFileTree( excludedDirs: string[], - options?: { useGitignore?: boolean }, -): Promise { + options?: FileTreeOptions, +): Promise { const workspaceFolders = vscode.workspace.workspaceFolders if (!workspaceFolders || workspaceFolders.length === 0) { vscode.window.showInformationMessage('No workspace folder open.') - return [] + return { roots: [], truncated: false } } - const allRootItems: VscodeTreeItem[] = [] - + const signal = options?.signal const useGitignore = options?.useGitignore !== false + const allRootItems: VscodeTreeItem[] = [] + const budget: ScanBudget = { nodesVisited: 0, truncated: false } for (const folder of workspaceFolders) { + throwIfAborted(signal) + if (!checkBudget(budget, signal)) break + const rootUri = folder.uri - const rootFsPath = rootUri.fsPath - - // Build a comprehensive ignore rule set similar to Git's behavior: - // - project .gitignore - // - .git/info/exclude - // - user's global excludes file (core.excludesFile), if available - // The logic is best-effort and silently skips files that aren't present. - const allIgnoreLines: string[] = [] - - if (useGitignore) { - // 1) Project .gitignore (root) - const gitignorePath = path.join(rootFsPath, '.gitignore') - try { - const bytes = await vscode.workspace.fs.readFile( - vscode.Uri.file(gitignorePath), - ) - const content = Buffer.from(bytes).toString('utf8') - allIgnoreLines.push(...content.split(/\r?\n/)) - } catch { - // No .gitignore at repo root; ignore quietly - } + const allIgnoreLines = await loadIgnoreRulesForRoot( + rootUri.fsPath, + useGitignore, + ) + const ctx = buildIgnoreContext(rootUri, excludedDirs, allIgnoreLines) - // 2) Repo-specific excludes: .git/info/exclude - const repoExcludePath = path.join(rootFsPath, '.git', 'info', 'exclude') - try { - const content = fs.readFileSync(repoExcludePath, 'utf8') - allIgnoreLines.push(...content.split(/\r?\n/)) - } catch { - // Not present or unreadable; ignore quietly - } + const subItems = await readDirectoryRecursiveForRoot( + rootUri, + rootUri, + ctx, + budget, + 0, + signal, + ) - // 3) Global excludes file as configured in Git (core.excludesFile) - // Attempt to query Git for the actual path; if not available, try common defaults. - try { - const excludesPath = execFileSync( - 'git', - ['config', '--get', 'core.excludesFile'], - { - cwd: rootFsPath, - encoding: 'utf8', - stdio: ['ignore', 'pipe', 'ignore'], - }, - ).trim() - if (excludesPath) { - const resolved = resolveExcludesPath(excludesPath) - if (resolved && fs.existsSync(resolved)) { - const content = fs.readFileSync(resolved, 'utf8') - allIgnoreLines.push(...content.split(/\r?\n/)) - } - } - } catch { - // Git not available or no core.excludesFile set; try common fallbacks - const candidates = [ - path.join(os.homedir(), '.config', 'git', 'ignore'), - path.join(os.homedir(), '.gitignore_global'), - path.join(os.homedir(), '.gitignore'), - ] - for (const p of candidates) { - try { - if (fs.existsSync(p)) { - const content = fs.readFileSync(p, 'utf8') - allIgnoreLines.push(...content.split(/\r?\n/)) - break - } - } catch { - // keep trying others - } - } - } + allRootItems.push({ + label: folder.name, + value: rootUri.toString(), + icons: FOLDER_ICONS, + subItems, + }) + } + + return { roots: allRootItems, truncated: budget.truncated } +} + +export interface ListFilesResult { + uris: string[] + truncated: boolean +} + +/** + * Lists file URIs under a folder URI (for select-all), respecting ignores and caps. + */ +export async function listFilesUnderUri( + targetUriString: string, + excludedDirs: string[], + options?: FileTreeOptions, +): Promise { + const signal = options?.signal + throwIfAborted(signal) + + const targetUri = vscode.Uri.parse(targetUriString) + const stat = await vscode.workspace.fs.stat(targetUri) + + const ctx = await createIgnoreContextForUri(targetUri, excludedDirs, options) + const uris: string[] = [] + let truncated = false + + const relativeSelf = path.relative(ctx.rootUri.fsPath, targetUri.fsPath) + + if (stat.type === vscode.FileType.File) { + if (!shouldIgnorePath(relativeSelf, ctx)) { + uris.push(targetUri.toString()) } + return { uris, truncated: false } + } - let ign = ignore() - if (allIgnoreLines.length > 0) { - ign = ignore().add(allIgnoreLines) + const queue: vscode.Uri[] = [targetUri] + let visited = 0 + + while (queue.length > 0) { + throwIfAborted(signal) + const currentUri = queue.shift()! + let entries: [string, vscode.FileType][] + try { + entries = await vscode.workspace.fs.readDirectory(currentUri) + } catch { + continue } - // Optional: add a few extremely common fallbacks to reduce noise/perf impact - // only when not already covered (safe, minimal defaults) - ign.add(['.git', '.hg', '.svn']) + for (const [name, type] of entries) { + throwIfAborted(signal) + if (uris.length >= MAX_LIST_FILES) { + truncated = true + return { uris, truncated } + } - // Create ignore object for user-defined excluded patterns - const userIgnore = ignore().add(excludedDirs) + const entryUri = vscode.Uri.joinPath(currentUri, name) + const relativePathForIgnore = path.relative( + ctx.rootUri.fsPath, + entryUri.fsPath, + ) - const subItems = await readDirectoryRecursiveForRoot( - rootUri, - rootUri, // rootUri itself is the base for relative ignore paths - excludedDirs, - ign, - userIgnore, - ) + if (shouldIgnorePath(relativePathForIgnore, ctx)) { + continue + } - const rootItem: VscodeTreeItem = { - label: folder.name, // Use workspace folder name as label - value: rootUri.toString(), // Use root URI string as value - icons: FOLDER_ICONS, // Root is a folder - subItems: subItems, - // Optionally, mark as open by default - // open: true, + visited++ + if (visited > MAX_TREE_NODES) { + truncated = true + return { uris, truncated } + } + + if (type === vscode.FileType.File) { + uris.push(entryUri.toString()) + } else if (type === vscode.FileType.Directory) { + queue.push(entryUri) + } } - allRootItems.push(rootItem) } - return allRootItems + return { uris, truncated } } diff --git a/src/webview-ui/src/App.tsx b/src/webview-ui/src/App.tsx index 5f799c6..f69c4d8 100644 --- a/src/webview-ui/src/App.tsx +++ b/src/webview-ui/src/App.tsx @@ -4,9 +4,40 @@ import type { VscodeTreeItem } from './types' // Import tree item type from loca import './App.css' import ApplyTab from './components/apply-tab/index' import ContextTab from './components/context-tab' +import { mergeChildren } from './components/context-tab/file-explorer/tree-merge' import SettingsTab from './components/settings-tab' import { getVsCodeApi } from './utils/vscode' // Import the new utility +interface UpdateFileTreePayload { + tree: VscodeTreeItem[] + truncated?: boolean +} + +interface UpdateDirectoryChildrenMessage { + command: 'updateDirectoryChildren' + parentUri: string + children: VscodeTreeItem[] + truncated?: boolean +} + +function parseUpdateFileTreePayload(payload: unknown): { + tree: VscodeTreeItem[] + truncated: boolean +} { + if (Array.isArray(payload)) { + return { tree: payload as VscodeTreeItem[], truncated: false } + } + if ( + payload && + typeof payload === 'object' && + Array.isArray((payload as UpdateFileTreePayload).tree) + ) { + const p = payload as UpdateFileTreePayload + return { tree: p.tree, truncated: Boolean(p.truncated) } + } + return { tree: [], truncated: false } +} + interface VsCodeMessage { command: string payload?: unknown // Use unknown instead of any for better type safety @@ -28,6 +59,10 @@ function App() { const [selectedUris, setSelectedUris] = useState>(new Set()) const [isLoading, setIsLoading] = useState(true) // For loading indicator const [, setErrorText] = useState(null) // Graceful error banner + const [treeTruncated, setTreeTruncated] = useState(false) + const [loadingFolderUris, setLoadingFolderUris] = useState>( + () => new Set(), + ) const [excludedFolders, setExcludedFolders] = useState('') // Persisted excluded folders const [readGitignore, setReadGitignore] = useState(true) @@ -119,13 +154,33 @@ function App() { const message = event.data switch (message.command) { - case 'updateFileTree': - // TODO: Add type check for payload - if (Array.isArray(message.payload)) { - setFileTreeData(message.payload as VscodeTreeItem[]) - } + case 'updateFileTree': { + const { tree, truncated } = parseUpdateFileTreePayload( + message.payload, + ) + setFileTreeData(tree) + setTreeTruncated(truncated) + setLoadingFolderUris(new Set()) setIsLoading(false) break + } + case 'updateDirectoryChildren': { + const msg = message as unknown as UpdateDirectoryChildrenMessage + if (msg.parentUri && Array.isArray(msg.children)) { + setFileTreeData((prev) => + mergeChildren(prev, msg.parentUri, msg.children), + ) + setLoadingFolderUris((prev) => { + const next = new Set(prev) + next.delete(msg.parentUri) + return next + }) + if (msg.truncated) setTreeTruncated(true) + } + break + } + case 'listFilesUnderUriResponse': + break case 'showError': // Display error message in a dismissible banner { @@ -203,6 +258,7 @@ function App() { const handleRefresh = useCallback( (excludedFoldersArg?: string) => { setIsLoading(true) + setTreeTruncated(false) sendMessage('getFileTree', { excludedFolders: excludedFoldersArg, readGitignore, @@ -211,6 +267,14 @@ function App() { [sendMessage, readGitignore], ) + const handleLoadChildren = useCallback( + (parentUri: string) => { + setLoadingFolderUris((prev) => new Set(prev).add(parentUri)) + sendMessage('getDirectoryChildren', { parentUri }) + }, + [sendMessage], + ) + // Save settings handler (excluded folders + readGitignore) const handleSaveSettings = useCallback( (payload: { excludedFolders: string; readGitignore: boolean }) => { @@ -295,15 +359,16 @@ function App() { diff --git a/src/webview-ui/src/components/context-tab/__tests__/context-tab.layout.test.tsx b/src/webview-ui/src/components/context-tab/__tests__/context-tab.layout.test.tsx index 806c1fe..8aa6b78 100644 --- a/src/webview-ui/src/components/context-tab/__tests__/context-tab.layout.test.tsx +++ b/src/webview-ui/src/components/context-tab/__tests__/context-tab.layout.test.tsx @@ -17,6 +17,8 @@ describe('ContextTab layout', () => { onSelect: noop, onRefresh: noop, isLoading: false, + loadingFolderUris: new Set(), + onLoadChildren: noop, } it('renders sticky footer with compact stats and dual copy buttons', () => { diff --git a/src/webview-ui/src/components/context-tab/__tests__/utils-descendants.test.ts b/src/webview-ui/src/components/context-tab/__tests__/utils-descendants.test.ts new file mode 100644 index 0000000..a806290 --- /dev/null +++ b/src/webview-ui/src/components/context-tab/__tests__/utils-descendants.test.ts @@ -0,0 +1,25 @@ +import { describe, expect, it } from 'vitest' +import type { VscodeTreeItem } from '../../../types' +import { getAllDescendantPaths } from '../utils' + +describe('getAllDescendantPaths', () => { + it('collects all descendant URIs iteratively without stack overflow', () => { + let leaf: VscodeTreeItem = { + label: 'deep', + value: 'deep', + icons: { branch: 'file', open: 'file', leaf: 'file' }, + } + for (let i = 0; i < 200; i++) { + leaf = { + label: `d${i}`, + value: `d${i}`, + icons: { branch: 'folder', open: 'folder-opened', leaf: 'file' }, + subItems: [leaf], + } + } + + const paths = getAllDescendantPaths(leaf) + expect(paths.length).toBe(201) + expect(paths[0]).toBe(leaf.value) + }) +}) diff --git a/src/webview-ui/src/components/context-tab/file-explorer/__tests__/index.test.tsx b/src/webview-ui/src/components/context-tab/file-explorer/__tests__/index.test.tsx index 0a882eb..48c581e 100644 --- a/src/webview-ui/src/components/context-tab/file-explorer/__tests__/index.test.tsx +++ b/src/webview-ui/src/components/context-tab/file-explorer/__tests__/index.test.tsx @@ -62,24 +62,48 @@ vi.mock('../row-decorations', () => ({ import FileExplorer from '../index' +const folderIcons = { + branch: 'folder', + open: 'folder-opened', + leaf: 'file', +} as const +const fileIcons = { branch: 'file', open: 'file', leaf: 'file' } as const + const mkTree = (): VscodeTreeItem[] => [ { label: 'workspace', value: 'ws', + icons: folderIcons, subItems: [ { label: 'src', value: 'src', + icons: folderIcons, subItems: [ - { label: 'a.ts', value: 'a' }, - { label: 'b.ts', value: 'b' }, + { label: 'a.ts', value: 'a', icons: fileIcons }, + { label: 'b.ts', value: 'b', icons: fileIcons }, ], }, - { label: 'README.md', value: 'r' }, + { label: 'README.md', value: 'r', icons: fileIcons }, ], }, ] +const defaultExplorerProps = { + loadingFolderUris: new Set(), + onLoadChildren: vi.fn(), + treeTruncated: false, +} + +vi.mock('../list-files-under-uri', () => ({ + listFilesUnderUriRemote: vi.fn(async (parentUri: string) => { + if (parentUri === 'src') { + return { uris: ['a', 'b'], truncated: false } + } + return { uris: [], truncated: false } + }), +})) + describe('FileExplorer (index.tsx)', () => { beforeEach(() => { vi.useFakeTimers() @@ -90,9 +114,78 @@ describe('FileExplorer (index.tsx)', () => { vi.useRealTimers() }) + it('auto-loads children for shallow roots on mount (lazy load)', () => { + const onLoadChildren = vi.fn() + // Shallow roots: folder icons but no subItems, matching getWorkspaceRoots() + const shallowRoots: VscodeTreeItem[] = [ + { label: 'rootA', value: 'file:///root-a', icons: folderIcons }, + { label: 'rootB', value: 'file:///root-b', icons: folderIcons }, + ] + + render( + {}} + isLoading={false} + searchQuery="" + actualTokenCounts={{}} + />, + ) + + expect(onLoadChildren).toHaveBeenCalledWith('file:///root-a') + expect(onLoadChildren).toHaveBeenCalledWith('file:///root-b') + expect(onLoadChildren).toHaveBeenCalledTimes(2) + }) + + it('does not re-request children for roots already loading', () => { + const onLoadChildren = vi.fn() + const shallowRoots: VscodeTreeItem[] = [ + { label: 'rootA', value: 'file:///root-a', icons: folderIcons }, + ] + + render( + {}} + isLoading={false} + searchQuery="" + actualTokenCounts={{}} + />, + ) + + expect(onLoadChildren).not.toHaveBeenCalled() + }) + + it('does not auto-load roots that already have children', () => { + const onLoadChildren = vi.fn() + + render( + {}} + isLoading={false} + searchQuery="" + actualTokenCounts={{}} + />, + ) + + expect(onLoadChildren).not.toHaveBeenCalled() + }) + it('renders tree labels', () => { render( {}} @@ -103,7 +196,6 @@ describe('FileExplorer (index.tsx)', () => { ) expect(screen.getByText('workspace')).toBeInTheDocument() - expect(screen.getByText('src')).toBeInTheDocument() expect(screen.getByText('README.md')).toBeInTheDocument() }) @@ -112,6 +204,7 @@ describe('FileExplorer (index.tsx)', () => { render( { render( { render( { it('double-clicking a file sends openFile message with fileUri', async () => { render( {}} @@ -221,6 +317,7 @@ describe('FileExplorer (index.tsx)', () => { postMessageSpy.mockClear() render( {}} diff --git a/src/webview-ui/src/components/context-tab/file-explorer/__tests__/tree-index.test.ts b/src/webview-ui/src/components/context-tab/file-explorer/__tests__/tree-index.test.ts index ecfc39d..8ec041e 100644 --- a/src/webview-ui/src/components/context-tab/file-explorer/__tests__/tree-index.test.ts +++ b/src/webview-ui/src/components/context-tab/file-explorer/__tests__/tree-index.test.ts @@ -2,20 +2,29 @@ import { describe, expect, it } from 'vitest' import type { VscodeTreeItem } from '../../../../types' import { buildTreeIndex } from '../tree-index' +const folderIcons = { + branch: 'folder', + open: 'folder-opened', + leaf: 'file', +} as const +const fileIcons = { branch: 'file', open: 'file', leaf: 'file' } as const + describe('buildTreeIndex', () => { const tree: VscodeTreeItem[] = [ { label: 'root', value: 'root', + icons: folderIcons, subItems: [ - { label: 'a.ts', value: 'a' }, - { label: 'b.ts', value: 'b' }, + { label: 'a.ts', value: 'a', icons: fileIcons }, + { label: 'b.ts', value: 'b', icons: fileIcons }, { label: 'folder', value: 'folder', + icons: folderIcons, subItems: [ - { label: 'c.ts', value: 'c' }, - { label: 'd.ts', value: 'd' }, + { label: 'c.ts', value: 'c', icons: fileIcons }, + { label: 'd.ts', value: 'd', icons: fileIcons }, ], }, ], @@ -34,4 +43,20 @@ describe('buildTreeIndex', () => { expect(idx.descendantFileCount.get('folder')).toBe(2) expect(idx.descendantFileCount.get('a')).toBe(1) }) + + it('handles lazily-loaded folders without subItems', () => { + // Shallow roots from getWorkspaceRoots(): folder icons, no subItems. + const shallow: VscodeTreeItem[] = [ + { label: 'rootA', value: 'file:///root-a', icons: folderIcons }, + { label: 'rootB', value: 'file:///root-b', icons: folderIcons }, + ] + + const idx = buildTreeIndex(shallow) + + expect(idx.nodes.get('file:///root-a')?.isFolder).toBe(true) + expect(idx.nodes.get('file:///root-a')?.children).toEqual([]) + // Unloaded folders contribute 0 files until their children are fetched. + expect(idx.descendantFileCount.get('file:///root-a')).toBe(0) + expect(idx.descendantFileCount.get('file:///root-b')).toBe(0) + }) }) diff --git a/src/webview-ui/src/components/context-tab/file-explorer/__tests__/tree-merge.test.ts b/src/webview-ui/src/components/context-tab/file-explorer/__tests__/tree-merge.test.ts new file mode 100644 index 0000000..b2ce9bf --- /dev/null +++ b/src/webview-ui/src/components/context-tab/file-explorer/__tests__/tree-merge.test.ts @@ -0,0 +1,39 @@ +import { describe, expect, it } from 'vitest' +import type { VscodeTreeItem } from '../../../../types' +import { folderNeedsLoad, isFolderItem, mergeChildren } from '../tree-merge' + +describe('tree-merge', () => { + const tree: VscodeTreeItem[] = [ + { + label: 'root', + value: 'root', + icons: { branch: 'folder', open: 'folder-opened', leaf: 'file' }, + }, + ] + + it('mergeChildren attaches children to the matching parent', () => { + const children: VscodeTreeItem[] = [ + { + label: 'child.ts', + value: 'child', + icons: { branch: 'file', open: 'file', leaf: 'file' }, + }, + ] + const merged = mergeChildren(tree, 'root', children) + expect(merged[0]?.subItems).toHaveLength(1) + expect(merged[0]?.subItems?.[0]?.value).toBe('child') + }) + + it('isFolderItem and folderNeedsLoad detect lazy folders', () => { + const folder: VscodeTreeItem = { + label: 'f', + value: 'f', + icons: { branch: 'folder', open: 'folder-opened', leaf: 'file' }, + } + expect(isFolderItem(folder)).toBe(true) + expect(folderNeedsLoad(folder)).toBe(true) + + const loaded: VscodeTreeItem = { ...folder, subItems: [] } + expect(folderNeedsLoad(loaded)).toBe(false) + }) +}) diff --git a/src/webview-ui/src/components/context-tab/file-explorer/__tests__/tree-node.test.tsx b/src/webview-ui/src/components/context-tab/file-explorer/__tests__/tree-node.test.tsx index 433d320..d597ca8 100644 --- a/src/webview-ui/src/components/context-tab/file-explorer/__tests__/tree-node.test.tsx +++ b/src/webview-ui/src/components/context-tab/file-explorer/__tests__/tree-node.test.tsx @@ -66,6 +66,7 @@ describe('TreeNode', () => { depth={0} isFolder isOpen + isLoadingChildren={false} totalDescendantFiles={1} selectedDescendantFiles={0} folderSelectionState="none" @@ -88,6 +89,30 @@ describe('TreeNode', () => { expect(onDeselectAll).toHaveBeenCalledTimes(1) }) + it('does not render children when folder is collapsed', () => { + render( + {}} + onSelectAllInSubtree={() => {}} + onDeselectAllInSubtree={() => {}} + renderChildren={() => child} + />, + ) + + expect(screen.queryByTestId('child')).not.toBeInTheDocument() + }) + it('renders file label and exposes toggle to callback', () => { const onToggleFile = vi.fn() @@ -97,6 +122,7 @@ describe('TreeNode', () => { depth={1} isFolder={false} isOpen + isLoadingChildren={false} totalDescendantFiles={0} selectedDescendantFiles={0} folderSelectionState="none" diff --git a/src/webview-ui/src/components/context-tab/file-explorer/index.tsx b/src/webview-ui/src/components/context-tab/file-explorer/index.tsx index 4436b6c..6c20c6b 100644 --- a/src/webview-ui/src/components/context-tab/file-explorer/index.tsx +++ b/src/webview-ui/src/components/context-tab/file-explorer/index.tsx @@ -11,10 +11,11 @@ import type { VscodeTreeItem } from '../../../types' import { getVsCodeApi } from '../../../utils/vscode' import { FileTreeSkeleton, LoadingOverlay } from '../../loading' import { filterTreeData, getAllDescendantPaths } from '../utils' +import { listFilesUnderUriRemote } from './list-files-under-uri' import type { FolderSelectionState } from './row-decorations' -import TreeNode from './tree-node' - import { buildTreeIndex } from './tree-index' +import { folderNeedsLoad, isFolderItem } from './tree-merge' +import TreeNode from './tree-node' interface FileExplorerProps { fileTreeData: VscodeTreeItem[] @@ -23,6 +24,9 @@ interface FileExplorerProps { isLoading: boolean searchQuery: string actualTokenCounts: Record + treeTruncated?: boolean + loadingFolderUris: Set + onLoadChildren: (parentUri: string) => void } type LoadingPhase = 'initial' | 'skeleton' | 'progressive' | 'complete' @@ -34,20 +38,43 @@ const FileExplorer: React.FC = ({ isLoading, searchQuery, actualTokenCounts, + treeTruncated = false, + loadingFolderUris, + onLoadChildren, }) => { - // Loading phase management const [loadingPhase, setLoadingPhase] = useState('initial') const [isRefreshing, setIsRefreshing] = useState(false) const [, setPrevFileTreeData] = useState([]) + const [expandedUris, setExpandedUris] = useState>(() => new Set()) - // Defer heavy recalculations when selection/token counts change massively const deferredSelectedUris = useDeferredValue(selectedUris) const deferredTokenCounts = useDeferredValue(actualTokenCounts) - // Loading phase effects + useEffect(() => { + if (fileTreeData.length > 0 && expandedUris.size === 0) { + setExpandedUris(new Set(fileTreeData.map((r) => r.value))) + } + }, [fileTreeData, expandedUris.size]) + + // Roots arrive shallow (no subItems) from the lazy `getWorkspaceRoots` API and + // are seeded as expanded above, but expanding via state alone never fetches + // their children. Auto-load children for any root that is expanded yet unloaded + // so the initial render (and refresh, where root URIs persist in expandedUris) + // populates without the user having to collapse and re-expand each root. + useEffect(() => { + for (const root of fileTreeData) { + if ( + expandedUris.has(root.value) && + folderNeedsLoad(root) && + !loadingFolderUris.has(root.value) + ) { + onLoadChildren(root.value) + } + } + }, [fileTreeData, expandedUris, loadingFolderUris, onLoadChildren]) + useEffect(() => { if (isLoading) { - // If we have existing data, this is a refresh if (fileTreeData.length > 0) { setIsRefreshing(true) setLoadingPhase('initial') @@ -56,7 +83,6 @@ const FileExplorer: React.FC = ({ setLoadingPhase('skeleton') } } else { - // Data has loaded if (fileTreeData.length > 0) { setLoadingPhase('progressive') setTimeout(() => { @@ -70,23 +96,54 @@ const FileExplorer: React.FC = ({ } }, [isLoading, fileTreeData.length]) - // Filtered items based on search const visibleItems = useMemo(() => { return searchQuery ? filterTreeData(fileTreeData, searchQuery) : fileTreeData }, [fileTreeData, searchQuery]) - // Build index for current visible tree const index = useMemo(() => buildTreeIndex(visibleItems), [visibleItems]) - // Stable refs to avoid function identity changes and stale closures const selectedUrisRef = useRef(selectedUris) const indexRef = useRef(index) + const fileTreeRef = useRef(fileTreeData) selectedUrisRef.current = selectedUris indexRef.current = index + fileTreeRef.current = fileTreeData + + const findItemByUri = useCallback( + (items: VscodeTreeItem[], uri: string): VscodeTreeItem | undefined => { + for (const item of items) { + if (item.value === uri) return item + if (item.subItems) { + const found = findItemByUri(item.subItems, uri) + if (found) return found + } + } + return undefined + }, + [], + ) + + const toggleFolderExpanded = useCallback( + (uri: string) => { + setExpandedUris((prev) => { + const next = new Set(prev) + if (next.has(uri)) { + next.delete(uri) + } else { + next.add(uri) + const item = findItemByUri(fileTreeRef.current, uri) + if (item && folderNeedsLoad(item)) { + onLoadChildren(uri) + } + } + return next + }) + }, + [findItemByUri, onLoadChildren], + ) - // Derived per-node metrics based on selection + tokens (single post-order pass) const { selectedCountMap, tokenTotalsMap } = useMemo(() => { const selectedCountMap = new Map() const tokenTotalsMap = new Map() @@ -110,7 +167,6 @@ const FileExplorer: React.FC = ({ return { selectedCountMap, tokenTotalsMap } }, [index, deferredSelectedUris, deferredTokenCounts]) - // Selection helpers const toggleFile = useCallback( (uri: string) => { const next = new Set(selectedUrisRef.current) @@ -123,13 +179,21 @@ const FileExplorer: React.FC = ({ const selectAllInSubtree = useCallback( (uri: string) => { - const node = indexRef.current.nodes.get(uri) - if (!node) return - // Yield to the browser, then perform heavy traversal in a task - setTimeout(() => { - const next = new Set(selectedUrisRef.current) - for (const u of getAllDescendantPaths(node.item)) next.add(u) - startTransition(() => onSelect(next)) + setTimeout(async () => { + try { + const { uris } = await listFilesUnderUriRemote(uri) + const next = new Set(selectedUrisRef.current) + next.add(uri) + for (const u of uris) next.add(u) + startTransition(() => onSelect(next)) + } catch (err) { + console.error('selectAllInSubtree failed:', err) + const node = indexRef.current.nodes.get(uri) + if (!node) return + const next = new Set(selectedUrisRef.current) + for (const u of getAllDescendantPaths(node.item)) next.add(u) + startTransition(() => onSelect(next)) + } }, 0) }, [onSelect], @@ -137,12 +201,24 @@ const FileExplorer: React.FC = ({ const deselectAllInSubtree = useCallback( (uri: string) => { - const node = indexRef.current.nodes.get(uri) - if (!node) return - setTimeout(() => { - const next = new Set(selectedUrisRef.current) - for (const u of getAllDescendantPaths(node.item)) next.delete(u) - startTransition(() => onSelect(next)) + setTimeout(async () => { + try { + const { uris } = await listFilesUnderUriRemote(uri) + const next = new Set(selectedUrisRef.current) + next.delete(uri) + for (const u of uris) next.delete(u) + const node = indexRef.current.nodes.get(uri) + if (node) { + for (const u of getAllDescendantPaths(node.item)) next.delete(u) + } + startTransition(() => onSelect(next)) + } catch { + const node = indexRef.current.nodes.get(uri) + if (!node) return + const next = new Set(selectedUrisRef.current) + for (const u of getAllDescendantPaths(node.item)) next.delete(u) + startTransition(() => onSelect(next)) + } }, 0) }, [onSelect], @@ -163,8 +239,8 @@ const FileExplorer: React.FC = ({ items: VscodeTreeItem[], depth = 0, ): React.ReactNode[] => { - return items.map((item, _itemIndex) => { - const isFolder = !!(item.subItems && item.subItems.length > 0) + return items.map((item) => { + const isFolder = isFolderItem(item) const totalDescFiles = index.descendantFileCount.get(item.value) || 0 const selectedDescFiles = selectedCountMap.get(item.value) || 0 const folderState = isFolder @@ -173,7 +249,7 @@ const FileExplorer: React.FC = ({ const folderTokens = isFolder ? tokenTotalsMap.get(item.value) || 0 : 0 const fileSelected = !isFolder && deferredSelectedUris.has(item.value) const fileTokens = !isFolder ? deferredTokenCounts[item.value] || 0 : 0 - const isOpen = item.open ?? depth === 0 + const isOpen = expandedUris.has(item.value) return ( = ({ depth={depth} isFolder={isFolder} isOpen={isOpen} + isLoadingChildren={loadingFolderUris.has(item.value)} totalDescendantFiles={totalDescFiles} selectedDescendantFiles={selectedDescFiles} folderSelectionState={folderState} @@ -197,7 +274,27 @@ const FileExplorer: React.FC = ({ }) } - // Handle actual double-clicks to open files in VS Code + const handleTreeClick = useCallback( + (e: React.MouseEvent) => { + const target = e.target as HTMLElement | null + if (!target) return + if (target.closest('button')) return + + const itemEl = target.closest( + 'vscode-tree-item[data-uri]', + ) as HTMLElement | null + if (!itemEl) return + const uri = itemEl.getAttribute('data-uri') + if (!uri) return + + const item = findItemByUri(fileTreeRef.current, uri) + if (!item || !isFolderItem(item)) return + + toggleFolderExpanded(uri) + }, + [findItemByUri, toggleFolderExpanded], + ) + const handleTreeDoubleClick = useCallback((e: React.MouseEvent) => { const target = e.target as HTMLElement | null if (!target) return @@ -205,14 +302,12 @@ const FileExplorer: React.FC = ({ if (!itemEl) return const uri = itemEl.getAttribute('data-uri') if (!uri) return - // Only open if this is a file (not a folder) const node = indexRef.current.nodes.get(uri) if (!node || node.isFolder) return const vscode = getVsCodeApi() vscode.postMessage({ command: 'openFile', payload: { fileUri: uri } }) }, []) - // Determine what content to show based on loading phase const renderContent = () => { switch (loadingPhase) { case 'initial': @@ -231,7 +326,19 @@ const FileExplorer: React.FC = ({
+ {treeTruncated ? ( +

+ File tree was truncated due to size limits. Expand folders to + load more, or add exclusions in Settings. +

+ ) : null} + {searchQuery ? ( +

+ Search only includes expanded and loaded folders. +

+ ) : null} = ({
{renderContent()} - {/* Refresh overlay - shows when refreshing existing data */} { + return new Promise((resolve, reject) => { + const requestId = `list-${Date.now()}-${Math.random().toString(36).slice(2)}` + const timeout = setTimeout(() => { + window.removeEventListener('message', handler) + reject(new Error('Listing files under folder timed out')) + }, LIST_FILES_TIMEOUT_MS) + + const handler = (event: MessageEvent) => { + const msg = event.data as { + command?: string + requestId?: string + uris?: string[] + truncated?: boolean + } + if ( + msg.command === 'listFilesUnderUriResponse' && + msg.requestId === requestId + ) { + clearTimeout(timeout) + window.removeEventListener('message', handler) + resolve({ + uris: Array.isArray(msg.uris) ? msg.uris : [], + truncated: Boolean(msg.truncated), + }) + } + } + + window.addEventListener('message', handler) + getVsCodeApi().postMessage({ + command: 'listFilesUnderUri', + payload: { parentUri, requestId }, + }) + }) +} diff --git a/src/webview-ui/src/components/context-tab/file-explorer/tree-index.ts b/src/webview-ui/src/components/context-tab/file-explorer/tree-index.ts index 2e5dade..0b057e7 100644 --- a/src/webview-ui/src/components/context-tab/file-explorer/tree-index.ts +++ b/src/webview-ui/src/components/context-tab/file-explorer/tree-index.ts @@ -19,12 +19,14 @@ export function buildTreeIndex(items: VscodeTreeItem[]): TreeIndex { const visit = (item: VscodeTreeItem): number => { const uri = item.value - const isFolder = !!(item.subItems && item.subItems.length > 0) + const isFolder = item.icons?.branch === 'folder' const children = item.subItems?.map((c) => c.value) ?? [] nodes.set(uri, { children, isFolder, item }) let fileCount = 0 if (isFolder) { - for (const child of item.subItems!) { + // subItems is undefined for lazily-loaded folders that haven't been + // expanded yet; treat as an empty (0-file) subtree until children load. + for (const child of item.subItems ?? []) { fileCount += visit(child) } } else { diff --git a/src/webview-ui/src/components/context-tab/file-explorer/tree-merge.ts b/src/webview-ui/src/components/context-tab/file-explorer/tree-merge.ts new file mode 100644 index 0000000..a7222a1 --- /dev/null +++ b/src/webview-ui/src/components/context-tab/file-explorer/tree-merge.ts @@ -0,0 +1,28 @@ +import type { VscodeTreeItem } from '../../../types' + +export function mergeChildren( + tree: VscodeTreeItem[], + parentUri: string, + children: VscodeTreeItem[], +): VscodeTreeItem[] { + const mergeInto = (items: VscodeTreeItem[]): VscodeTreeItem[] => { + return items.map((item) => { + if (item.value === parentUri) { + return { ...item, subItems: children } + } + if (item.subItems && item.subItems.length > 0) { + return { ...item, subItems: mergeInto(item.subItems) } + } + return item + }) + } + return mergeInto(tree) +} + +export function isFolderItem(item: VscodeTreeItem): boolean { + return item.icons?.branch === 'folder' +} + +export function folderNeedsLoad(item: VscodeTreeItem): boolean { + return isFolderItem(item) && item.subItems === undefined +} diff --git a/src/webview-ui/src/components/context-tab/file-explorer/tree-node.tsx b/src/webview-ui/src/components/context-tab/file-explorer/tree-node.tsx index 828577d..9daaae7 100644 --- a/src/webview-ui/src/components/context-tab/file-explorer/tree-node.tsx +++ b/src/webview-ui/src/components/context-tab/file-explorer/tree-node.tsx @@ -7,7 +7,8 @@ interface TreeNodeProps { item: VscodeTreeItem depth: number isFolder: boolean - isOpen?: boolean + isOpen: boolean + isLoadingChildren: boolean totalDescendantFiles: number selectedDescendantFiles: number folderSelectionState: FolderSelectionState @@ -24,7 +25,8 @@ const TreeNode: React.FC = ({ item, depth, isFolder, - isOpen = depth === 0, + isOpen, + isLoadingChildren, totalDescendantFiles, selectedDescendantFiles, folderSelectionState, @@ -36,6 +38,9 @@ const TreeNode: React.FC = ({ onDeselectAllInSubtree, renderChildren, }) => { + const showChildren = + isFolder && isOpen && item.subItems !== undefined && !isLoadingChildren + return ( = ({
{item.label} + {isLoadingChildren ? ( + Loading… + ) : null}
= ({ fileTokenCount={fileTokenCount} />
- {isFolder ? renderChildren(item.subItems || [], depth + 1) : null} + {showChildren ? renderChildren(item.subItems || [], depth + 1) : null} ) } @@ -79,13 +87,15 @@ function areEqual(prev: TreeNodeProps, next: TreeNodeProps): boolean { prev.item.value === next.item.value && prev.isFolder === next.isFolder && prev.isOpen === next.isOpen && + prev.isLoadingChildren === next.isLoadingChildren && prev.totalDescendantFiles === next.totalDescendantFiles && prev.selectedDescendantFiles === next.selectedDescendantFiles && prev.folderSelectionState === next.folderSelectionState && prev.folderTokenTotal === next.folderTokenTotal && prev.fileIsSelected === next.fileIsSelected && prev.fileTokenCount === next.fileTokenCount && - prev.depth === next.depth + prev.depth === next.depth && + prev.item.subItems === next.item.subItems ) } diff --git a/src/webview-ui/src/components/context-tab/index.tsx b/src/webview-ui/src/components/context-tab/index.tsx index dff9dd9..b16b963 100644 --- a/src/webview-ui/src/components/context-tab/index.tsx +++ b/src/webview-ui/src/components/context-tab/index.tsx @@ -21,6 +21,9 @@ interface ContextTabProps { onSelect: (uris: Set) => void onRefresh: (excludedFolders?: string) => void isLoading: boolean + treeTruncated?: boolean + loadingFolderUris: Set + onLoadChildren: (parentUri: string) => void } const ContextTab: React.FC = ({ @@ -31,6 +34,9 @@ const ContextTab: React.FC = ({ onSelect, onRefresh, isLoading, + treeTruncated, + loadingFolderUris, + onLoadChildren, }) => { const [userInstructions, setUserInstructions] = useState('') const [searchQuery, setSearchQuery] = useState('') @@ -208,6 +214,9 @@ const ContextTab: React.FC = ({ isLoading={isLoading} searchQuery={searchQuery} actualTokenCounts={actualTokenCounts} + treeTruncated={treeTruncated} + loadingFolderUris={loadingFolderUris} + onLoadChildren={onLoadChildren} />
diff --git a/src/webview-ui/src/components/context-tab/utils.ts b/src/webview-ui/src/components/context-tab/utils.ts index 058f3b5..3928cff 100644 --- a/src/webview-ui/src/components/context-tab/utils.ts +++ b/src/webview-ui/src/components/context-tab/utils.ts @@ -48,12 +48,17 @@ export function formatTokenCount(count: number): string { return `${(count / 1000).toFixed(1)}k` } -// Helper function to recursively gather all descendant paths +// Iteratively gather descendant paths (avoids deep recursion stack overflow) export const getAllDescendantPaths = (item: VscodeTreeItem): string[] => { - const paths = [item.value] - if (item.subItems) { - for (const sub of item.subItems) { - paths.push(...getAllDescendantPaths(sub)) + const paths: string[] = [] + const stack: VscodeTreeItem[] = [item] + while (stack.length > 0) { + const current = stack.pop()! + paths.push(current.value) + if (current.subItems) { + for (let i = current.subItems.length - 1; i >= 0; i--) { + stack.push(current.subItems[i]!) + } } } return paths diff --git a/src/webview-ui/src/utils/mock.ts b/src/webview-ui/src/utils/mock.ts index 18a7197..2f843dc 100644 --- a/src/webview-ui/src/utils/mock.ts +++ b/src/webview-ui/src/utils/mock.ts @@ -349,6 +349,65 @@ function buildMockFileTree(): VscodeTreeItem[] { return tree } +function isMockFolder(item: VscodeTreeItem): boolean { + return item.icons?.branch === 'folder' +} + +function findItemByUri( + items: VscodeTreeItem[], + uri: string, +): VscodeTreeItem | undefined { + for (const item of items) { + if (item.value === uri) return item + if (item.subItems) { + const found = findItemByUri(item.subItems, uri) + if (found) return found + } + } + return undefined +} + +function buildMockShallowRoots(full: VscodeTreeItem[]): VscodeTreeItem[] { + return full.map((root) => ({ + label: root.label, + value: root.value, + icons: root.icons, + })) +} + +function toLazyChildren( + subItems: VscodeTreeItem[] | undefined, +): VscodeTreeItem[] { + if (!subItems) return [] + return subItems.map((item) => { + if (isMockFolder(item)) { + const { subItems: _nested, ...rest } = item + return rest + } + return { ...item } + }) +} + +function collectFileUrisUnder(root: VscodeTreeItem): string[] { + const uris: string[] = [] + const walk = (items: VscodeTreeItem[] | undefined) => { + if (!items) return + for (const item of items) { + if (!isMockFolder(item)) { + uris.push(item.value) + } else { + walk(item.subItems) + } + } + } + if (isMockFolder(root)) { + walk(root.subItems) + } else { + uris.push(root.value) + } + return uris +} + function estimateTokensFromText(text: string): number { return Math.ceil((text || '').length / 4) } @@ -384,7 +443,44 @@ export function createMockVsCodeApi(): VsCodeApi { readGitignore = (payload as { readGitignore?: boolean }) .readGitignore as boolean } - sendToWebview({ command: 'updateFileTree', payload: fileTree }) + sendToWebview({ + command: 'updateFileTree', + payload: { + tree: buildMockShallowRoots(fileTree), + truncated: false, + }, + }) + break + } + case 'getDirectoryChildren': { + const parentUri = + isObject(payload) && + typeof (payload as { parentUri?: string }).parentUri === 'string' + ? (payload as { parentUri: string }).parentUri + : '' + const item = findItemByUri(fileTree, parentUri) + sendToWebview({ + command: 'updateDirectoryChildren', + parentUri, + children: toLazyChildren(item?.subItems), + truncated: false, + }) + break + } + case 'listFilesUnderUri': { + const p = payload as { + parentUri?: string + requestId?: string + } + const parentUri = p?.parentUri ?? '' + const requestId = p?.requestId ?? '' + const item = findItemByUri(fileTree, parentUri) + sendToWebview({ + command: 'listFilesUnderUriResponse', + requestId, + uris: item ? collectFileUrisUnder(item) : [], + truncated: false, + }) break } case 'getSettings': {