Merge pull request #10648 from Kilo-Org/serene-run

fix(cli): isolate semantic indexing in a worker

.changeset/isolate-indexing-worker.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/cli": patch
+---
+
+Keep the extension responsive while semantic indexing processes large workspaces.

packages/opencode/script/build.ts

@@ -214,6 +214,7 @@ for (const item of targets) {
   const rootPath = path.resolve(dir, "../../node_modules/@opentui/core/parser.worker.js")
   const parserWorker = fs.realpathSync(fs.existsSync(localPath) ? localPath : rootPath)
   const workerPath = "./src/cli/cmd/tui/worker.ts"
+  const indexingWorkerPath = "./src/kilocode/indexing-worker.ts" // kilocode_change
 
   // Use platform-specific bunfs root path based on target OS
   const bunfsRoot = item.os === "win32" ? "B:/~BUN/root/" : "/$bunfs/root/"
@@ -242,13 +243,14 @@ for (const item of targets) {
     },
     // kilocode_change start - packages/app was removed; no embedded web UI
     files: {},
-    entrypoints: ["./src/index.ts", parserWorker, workerPath],
+    entrypoints: ["./src/index.ts", parserWorker, workerPath, indexingWorkerPath],
     // kilocode_change end
     define: {
       KILO_VERSION: `'${Script.version}'`,
       KILO_MIGRATIONS: JSON.stringify(migrations),
       OTUI_TREE_SITTER_WORKER_PATH: bunfsRoot + workerRelativePath,
       KILO_WORKER_PATH: workerPath,
+      KILO_INDEXING_WORKER_PATH: indexingWorkerPath, // kilocode_change
       KILO_CHANNEL: `'${Script.channel}'`,
       KILO_LIBC: item.os === "linux" ? `'${item.abi ?? "glibc"}'` : "",
       KILO_BUILD_KIND: Script.release ? `'release'` : `'source'`, // kilocode_change

packages/opencode/src/kilocode/indexing-worker-client.ts

@@ -0,0 +1,150 @@
+import type {
+  IndexingConfigInput,
+  IndexingTelemetryEvent,
+  VectorStoreSearchResult,
+} from "@kilocode/kilo-indexing/engine"
+import type { IndexingStatus } from "@kilocode/kilo-indexing/status"
+import { withTimeout } from "@/util/timeout"
+import type { Message, Request, Result } from "./indexing-worker-protocol"
+
+declare global {
+  const KILO_INDEXING_WORKER_PATH: string
+}
+
+export namespace IndexingWorker {
+  export type Hooks = {
+    status(status: IndexingStatus): void
+    telemetry(event: IndexingTelemetryEvent): void
+    failure(err: unknown): void
+  }
+
+  export type Driver = {
+    init(input: IndexingConfigInput): Promise<IndexingStatus>
+    search(query: string, directoryPrefix?: string): Promise<VectorStoreSearchResult[]>
+    dispose(): Promise<void>
+  }
+
+  export type Factory = (directory: string, root: string, hooks: Hooks) => Driver
+
+  const worker = (directory: string, root: string, hooks: Hooks): Driver => {
+    const file =
+      typeof KILO_INDEXING_WORKER_PATH !== "undefined"
+        ? KILO_INDEXING_WORKER_PATH
+        : new URL("./indexing-worker.ts", import.meta.url)
+    const task = new Worker(file)
+    const pending = new Map<number, { resolve(message: Result): void; reject(err: unknown): void }>()
+    let id = 0
+    let stopped = false
+    let stopping = false
+
+    const reject = (err: unknown) => {
+      for (const item of pending.values()) item.reject(err)
+      pending.clear()
+    }
+
+    const fail = (err: unknown) => {
+      if (stopped || stopping) return
+      stopped = true
+      reject(err)
+      task.terminate()
+      hooks.failure(err)
+    }
+
+    task.onmessage = (event: MessageEvent<Message>) => {
+      const message = event.data
+      if (message.type === "event") {
+        if (stopping || stopped) return
+        if (message.event === "status") hooks.status(message.data)
+        if (message.event === "telemetry") hooks.telemetry(message.data)
+        return
+      }
+
+      const request = pending.get(message.id)
+      if (!request) return
+      pending.delete(message.id)
+      if (message.ok) {
+        request.resolve(message)
+        return
+      }
+      request.reject(new Error(message.error))
+    }
+
+    task.onerror = (event) => {
+      fail(event.error ?? new Error(event.message))
+    }
+
+    const call = <T>(request: Request, read: (message: Result) => T, allowStopping = false) => {
+      if (stopped || (stopping && !allowStopping)) return Promise.reject(new Error("Indexing worker is disposed."))
+      return new Promise<T>((resolve, reject) => {
+        pending.set(request.id, {
+          resolve(message) {
+            try {
+              resolve(read(message))
+            } catch (err) {
+              reject(err)
+            }
+          },
+          reject,
+        })
+        task.postMessage(request)
+      })
+    }
+
+    return {
+      init(config) {
+        const request: Request = {
+          type: "request",
+          id: id++,
+          method: "init",
+          input: {
+            directory,
+            root,
+            config,
+            lancedbPath: process.env.KILO_LANCEDB_PATH,
+          },
+        }
+        return call(request, (message) => {
+          if (message.ok && message.method === "init") return message.value
+          throw new Error("Unexpected indexing worker init response.")
+        })
+      },
+      search(query, directoryPrefix) {
+        const request: Request = { type: "request", id: id++, method: "search", input: { query, directoryPrefix } }
+        return call(request, (message) => {
+          if (message.ok && message.method === "search") return message.value
+          throw new Error("Unexpected indexing worker search response.")
+        })
+      },
+      async dispose() {
+        if (stopped || stopping) return
+        stopping = true
+        const request: Request = { type: "request", id: id++, method: "dispose", input: undefined }
+        await withTimeout(
+          call(
+            request,
+            (message) => {
+              if (message.ok && message.method === "dispose") return message.value
+              throw new Error("Unexpected indexing worker dispose response.")
+            },
+            true,
+          ),
+          1000,
+          "Indexing worker shutdown timed out",
+        ).catch(() => undefined)
+        stopped = true
+        reject(new Error("Indexing worker is disposed."))
+        task.terminate()
+      },
+    }
+  }
+
+  let factory: Factory = worker
+
+  export function create(directory: string, root: string, hooks: Hooks) {
+    return factory(directory, root, hooks)
+  }
+
+  export function override(next?: Factory) {
+    factory = next ?? worker
+  }
+}

packages/opencode/src/kilocode/indexing-worker-protocol.ts

@@ -0,0 +1,30 @@
+import type {
+  IndexingConfigInput,
+  IndexingTelemetryEvent,
+  VectorStoreSearchResult,
+} from "@kilocode/kilo-indexing/engine"
+import type { IndexingStatus } from "@kilocode/kilo-indexing/status"
+
+export type InitInput = {
+  directory: string
+  root: string
+  config: IndexingConfigInput
+  lancedbPath?: string
+}
+
+export type Request =
+  | { type: "request"; id: number; method: "init"; input: InitInput }
+  | { type: "request"; id: number; method: "search"; input: { query: string; directoryPrefix?: string } }
+  | { type: "request"; id: number; method: "dispose"; input: undefined }
+
+export type Result =
+  | { type: "result"; id: number; method: "init"; ok: true; value: IndexingStatus }
+  | { type: "result"; id: number; method: "search"; ok: true; value: VectorStoreSearchResult[] }
+  | { type: "result"; id: number; method: "dispose"; ok: true; value: undefined }
+  | { type: "result"; id: number; method: Request["method"]; ok: false; error: string }
+
+export type Event =
+  | { type: "event"; event: "status"; data: IndexingStatus }
+  | { type: "event"; event: "telemetry"; data: IndexingTelemetryEvent }
+
+export type Message = Result | Event

packages/opencode/src/kilocode/indexing-worker.ts

@@ -0,0 +1,57 @@
+import { CodeIndexManager } from "@kilocode/kilo-indexing/engine"
+import { normalizeIndexingStatus } from "@kilocode/kilo-indexing/status"
+import type { Request, Result, Event } from "./indexing-worker-protocol"
+
+let manager: CodeIndexManager | undefined
+let progress: { dispose(): void } | undefined
+let telemetry: { dispose(): void } | undefined
+
+function send(message: Result | Event) {
+  postMessage(message)
+}
+
+function dispose() {
+  progress?.dispose()
+  telemetry?.dispose()
+  progress = undefined
+  telemetry = undefined
+  manager?.dispose()
+  manager = undefined
+}
+
+async function init(request: Extract<Request, { method: "init" }>) {
+  dispose()
+  if (request.input.lancedbPath) process.env.KILO_LANCEDB_PATH = request.input.lancedbPath
+  const next = new CodeIndexManager(request.input.directory, request.input.root)
+  manager = next
+  progress = next.onProgressUpdate.on(() => {
+    send({ type: "event", event: "status", data: normalizeIndexingStatus(next) })
+  })
+  telemetry = next.onTelemetry.on((data) => {
+    send({ type: "event", event: "telemetry", data })
+  })
+  await next.initialize(request.input.config)
+  send({ type: "result", id: request.id, method: "init", ok: true, value: normalizeIndexingStatus(next) })
+}
+
+onmessage = async (event: MessageEvent<Request>) => {
+  const request = event.data
+  try {
+    if (request.method === "dispose") {
+      dispose()
+      send({ type: "result", id: request.id, method: "dispose", ok: true, value: undefined })
+      return
+    }
+
+    if (request.method === "search") {
+      const value = manager ? await manager.searchIndex(request.input.query, request.input.directoryPrefix) : []
+      send({ type: "result", id: request.id, method: "search", ok: true, value })
+      return
+    }
+
+    await init(request)
+  } catch (err) {
+    const error = err instanceof Error ? err.message : String(err)
+    send({ type: "result", id: request.id, method: request.method, ok: false, error })
+  }
+}

packages/opencode/src/kilocode/indexing.ts

@@ -1,13 +1,9 @@
 import z from "zod"
 import path from "path"
-import {
-  CodeIndexManager,
-  type IndexingTelemetryEvent,
-  type VectorStoreSearchResult,
-} from "@kilocode/kilo-indexing/engine"
+import { type IndexingTelemetryEvent, type VectorStoreSearchResult } from "@kilocode/kilo-indexing/engine"
 import { toIndexingConfigInput, type IndexingConfig } from "@kilocode/kilo-indexing/config"
 import { hasIndexingPlugin } from "@kilocode/kilo-indexing/detect"
-import { IndexingStatus, disabledIndexingStatus, normalizeIndexingStatus } from "@kilocode/kilo-indexing/status"
+import { IndexingStatus, disabledIndexingStatus } from "@kilocode/kilo-indexing/status"
 import { Telemetry } from "@kilocode/kilo-telemetry"
 import { fetchKiloEmbeddingModelCatalog } from "@kilocode/kilo-gateway"
 import { Instance } from "@/project/instance"
@@ -20,6 +16,7 @@ import { registerDisposer } from "@/effect/instance-registry"
 import { Global } from "@opencode-ai/core/global"
 import * as Log from "@opencode-ai/core/util/log"
 import { Event as IndexingEvent } from "./indexing-event"
+import { IndexingWorker } from "./indexing-worker-client"
 import { LanceDBRuntime } from "./lancedb" // kilocode_change
 import { indexingWithKiloDefault, resolveKiloIndexingAuth, type KiloIndexingAuth } from "./indexing-auth" // kilocode_change
 
@@ -196,10 +193,11 @@ export namespace KiloIndexing {
   }
 
   type Entry = {
-    manager?: CodeIndexManager
+    engine?: IndexingWorker.Driver
+    initialized?: boolean
     current(): Status
     publish(): Promise<void>
-    dispose(): void
+    dispose(): Promise<void>
   }
 
   type Cache = {
@@ -223,14 +221,14 @@ export namespace KiloIndexing {
     return {
       current,
       publish,
-      dispose() {},
+      async dispose() {},
     }
   }
 
   function track(hit: Cache, entry: Entry) {
     if (!hit.entry) hit.resolve(entry)
     hit.entry = entry
-    if (hit.disposed) entry.dispose()
+    if (hit.disposed) void entry.dispose()
     return entry
   }
 
@@ -259,61 +257,83 @@ export namespace KiloIndexing {
 
     log.info("initializing project indexing", { workspacePath: dir })
     const root = path.join(Global.Path.state, "indexing")
-    const manager = new CodeIndexManager(dir, root)
     const auth = await kiloAuth(cfg)
     const globalConfig = await AppRuntime.runPromise(Config.Service.use((svc) => svc.getGlobal()))
     const global = globalConfig.indexing
     const merged = indexingWithKiloDefault({ ...global, ...cfg.indexing }, auth)
     const cfgInput = await model(enrichKilo(input(merged, global), auth), auth)
-    const box = { status: pending() as Status | undefined }
-    const current = () => box.status ?? normalizeIndexingStatus(manager)
+    const box = { status: pending() }
+    const current = () => box.status
     let disposed = false
 
     const publish = async () => {
       await Bus.publish(Event, { status: current() })
     }
-    const report = async () => {
+    const report = Instance.bind(async () => {
       try {
         return await publish()
       } catch (err) {
         log.error("failed to publish indexing status", { err })
       }
-    }
-
-    const unsub = manager.onProgressUpdate.on(() => {
+    })
+    const status = Instance.bind((next: Status) => {
+      if (disposed) return
+      box.status = next
       void report()
     })
-    const telemetrySub = manager.onTelemetry.on((event) => {
+    const telemetry = Instance.bind((event: IndexingTelemetryEvent) => {
+      if (disposed) return
       trackTelemetry(event)
     })
-
     const base: Entry = {
       current,
       publish,
-      dispose() {
+      async dispose() {
         if (disposed) return
         disposed = true
-        unsub.dispose()
-        telemetrySub.dispose()
-        manager.dispose()
+        base.initialized = false
+        await base.engine?.dispose().catch((err) => {
+          log.warn("failed to dispose project indexing worker", { err, workspacePath: dir })
+        })
       },
     }
+    const failure = Instance.bind((err: unknown) => {
+      if (disposed) return
+      base.initialized = false
+      box.status = failed(err)
+      log.error("project indexing worker failed", { err, workspacePath: dir })
+      void report()
+    })
     track(hit, base)
     await report()
 
     if (hit.disposed) return base
 
-    // kilocode_change start
+    if (!cfgInput.enabled) {
+      box.status = disabledIndexingStatus()
+      await report()
+      return base
+    }
+
     const err = await LanceDBRuntime.ensure(cfgInput.vectorStoreProvider)
-      .then(() => manager.initialize(cfgInput))
+      .then(async () => {
+        if (hit.disposed) return
+        const engine = IndexingWorker.create(dir, root, { status, telemetry, failure })
+        base.engine = engine
+        box.status = await engine.init(cfgInput)
+        base.initialized = true
+      })
       .then(
         () => undefined,
         (err) => err,
       )
-    // kilocode_change end
     if (hit.disposed) return base
 
     if (err) {
+      await base.engine?.dispose().catch((disposeErr) => {
+        log.warn("failed to dispose failed project indexing worker", { err: disposeErr, workspacePath: dir })
+      })
+      base.engine = undefined
       box.status = failed(err)
       log.error("project indexing initialization failed", {
         err,
@@ -322,14 +342,10 @@ export namespace KiloIndexing {
       await report()
       return base
     }
-    box.status = undefined
-    base.manager = manager
 
     log.info("project indexing initialized", {
       workspacePath: dir,
-      featureEnabled: manager.isFeatureEnabled,
-      featureConfigured: manager.isFeatureConfigured,
-      state: manager.getCurrentStatus().systemStatus,
+      state: current().state,
     })
     await report()
 
@@ -348,9 +364,9 @@ export namespace KiloIndexing {
       reject: gate.reject,
     } as Cache
     next.promise = boot(next)
-      .then((entry) => {
+      .then(async (entry) => {
         if (next.disposed) {
-          entry.dispose()
+          await entry.dispose()
           return entry
         }
         next.entry = entry
@@ -370,7 +386,7 @@ export namespace KiloIndexing {
     cache.delete(dir)
     if (hit) hit.disposed = true
     if (hit?.entry) {
-      hit.entry.dispose()
+      await hit.entry.dispose()
       return
     }
   })
@@ -389,19 +405,19 @@ export namespace KiloIndexing {
 
   export function ready(): boolean {
     const entry = cache.get(Instance.directory)?.entry
-    if (!entry?.manager) return false
+    if (!entry?.initialized) return false
     return entry.current().state !== "Disabled"
   }
 
   export async function available(): Promise<boolean> {
     const entry = await hit().ready
-    if (!entry.manager) return false
+    if (!entry.initialized) return false
     return entry.current().state !== "Disabled"
   }
 
   export async function search(query: string, directoryPrefix?: string): Promise<VectorStoreSearchResult[]> {
     const entry = await hit().ready
-    if (!entry.manager) return []
-    return entry.manager.searchIndex(query, directoryPrefix)
+    if (!entry.initialized || entry.current().state === "Disabled" || !entry.engine) return []
+    return entry.engine.search(query, directoryPrefix)
   }
 }

packages/opencode/test/kilocode/indexing-startup.test.ts

@@ -1,8 +1,10 @@
-import { afterEach, describe, expect, spyOn, test } from "bun:test"
+import { afterEach, beforeEach, describe, expect, spyOn, test } from "bun:test"
 import { CodeIndexManager } from "@kilocode/kilo-indexing/engine"
+import { normalizeIndexingStatus } from "@kilocode/kilo-indexing/status"
 import type { Config } from "../../src/config/config"
 import { GlobalBus } from "../../src/bus/global"
 import { KiloIndexing } from "../../src/kilocode/indexing"
+import { IndexingWorker } from "../../src/kilocode/indexing-worker-client"
 import { WithInstance } from "../../src/project/with-instance"
 import { Server } from "../../src/server/server"
 import * as Log from "@opencode-ai/core/util/log"
@@ -41,6 +43,17 @@ const off: Partial<Config.Info> = {
     },
   },
 }
+const inactive: Partial<Config.Info> = {
+  plugin: ["@kilocode/kilo-indexing"],
+  experimental: {
+    semantic_indexing: true,
+  },
+  indexing: {
+    enabled: false,
+    provider: "ollama",
+    vectorStore: "qdrant",
+  },
+}
 const kilo: Partial<Config.Info> = {
   plugin: ["@kilocode/kilo-indexing"],
   experimental: {
@@ -81,6 +94,25 @@ const configDir = process.env["KILO_CONFIG_DIR"]
 const disabled = process.env["KILO_DISABLE_CODEBASE_INDEXING"]
 const error = new Error("test indexing initialization failed")
 
+function inline(directory: string, root: string, hooks: IndexingWorker.Hooks): IndexingWorker.Driver {
+  const manager = new CodeIndexManager(directory, root)
+  const progress = manager.onProgressUpdate.on(() => hooks.status(normalizeIndexingStatus(manager)))
+  const telemetry = manager.onTelemetry.on(hooks.telemetry)
+
+  return {
+    async init(input) {
+      await manager.initialize(input)
+      return normalizeIndexingStatus(manager)
+    },
+    search: (query, directoryPrefix) => manager.searchIndex(query, directoryPrefix),
+    async dispose() {
+      progress.dispose()
+      telemetry.dispose()
+      manager.dispose()
+    },
+  }
+}
+
 async function wait(read: () => Promise<KiloIndexing.Status>, state: KiloIndexing.Status["state"]) {
   for (const _ of Array.from({ length: 100 })) {
     const status = await read()
@@ -98,7 +130,12 @@ async function called(init: ReturnType<typeof spyOn<CodeIndexManager, "initializ
   throw new Error("indexing initialization did not start")
 }
 
+beforeEach(() => {
+  IndexingWorker.override(inline)
+})
+
 afterEach(async () => {
+  IndexingWorker.override()
   if (configDir === undefined) delete process.env["KILO_CONFIG_DIR"]
   else process.env["KILO_CONFIG_DIR"] = configDir
   if (disabled === undefined) delete process.env["KILO_DISABLE_CODEBASE_INDEXING"]
@@ -215,8 +252,9 @@ describe("indexing startup degradation", () => {
     }
   })
 
-  test("keeps degraded indexing queryable but unavailable", async () => {
+  test("keeps degraded indexing queryable but releases its failed engine", async () => {
     const init = spyOn(CodeIndexManager.prototype, "initialize").mockRejectedValue(error)
+    const dispose = spyOn(CodeIndexManager.prototype, "dispose")
 
     await using tmp = await tmpdir({ git: true, config: cfg })
     process.env["KILO_CONFIG_DIR"] = tmp.path
@@ -232,9 +270,11 @@ describe("indexing startup degradation", () => {
           expect(await KiloIndexing.available()).toBe(false)
           expect(KiloIndexing.ready()).toBe(false)
           expect(await KiloIndexing.search("boot failure")).toEqual([])
+          expect(dispose).toHaveBeenCalledTimes(1)
         },
       })
     } finally {
+      dispose.mockRestore()
       init.mockRestore()
     }
   })
@@ -285,6 +325,33 @@ describe("indexing startup degradation", () => {
     })
   })
 
+  test("does not allocate an engine when indexing configuration is disabled", async () => {
+    const created: string[] = []
+    IndexingWorker.override((directory, root, hooks) => {
+      created.push(directory)
+      return inline(directory, root, hooks)
+    })
+
+    await using tmp = await tmpdir({ git: true, config: inactive })
+    process.env["KILO_CONFIG_DIR"] = tmp.path
+
+    await WithInstance.provide({
+      directory: tmp.path,
+      fn: async () => {
+        const status = await wait(() => KiloIndexing.current(), "Disabled")
+
+        expect(status).toMatchObject({
+          state: "Disabled",
+          message: "Indexing disabled.",
+        })
+        expect(await KiloIndexing.available()).toBe(false)
+        expect(KiloIndexing.ready()).toBe(false)
+        expect(await KiloIndexing.search("disabled")).toEqual([])
+        expect(created).toEqual([])
+      },
+    })
+  })
+
   test("enriches Kilo provider config from env auth", async () => {
     global.fetch = (() =>
       Promise.resolve(

packages/opencode/test/kilocode/indexing-worker.test.ts

@@ -0,0 +1,24 @@
+import { expect, test } from "bun:test"
+import { IndexingWorker } from "../../src/kilocode/indexing-worker-client"
+import { tmpdir } from "../fixture/fixture"
+
+test("runs indexing engine requests in its worker", async () => {
+  await using tmp = await tmpdir()
+  const failures: unknown[] = []
+  const engine = IndexingWorker.create(tmp.path, tmp.path, {
+    status() {},
+    telemetry() {},
+    failure(err) {
+      failures.push(err)
+    },
+  })
+
+  try {
+    const status = await engine.init({ enabled: false, embedderProvider: "openai" })
+    expect(status.state).toBe("Disabled")
+  } finally {
+    await engine.dispose()
+  }
+
+  expect(failures).toEqual([])
+})