From c4dcb3d225feb571c71b08c85f129aba5600dc3e Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Mon, 17 Aug 2026 10:06:45 +0200 Subject: [PATCH] fix(vscode): keep file mentions fresh --- .changeset/fresh-file-mentions.md | 5 + packages/kilo-vscode/src/KiloProvider.ts | 3 +- .../src/kilo-provider/file-search.ts | 37 +-- .../tests/unit/file-search.test.ts | 74 ++++++ .../tests/unit/use-file-mention.test.ts | 236 +++++++++++++++++- .../webview-ui/src/hooks/useFileMention.ts | 137 ++++++++-- 6 files changed, 426 insertions(+), 66 deletions(-) create mode 100644 .changeset/fresh-file-mentions.md create mode 100644 packages/kilo-vscode/tests/unit/file-search.test.ts diff --git a/.changeset/fresh-file-mentions.md b/.changeset/fresh-file-mentions.md new file mode 100644 index 0000000000..9aebe9a418 --- /dev/null +++ b/.changeset/fresh-file-mentions.md @@ -0,0 +1,5 @@ +--- +"kilo-code": patch +--- + +Keep file mention suggestions current and scoped to the active workspace while preserving instant cached results. diff --git a/packages/kilo-vscode/src/KiloProvider.ts b/packages/kilo-vscode/src/KiloProvider.ts index c61671bc46..144b7fdd61 100644 --- a/packages/kilo-vscode/src/KiloProvider.ts +++ b/packages/kilo-vscode/src/KiloProvider.ts @@ -63,7 +63,7 @@ import { slimInfo, slimPart, slimParts } from "./kilo-provider/slim-metadata" import { handleSidebarWorktreeMessage } from "./kilo-provider/sidebar-worktree" import { parseMessageFiles, type MessageFile } from "./kilo-provider/message-files" import { renameSession } from "./kilo-provider/rename-session" -import { handleFileSearch, prewarmFileSearch } from "./kilo-provider/file-search" +import { handleFileSearch } from "./kilo-provider/file-search" import { handleSessionSearch } from "./kilo-provider/session-search" import { handleFilePicker } from "./kilo-provider/file-picker" import { watchFontSizeConfig } from "./kilo-provider/font-size" @@ -1853,7 +1853,6 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper this.seedSessionStatusMap(), ]) await this.refreshGitStatus(this.getWorkspaceDirectory()) - prewarmFileSearch(this.client, this.getWorkspaceDirectory()) this.sendNotificationSettings() this.sendTimelineSetting() this.postMessage(buildThroughputSettingMessage()) diff --git a/packages/kilo-vscode/src/kilo-provider/file-search.ts b/packages/kilo-vscode/src/kilo-provider/file-search.ts index ee51ee88d0..f709361099 100644 --- a/packages/kilo-vscode/src/kilo-provider/file-search.ts +++ b/packages/kilo-vscode/src/kilo-provider/file-search.ts @@ -20,23 +20,6 @@ type Input = { post: (message: unknown) => void } -type Cache = { - files: string[] - folders: string[] - updated: number -} - -const cache = new Map() - -export function prewarmFileSearch(client: KiloClient | null, dir: string): void { - if (!client || !dir) return - void fetchBackend(client, dir, "").then(([files, folders]) => { - if (files.length || folders.length) { - cache.set(dir, { files, folders, updated: Date.now() }) - } - }) -} - async function fetchBackend(client: KiloClient, dir: string, query: string): Promise<[string[], string[]]> { if (!client?.find?.files) return [[], []] const [fileRes, folderRes] = await Promise.allSettled([ @@ -76,22 +59,10 @@ export async function handleFileSearch(input: Input): Promise { const id = input.message.sessionID ?? input.current ?? input.context const dir = input.dir(id) const query = input.message.query - const entry = !query && dir ? cache.get(dir) : undefined - - if (entry) { - const open = await input.open(dir) - const { paths, items } = assemble(query, dir, entry.files, entry.folders, open) - input.post({ type: "fileSearchResult", paths, items, dir, requestId: input.message.requestId }) - } - - void fetchBackend(client, dir, query).then(async ([files, folders]) => { - if (!query && dir && (files.length || folders.length)) { - cache.set(dir, { files, folders, updated: Date.now() }) - } - const open = dir ? await input.open(dir) : new Set() - const { paths, items } = assemble(query, dir, files, folders, open) - input.post({ type: "fileSearchResult", paths, items, dir, requestId: input.message.requestId }) - }) + const [files, folders] = await fetchBackend(client, dir, query) + const open = dir ? await input.open(dir) : new Set() + const { paths, items } = assemble(query, dir, files, folders, open) + input.post({ type: "fileSearchResult", paths, items, dir, requestId: input.message.requestId }) } function settled(result: PromiseSettledResult<{ data: string[] }>, kind: "file" | "folder"): string[] { diff --git a/packages/kilo-vscode/tests/unit/file-search.test.ts b/packages/kilo-vscode/tests/unit/file-search.test.ts new file mode 100644 index 0000000000..742735d270 --- /dev/null +++ b/packages/kilo-vscode/tests/unit/file-search.test.ts @@ -0,0 +1,74 @@ +import { describe, expect, it } from "bun:test" +import { handleFileSearch } from "../../src/kilo-provider/file-search" + +type Query = { query: string; directory: string; type: "file" | "directory"; limit: number } + +function client(data: { files: string[]; folders: string[] }) { + const calls: Query[] = [] + return { + calls, + value: { + find: { + files: async (query: Query) => { + calls.push(query) + return { data: query.type === "file" ? data.files : data.folders } + }, + }, + }, + } +} + +describe("handleFileSearch", () => { + it("posts one fresh response for each request", async () => { + const api = client({ files: ["src/a.ts"], folders: ["src"] }) + const posted: unknown[] = [] + + await handleFileSearch({ + client: api.value as never, + message: { query: "", requestId: "request-1", sessionID: "session-1" }, + dir: (id) => (id === "session-1" ? "/repo" : ""), + open: async () => new Set(["src/open.ts"]), + post: (message) => posted.push(message), + }) + + expect(api.calls).toEqual([ + { query: "", directory: "/repo", type: "file", limit: 50 }, + { query: "", directory: "/repo", type: "directory", limit: 50 }, + ]) + expect(posted).toHaveLength(1) + expect(posted[0]).toEqual({ + type: "fileSearchResult", + requestId: "request-1", + dir: "/repo", + paths: ["src/open.ts", "src/a.ts"], + items: [ + { path: "src/open.ts", type: "opened-file" }, + { path: "src/a.ts", type: "file" }, + { path: "src", type: "folder" }, + ], + }) + }) + + it("returns an empty fresh response when files were deleted", async () => { + const api = client({ files: [], folders: [] }) + const posted: unknown[] = [] + + await handleFileSearch({ + client: api.value as never, + message: { query: "", requestId: "request-empty" }, + dir: () => "/repo", + open: async () => new Set(), + post: (message) => posted.push(message), + }) + + expect(posted).toEqual([ + { + type: "fileSearchResult", + requestId: "request-empty", + dir: "/repo", + paths: [], + items: [], + }, + ]) + }) +}) diff --git a/packages/kilo-vscode/tests/unit/use-file-mention.test.ts b/packages/kilo-vscode/tests/unit/use-file-mention.test.ts index 59445fa01d..9b9d5ec974 100644 --- a/packages/kilo-vscode/tests/unit/use-file-mention.test.ts +++ b/packages/kilo-vscode/tests/unit/use-file-mention.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "bun:test" -import { createRoot } from "solid-js" +import { createRoot, createSignal } from "solid-js" import { useFileMention } from "../../webview-ui/src/hooks/useFileMention" import { FILE_PICKER_RESULT } from "../../webview-ui/src/hooks/file-mention-utils" import type { ExtensionMessage, WebviewMessage } from "../../webview-ui/src/types/messages" @@ -751,7 +751,7 @@ describe("useFileMention", () => { dispose.fn?.() }) - it("renders cached files instantly when opening @ with empty query", async () => { + it("renders cached files instantly when opening @ with empty query", () => { const posted: WebviewMessage[] = [] const handlers = new Set<(message: ExtensionMessage) => void>() const ctx = { @@ -767,12 +767,13 @@ describe("useFileMention", () => { dispose.fn = root return useFileMention(ctx, undefined, () => false) }) - - // Simulate prewarm response + mention.onInput("@", 1) + const refresh = posted.at(-1) + expect(refresh?.type).toBe("requestFileSearch") for (const handler of handlers) { handler({ type: "fileSearchResult", - requestId: "file-search-prewarm", + requestId: refresh?.type === "requestFileSearch" ? refresh.requestId : "", dir: "/repo", paths: ["src/index.ts", "package.json"], items: [ @@ -782,10 +783,6 @@ describe("useFileMention", () => { }) } - // Now user types @ - mention.onInput("@", 1) - - // Mention results must immediately contain the cached files without waiting for a debounce/IPC round-trip expect(mention.mentionResults()).toEqual([ { type: "terminal", value: "terminal", label: "Terminal", description: "Active terminal output" }, { type: "past-chats", value: "past-chats", label: "Past chats", description: "Search previous sessions" }, @@ -809,4 +806,225 @@ describe("useFileMention", () => { dispose.fn?.() }) + + it("replaces deleted cached files when an empty refresh completes", () => { + const posted: WebviewMessage[] = [] + const handlers = new Set<(message: ExtensionMessage) => void>() + const ctx = { + postMessage: (message: WebviewMessage) => posted.push(message), + onMessage: (handler: (message: ExtensionMessage) => void) => { + handlers.add(handler) + return () => handlers.delete(handler) + }, + } + + const dispose: { fn?: () => void } = {} + const mention = createRoot((root) => { + dispose.fn = root + return useFileMention(ctx, undefined, () => false) + }) + mention.onInput("@", 1) + const prewarm = posted.at(-1) + expect(prewarm?.type).toBe("requestFileSearch") + for (const handler of handlers) { + handler({ + type: "fileSearchResult", + requestId: prewarm?.type === "requestFileSearch" ? prewarm.requestId : "", + dir: "/repo", + paths: ["deleted.ts"], + items: [{ path: "deleted.ts", type: "file" }], + }) + } + + expect(mention.mentionResults()).toContainEqual({ type: "file", value: "deleted.ts" }) + mention.closeMention() + mention.onInput("@", 1) + expect(mention.mentionResults()).toContainEqual({ type: "file", value: "deleted.ts" }) + const refresh = posted.at(-1) + expect(refresh?.type).toBe("requestFileSearch") + for (const handler of handlers) { + handler({ + type: "fileSearchResult", + requestId: refresh?.type === "requestFileSearch" ? refresh.requestId : "", + dir: "/repo", + paths: [], + items: [], + }) + } + + expect(mention.mentionResults()).not.toContainEqual({ type: "file", value: "deleted.ts" }) + mention.closeMention() + mention.onInput("@", 1) + expect(mention.mentionResults()).not.toContainEqual({ type: "file", value: "deleted.ts" }) + + dispose.fn?.() + }) + + it("does not reuse cached files after switching sessions", () => { + const posted: WebviewMessage[] = [] + const handlers = new Set<(message: ExtensionMessage) => void>() + const ctx = { + postMessage: (message: WebviewMessage) => posted.push(message), + onMessage: (handler: (message: ExtensionMessage) => void) => { + handlers.add(handler) + return () => handlers.delete(handler) + }, + } + + const dispose: { fn?: () => void } = {} + const state = createRoot((root) => { + dispose.fn = root + const [session, setSession] = createSignal("session-a") + return { mention: useFileMention(ctx, session, () => false), setSession } + }) + state.mention.onInput("@", 1) + const first = posted.at(-1) + expect(first).toMatchObject({ type: "requestFileSearch", sessionID: "session-a" }) + for (const handler of handlers) { + handler({ + type: "fileSearchResult", + requestId: first?.type === "requestFileSearch" ? first.requestId : "", + dir: "/repo-a", + paths: ["only-a.ts"], + items: [{ path: "only-a.ts", type: "file" }], + }) + } + + state.mention.closeMention() + state.setSession("session-b") + state.mention.onInput("@", 1) + + expect(state.mention.mentionResults()).not.toContainEqual({ type: "file", value: "only-a.ts" }) + expect(posted.at(-1)).toMatchObject({ type: "requestFileSearch", sessionID: "session-b", query: "" }) + + dispose.fn?.() + }) + + it("preserves the highlighted file when fresh results replace cached results", () => { + const posted: WebviewMessage[] = [] + const handlers = new Set<(message: ExtensionMessage) => void>() + const ctx = { + postMessage: (message: WebviewMessage) => posted.push(message), + onMessage: (handler: (message: ExtensionMessage) => void) => { + handlers.add(handler) + return () => handlers.delete(handler) + }, + } + + const dispose: { fn?: () => void } = {} + const mention = createRoot((root) => { + dispose.fn = root + return useFileMention(ctx, undefined, () => false) + }) + mention.onInput("@", 1) + const prewarm = posted.at(-1) + for (const handler of handlers) { + handler({ + type: "fileSearchResult", + requestId: prewarm?.type === "requestFileSearch" ? prewarm.requestId : "", + dir: "/repo", + paths: ["a.ts", "b.ts"], + items: [ + { path: "a.ts", type: "file" }, + { path: "b.ts", type: "file" }, + ], + }) + } + mention.closeMention() + mention.onInput("@", 1) + const selected = mention.mentionResults().findIndex((item) => item.type === "file" && item.value === "b.ts") + mention.setMentionIndex(selected) + + const refresh = posted.at(-1) + for (const handler of handlers) { + handler({ + type: "fileSearchResult", + requestId: refresh?.type === "requestFileSearch" ? refresh.requestId : "", + dir: "/repo", + paths: ["new.ts", "b.ts"], + items: [ + { path: "new.ts", type: "file" }, + { path: "b.ts", type: "file" }, + ], + }) + } + + expect(mention.mentionResults()[mention.mentionIndex()]).toEqual({ type: "file", value: "b.ts" }) + + dispose.fn?.() + }) + + it("ignores a response after the query changes", async () => { + const posted: WebviewMessage[] = [] + const handlers = new Set<(message: ExtensionMessage) => void>() + const ctx = { + postMessage: (message: WebviewMessage) => posted.push(message), + onMessage: (handler: (message: ExtensionMessage) => void) => { + handlers.add(handler) + return () => handlers.delete(handler) + }, + } + + const dispose: { fn?: () => void } = {} + const mention = createRoot((root) => { + dispose.fn = root + return useFileMention(ctx, undefined, () => false) + }) + + mention.onInput("@old", 4) + await wait(170) + const old = posted.at(-1) + expect(old).toMatchObject({ type: "requestFileSearch", query: "old" }) + mention.onInput("@new", 4) + + for (const handler of handlers) { + handler({ + type: "fileSearchResult", + requestId: old?.type === "requestFileSearch" ? old.requestId : "", + dir: "/repo", + paths: ["old.ts"], + items: [{ path: "old.ts", type: "file" }], + }) + } + + expect(mention.mentionResults()).not.toContainEqual({ type: "file", value: "old.ts" }) + + dispose.fn?.() + }) + + it("ignores a response after closing the mention menu", () => { + const posted: WebviewMessage[] = [] + const handlers = new Set<(message: ExtensionMessage) => void>() + const ctx = { + postMessage: (message: WebviewMessage) => posted.push(message), + onMessage: (handler: (message: ExtensionMessage) => void) => { + handlers.add(handler) + return () => handlers.delete(handler) + }, + } + + const dispose: { fn?: () => void } = {} + const mention = createRoot((root) => { + dispose.fn = root + return useFileMention(ctx, undefined, () => false) + }) + + mention.onInput("@", 1) + const request = posted.at(-1) + mention.closeMention() + for (const handler of handlers) { + handler({ + type: "fileSearchResult", + requestId: request?.type === "requestFileSearch" ? request.requestId : "", + dir: "/repo", + paths: ["late.ts"], + items: [{ path: "late.ts", type: "file" }], + }) + } + mention.onInput("@", 1) + + expect(mention.mentionResults()).not.toContainEqual({ type: "file", value: "late.ts" }) + + dispose.fn?.() + }) }) diff --git a/packages/kilo-vscode/webview-ui/src/hooks/useFileMention.ts b/packages/kilo-vscode/webview-ui/src/hooks/useFileMention.ts index 051254048a..80c9b01927 100644 --- a/packages/kilo-vscode/webview-ui/src/hooks/useFileMention.ts +++ b/packages/kilo-vscode/webview-ui/src/hooks/useFileMention.ts @@ -25,6 +25,21 @@ import { } from "./file-mention-utils" const FILE_SEARCH_DEBOUNCE_MS = 150 +const FILE_SEARCH_CACHE_MS = 5000 +const FILE_SEARCH_CACHE_LIMIT = 8 + +type FileSearchCache = { + items: Array + updated: number + revision: number +} + +type FileSearchRequest = { + id: string + query: string + scope: string + revision: number +} interface VSCodeContext { postMessage: (message: WebviewMessage) => void @@ -118,8 +133,8 @@ export function useFileMention( const [sessionPicker, setSessionPicker] = createSignal(false) const [sessionCandidates, setSessionCandidates] = createSignal([]) let workspaceDir = "" - let cached: Array = [] - const cache = new Map>() + const cache = new Map() + const dirs = new Map() // Accumulates every path ever mentioned so syncMentionedPaths can // rediscover them after a native undo restores the text. const knownPaths = new Set() @@ -129,6 +144,9 @@ export function useFileMention( let fileSearchTimer: ReturnType | undefined let fileSearchCounter = 0 + let fileSearchRevision = 0 + let fileSearchRequest: FileSearchRequest | undefined + let prewarmRequest: FileSearchRequest | undefined let filePickerCounter = 0 let sessionSearchCounter = 0 let pickerState: { @@ -142,17 +160,75 @@ export function useFileMention( let pendingArrowSnap: { timer: ReturnType; prevValue: string; prevPosition: number } | undefined const showMention = () => mentionQuery() !== null + const scope = () => sessionID?.() ?? "" + let activeScope = scope() + + const syncScope = () => { + const value = scope() + if (value === activeScope) return value + activeScope = value + if (fileSearchTimer) clearTimeout(fileSearchTimer) + fileSearchRevision++ + fileSearchRequest = undefined + prewarmRequest = undefined + workspaceDir = dirs.get(value) ?? "" + setMentionResults([]) + setMentionIndex(0) + return value + } + + const readCache = (dir: string): Array => { + if (!dir) return [] + const entry = cache.get(dir) + if (!entry) return [] + if (Date.now() - entry.updated <= FILE_SEARCH_CACHE_MS) return entry.items + cache.delete(dir) + return [] + } + + const writeCache = (dir: string, items: Array, revision: number) => { + if (!dir) return + const entry = cache.get(dir) + if (entry && entry.revision > revision) return + cache.delete(dir) + cache.set(dir, { items, updated: Date.now(), revision }) + while (cache.size > FILE_SEARCH_CACHE_LIMIT) { + const oldest = cache.keys().next().value + if (!oldest) return + cache.delete(oldest) + } + } + + const replaceResults = (items: MentionResult[]) => { + const index = mentionIndex() + const selected = mentionResults()[index] + setMentionResults(items) + if (!selected) { + setMentionIndex(0) + return + } + const next = items.findIndex((item) => item.type === selected.type && item.value === selected.value) + setMentionIndex(next >= 0 ? next : Math.min(index, Math.max(items.length - 1, 0))) + } createEffect(() => { if (!showMention()) setMentionIndex(0) }) createEffect(() => { - const id = sessionID?.() + const id = syncScope() + if (fileSearchTimer) clearTimeout(fileSearchTimer) + fileSearchRequest = undefined + setMentionQuery(null) + setMentionResults([]) + setMentionIndex(0) + const revision = ++fileSearchRevision + const requestId = `file-search-prewarm-${revision}` + prewarmRequest = { id: requestId, query: "", scope: id, revision } vscode.postMessage({ type: "requestFileSearch", query: "", - requestId: "file-search-prewarm", + requestId, ...(id ? { sessionID: id } : {}), }) }) @@ -170,18 +246,25 @@ export function useFileMention( return } if (message.type !== "fileSearchResult") return - if (message.requestId === `file-search-${fileSearchCounter}` || message.requestId === "file-search-prewarm") { - const items = message.items ?? message.paths.map((path) => ({ path, type: "file" as const })) + const request = + message.requestId === fileSearchRequest?.id + ? fileSearchRequest + : message.requestId === prewarmRequest?.id + ? prewarmRequest + : undefined + if (!request || request.scope !== scope()) return + if (request === fileSearchRequest) fileSearchRequest = undefined + if (request === prewarmRequest) prewarmRequest = undefined + if (request.revision < fileSearchRevision) return + + const items = message.items ?? message.paths.map((path) => ({ path, type: "file" as const })) + if (message.dir) { + dirs.set(request.scope, message.dir) workspaceDir = message.dir - if (message.dir) cache.set(message.dir, items) - if (!mentionQuery()) { - cached = items - } - if (showMention() && message.requestId === `file-search-${fileSearchCounter}`) { - setMentionResults(buildMentionResults(mentionQuery() ?? "", items, git?.() ?? true)) - setMentionIndex(0) - } } + if (!request.query) writeCache(message.dir, items, request.revision) + if (!showMention() || request.query !== mentionQuery()) return + replaceResults(buildMentionResults(request.query, items, git?.() ?? true)) }) onCleanup(() => { @@ -190,19 +273,25 @@ export function useFileMention( if (pendingArrowSnap) clearTimeout(pendingArrowSnap.timer) }) - const requestFileSearch = (query: string, immediate = false) => { + const requestFileSearch = (query: string) => { if (fileSearchTimer) clearTimeout(fileSearchTimer) + const revision = ++fileSearchRevision + const request = { + id: `file-search-${++fileSearchCounter}`, + query, + scope: syncScope(), + revision, + } + fileSearchRequest = request const send = () => { - fileSearchCounter++ - const id = sessionID?.() vscode.postMessage({ type: "requestFileSearch", query, - requestId: `file-search-${fileSearchCounter}`, - ...(id ? { sessionID: id } : {}), + requestId: request.id, + ...(request.scope ? { sessionID: request.scope } : {}), }) } - if (immediate || !query) { + if (!query) { send() return } @@ -210,6 +299,9 @@ export function useFileMention( } const closeMention = () => { + if (fileSearchTimer) clearTimeout(fileSearchTimer) + fileSearchRevision++ + fileSearchRequest = undefined setMentionQuery(null) setMentionResults([]) setSessionPicker(false) @@ -319,6 +411,7 @@ export function useFileMention( let suppress = false const onInput = (val: string, cursor: number) => { + syncScope() syncMentionedPaths(val) if (suppress) return closeSessionPicker() @@ -327,11 +420,11 @@ export function useFileMention( if (match) { const query = match[1] ?? "" setMentionQuery(query) - const items = (workspaceDir && cache.get(workspaceDir)) || cached + const items = readCache(workspaceDir) if (!query) { setMentionResults(buildMentionResults("", items, git?.() ?? true)) setMentionIndex(0) - requestFileSearch("", true) + requestFileSearch("") return } setMentionResults((prev) => {