fix(vscode): keep file mentions fresh

.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.

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

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<string, Cache>()
-
-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<void> {
   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<string>()
-    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<string>()
+  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[] {

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: [],
+      },
+    ])
+  })
+})

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?.()
+  })
 })

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<FileSearchItem | string>
+  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<SessionSearchItem[]>([])
   let workspaceDir = ""
-  let cached: Array<FileSearchItem | string> = []
-  const cache = new Map<string, Array<FileSearchItem | string>>()
+  const cache = new Map<string, FileSearchCache>()
+  const dirs = new Map<string, string>()
   // Accumulates every path ever mentioned so syncMentionedPaths can
   // rediscover them after a native undo restores the text.
   const knownPaths = new Set<string>()
@@ -129,6 +144,9 @@ export function useFileMention(
 
   let fileSearchTimer: ReturnType<typeof setTimeout> | 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<typeof setTimeout>; 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<FileSearchItem | string> => {
+    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<FileSearchItem | string>, 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) => {