diff --git a/src/providers/file-explorer/index.ts b/src/providers/file-explorer/index.ts index 18b3b98..5578e06 100644 --- a/src/providers/file-explorer/index.ts +++ b/src/providers/file-explorer/index.ts @@ -9,6 +9,8 @@ import { import { telemetry } from '../../services/telemetry' import { countManyWithInfo, encodeText } from '../../services/token-counter' import { + type ListFilesProgress, + type ListFilesResult, getDirectoryChildren, getWorkspaceRoots, listFilesUnderUri, @@ -28,6 +30,21 @@ import type { UpdateSettingsPayload, } from './types' +const LIST_FILES_TIMEOUT_MS = 45_000 + +interface InFlightListFilesRequest { + key: string + controller: AbortController + timeout: NodeJS.Timeout + timedOut: boolean + requestIds: Set + promise: Promise +} + +function isAbortError(error: unknown): boolean { + return error instanceof Error && error.name === 'AbortError' +} + export class FileExplorerWebviewProvider implements vscode.WebviewViewProvider { public static readonly viewType = 'overwriteFilesWebview' @@ -37,6 +54,10 @@ export class FileExplorerWebviewProvider implements vscode.WebviewViewProvider { private _treeBuildAbort: AbortController | null = null private _treeBuildTimeout: NodeJS.Timeout | null = null private _pendingFileTreePayload: GetFileTreePayload | undefined + private readonly _listFilesByParent = new Map< + string, + InFlightListFilesRequest + >() private readonly _excludedDirs = [ '.git', '.hg', @@ -558,34 +579,115 @@ export class FileExplorerWebviewProvider implements vscode.WebviewViewProvider { ): 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), - }, - ) + const parentUri = payload.parentUri + const requestId = payload.requestId + const allExcludedDirs = this._resolveExcludedDirs(payload) + const useGitignore = this._readGitignoreFromPayload(payload) + const requestKey = JSON.stringify({ + parentUri, + excludedDirs: allExcludedDirs, + useGitignore, + }) + const existing = this._listFilesByParent.get(parentUri) - this._view.webview.postMessage({ - command: 'listFilesUnderUriResponse', - requestId: payload.requestId, - uris: result.uris, - truncated: result.truncated, + if (existing && existing.key === requestKey) { + existing.requestIds.add(requestId) + this._postListFilesProgress(existing, { + filesFound: 0, + nodesVisited: 0, + truncated: false, }) + try { + await this._respondWhenListFilesCompletes(existing, requestId) + } catch (error) { + this._view?.webview.postMessage({ + command: 'listFilesUnderUriResponse', + requestId, + uris: [], + truncated: false, + aborted: isAbortError(error) && !existing.timedOut, + error: error instanceof Error ? error.message : String(error), + }) + } + return + } + + if (existing) { + existing.controller.abort() + clearTimeout(existing.timeout) + this._listFilesByParent.delete(parentUri) + } + + const controller = new AbortController() + let inFlight: InFlightListFilesRequest + const timeout = setTimeout(() => { + if (inFlight) inFlight.timedOut = true + controller.abort() + }, LIST_FILES_TIMEOUT_MS) + const promise = listFilesUnderUri(parentUri, allExcludedDirs, { + useGitignore, + signal: controller.signal, + onProgress: (progress) => this._postListFilesProgress(inFlight, progress), + }) + inFlight = { + key: requestKey, + controller, + timeout, + timedOut: false, + requestIds: new Set([requestId]), + promise, + } + this._listFilesByParent.set(parentUri, inFlight) + + try { + await this._respondWhenListFilesCompletes(inFlight, requestId) } catch (error) { - this._handleError(error, 'Error listing files under folder') - this._view.webview.postMessage({ + const aborted = isAbortError(error) && !inFlight.timedOut + if (!aborted) { + this._handleError(error, 'Error listing files under folder') + } + this._view?.webview.postMessage({ command: 'listFilesUnderUriResponse', - requestId: payload.requestId, + requestId, uris: [], truncated: false, + aborted, error: error instanceof Error ? error.message : String(error), }) + } finally { + if (this._listFilesByParent.get(parentUri) === inFlight) { + this._listFilesByParent.delete(parentUri) + } + clearTimeout(timeout) } } + private _postListFilesProgress( + request: InFlightListFilesRequest, + progress: ListFilesProgress, + ): void { + for (const requestId of request.requestIds) { + this._view?.webview.postMessage({ + command: 'listFilesUnderUriProgress', + requestId, + ...progress, + }) + } + } + + private async _respondWhenListFilesCompletes( + request: InFlightListFilesRequest, + requestId: string, + ): Promise { + const result = await request.promise + this._view?.webview.postMessage({ + command: 'listFilesUnderUriResponse', + requestId, + uris: result.uris, + truncated: result.truncated, + }) + } + /** * Opens the specified file in the VS Code editor. */ @@ -1112,21 +1214,12 @@ export class FileExplorerWebviewProvider implements vscode.WebviewViewProvider { const urisToCount = Array.isArray(payload?.selectedUris) ? payload.selectedUris : [] + const tokenRequestId = payload?.requestId const uris = urisToCount.map((uriString) => vscode.Uri.parse(uriString)) - // Calculate total selected size for bucketing - let totalSizeBytes = 0 - for (const uri of uris) { - try { - const stat = await vscode.workspace.fs.stat(uri) - if (stat.type === vscode.FileType.File) { - totalSizeBytes += stat.size - } - } catch { - // Skip files that can't be accessed - } - } + const { tokenCounts, skippedFiles, totalSizeBytes } = + await countManyWithInfo(uris) // Track token count started try { @@ -1138,8 +1231,6 @@ export class FileExplorerWebviewProvider implements vscode.WebviewViewProvider { console.warn('[telemetry] failed to capture token_count_started', e) } - const { tokenCounts, skippedFiles } = await countManyWithInfo(uris) - const duration = Date.now() - startTime const totalTokens = Object.values(tokenCounts).reduce( (sum, count) => sum + count, @@ -1160,7 +1251,7 @@ export class FileExplorerWebviewProvider implements vscode.WebviewViewProvider { // Send the results back to the webview this._view.webview.postMessage({ command: 'updateTokenCounts', - payload: { tokenCounts, skippedFiles }, + payload: { requestId: tokenRequestId, tokenCounts, skippedFiles }, }) // Log skipped files for debugging @@ -1184,7 +1275,11 @@ export class FileExplorerWebviewProvider implements vscode.WebviewViewProvider { // Send empty counts back on error this._view.webview.postMessage({ command: 'updateTokenCounts', - payload: { tokenCounts: {}, skippedFiles: [] }, + payload: { + requestId: payload?.requestId, + tokenCounts: {}, + skippedFiles: [], + }, }) } } diff --git a/src/providers/file-explorer/types.ts b/src/providers/file-explorer/types.ts index a320a8e..2f442e0 100644 --- a/src/providers/file-explorer/types.ts +++ b/src/providers/file-explorer/types.ts @@ -13,6 +13,7 @@ export interface CopyContextPayload { // defining it here for clarity or future use if refactored. export interface GetTokenCountsPayload { selectedUris: string[] // Array of URI strings + requestId?: string } export interface GetFileTreePayload { diff --git a/src/services/token-counter.ts b/src/services/token-counter.ts index ae1e37d..8d5d166 100644 --- a/src/services/token-counter.ts +++ b/src/services/token-counter.ts @@ -133,6 +133,7 @@ export async function countTokens(uri: vscode.Uri): Promise { */ export async function countTokensWithInfo(uri: vscode.Uri): Promise<{ tokens: number + sizeBytes: number skipped?: SkippedFileInfo }> { try { @@ -140,24 +141,24 @@ export async function countTokensWithInfo(uri: vscode.Uri): Promise<{ // Only process files if (stats.type !== vscode.FileType.File) { - return { tokens: 0 } + return { tokens: 0, sizeBytes: 0 } } - const stats2 = await stat(uri.fsPath) - const mtime = Math.floor(stats2.mtime.getTime() / 1000) - const size = stats2.size + const mtime = Math.floor(stats.mtime / 1000) + const size = stats.size // 1️⃣ Cache hit? const key = uri.fsPath const entry = cache.get(key) if (entry && entry.mtime === mtime && entry.size === size) { - return { tokens: entry.tokens } + return { tokens: entry.tokens, sizeBytes: size } } // 2️⃣ Safeguards if (size > MAX_BYTES) { return { tokens: 0, + sizeBytes: size, skipped: { uri: uri.toString(), reason: 'too-large', @@ -169,70 +170,74 @@ export async function countTokensWithInfo(uri: vscode.Uri): Promise<{ const encoder = await getEncoder() // 3️⃣ Stream + incremental encode → memory stays flat - return new Promise<{ tokens: number; skipped?: SkippedFileInfo }>( - (resolve, reject) => { - let tokens = 0 - let firstChunkProcessed = false - let hasError = false - - const stream = createReadStream(uri.fsPath, { highWaterMark: 64_000 }) - - stream.on('data', (chunk: Buffer | string) => { - if (hasError) return - - const buf = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk) - - if (!firstChunkProcessed) { - firstChunkProcessed = true - if (looksBinary(buf)) { - // early out for binaries - hasError = true - stream.destroy() - resolve({ - tokens: 0, - skipped: { - uri: uri.toString(), - reason: 'binary', - message: 'Binary file detected', - }, - }) - return - } - } + return new Promise<{ + tokens: number + sizeBytes: number + skipped?: SkippedFileInfo + }>((resolve, reject) => { + let tokens = 0 + let firstChunkProcessed = false + let hasError = false + + const stream = createReadStream(uri.fsPath, { highWaterMark: 64_000 }) + + stream.on('data', (chunk: Buffer | string) => { + if (hasError) return + + const buf = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk) - try { - const text = buf.toString('utf8') - tokens += encoder.encode(text).length - } catch (error) { + if (!firstChunkProcessed) { + firstChunkProcessed = true + if (looksBinary(buf)) { + // early out for binaries hasError = true stream.destroy() - reject(error) + resolve({ + tokens: 0, + sizeBytes: size, + skipped: { + uri: uri.toString(), + reason: 'binary', + message: 'Binary file detected', + }, + }) + return } - }) + } - stream.on('error', (error) => { + try { + const text = buf.toString('utf8') + tokens += encoder.encode(text).length + } catch (error) { hasError = true - // Track token counting errors - try { - telemetry.trackUnhandled('backend', error) - } catch (e) { - console.warn('[telemetry] failed to track token counter error', e) - } + stream.destroy() reject(error) - }) + } + }) - stream.on('end', () => { - if (!hasError) { - cache.set(key, { mtime, size, tokens }) - resolve({ tokens }) - } - }) - }, - ) + stream.on('error', (error) => { + hasError = true + // Track token counting errors + try { + telemetry.trackUnhandled('backend', error) + } catch (e) { + console.warn('[telemetry] failed to track token counter error', e) + } + reject(error) + }) + + stream.on('end', () => { + if (!hasError) { + cache.set(key, { mtime, size, tokens }) + resolve({ tokens, sizeBytes: size }) + } + }) + }) } catch (error) { // File doesn't exist or can't be accessed return { tokens: 0, + sizeBytes: 0, skipped: { uri: uri.toString(), reason: 'error', @@ -261,6 +266,7 @@ export async function countMany( export async function countManyWithInfo(uris: vscode.Uri[]): Promise<{ tokenCounts: Record skippedFiles: SkippedFileInfo[] + totalSizeBytes: number }> { const limit = await getLimit() const results = await Promise.all( @@ -269,19 +275,21 @@ export async function countManyWithInfo(uris: vscode.Uri[]): Promise<{ const tokenCounts: Record = {} const skippedFiles: SkippedFileInfo[] = [] + let totalSizeBytes = 0 for (let i = 0; i < uris.length; i++) { const uriString = uris[i].toString() const result = results[i] tokenCounts[uriString] = result.tokens + totalSizeBytes += result.sizeBytes if (result.skipped) { skippedFiles.push(result.skipped) } } - return { tokenCounts, skippedFiles } + return { tokenCounts, skippedFiles, totalSizeBytes } } /** diff --git a/src/utils/file-system.ts b/src/utils/file-system.ts index 5581e51..39c6c27 100644 --- a/src/utils/file-system.ts +++ b/src/utils/file-system.ts @@ -119,6 +119,7 @@ const FILE_ICONS = { export interface FileTreeOptions { useGitignore?: boolean signal?: AbortSignal + onProgress?: (progress: ListFilesProgress) => void } export interface FileTreeResult { @@ -730,6 +731,12 @@ export interface ListFilesResult { truncated: boolean } +export interface ListFilesProgress { + filesFound: number + nodesVisited: number + truncated: boolean +} + /** * Lists file URIs under a folder URI (for select-all), respecting ignores and caps. */ @@ -758,14 +765,31 @@ export async function listFilesUnderUri( } const queue: vscode.Uri[] = [targetUri] + let cursor = 0 let visited = 0 + let lastProgressAt = Date.now() + + const reportProgress = (force = false) => { + if (!options?.onProgress) return + const now = Date.now() + if (!force && now - lastProgressAt < 250) return + lastProgressAt = now + options.onProgress({ + filesFound: uris.length, + nodesVisited: visited, + truncated, + }) + } + + reportProgress(true) - while (queue.length > 0) { + while (cursor < queue.length) { throwIfAborted(signal) - const currentUri = queue.shift()! + const currentUri = queue[cursor++] let entries: [string, vscode.FileType][] try { entries = await vscode.workspace.fs.readDirectory(currentUri) + reportProgress() } catch { continue } @@ -774,6 +798,7 @@ export async function listFilesUnderUri( throwIfAborted(signal) if (uris.length >= MAX_LIST_FILES) { truncated = true + reportProgress(true) return { uris, truncated } } @@ -790,16 +815,19 @@ export async function listFilesUnderUri( visited++ if (visited > MAX_TREE_NODES) { truncated = true + reportProgress(true) return { uris, truncated } } if (type === vscode.FileType.File) { uris.push(entryUri.toString()) + reportProgress() } else if (type === vscode.FileType.Directory) { queue.push(entryUri) } } } + reportProgress(true) return { uris, truncated } } diff --git a/src/webview-ui/src/App.tsx b/src/webview-ui/src/App.tsx index f69c4d8..650f266 100644 --- a/src/webview-ui/src/App.tsx +++ b/src/webview-ui/src/App.tsx @@ -1,10 +1,13 @@ import type { VscTabsSelectEvent } from '@vscode-elements/elements/dist/vscode-tabs/vscode-tabs' -import { useCallback, useEffect, useState } from 'react' +import { useCallback, useEffect, useRef, useState } from 'react' import type { VscodeTreeItem } from './types' // Import tree item type from local types 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 { + buildTreeIndex, + mergeChildrenIntoIndexedTree, +} from './components/context-tab/file-explorer/tree-index' import SettingsTab from './components/settings-tab' import { getVsCodeApi } from './utils/vscode' // Import the new utility @@ -65,6 +68,7 @@ function App() { ) const [excludedFolders, setExcludedFolders] = useState('') // Persisted excluded folders const [readGitignore, setReadGitignore] = useState(true) + const treeIndexRef = useRef(buildTreeIndex([])) // Send message to extension using the utility const sendMessage = useCallback((command: string, payload?: unknown) => { @@ -158,6 +162,7 @@ function App() { const { tree, truncated } = parseUpdateFileTreePayload( message.payload, ) + treeIndexRef.current = buildTreeIndex(tree) setFileTreeData(tree) setTreeTruncated(truncated) setLoadingFolderUris(new Set()) @@ -167,9 +172,16 @@ function App() { case 'updateDirectoryChildren': { const msg = message as unknown as UpdateDirectoryChildrenMessage if (msg.parentUri && Array.isArray(msg.children)) { - setFileTreeData((prev) => - mergeChildren(prev, msg.parentUri, msg.children), - ) + setFileTreeData((prev) => { + const merged = mergeChildrenIntoIndexedTree( + prev, + treeIndexRef.current, + msg.parentUri, + msg.children, + ) + treeIndexRef.current = merged.index + return merged.tree + }) setLoadingFolderUris((prev) => { const next = new Set(prev) next.delete(msg.parentUri) @@ -181,6 +193,8 @@ function App() { } case 'listFilesUnderUriResponse': break + case 'listFilesUnderUriProgress': + break case 'showError': // Display error message in a dismissible banner { 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 6c20c6b..002ee3b 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 @@ -46,6 +46,12 @@ const FileExplorer: React.FC = ({ const [isRefreshing, setIsRefreshing] = useState(false) const [, setPrevFileTreeData] = useState([]) const [expandedUris, setExpandedUris] = useState>(() => new Set()) + const [listingFolderUris, setListingFolderUris] = useState>( + () => new Set(), + ) + const [listingFileCounts, setListingFileCounts] = useState< + Map + >(() => new Map()) const deferredSelectedUris = useDeferredValue(selectedUris) const deferredTokenCounts = useDeferredValue(actualTokenCounts) @@ -103,27 +109,17 @@ const FileExplorer: React.FC = ({ }, [fileTreeData, searchQuery]) const index = useMemo(() => buildTreeIndex(visibleItems), [visibleItems]) + const fullTreeIndex = useMemo( + () => buildTreeIndex(fileTreeData), + [fileTreeData], + ) const selectedUrisRef = useRef(selectedUris) const indexRef = useRef(index) - const fileTreeRef = useRef(fileTreeData) + const fullTreeIndexRef = useRef(fullTreeIndex) 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 - }, - [], - ) + fullTreeIndexRef.current = fullTreeIndex const toggleFolderExpanded = useCallback( (uri: string) => { @@ -133,7 +129,7 @@ const FileExplorer: React.FC = ({ next.delete(uri) } else { next.add(uri) - const item = findItemByUri(fileTreeRef.current, uri) + const item = fullTreeIndexRef.current.nodes.get(uri)?.item if (item && folderNeedsLoad(item)) { onLoadChildren(uri) } @@ -141,7 +137,7 @@ const FileExplorer: React.FC = ({ return next }) }, - [findItemByUri, onLoadChildren], + [onLoadChildren], ) const { selectedCountMap, tokenTotalsMap } = useMemo(() => { @@ -180,19 +176,38 @@ const FileExplorer: React.FC = ({ const selectAllInSubtree = useCallback( (uri: string) => { setTimeout(async () => { + setListingFolderUris((prev) => new Set(prev).add(uri)) try { - const { uris } = await listFilesUnderUriRemote(uri) + const { uris } = await listFilesUnderUriRemote(uri, (progress) => { + setListingFileCounts((prev) => { + const next = new Map(prev) + next.set(uri, progress.filesFound) + return next + }) + }) const next = new Set(selectedUrisRef.current) next.add(uri) for (const u of uris) next.add(u) startTransition(() => onSelect(next)) } catch (err) { + if (err instanceof Error && err.name === 'AbortError') return 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)) + } finally { + setListingFolderUris((prev) => { + const next = new Set(prev) + next.delete(uri) + return next + }) + setListingFileCounts((prev) => { + const next = new Map(prev) + next.delete(uri) + return next + }) } }, 0) }, @@ -202,8 +217,15 @@ const FileExplorer: React.FC = ({ const deselectAllInSubtree = useCallback( (uri: string) => { setTimeout(async () => { + setListingFolderUris((prev) => new Set(prev).add(uri)) try { - const { uris } = await listFilesUnderUriRemote(uri) + const { uris } = await listFilesUnderUriRemote(uri, (progress) => { + setListingFileCounts((prev) => { + const next = new Map(prev) + next.set(uri, progress.filesFound) + return next + }) + }) const next = new Set(selectedUrisRef.current) next.delete(uri) for (const u of uris) next.delete(u) @@ -212,12 +234,24 @@ const FileExplorer: React.FC = ({ for (const u of getAllDescendantPaths(node.item)) next.delete(u) } startTransition(() => onSelect(next)) - } catch { + } catch (err) { + if (err instanceof Error && err.name === 'AbortError') return 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)) + } finally { + setListingFolderUris((prev) => { + const next = new Set(prev) + next.delete(uri) + return next + }) + setListingFileCounts((prev) => { + const next = new Map(prev) + next.delete(uri) + return next + }) } }, 0) }, @@ -250,6 +284,8 @@ const FileExplorer: React.FC = ({ const fileSelected = !isFolder && deferredSelectedUris.has(item.value) const fileTokens = !isFolder ? deferredTokenCounts[item.value] || 0 : 0 const isOpen = expandedUris.has(item.value) + const isListingFiles = listingFolderUris.has(item.value) + const listingFileCount = listingFileCounts.get(item.value) ?? 0 return ( = ({ isFolder={isFolder} isOpen={isOpen} isLoadingChildren={loadingFolderUris.has(item.value)} + isListingFiles={isListingFiles} + listingFileCount={listingFileCount} totalDescendantFiles={totalDescFiles} selectedDescendantFiles={selectedDescFiles} folderSelectionState={folderState} @@ -287,12 +325,12 @@ const FileExplorer: React.FC = ({ const uri = itemEl.getAttribute('data-uri') if (!uri) return - const item = findItemByUri(fileTreeRef.current, uri) - if (!item || !isFolderItem(item)) return + const node = fullTreeIndexRef.current.nodes.get(uri) + if (!node?.isFolder) return toggleFolderExpanded(uri) }, - [findItemByUri, toggleFolderExpanded], + [toggleFolderExpanded], ) const handleTreeDoubleClick = useCallback((e: React.MouseEvent) => { diff --git a/src/webview-ui/src/components/context-tab/file-explorer/list-files-under-uri.ts b/src/webview-ui/src/components/context-tab/file-explorer/list-files-under-uri.ts index 5afcb87..b3c5653 100644 --- a/src/webview-ui/src/components/context-tab/file-explorer/list-files-under-uri.ts +++ b/src/webview-ui/src/components/context-tab/file-explorer/list-files-under-uri.ts @@ -4,6 +4,11 @@ const LIST_FILES_TIMEOUT_MS = 60_000 export function listFilesUnderUriRemote( parentUri: string, + onProgress?: (progress: { + filesFound: number + nodesVisited: number + truncated: boolean + }) => void, ): Promise<{ uris: string[]; truncated: boolean }> { return new Promise((resolve, reject) => { const requestId = `list-${Date.now()}-${Math.random().toString(36).slice(2)}` @@ -18,6 +23,22 @@ export function listFilesUnderUriRemote( requestId?: string uris?: string[] truncated?: boolean + filesFound?: number + nodesVisited?: number + error?: string + aborted?: boolean + } + if ( + msg.command === 'listFilesUnderUriProgress' && + msg.requestId === requestId + ) { + onProgress?.({ + filesFound: typeof msg.filesFound === 'number' ? msg.filesFound : 0, + nodesVisited: + typeof msg.nodesVisited === 'number' ? msg.nodesVisited : 0, + truncated: Boolean(msg.truncated), + }) + return } if ( msg.command === 'listFilesUnderUriResponse' && @@ -25,6 +46,12 @@ export function listFilesUnderUriRemote( ) { clearTimeout(timeout) window.removeEventListener('message', handler) + if (msg.error) { + const error = new Error(msg.error) + if (msg.aborted) error.name = 'AbortError' + reject(error) + return + } resolve({ uris: Array.isArray(msg.uris) ? msg.uris : [], truncated: Boolean(msg.truncated), 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 0b057e7..0eae542 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 @@ -7,6 +7,9 @@ export interface IndexedNode { } export interface TreeIndex { nodes: Map + parentByUri: Map + childIndexByUri: Map + rootIndexByUri: Map postOrder: string[] roots: VscodeTreeItem[] descendantFileCount: Map @@ -14,20 +17,30 @@ export interface TreeIndex { export function buildTreeIndex(items: VscodeTreeItem[]): TreeIndex { const nodes = new Map() + const parentByUri = new Map() + const childIndexByUri = new Map() + const rootIndexByUri = new Map() const descendantFileCount = new Map() const postOrder: string[] = [] - const visit = (item: VscodeTreeItem): number => { + const visit = ( + item: VscodeTreeItem, + parentUri?: string, + childIndex?: number, + ): number => { const uri = item.value const isFolder = item.icons?.branch === 'folder' const children = item.subItems?.map((c) => c.value) ?? [] nodes.set(uri, { children, isFolder, item }) + if (parentUri) parentByUri.set(uri, parentUri) + if (childIndex !== undefined) childIndexByUri.set(uri, childIndex) let fileCount = 0 if (isFolder) { // 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) + for (let i = 0; i < (item.subItems ?? []).length; i++) { + const child = item.subItems![i] + fileCount += visit(child, uri, i) } } else { fileCount = 1 @@ -36,6 +49,145 @@ export function buildTreeIndex(items: VscodeTreeItem[]): TreeIndex { postOrder.push(uri) return fileCount } - for (const root of items) visit(root) - return { nodes, postOrder, roots: items, descendantFileCount } + for (let i = 0; i < items.length; i++) { + rootIndexByUri.set(items[i].value, i) + visit(items[i], undefined, i) + } + return { + nodes, + parentByUri, + childIndexByUri, + rootIndexByUri, + postOrder, + roots: items, + descendantFileCount, + } +} + +function removeIndexedSubtree(index: TreeIndex, uri: string): void { + const node = index.nodes.get(uri) + if (!node) return + for (const childUri of node.children) { + removeIndexedSubtree(index, childUri) + } + index.nodes.delete(uri) + index.parentByUri.delete(uri) + index.childIndexByUri.delete(uri) + index.rootIndexByUri.delete(uri) + index.descendantFileCount.delete(uri) +} + +function addIndexedSubtree( + index: TreeIndex, + item: VscodeTreeItem, + parentUri: string, + childIndex: number, +): number { + const isFolder = item.icons?.branch === 'folder' + const children = item.subItems?.map((child) => child.value) ?? [] + index.nodes.set(item.value, { children, isFolder, item }) + index.parentByUri.set(item.value, parentUri) + index.childIndexByUri.set(item.value, childIndex) + + let fileCount = isFolder ? 0 : 1 + if (isFolder) { + for (let i = 0; i < (item.subItems ?? []).length; i++) { + fileCount += addIndexedSubtree(index, item.subItems![i], item.value, i) + } + } + index.descendantFileCount.set(item.value, fileCount) + return fileCount +} + +function buildPathToRoot(index: TreeIndex, uri: string): string[] { + const path = [uri] + let current = uri + while (index.parentByUri.has(current)) { + current = index.parentByUri.get(current)! + path.push(current) + } + return path.reverse() +} + +function cloneWithChildrenAtPath( + item: VscodeTreeItem, + path: string[], + pathIndex: number, + children: VscodeTreeItem[], + index: TreeIndex, +): VscodeTreeItem { + if (pathIndex === path.length - 1) { + return { ...item, subItems: children } + } + + const nextUri = path[pathIndex + 1] + const childIndex = index.childIndexByUri.get(nextUri) + if (childIndex === undefined || !item.subItems) return item + + const nextSubItems = item.subItems.slice() + nextSubItems[childIndex] = cloneWithChildrenAtPath( + nextSubItems[childIndex], + path, + pathIndex + 1, + children, + index, + ) + return { ...item, subItems: nextSubItems } +} + +export function mergeChildrenIntoIndexedTree( + tree: VscodeTreeItem[], + index: TreeIndex, + parentUri: string, + children: VscodeTreeItem[], +): { tree: VscodeTreeItem[]; index: TreeIndex } { + if (!index.nodes.has(parentUri)) return { tree, index } + + const path = buildPathToRoot(index, parentUri) + const rootUri = path[0] + const rootIndex = index.rootIndexByUri.get(rootUri) + if (rootIndex === undefined) return { tree, index } + + const nextTree = tree.slice() + const updatedRoot = cloneWithChildrenAtPath( + nextTree[rootIndex], + path, + 0, + children, + index, + ) + nextTree[rootIndex] = updatedRoot + + const nextIndex: TreeIndex = { + ...index, + nodes: new Map(index.nodes), + parentByUri: new Map(index.parentByUri), + childIndexByUri: new Map(index.childIndexByUri), + rootIndexByUri: new Map(index.rootIndexByUri), + descendantFileCount: new Map(index.descendantFileCount), + roots: nextTree, + } + + const parentNode = nextIndex.nodes.get(parentUri) + if (parentNode) { + for (const childUri of parentNode.children) { + removeIndexedSubtree(nextIndex, childUri) + } + const nextChildren = children.map((child) => child.value) + nextIndex.nodes.set(parentUri, { + ...parentNode, + children: nextChildren, + item: + path.length === 1 + ? updatedRoot + : { ...parentNode.item, subItems: children }, + }) + let fileCount = 0 + for (let i = 0; i < children.length; i++) { + fileCount += addIndexedSubtree(nextIndex, children[i], parentUri, i) + } + nextIndex.descendantFileCount.set(parentUri, fileCount) + } + + return { tree: nextTree, index: nextIndex } } 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 9daaae7..99ee7c9 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 @@ -9,6 +9,8 @@ interface TreeNodeProps { isFolder: boolean isOpen: boolean isLoadingChildren: boolean + isListingFiles: boolean + listingFileCount: number totalDescendantFiles: number selectedDescendantFiles: number folderSelectionState: FolderSelectionState @@ -27,6 +29,8 @@ const TreeNode: React.FC = ({ isFolder, isOpen, isLoadingChildren, + isListingFiles, + listingFileCount, totalDescendantFiles, selectedDescendantFiles, folderSelectionState, @@ -57,6 +61,13 @@ const TreeNode: React.FC = ({ {isLoadingChildren ? ( Loading… ) : null} + {isListingFiles ? ( + + {listingFileCount > 0 + ? `Selecting ${listingFileCount}…` + : 'Selecting…'} + + ) : null}
= ({ // Debounce timer for user instructions token counting (use ref to avoid re-renders) const debounceRef = useRef | null>(null) + const latestTokenRequestIdRef = useRef(null) // Constant for XML formatting instructions const XML_INSTRUCTIONS_TOKENS = 5000 // This is an approximation @@ -110,14 +111,19 @@ const ContextTab: React.FC = ({ const vscode = getVsCodeApi() const urisArray = Array.from(selectedUris) if (urisArray.length === 0) { + latestTokenRequestIdRef.current = null setActualTokenCounts({}) setSkippedFiles([]) return } + const requestId = `tokens-${Date.now()}-${Math.random() + .toString(36) + .slice(2)}` + latestTokenRequestIdRef.current = requestId const handle = setTimeout(() => { vscode.postMessage({ command: 'getTokenCounts', - payload: { selectedUris: urisArray }, + payload: { selectedUris: urisArray, requestId }, }) }, 200) return () => clearTimeout(handle) @@ -128,31 +134,40 @@ const ContextTab: React.FC = ({ const handleMessage = (event: MessageEvent) => { const message = event.data if (message.command === 'updateTokenCounts') { + const requestId = message.payload?.requestId + if ( + typeof requestId === 'string' && + requestId !== latestTokenRequestIdRef.current + ) { + return + } const incoming: Record = message.payload.tokenCounts || {} // Shallow diff and only update if changed - let changed = false - const next: Record = { ...actualTokenCounts } - for (const [k, v] of Object.entries(incoming)) { - if (next[k] !== v) { - next[k] = v - changed = true + setActualTokenCounts((current) => { + let changed = false + const next: Record = { ...current } + for (const [k, v] of Object.entries(incoming)) { + if (next[k] !== v) { + next[k] = v + changed = true + } } - } - // Remove keys that no longer exist - for (const k of Object.keys(next)) { - if (!(k in incoming)) { - delete next[k] - changed = true + // Remove keys that no longer exist + for (const k of Object.keys(next)) { + if (!(k in incoming)) { + delete next[k] + changed = true + } } - } - if (changed) setActualTokenCounts(next) + return changed ? next : current + }) setSkippedFiles(message.payload.skippedFiles || []) } } window.addEventListener('message', handleMessage) return () => window.removeEventListener('message', handleMessage) - }, [actualTokenCounts]) + }, []) const handleRefreshClick = useCallback(() => { // Reset skipped files and token counts when refreshing to clear any deleted files diff --git a/src/webview-ui/src/utils/mock.ts b/src/webview-ui/src/utils/mock.ts index 2f843dc..1a9a21b 100644 --- a/src/webview-ui/src/utils/mock.ts +++ b/src/webview-ui/src/utils/mock.ts @@ -18,7 +18,7 @@ import type { VscodeTreeItem } from '../types' type ExcludedFoldersPayload = { excludedFolders: string } type SaveSettingsPayload = { excludedFolders: string; readGitignore: boolean } -type TokenCountsPayload = { selectedUris: string[] } +type TokenCountsPayload = { selectedUris: string[]; requestId?: string } type TokenCountPayload = { text: string; requestId: string } type OpenFilePayload = { fileUri: string } @@ -475,10 +475,18 @@ export function createMockVsCodeApi(): VsCodeApi { const parentUri = p?.parentUri ?? '' const requestId = p?.requestId ?? '' const item = findItemByUri(fileTree, parentUri) + const uris = item ? collectFileUrisUnder(item) : [] + sendToWebview({ + command: 'listFilesUnderUriProgress', + requestId, + filesFound: uris.length, + nodesVisited: uris.length, + truncated: false, + }) sendToWebview({ command: 'listFilesUnderUriResponse', requestId, - uris: item ? collectFileUrisUnder(item) : [], + uris, truncated: false, }) break @@ -524,13 +532,16 @@ export function createMockVsCodeApi(): VsCodeApi { const selectedUris: string[] = isTokenCountsPayload(payload) ? payload.selectedUris : [] + const requestId = isTokenCountsPayload(payload) + ? payload.requestId + : undefined const tokenCounts: Record = {} for (const uri of selectedUris) { tokenCounts[uri] = Math.max(10, estimateTokensFromText(uri) * 3) } sendToWebview({ command: 'updateTokenCounts', - payload: { tokenCounts, skippedFiles: [] }, + payload: { requestId, tokenCounts, skippedFiles: [] }, }) break }