From 9fbd5479b09739b21ca636612a85501f0d0f548f Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Thu, 28 May 2026 10:47:41 +0200 Subject: [PATCH] fix(cli): isolate semantic indexing in a worker --- .changeset/isolate-indexing-worker.md | 5 + packages/opencode/script/build.ts | 4 +- .../src/kilocode/indexing-worker-client.ts | 150 ++++++++++++++++++ .../src/kilocode/indexing-worker-protocol.ts | 30 ++++ .../opencode/src/kilocode/indexing-worker.ts | 57 +++++++ packages/opencode/src/kilocode/indexing.ts | 92 ++++++----- .../test/kilocode/indexing-startup.test.ts | 71 ++++++++- .../test/kilocode/indexing-worker.test.ts | 24 +++ 8 files changed, 392 insertions(+), 41 deletions(-) create mode 100644 .changeset/isolate-indexing-worker.md create mode 100644 packages/opencode/src/kilocode/indexing-worker-client.ts create mode 100644 packages/opencode/src/kilocode/indexing-worker-protocol.ts create mode 100644 packages/opencode/src/kilocode/indexing-worker.ts create mode 100644 packages/opencode/test/kilocode/indexing-worker.test.ts diff --git a/.changeset/isolate-indexing-worker.md b/.changeset/isolate-indexing-worker.md new file mode 100644 index 0000000000..34f2c9fb10 --- /dev/null +++ b/.changeset/isolate-indexing-worker.md @@ -0,0 +1,5 @@ +--- +"@kilocode/cli": patch +--- + +Keep the extension responsive while semantic indexing processes large workspaces. diff --git a/packages/opencode/script/build.ts b/packages/opencode/script/build.ts index a9756daac4..d1bc8d2074 100755 --- a/packages/opencode/script/build.ts +++ b/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 diff --git a/packages/opencode/src/kilocode/indexing-worker-client.ts b/packages/opencode/src/kilocode/indexing-worker-client.ts new file mode 100644 index 0000000000..8b7f6e2a96 --- /dev/null +++ b/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 + search(query: string, directoryPrefix?: string): Promise + dispose(): Promise + } + + 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() + 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) => { + 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 = (request: Request, read: (message: Result) => T, allowStopping = false) => { + if (stopped || (stopping && !allowStopping)) return Promise.reject(new Error("Indexing worker is disposed.")) + return new Promise((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 + } +} diff --git a/packages/opencode/src/kilocode/indexing-worker-protocol.ts b/packages/opencode/src/kilocode/indexing-worker-protocol.ts new file mode 100644 index 0000000000..01726dee8f --- /dev/null +++ b/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 diff --git a/packages/opencode/src/kilocode/indexing-worker.ts b/packages/opencode/src/kilocode/indexing-worker.ts new file mode 100644 index 0000000000..1c07845381 --- /dev/null +++ b/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) { + 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) => { + 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 }) + } +} diff --git a/packages/opencode/src/kilocode/indexing.ts b/packages/opencode/src/kilocode/indexing.ts index b0bc479ba5..19a1b1a891 100644 --- a/packages/opencode/src/kilocode/indexing.ts +++ b/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 @@ -182,10 +179,11 @@ export namespace KiloIndexing { } type Entry = { - manager?: CodeIndexManager + engine?: IndexingWorker.Driver + initialized?: boolean current(): Status publish(): Promise - dispose(): void + dispose(): Promise } type Cache = { @@ -209,14 +207,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 } @@ -245,61 +243,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, @@ -308,14 +328,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() @@ -334,9 +350,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 @@ -356,7 +372,7 @@ export namespace KiloIndexing { cache.delete(dir) if (hit) hit.disposed = true if (hit?.entry) { - hit.entry.dispose() + await hit.entry.dispose() return } }) @@ -375,19 +391,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 { 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 { 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) } } diff --git a/packages/opencode/test/kilocode/indexing-startup.test.ts b/packages/opencode/test/kilocode/indexing-startup.test.ts index 2a6522339c..b369b82790 100644 --- a/packages/opencode/test/kilocode/indexing-startup.test.ts +++ b/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 = { }, }, } +const inactive: Partial = { + plugin: ["@kilocode/kilo-indexing"], + experimental: { + semantic_indexing: true, + }, + indexing: { + enabled: false, + provider: "ollama", + vectorStore: "qdrant", + }, +} const kilo: Partial = { plugin: ["@kilocode/kilo-indexing"], experimental: { @@ -68,6 +81,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, state: KiloIndexing.Status["state"]) { for (const _ of Array.from({ length: 100 })) { const status = await read() @@ -85,7 +117,12 @@ async function called(init: ReturnType { + 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"] @@ -202,8 +239,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 @@ -219,9 +257,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() } }) @@ -272,6 +312,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( diff --git a/packages/opencode/test/kilocode/indexing-worker.test.ts b/packages/opencode/test/kilocode/indexing-worker.test.ts new file mode 100644 index 0000000000..235c82b4dd --- /dev/null +++ b/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([]) +})