diff --git a/src/__tests__/extension.spec.ts b/src/__tests__/extension.spec.ts index 56ccd52588..05bc327230 100644 --- a/src/__tests__/extension.spec.ts +++ b/src/__tests__/extension.spec.ts @@ -1,6 +1,7 @@ // npx vitest run __tests__/extension.spec.ts import type * as vscode from "vscode" +import { makeUri } from "../test-utils/vscode" vi.mock("vscode", () => ({ window: { @@ -15,6 +16,7 @@ vi.mock("vscode", () => ({ onDidChangeActiveTextEditor: vi.fn(), }, workspace: { + workspaceFolders: undefined, registerTextDocumentContentProvider: vi.fn(), getConfiguration: vi.fn().mockReturnValue({ get: vi.fn().mockReturnValue([]), @@ -139,11 +141,10 @@ vi.mock("../services/mcp/McpServerManager", () => ({ }, })) -vi.mock("../services/code-index/code-index-manager-registry", () => ({ - CodeIndexManagerRegistry: { - getOrCreate: vi.fn().mockReturnValue(null), - disposeAll: vi.fn(), - }, +vi.mock("../services/code-index/manager", () => ({ + CodeIndexManager: vi.fn().mockImplementation(function () { + return { initialize: vi.fn().mockResolvedValue({ requiresRestart: false }), dispose: vi.fn() } + }), })) vi.mock("../services/mdm/MdmService", () => ({ @@ -268,6 +269,190 @@ describe("extension.ts", () => { expect(dotenv.config).toHaveBeenCalledTimes(1) }) + describe("code index workspace activation", () => { + beforeEach(() => { + vi.resetModules() + }) + + afterEach(async () => { + const vscode = await import("vscode") + vi.mocked(vscode.workspace).workspaceFolders = undefined + const { codeIndexWorkspaceScopeRegistry } = + await import("../services/code-index/code-index-workspace-scope-registry") + codeIndexWorkspaceScopeRegistry.disposeAll() + }) + + test("initializes every workspace once with the shared context without blocking activation", async () => { + const vscode = await import("vscode") + const folders = ["/first", "/second"].map((name, index) => ({ name, index, uri: makeUri(name) })) + vi.mocked(vscode.workspace).workspaceFolders = folders + const { codeIndexWorkspaceScopeRegistry: registry } = + await import("../services/code-index/code-index-workspace-scope-registry") + const scopes = folders.map((folder) => registry.getScope(mockContext, folder.uri.fsPath)!) + let release!: () => void + const pending = new Promise((resolve) => { + release = resolve + }) + for (const scope of scopes) { + vi.mocked(scope.codeIndexManager.initialize).mockImplementationOnce(async () => { + await pending + return { requiresRestart: false } + }) + } + const { ContextProxy } = await import("../core/config/ContextProxy") + const { activate } = await import("../extension") + let activated = false + const activation = activate(mockContext).then(() => { + activated = true + }) + try { + await vi.waitFor(() => expect(activated).toBe(true)) + const contextProxy = await ContextProxy.getInstance(mockContext) + for (const scope of scopes) { + expect(scope.codeIndexManager.initialize).toHaveBeenCalledExactlyOnceWith(contextProxy) + expect(scope.codeIndexManager.dispose).not.toHaveBeenCalled() + } + expect(vscode.commands.executeCommand).toHaveBeenCalledWith("test-extension.activationCompleted") + } finally { + release() + await activation + } + }) + + test.each([new Error("configuration failed"), "configuration failed"])( + "logs background rejection %s with its workspace and continues other initialization and cleanup", + async (error) => { + const vscode = await import("vscode") + vi.mocked(vscode.workspace).workspaceFolders = ["/broken", "/healthy"].map((name, index) => ({ + name, + index, + uri: makeUri(name), + })) + const { codeIndexWorkspaceScopeRegistry: registry } = + await import("../services/code-index/code-index-workspace-scope-registry") + const broken = registry.getScope(mockContext, "/broken")! + const healthy = registry.getScope(mockContext, "/healthy")! + vi.mocked(broken.codeIndexManager.initialize).mockRejectedValueOnce(error) + const { activate } = await import("../extension") + await expect(activate(mockContext)).resolves.toBeDefined() + + expect(healthy.codeIndexManager.initialize).toHaveBeenCalledTimes(1) + const channel = vi.mocked(vscode.window.createOutputChannel).mock.results.at(-1)?.value + expect(channel.appendLine).toHaveBeenCalledWith( + "[CodeIndexManager] Error during background CodeIndexManager configuration/indexing for /broken: configuration failed", + ) + await Promise.all(mockContext.subscriptions.map((subscription) => subscription?.dispose?.())) + expect(broken.codeIndexManager.dispose).toHaveBeenCalledTimes(1) + expect(healthy.codeIndexManager.dispose).toHaveBeenCalledTimes(1) + expect(registry.getAllScopes()).toEqual([]) + }, + ) + + test.each([undefined, []])( + "activation without workspace folders (%s) still owns lazy cleanup", + async (folders) => { + const vscode = await import("vscode") + vi.mocked(vscode.workspace).workspaceFolders = folders + const { CodeIndexManager } = await import("../services/code-index/manager") + const { activate } = await import("../extension") + await expect(activate(mockContext)).resolves.toBeDefined() + expect(CodeIndexManager).not.toHaveBeenCalled() + + vi.mocked(vscode.workspace).workspaceFolders = [{ name: "late", index: 0, uri: makeUri("/late") }] + const { codeIndexWorkspaceScopeRegistry: registry } = + await import("../services/code-index/code-index-workspace-scope-registry") + const lazy = registry.getScope(mockContext, "/late")! + await Promise.all(mockContext.subscriptions.map((subscription) => subscription?.dispose?.())) + expect(lazy.codeIndexManager.dispose).toHaveBeenCalledTimes(1) + expect(registry.getAllScopes()).toEqual([]) + }, + ) + + test("one registry cleanup owner disposes startup and lazy scopes exactly once", async () => { + const vscode = await import("vscode") + const first = { name: "first", index: 0, uri: makeUri("/first") } + const late = { name: "late", index: 1, uri: makeUri("/late") } + vi.mocked(vscode.workspace).workspaceFolders = [first] + const { codeIndexWorkspaceScopeRegistry: registry } = + await import("../services/code-index/code-index-workspace-scope-registry") + const { activate, deactivate } = await import("../extension") + await activate(mockContext) + const startup = registry.getScope(mockContext, first.uri.fsPath)! + vi.mocked(vscode.workspace).workspaceFolders = [first, late] + const lazy = registry.getScope(mockContext, late.uri.fsPath)! + + await deactivate() + for (const subscription of mockContext.subscriptions) { + // Unrelated activation mocks do not all return a disposable. + await subscription?.dispose?.() + } + + expect(startup.codeIndexManager.dispose).toHaveBeenCalledTimes(1) + expect(lazy.codeIndexManager.dispose).toHaveBeenCalledTimes(1) + expect(registry.getAllScopes()).toEqual([]) + expect(mockContext.subscriptions).not.toContain(startup) + expect(mockContext.subscriptions).not.toContain(lazy) + registry.disposeAll() + expect(startup.codeIndexManager.dispose).toHaveBeenCalledTimes(1) + expect(lazy.codeIndexManager.dispose).toHaveBeenCalledTimes(1) + }) + + test.each([false, true])( + "cleanup waits for pending initialization (rejects=%s) before disposal", + async (rejects) => { + const vscode = await import("vscode") + vi.mocked(vscode.workspace).workspaceFolders = [{ name: "first", index: 0, uri: makeUri("/first") }] + const { codeIndexWorkspaceScopeRegistry: registry } = + await import("../services/code-index/code-index-workspace-scope-registry") + const scope = registry.getScope(mockContext, "/first")! + let release!: () => void + const pending = new Promise((resolve) => { + release = resolve + }) + vi.mocked(scope.codeIndexManager.initialize).mockImplementationOnce(async () => { + await pending + if (rejects) throw new Error("late initialization failure") + return { requiresRestart: false } + }) + const { activate, deactivate } = await import("../extension") + await activate(mockContext) + const cleanup = Promise.all(mockContext.subscriptions.map((subscription) => subscription?.dispose?.())) + const disposedBeforeInitialization = vi.mocked(scope.codeIndexManager.dispose).mock.calls.length + release() + await cleanup + await deactivate() + + expect(disposedBeforeInitialization).toBe(0) + expect(scope.codeIndexManager.dispose).toHaveBeenCalledTimes(1) + expect(registry.getAllScopes()).toEqual([]) + }, + ) + + test("skips an unavailable scope without skipping later workspace initialization", async () => { + const vscode = await import("vscode") + const folders = ["/unavailable", "/healthy"].map((name, index) => ({ + name, + index, + uri: makeUri(name), + })) + vi.mocked(vscode.workspace).workspaceFolders = folders + const { codeIndexWorkspaceScopeRegistry: registry } = + await import("../services/code-index/code-index-workspace-scope-registry") + const getScope = vi.spyOn(registry, "getScope").mockReturnValueOnce(undefined) + try { + const { activate } = await import("../extension") + await expect(activate(mockContext)).resolves.toBeDefined() + expect(getScope).toHaveBeenCalledWith(mockContext, folders[0]) + expect(getScope).toHaveBeenCalledWith(mockContext, folders[1]) + const scopes = registry.getAllScopes() + expect(scopes).toHaveLength(1) + expect(scopes[0].codeIndexManager.initialize).toHaveBeenCalledTimes(1) + } finally { + getScope.mockRestore() + } + }) + }) + describe("cloud organization settings handling", () => { beforeEach(() => { vi.resetModules() @@ -464,7 +649,9 @@ describe("extension.ts", () => { const { TelemetryService } = await import("@roo-code/telemetry") const { Terminal } = await import("../integrations/terminal/Terminal") const { TerminalRegistry } = await import("../integrations/terminal/TerminalRegistry") - const { CodeIndexManagerRegistry } = await import("../services/code-index/code-index-manager-registry") + const { codeIndexWorkspaceScopeRegistry } = + await import("../services/code-index/code-index-workspace-scope-registry") + const disposeAll = vi.spyOn(codeIndexWorkspaceScopeRegistry, "disposeAll") vi.mocked(TelemetryService.instance.shutdown).mockRejectedValue(new Error("shutdown failed")) const setTerminalProfileSpy = vi.spyOn(Terminal, "setTerminalProfile") @@ -476,7 +663,7 @@ describe("extension.ts", () => { expect(setTerminalProfileSpy).toHaveBeenCalledWith(undefined) expect(TerminalRegistry.cleanup).toHaveBeenCalledTimes(1) - expect(CodeIndexManagerRegistry.disposeAll).toHaveBeenCalledTimes(1) + expect(disposeAll).toHaveBeenCalledTimes(1) setTerminalProfileSpy.mockRestore() }) @@ -489,7 +676,9 @@ describe("extension.ts", () => { const { TelemetryService } = await import("@roo-code/telemetry") const { Terminal } = await import("../integrations/terminal/Terminal") const { TerminalRegistry } = await import("../integrations/terminal/TerminalRegistry") - const { CodeIndexManagerRegistry } = await import("../services/code-index/code-index-manager-registry") + const { codeIndexWorkspaceScopeRegistry } = + await import("../services/code-index/code-index-workspace-scope-registry") + const disposeAll = vi.spyOn(codeIndexWorkspaceScopeRegistry, "disposeAll") const setTerminalProfileSpy = vi.spyOn(Terminal, "setTerminalProfile") @@ -513,7 +702,7 @@ describe("extension.ts", () => { expect(mockTelemetryServiceInstance.shutdown).not.toHaveBeenCalled() expect(setTerminalProfileSpy).toHaveBeenCalledWith(undefined) expect(TerminalRegistry.cleanup).toHaveBeenCalledTimes(1) - expect(CodeIndexManagerRegistry.disposeAll).toHaveBeenCalledTimes(1) + expect(disposeAll).toHaveBeenCalledTimes(1) instanceGetterSpy.mockRestore() setTerminalProfileSpy.mockRestore() diff --git a/src/activate/__tests__/registerCommands.spec.ts b/src/activate/__tests__/registerCommands.spec.ts index 7088560700..8f480e1bbb 100644 --- a/src/activate/__tests__/registerCommands.spec.ts +++ b/src/activate/__tests__/registerCommands.spec.ts @@ -67,9 +67,9 @@ vi.mock("../../core/config/importExport", () => ({ importSettingsWithFeedback: vi.fn(), })) -vi.mock("../../services/code-index/code-index-manager-registry", () => ({ - CodeIndexManagerRegistry: { - getOrCreate: vi.fn(), +vi.mock("../../services/code-index/code-index-workspace-scope-registry", () => ({ + codeIndexWorkspaceScopeRegistry: { + getScope: vi.fn(), }, })) diff --git a/src/core/prompts/system.ts b/src/core/prompts/system.ts index 38283087bf..3c6ea068e5 100644 --- a/src/core/prompts/system.ts +++ b/src/core/prompts/system.ts @@ -14,7 +14,7 @@ import { formatLanguage } from "../../shared/language" import { isEmpty } from "../../utils/object" import { McpHub } from "../../services/mcp/McpHub" -import { CodeIndexManagerRegistry } from "../../services/code-index/code-index-manager-registry" +import type { CodeIndexWorkspaceScope } from "../../services/code-index/code-index-workspace-scope" import { SkillsManager } from "../../services/skills/SkillsManager" import type { SystemPromptSettings } from "./types" @@ -65,6 +65,7 @@ async function generatePrompt( skillsManager?: SkillsManager, disabledTools?: string[], modelInfo?: ModelInfo, + codeIndexWorkspaceScope?: CodeIndexWorkspaceScope, ): Promise { if (!context) { throw new Error("Extension context is required for generating system prompt") @@ -74,8 +75,6 @@ async function generatePrompt( const modeConfig = getModeBySlug(mode, customModeConfigs) || modes.find((m) => m.slug === mode) || modes[0] const { roleDefinition, baseInstructions } = getModeSelection(mode, promptComponent, customModeConfigs) - const codeIndexManager = CodeIndexManagerRegistry.getOrCreate(context, cwd) - // Resolve the single, request-scoped effective tool policy ONCE, then have every // prompt section and the MCP short-circuit derive from it. This is the one source of // truth shared by prompt generation, API tool construction, runtime validation, and @@ -88,7 +87,7 @@ async function generatePrompt( modelInfo, experiments, todoListEnabled: settings?.todoListEnabled, - codeIndexManager, + codeIndexManager: codeIndexWorkspaceScope?.codeIndexManager, }) // Tool calling is native-only. @@ -148,6 +147,7 @@ export const SYSTEM_PROMPT = async ( skillsManager?: SkillsManager, disabledTools?: string[], modelInfo?: ModelInfo, + codeIndexWorkspaceScope?: CodeIndexWorkspaceScope, ): Promise => { if (!context) { throw new Error("Extension context is required for generating system prompt") @@ -178,5 +178,6 @@ export const SYSTEM_PROMPT = async ( skillsManager, disabledTools, modelInfo, + codeIndexWorkspaceScope, ) } diff --git a/src/core/prompts/tools/__tests__/codebase-search-readiness.spec.ts b/src/core/prompts/tools/__tests__/codebase-search-readiness.spec.ts index 2e05f6c6cd..a00258a24a 100644 --- a/src/core/prompts/tools/__tests__/codebase-search-readiness.spec.ts +++ b/src/core/prompts/tools/__tests__/codebase-search-readiness.spec.ts @@ -1,6 +1,7 @@ import type OpenAI from "openai" import { toolNamesSchema, type ModeConfig } from "@roo-code/types" import type { CodeIndexManager } from "../../../../services/code-index/manager" +import type { CodeIndexWorkspaceScope } from "../../../../services/code-index/code-index-workspace-scope" import { filterNativeToolsForMode } from "../filter-tools-for-mode" import { resolveEffectiveToolPolicy } from "../effective-tool-policy" import { getNativeTools } from "../native-tools" @@ -14,6 +15,10 @@ function makeManager(flags: Readiness): CodeIndexManager { return flags as CodeIndexManager } +function makeScope(flags: Readiness): CodeIndexWorkspaceScope { + return { codeIndexManager: makeManager(flags) } as CodeIndexWorkspaceScope +} + function toolNames(definitions: OpenAI.Chat.ChatCompletionTool[]) { return definitions.flatMap((tool) => ("function" in tool ? [tool.function.name] : [])) } @@ -44,7 +49,7 @@ describe("codebase_search readiness", () => { (flags) => { const manager = makeManager(flags) const policy = resolveEffectiveToolPolicy({ mode: "code", codeIndexManager: manager }) - const filtered = filterNativeToolsForMode(getNativeTools(), "code", [], {}, manager) + const filtered = filterNativeToolsForMode(getNativeTools(), "code", [], {}, makeScope(flags)) expect(policy.tools).not.toContain(tools.codebase_search) expect(toolNames(filtered)).not.toContain(tools.codebase_search) @@ -58,7 +63,13 @@ describe("codebase_search readiness", () => { it("includes search when all three readiness conditions are met", () => { const manager = makeManager({ isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true }) const policy = resolveEffectiveToolPolicy({ mode: "code", codeIndexManager: manager }) - const filtered = filterNativeToolsForMode(getNativeTools(), "code", [], {}, manager) + const filtered = filterNativeToolsForMode( + getNativeTools(), + "code", + [], + {}, + makeScope({ isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true }), + ) expect(policy.tools).toContain(tools.codebase_search) expect(toolNames(filtered)).toContain(tools.codebase_search) @@ -77,7 +88,7 @@ describe("codebase_search readiness", () => { expect(resolveEffectiveToolPolicy({ mode: "code", codeIndexManager: manager }).tools).toContain( tools.codebase_search, ) - expect(toolNames(filterNativeToolsForMode(getNativeTools(), "code", [], {}, manager))).toContain( + expect(toolNames(filterNativeToolsForMode(getNativeTools(), "code", [], {}, makeScope(flags)))).toContain( tools.codebase_search, ) @@ -86,16 +97,16 @@ describe("codebase_search readiness", () => { expect(resolveEffectiveToolPolicy({ mode: "code", codeIndexManager: manager }).tools).not.toContain( tools.codebase_search, ) - expect(toolNames(filterNativeToolsForMode(getNativeTools(), "code", [], {}, manager))).not.toContain( - tools.codebase_search, - ) + expect( + toolNames(filterNativeToolsForMode(getNativeTools(), "code", [], {}, makeScope(flags))), + ).not.toContain(tools.codebase_search) flags[flag] = true expect(resolveEffectiveToolPolicy({ mode: "code", codeIndexManager: manager }).tools).toContain( tools.codebase_search, ) - expect(toolNames(filterNativeToolsForMode(getNativeTools(), "code", [], {}, manager))).toContain( + expect(toolNames(filterNativeToolsForMode(getNativeTools(), "code", [], {}, makeScope(flags)))).toContain( tools.codebase_search, ) }, @@ -108,15 +119,24 @@ describe("codebase_search readiness", () => { expect(resolveEffectiveToolPolicy({ mode: "code", codeIndexManager: ready }).tools).toContain( tools.codebase_search, ) - expect(toolNames(filterNativeToolsForMode(getNativeTools(), "code", [], {}, ready))).toContain( - tools.codebase_search, - ) + expect( + toolNames( + filterNativeToolsForMode( + getNativeTools(), + "code", + [], + {}, + makeScope({ isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true }), + ), + ), + ).toContain(tools.codebase_search) for (const manager of [disabled, undefined]) { expect(resolveEffectiveToolPolicy({ mode: "code", codeIndexManager: manager }).tools).not.toContain( tools.codebase_search, ) - expect(toolNames(filterNativeToolsForMode(getNativeTools(), "code", [], {}, manager))).not.toContain( + const scope = manager ? ({ codeIndexManager: manager } as CodeIndexWorkspaceScope) : undefined + expect(toolNames(filterNativeToolsForMode(getNativeTools(), "code", [], {}, scope))).not.toContain( tools.codebase_search, ) } @@ -124,16 +144,30 @@ describe("codebase_search readiness", () => { expect(resolveEffectiveToolPolicy({ mode: "code", codeIndexManager: ready }).tools).toContain( tools.codebase_search, ) - expect(toolNames(filterNativeToolsForMode(getNativeTools(), "code", [], {}, ready))).toContain( - tools.codebase_search, - ) + expect( + toolNames( + filterNativeToolsForMode( + getNativeTools(), + "code", + [], + {}, + makeScope({ isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true }), + ), + ), + ).toContain(tools.codebase_search) }) it("does not grant read permissions merely because the manager is ready", () => { const manager = makeManager({ isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true }) const mode: ModeConfig = { slug: "no-read", name: "No read", roleDefinition: "No reading", groups: ["command"] } const policy = resolveEffectiveToolPolicy({ mode: mode.slug, customModes: [mode], codeIndexManager: manager }) - const filtered = filterNativeToolsForMode(getNativeTools(), mode.slug, [mode], {}, manager) + const filtered = filterNativeToolsForMode( + getNativeTools(), + mode.slug, + [mode], + {}, + makeScope({ isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true }), + ) for (const tool of [tools.codebase_search, ...ordinaryReadTools]) { expect(policy.tools).not.toContain(tool) @@ -148,7 +182,14 @@ describe("codebase_search readiness", () => { const manager = makeManager({ isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true }) const disabledTools = [tools.codebase_search] const policy = resolveEffectiveToolPolicy({ mode: "code", codeIndexManager: manager, disabledTools }) - const filtered = filterNativeToolsForMode(getNativeTools(), "code", [], {}, manager, { disabledTools }) + const filtered = filterNativeToolsForMode( + getNativeTools(), + "code", + [], + {}, + makeScope({ isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true }), + { disabledTools }, + ) expect(policy.tools).not.toContain(tools.codebase_search) expect(toolNames(filtered)).not.toContain(tools.codebase_search) diff --git a/src/core/prompts/tools/__tests__/filter-tools-for-mode.spec.ts b/src/core/prompts/tools/__tests__/filter-tools-for-mode.spec.ts index 11198caff9..f9be9d4c93 100644 --- a/src/core/prompts/tools/__tests__/filter-tools-for-mode.spec.ts +++ b/src/core/prompts/tools/__tests__/filter-tools-for-mode.spec.ts @@ -3,7 +3,21 @@ import type OpenAI from "openai" import type { ModeConfig } from "@roo-code/types" +import type { CodeIndexManager } from "../../../../services/code-index/manager" +import type { CodeIndexWorkspaceScope } from "../../../../services/code-index/code-index-workspace-scope" import { filterMcpToolsForMode, filterNativeToolsForMode } from "../filter-tools-for-mode" +import { resolveEffectiveToolPolicy } from "../effective-tool-policy" + +type Readiness = Pick + +function makeWorkspaceScope(readiness: Readiness): CodeIndexWorkspaceScope { + return { + // Filtering only reads readiness; keep the same object so tests can change live state. + codeIndexManager: readiness as CodeIndexManager, + initialize: vi.fn(), + dispose: vi.fn(), + } +} function makeTool(name: string): OpenAI.Chat.ChatCompletionTool { return { @@ -16,6 +30,121 @@ function makeTool(name: string): OpenAI.Chat.ChatCompletionTool { } as OpenAI.Chat.ChatCompletionTool } +describe("workspace-scoped codebase_search filtering", () => { + it("retains search in native filtering and policy when the supplied scope is ready", () => { + const scope = makeWorkspaceScope({ isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true }) + const search = makeTool("codebase_search") + const read = makeTool("read_file") + + expect(filterNativeToolsForMode([read, search], "code", undefined, undefined, scope)).toEqual([read, search]) + expect( + resolveEffectiveToolPolicy({ mode: "code", codeIndexManager: scope.codeIndexManager }).tools.has( + "codebase_search", + ), + ).toBe(true) + }) + + const consumers = [ + { + name: "filterNativeToolsForMode", + available: (scope?: CodeIndexWorkspaceScope) => + filterNativeToolsForMode( + [makeTool("read_file"), makeTool("codebase_search")], + "code", + undefined, + undefined, + scope, + ).flatMap((tool) => (tool.type === "function" ? [tool.function.name] : [])), + }, + { + name: "resolveEffectiveToolPolicy", + available: (scope?: CodeIndexWorkspaceScope) => + Array.from( + resolveEffectiveToolPolicy({ mode: "code", codeIndexManager: scope?.codeIndexManager }).tools, + ), + }, + ] + + describe.each(consumers)("$name", ({ available }) => { + it.each([ + [false, false, false, false], + [false, false, true, false], + [false, true, false, false], + [false, true, true, false], + [true, false, false, false], + [true, false, true, false], + [true, true, false, false], + [true, true, true, true], + ])( + "enabled=%s configured=%s initialized=%s exposes search=%s", + (isFeatureEnabled, isFeatureConfigured, isInitialized, expected) => { + const scope = makeWorkspaceScope({ isFeatureEnabled, isFeatureConfigured, isInitialized }) + const tools = available(scope) + + expect(tools.includes("codebase_search")).toBe(expected) + expect(tools).toContain("read_file") + }, + ) + + it("hides search without a workspace scope but preserves ordinary read tools", () => { + const tools = available(undefined) + + expect(tools).not.toContain("codebase_search") + expect(tools).toContain("read_file") + }) + + it.each(["isFeatureEnabled", "isFeatureConfigured", "isInitialized"] as const)( + "rereads %s from the same manager on every call", + (flag) => { + const readiness = { isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true } + const scope = makeWorkspaceScope(readiness) + + expect(available(scope)).toContain("codebase_search") + readiness[flag] = false + expect(available(scope)).not.toContain("codebase_search") + readiness[flag] = true + expect(available(scope)).toContain("codebase_search") + }, + ) + + it("uses the supplied workspace rather than readiness from a previous workspace", () => { + const ready = makeWorkspaceScope({ isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true }) + const unready = makeWorkspaceScope({ + isFeatureEnabled: true, + isFeatureConfigured: true, + isInitialized: false, + }) + + expect(available(ready)).toContain("codebase_search") + expect(available(unready)).not.toContain("codebase_search") + expect(available(undefined)).not.toContain("codebase_search") + expect(available(ready)).toContain("codebase_search") + }) + }) + + it("does not let a ready workspace bypass a custom mode without the read group", () => { + const scope = makeWorkspaceScope({ isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true }) + const mode: ModeConfig = { + slug: "command-only", + name: "Command only", + roleDefinition: "Run commands only", + groups: ["command"], + } + const command = makeTool("execute_command") + + expect( + filterNativeToolsForMode([command, makeTool("codebase_search")], mode.slug, [mode], undefined, scope), + ).toEqual([command]) + expect( + resolveEffectiveToolPolicy({ + mode: mode.slug, + customModes: [mode], + codeIndexManager: scope.codeIndexManager, + }).tools.has("codebase_search"), + ).toBe(false) + }) +}) + describe("filterNativeToolsForMode - disabledTools", () => { const nativeTools: OpenAI.Chat.ChatCompletionTool[] = [ makeTool("execute_command"), diff --git a/src/core/prompts/tools/filter-tools-for-mode.ts b/src/core/prompts/tools/filter-tools-for-mode.ts index 45ccb39c5d..462f3e2ef8 100644 --- a/src/core/prompts/tools/filter-tools-for-mode.ts +++ b/src/core/prompts/tools/filter-tools-for-mode.ts @@ -1,7 +1,7 @@ import type OpenAI from "openai" import type { ModeConfig, ModelInfo } from "@roo-code/types" import { defaultModeSlug } from "../../../shared/modes" -import type { CodeIndexManager } from "../../../services/code-index/manager" +import type { CodeIndexWorkspaceScope } from "../../../services/code-index/code-index-workspace-scope" import type { McpHub } from "../../../services/mcp/McpHub" import { resolveEffectiveToolPolicy, resolveToolAlias, isToolDisabledOrExcluded } from "./effective-tool-policy" import { isToolAllowedForMode } from "../../../core/tools/validateToolUse" @@ -58,7 +58,7 @@ function getOrCreateRenamedTool( * @param mode - Current mode slug * @param customModes - Custom mode configurations * @param experiments - Experiment flags - * @param codeIndexManager - Code index manager for codebase_search feature check + * @param codeIndexWorkspaceScope - Workspace scope for codebase_search feature check * @param settings - Additional settings for tool filtering (includes modelInfo for model-specific customization) * @param mcpHub - MCP hub for checking available resources * @param allowedMcpServers - Optional allowlist of MCP server names for the current mode. When @@ -71,7 +71,7 @@ export function filterNativeToolsForMode( mode: string | undefined, customModes: ModeConfig[] | undefined, experiments: Record | undefined, - codeIndexManager?: CodeIndexManager, + codeIndexWorkspaceScope?: CodeIndexWorkspaceScope, settings?: Record, mcpHub?: McpHub, allowedMcpServers?: string[], @@ -91,7 +91,7 @@ export function filterNativeToolsForMode( modelInfo, experiments, todoListEnabled: settings?.todoListEnabled, - codeIndexManager, + codeIndexManager: codeIndexWorkspaceScope?.codeIndexManager, allowedMcpServers, }) diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index 55798437c3..0e8bb973a6 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -80,6 +80,8 @@ import { getModelMaxOutputTokens } from "../../shared/api" import { McpHub } from "../../services/mcp/McpHub" import { McpServerManager } from "../../services/mcp/McpServerManager" import { RepoPerTaskCheckpointService } from "../../services/checkpoints" +import type { CodeIndexWorkspaceScope } from "../../services/code-index/code-index-workspace-scope" +import { codeIndexWorkspaceScopeRegistry } from "../../services/code-index/code-index-workspace-scope-registry" // integrations import { DiffViewProvider } from "../../integrations/editor/DiffViewProvider" @@ -1875,6 +1877,7 @@ export class Task extends EventEmitter implements TaskLike { // from one snapshot. const state = await this.providerRef.deref()?.getState() const requestModelInfo = await this.safeEnsureModelFetched() + const codeIndexWorkspaceScope = await this.initializeCodeIndexWorkspaceScope() // A cancellation landing during the bounded metadata wait must stop // manual condensation before any prompt build or summarization request. @@ -1882,7 +1885,7 @@ export class Task extends EventEmitter implements TaskLike { return } - const systemPrompt = await this.getSystemPrompt(state, requestModelInfo) + const systemPrompt = await this.getSystemPrompt(state, requestModelInfo, codeIndexWorkspaceScope) // A cancellation landing during the prompt build's bounded MCP wait must // stop manual condensation before any summarization request is issued. @@ -1911,6 +1914,7 @@ export class Task extends EventEmitter implements TaskLike { apiConfiguration, disabledTools: state?.disabledTools, modelInfo: requestModelInfo, + codeIndexWorkspaceScope: codeIndexWorkspaceScope ?? null, includeAllToolsWithRestrictions: false, }) allTools = toolsResult.tools @@ -4242,6 +4246,7 @@ export class Task extends EventEmitter implements TaskLike { private async getSystemPrompt( requestState: Awaited> | undefined, requestModelInfo?: ModelInfo, + codeIndexWorkspaceScope?: CodeIndexWorkspaceScope, ): Promise { const { mcpEnabled } = requestState ?? {} let mcpHub: McpHub | undefined @@ -4313,6 +4318,7 @@ export class Task extends EventEmitter implements TaskLike { provider.getSkillsManager(), requestState?.disabledTools, modelInfo, + codeIndexWorkspaceScope, ) })() } @@ -4387,7 +4393,26 @@ export class Task extends EventEmitter implements TaskLike { return this.api.getModel().info } - private async handleContextWindowExceededError(requestModelInfo: ModelInfo): Promise { + private async initializeCodeIndexWorkspaceScope(): Promise { + const provider = this.providerRef.deref() + if (!provider) { + return undefined + } + + const scope = codeIndexWorkspaceScopeRegistry.getScope(provider.context, this.cwd) + try { + await scope?.initialize(provider.contextProxy) + return scope + } catch (error) { + console.error(`[Task#${this.taskId}] Failed to initialize code index workspace scope:`, error) + return undefined + } + } + + private async handleContextWindowExceededError( + requestModelInfo: ModelInfo, + codeIndexWorkspaceScope?: CodeIndexWorkspaceScope, + ): Promise { const state = await this.providerRef.deref()?.getState() const { profileThresholds = {} } = state ?? {} // Use task-local values, not provider state, to prevent cross-task configuration leaks. @@ -4437,6 +4462,7 @@ export class Task extends EventEmitter implements TaskLike { apiConfiguration, disabledTools: state?.disabledTools, modelInfo, + codeIndexWorkspaceScope: codeIndexWorkspaceScope ?? null, includeAllToolsWithRestrictions: false, }) allTools = toolsResult.tools @@ -4473,7 +4499,7 @@ export class Task extends EventEmitter implements TaskLike { apiHandler: this.api, autoCondenseContext: true, autoCondenseContextPercent: FORCED_CONTEXT_REDUCTION_PERCENT, - systemPrompt: await this.getSystemPrompt(state, modelInfo), + systemPrompt: await this.getSystemPrompt(state, modelInfo, codeIndexWorkspaceScope), taskId: this.taskId, profileThresholds, currentProfileId, @@ -4601,13 +4627,14 @@ export class Task extends EventEmitter implements TaskLike { // prompt and every tool array built below; prefer the caller's snapshot // when one was threaded. const requestModelInfo = options.requestModelInfo ?? (await this.safeEnsureModelFetched()) + const codeIndexWorkspaceScope = await this.initializeCodeIndexWorkspaceScope() // Retry recursions must reuse this snapshot instead of re-deriving it: a // metadata fetch landing between attempts would otherwise move // model-specific tool policy or `preserveReasoning` mid-request. When the // caller threaded a snapshot its options object is forwarded unchanged — // same reference, and never mutated. const retryOptions = options.requestModelInfo === undefined ? { ...options, requestModelInfo } : options - const systemPrompt = await this.getSystemPrompt(state, requestModelInfo) + const systemPrompt = await this.getSystemPrompt(state, requestModelInfo, codeIndexWorkspaceScope) // A cancellation landing during the rate-limit countdown, the bounded metadata // wait, or the MCP wait inside getSystemPrompt must stop this request before any @@ -4686,6 +4713,7 @@ export class Task extends EventEmitter implements TaskLike { apiConfiguration, disabledTools: state?.disabledTools, modelInfo: requestModelInfo, + codeIndexWorkspaceScope: codeIndexWorkspaceScope ?? null, includeAllToolsWithRestrictions: false, }) contextMgmtTools = toolsResult.tools @@ -4860,6 +4888,7 @@ export class Task extends EventEmitter implements TaskLike { apiConfiguration, disabledTools: state?.disabledTools, modelInfo, + codeIndexWorkspaceScope: codeIndexWorkspaceScope ?? null, includeAllToolsWithRestrictions: supportsAllowedFunctionNames, }) allTools = toolsResult.tools @@ -4948,7 +4977,7 @@ export class Task extends EventEmitter implements TaskLike { `Retry attempt ${retryAttempt + 1}/${MAX_CONTEXT_WINDOW_RETRIES}. ` + `Attempting automatic truncation...`, ) - await this.handleContextWindowExceededError(requestModelInfo) + await this.handleContextWindowExceededError(requestModelInfo, codeIndexWorkspaceScope) // Retry the request after handling the context window error yield* this.attemptApiRequest(retryAttempt + 1, retryOptions) return diff --git a/src/core/task/__tests__/Task.spec.ts b/src/core/task/__tests__/Task.spec.ts index c35acf864d..1731410e3c 100644 --- a/src/core/task/__tests__/Task.spec.ts +++ b/src/core/task/__tests__/Task.spec.ts @@ -141,11 +141,9 @@ vi.mock("p-wait-for", () => ({ })) // Task tests do not exercise indexing; keep workspace resolution and its cache out of this suite. -vi.mock("../../../services/code-index/code-index-manager-registry", () => ({ - CodeIndexManagerRegistry: { - getOrCreate: vi.fn().mockReturnValue(undefined), - getAllInstances: vi.fn().mockReturnValue([]), - disposeAll: vi.fn(), +vi.mock("../../../services/code-index/code-index-workspace-scope-registry", () => ({ + codeIndexWorkspaceScopeRegistry: { + getScope: vi.fn().mockReturnValue(undefined), }, })) @@ -1149,7 +1147,7 @@ describe("Cline", () => { // undefined and getSystemPrompt would re-read the divergent state and // re-guard the model metadata. The second argument must be the handler's // own settled snapshot, not just any object. - expect(getSystemPromptSpy).toHaveBeenCalledWith(snapshot, task.api.getModel().info) + expect(getSystemPromptSpy).toHaveBeenCalledWith(snapshot, task.api.getModel().info, undefined) }) it("threads the captured state snapshot into the system prompt when the context window is exceeded", async () => { @@ -1200,7 +1198,7 @@ describe("Cline", () => { // undefined and getSystemPrompt would re-read the divergent state and // the model info; the second argument must be the snapshot threaded // into the handler, not a fresh re-read. - expect(getSystemPromptSpy).toHaveBeenCalledWith(snapshot, ctxModelInfo) + expect(getSystemPromptSpy).toHaveBeenCalledWith(snapshot, ctxModelInfo, undefined) }) it("uses the task mode when manually condensing after focused state changes", async () => { @@ -1289,7 +1287,7 @@ describe("Cline", () => { // The state snapshot stays undefined for a gone provider; the model-info // snapshot is still captured from the task's own api handler. - expect(getSystemPromptSpy).toHaveBeenCalledWith(undefined, task.api.getModel().info) + expect(getSystemPromptSpy).toHaveBeenCalledWith(undefined, task.api.getModel().info, undefined) expect(overwriteSpy).toHaveBeenCalledTimes(1) }) diff --git a/src/core/task/__tests__/build-tools-readiness.integration.spec.ts b/src/core/task/__tests__/build-tools-readiness.integration.spec.ts index e1f8888c9d..1f48d9afad 100644 --- a/src/core/task/__tests__/build-tools-readiness.integration.spec.ts +++ b/src/core/task/__tests__/build-tools-readiness.integration.spec.ts @@ -1,13 +1,15 @@ import type OpenAI from "openai" import { toolNamesSchema } from "@roo-code/types" import type { CodeIndexManager } from "../../../services/code-index/manager" -import { CodeIndexManagerRegistry } from "../../../services/code-index/code-index-manager-registry" +import type { CodeIndexWorkspaceScope } from "../../../services/code-index/code-index-workspace-scope" +import { codeIndexWorkspaceScopeRegistry } from "../../../services/code-index/code-index-workspace-scope-registry" +import type { ContextProxy } from "../../config/ContextProxy" import { makeExtensionContext } from "../../../test-utils/vscode" import type { ClineProvider } from "../../webview/ClineProvider" import { buildNativeToolsArrayWithRestrictions } from "../build-tools" -vi.mock("../../../services/code-index/code-index-manager-registry", () => ({ - CodeIndexManagerRegistry: { getOrCreate: vi.fn() }, +vi.mock("../../../services/code-index/code-index-workspace-scope-registry", () => ({ + codeIndexWorkspaceScopeRegistry: { getScope: vi.fn() }, })) const tools = toolNamesSchema.enum @@ -26,12 +28,12 @@ describe.each([ { strategy: "filtered definitions", includeAllToolsWithRestrictions: false }, { strategy: "all definitions with an allowlist", includeAllToolsWithRestrictions: true }, ])("task readiness with $strategy", ({ includeAllToolsWithRestrictions }) => { - beforeEach(() => vi.mocked(CodeIndexManagerRegistry.getOrCreate).mockReset()) + beforeEach(() => vi.mocked(codeIndexWorkspaceScopeRegistry.getScope).mockReset()) function makeOptions() { const context = makeExtensionContext() // Only context and getMcpHub are needed; constructing a webview provider is unrelated to this test. - const provider = { context, getMcpHub: () => undefined } as ClineProvider + const provider = { context, contextProxy: {} as ContextProxy, getMcpHub: () => undefined } as ClineProvider return { provider, cwd: "/tasks/ready", @@ -47,22 +49,34 @@ describe.each([ return includeAllToolsWithRestrictions ? result.allowedFunctionNames : toolNames(result.tools) } + function makeScope(manager: CodeIndexManager): CodeIndexWorkspaceScope { + return { + codeIndexManager: manager, + initialize: vi.fn().mockResolvedValue({ requiresRestart: false }), + } as unknown as CodeIndexWorkspaceScope + } + it("uses the task context and cwd without leaking readiness between workspaces", async () => { const options = makeOptions() const ready = makeManager({ isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true }) const unready = makeManager({ isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: false }) - const managers = new Map([ - ["/tasks/ready", ready], - ["/tasks/unready", unready], + const scopes = new Map([ + ["/tasks/ready", makeScope(ready)], + ["/tasks/unready", makeScope(unready)], ]) - vi.mocked(CodeIndexManagerRegistry.getOrCreate).mockImplementation((_context, cwd) => managers.get(cwd ?? "")) + vi.mocked(codeIndexWorkspaceScopeRegistry.getScope).mockImplementation((_context, cwd) => + typeof cwd === "string" ? scopes.get(cwd) : undefined, + ) const first = await buildNativeToolsArrayWithRestrictions(options) - expect(CodeIndexManagerRegistry.getOrCreate).toHaveBeenLastCalledWith(options.provider.context, "/tasks/ready") + expect(codeIndexWorkspaceScopeRegistry.getScope).toHaveBeenLastCalledWith( + options.provider.context, + "/tasks/ready", + ) expect(callable(first)).toContain(tools.codebase_search) const other = await buildNativeToolsArrayWithRestrictions({ ...options, cwd: "/tasks/unready" }) - expect(CodeIndexManagerRegistry.getOrCreate).toHaveBeenLastCalledWith( + expect(codeIndexWorkspaceScopeRegistry.getScope).toHaveBeenLastCalledWith( options.provider.context, "/tasks/unready", ) @@ -80,13 +94,16 @@ describe.each([ } const restored = await buildNativeToolsArrayWithRestrictions(options) - expect(CodeIndexManagerRegistry.getOrCreate).toHaveBeenLastCalledWith(options.provider.context, "/tasks/ready") + expect(codeIndexWorkspaceScopeRegistry.getScope).toHaveBeenLastCalledWith( + options.provider.context, + "/tasks/ready", + ) expect(callable(restored)).toContain(tools.codebase_search) - expect(CodeIndexManagerRegistry.getOrCreate).toHaveBeenCalledTimes(3) + expect(codeIndexWorkspaceScopeRegistry.getScope).toHaveBeenCalledTimes(3) }) it("omits search without a manager while retaining ordinary read tools", async () => { - vi.mocked(CodeIndexManagerRegistry.getOrCreate).mockReturnValue(undefined) + vi.mocked(codeIndexWorkspaceScopeRegistry.getScope).mockReturnValue(undefined) const result = await buildNativeToolsArrayWithRestrictions({ ...makeOptions(), cwd: "/tasks/missing" }) if (includeAllToolsWithRestrictions) { @@ -102,12 +119,24 @@ describe.each([ } }) + it("isolates initialization failure while retaining ordinary read tools", async () => { + const scope = makeScope(makeManager({ isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true })) + vi.mocked(scope.initialize).mockRejectedValueOnce(new Error("initialization failed")) + vi.mocked(codeIndexWorkspaceScopeRegistry.getScope).mockReturnValue(scope) + + const result = await buildNativeToolsArrayWithRestrictions(makeOptions()) + expect(callable(result)).not.toContain(tools.codebase_search) + for (const tool of ordinaryReadTools) { + expect(callable(result)).toContain(tool) + } + }) + it.each(["isFeatureEnabled", "isFeatureConfigured", "isInitialized"] as const)( "rereads %s on subsequent builds with the same manager", async (flag) => { const options = makeOptions() const flags = { isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true } - vi.mocked(CodeIndexManagerRegistry.getOrCreate).mockReturnValue(makeManager(flags)) + vi.mocked(codeIndexWorkspaceScopeRegistry.getScope).mockReturnValue(makeScope(makeManager(flags))) const initial = await buildNativeToolsArrayWithRestrictions(options) expect(callable(initial)).toContain(tools.codebase_search) @@ -133,8 +162,8 @@ describe.each([ ) it("does not grant read tools to a command-only mode even with a ready manager", async () => { - vi.mocked(CodeIndexManagerRegistry.getOrCreate).mockReturnValue( - makeManager({ isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true }), + vi.mocked(codeIndexWorkspaceScopeRegistry.getScope).mockReturnValue( + makeScope(makeManager({ isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true })), ) const result = await buildNativeToolsArrayWithRestrictions({ ...makeOptions(), @@ -156,8 +185,8 @@ describe.each([ }) it("honors disabledTools with a ready manager without disabling ordinary read tools", async () => { - vi.mocked(CodeIndexManagerRegistry.getOrCreate).mockReturnValue( - makeManager({ isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true }), + vi.mocked(codeIndexWorkspaceScopeRegistry.getScope).mockReturnValue( + makeScope(makeManager({ isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true })), ) const result = await buildNativeToolsArrayWithRestrictions({ ...makeOptions(), diff --git a/src/core/task/__tests__/build-tools-workspace-scope.spec.ts b/src/core/task/__tests__/build-tools-workspace-scope.spec.ts new file mode 100644 index 0000000000..6dd59e33a9 --- /dev/null +++ b/src/core/task/__tests__/build-tools-workspace-scope.spec.ts @@ -0,0 +1,103 @@ +import type OpenAI from "openai" + +import type { ClineProvider } from "../../webview/ClineProvider" +import type { CodeIndexManager } from "../../../services/code-index/manager" +import type { CodeIndexWorkspaceScope } from "../../../services/code-index/code-index-workspace-scope" +import { makeExtensionContext } from "../../../test-utils/vscode" +import { codeIndexWorkspaceScopeRegistry } from "../../../services/code-index/code-index-workspace-scope-registry" +import * as filtering from "../../prompts/tools/filter-tools-for-mode" +import { buildNativeToolsArrayWithRestrictions } from "../build-tools" + +vi.mock("../../../services/code-index/code-index-workspace-scope-registry", () => ({ + codeIndexWorkspaceScopeRegistry: { getScope: vi.fn() }, +})) +vi.mock("@roo-code/core", () => ({ customToolRegistry: {}, formatNative: vi.fn() })) +vi.mock("../../../services/roo-config/index.js", () => ({ getRooDirectoriesForCwd: vi.fn() })) +vi.mock("../../prompts/tools/native-tools", () => ({ + getNativeTools: () => + ["read_file", "codebase_search"].map((name) => ({ + type: "function", + function: { name, description: name, parameters: { type: "object", properties: {} } }, + })), + getMcpServerTools: () => [], +})) + +function toolNames(tools: OpenAI.Chat.ChatCompletionTool[]): string[] { + return tools.flatMap((tool) => (tool.type === "function" ? [tool.function.name] : [])) +} + +describe("build tools workspace scope", () => { + afterEach(() => vi.restoreAllMocks()) + + it("forwards the full workspace scope and request cwd to real native filtering", async () => { + const context = makeExtensionContext() + // Building tools needs only the context and MCP accessor, not a webview host. + const provider = { context, getMcpHub: () => undefined } as ClineProvider + const readiness = { isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: true } + const scope: CodeIndexWorkspaceScope = { + // The real filter consumes only readiness; no indexing services are started. + codeIndexManager: readiness as CodeIndexManager, + initialize: vi.fn(), + dispose: vi.fn(), + } + vi.mocked(codeIndexWorkspaceScopeRegistry.getScope).mockReturnValue(scope) + const filter = vi.spyOn(filtering, "filterNativeToolsForMode") + const options = { + provider, + cwd: "/task-workspace", + mode: "code", + customModes: undefined, + experiments: undefined, + apiConfiguration: undefined, + } + + const result = await buildNativeToolsArrayWithRestrictions(options) + + expect(codeIndexWorkspaceScopeRegistry.getScope).toHaveBeenCalledWith(context, options.cwd) + expect(filter.mock.calls[0][4]).toBe(scope) + expect(toolNames(result.tools)).toEqual(["read_file", "codebase_search"]) + expect(result.allowedFunctionNames).toBeUndefined() + }) + + it.each([false, true])( + "updates search availability across scope changes with restrictions=%s", + async (includeAllToolsWithRestrictions) => { + const context = makeExtensionContext() + // Only the provider members read by the builder are needed at this boundary. + const provider = { context, getMcpHub: () => undefined } as ClineProvider + const readiness = { isFeatureEnabled: true, isFeatureConfigured: true, isInitialized: false } + const scope: CodeIndexWorkspaceScope = { + // The filter reads these live flags; the remaining manager services are irrelevant here. + codeIndexManager: readiness as CodeIndexManager, + initialize: vi.fn(), + dispose: vi.fn(), + } + const options = { + provider, + cwd: "/other-workspace", + mode: "code", + customModes: undefined, + experiments: undefined, + apiConfiguration: undefined, + includeAllToolsWithRestrictions, + } + const expectSearch = async (expected: boolean) => { + const result = await buildNativeToolsArrayWithRestrictions(options) + const allowed = includeAllToolsWithRestrictions ? result.allowedFunctionNames : toolNames(result.tools) + expect(allowed).toEqual(expected ? ["read_file", "codebase_search"] : ["read_file"]) + if (includeAllToolsWithRestrictions) { + expect(toolNames(result.tools)).toEqual(["read_file", "codebase_search"]) + } + } + + vi.mocked(codeIndexWorkspaceScopeRegistry.getScope).mockReturnValue(undefined) + await expectSearch(false) + vi.mocked(codeIndexWorkspaceScopeRegistry.getScope).mockReturnValue(scope) + await expectSearch(false) + readiness.isInitialized = true + await expectSearch(true) + readiness.isFeatureEnabled = false + await expectSearch(false) + }, + ) +}) diff --git a/src/core/task/__tests__/build-tools.spec.ts b/src/core/task/__tests__/build-tools.spec.ts index 65990932a2..6803295d53 100644 --- a/src/core/task/__tests__/build-tools.spec.ts +++ b/src/core/task/__tests__/build-tools.spec.ts @@ -18,9 +18,9 @@ import type { McpHub } from "../../../services/mcp/McpHub" // real, getOrCreate would construct a live manager from the stubbed context. // The all-false flags keep codebase_search out of every filter result, matching // the disabled-index baseline the assertions below assume. -vi.mock("../../../services/code-index/code-index-manager-registry", () => ({ - CodeIndexManagerRegistry: { - getOrCreate: () => ({ isFeatureEnabled: false, isFeatureConfigured: false, isInitialized: false }), +vi.mock("../../../services/code-index/code-index-workspace-scope-registry", () => ({ + codeIndexWorkspaceScopeRegistry: { + getScope: () => undefined, }, })) diff --git a/src/core/task/build-tools.ts b/src/core/task/build-tools.ts index 9d395eaa21..44093bbad2 100644 --- a/src/core/task/build-tools.ts +++ b/src/core/task/build-tools.ts @@ -6,6 +6,7 @@ import type { ProviderSettings, ModeConfig, ModelInfo } from "@roo-code/types" import { customToolRegistry, formatNative } from "@roo-code/core" import type { ClineProvider } from "../webview/ClineProvider" +import type { CodeIndexWorkspaceScope } from "../../services/code-index/code-index-workspace-scope" import { getRooDirectoriesForCwd } from "../../services/roo-config/index.js" import { getModeBySlug, defaultModeSlug } from "../../shared/modes" @@ -25,6 +26,8 @@ interface BuildToolsOptions { apiConfiguration: ProviderSettings | undefined disabledTools?: string[] modelInfo?: ModelInfo + /** A request-scoped, initialized scope. Null explicitly means initialization was unavailable. */ + codeIndexWorkspaceScope?: CodeIndexWorkspaceScope | null /** * If true, returns all tools without mode filtering, but also includes * the list of allowed tool names for use with allowedFunctionNames. @@ -93,14 +96,25 @@ export async function buildNativeToolsArrayWithRestrictions(options: BuildToolsO apiConfiguration, disabledTools, modelInfo, + codeIndexWorkspaceScope: suppliedCodeIndexWorkspaceScope, includeAllToolsWithRestrictions, } = options const mcpHub = provider.getMcpHub() - // Get CodeIndexManager for feature checking. - const { CodeIndexManagerRegistry } = await import("../../services/code-index/code-index-manager-registry") - const codeIndexManager = CodeIndexManagerRegistry.getOrCreate(provider.context, cwd) + // Get the workspace scope for code-index feature checking. + const { codeIndexWorkspaceScopeRegistry } = + await import("../../services/code-index/code-index-workspace-scope-registry") + let codeIndexWorkspaceScope = suppliedCodeIndexWorkspaceScope ?? undefined + if (suppliedCodeIndexWorkspaceScope === undefined) { + codeIndexWorkspaceScope = codeIndexWorkspaceScopeRegistry.getScope(provider.context, cwd) + try { + await codeIndexWorkspaceScope?.initialize(provider.contextProxy) + } catch (error) { + console.error("Failed to initialize code index workspace scope while building tools:", error) + codeIndexWorkspaceScope = undefined + } + } // Build settings object for tool filtering. const filterSettings = { @@ -129,7 +143,7 @@ export async function buildNativeToolsArrayWithRestrictions(options: BuildToolsO mode, customModes, experiments, - codeIndexManager, + codeIndexWorkspaceScope, filterSettings, mcpHub, allowedMcpServers, diff --git a/src/core/tools/CodebaseSearchTool.ts b/src/core/tools/CodebaseSearchTool.ts index ba1eb9bf75..e086d3b270 100644 --- a/src/core/tools/CodebaseSearchTool.ts +++ b/src/core/tools/CodebaseSearchTool.ts @@ -2,7 +2,7 @@ import * as vscode from "vscode" import path from "path" import { Task } from "../task/Task" -import { CodeIndexManagerRegistry } from "../../services/code-index/code-index-manager-registry" +import { codeIndexWorkspaceScopeRegistry } from "../../services/code-index/code-index-workspace-scope-registry" import { getWorkspacePath } from "../../utils/path" import { formatResponse } from "../prompts/responses" import { VectorStoreSearchResult } from "../../services/code-index/interfaces" @@ -57,7 +57,7 @@ export class CodebaseSearchTool extends BaseTool<"codebase_search"> { throw new Error("Extension context is not available.") } - const manager = CodeIndexManagerRegistry.getOrCreate(context) + const manager = codeIndexWorkspaceScopeRegistry.getScope(context, workspacePath)?.codeIndexManager if (!manager) { throw new Error("CodeIndexManager is not available.") diff --git a/src/core/tools/__tests__/CodebaseSearchTool.workspace-scope.spec.ts b/src/core/tools/__tests__/CodebaseSearchTool.workspace-scope.spec.ts new file mode 100644 index 0000000000..301c495f75 --- /dev/null +++ b/src/core/tools/__tests__/CodebaseSearchTool.workspace-scope.spec.ts @@ -0,0 +1,159 @@ +import * as vscode from "vscode" + +import { CodebaseSearchTool } from "../CodebaseSearchTool" +import type { ToolCallbacks } from "../BaseTool" +import type { Task } from "../../task/Task" +import type { ClineProvider } from "../../webview/ClineProvider" +import { CodeIndexManager } from "../../../services/code-index/manager" +import { codeIndexWorkspaceScopeRegistry as registry } from "../../../services/code-index/code-index-workspace-scope-registry" +import { getWorkspacePath } from "../../../utils/path" +import { makeExtensionContext, makeTextEditor, makeUri } from "../../../test-utils/vscode" + +vi.mock("../../../services/code-index/manager", () => ({ CodeIndexManager: vi.fn() })) +vi.mock("../../../utils/path", () => ({ getWorkspacePath: vi.fn() })) +vi.mock("vscode", () => ({ + workspace: { getWorkspaceFolder: vi.fn(), asRelativePath: vi.fn() }, + window: {}, +})) + +describe("CodebaseSearchTool workspace-scope consumer", () => { + const first = { name: "first", index: 0, uri: makeUri("/first") } + const second = { name: "second", index: 1, uri: makeUri("/second") } + const context = makeExtensionContext() + let task: { cwd: string } & Pick + let callbacks: ToolCallbacks + let manager: Pick + + beforeEach(() => { + vi.clearAllMocks() + manager = { + isFeatureEnabled: true, + isFeatureConfigured: true, + searchIndex: vi.fn().mockResolvedValue([]), + dispose: vi.fn(), + } + vi.mocked(CodeIndexManager).mockImplementation(function () { + // Only the search/disposal boundary is exercised; no indexing infrastructure is constructed. + return manager as CodeIndexManager + }) + Object.defineProperty(vscode.workspace, "workspaceFolders", { configurable: true, value: [first, second] }) + Object.defineProperty(vscode.window, "activeTextEditor", { configurable: true, value: makeTextEditor() }) + vi.spyOn(vscode.workspace, "getWorkspaceFolder").mockReturnValue(first) + vi.mocked(getWorkspacePath).mockReturnValue(first.uri.fsPath) + // The tool needs only task cwd, provider context, mistake count and result reporting. + task = { + cwd: second.uri.fsPath, + providerRef: new WeakRef({ context } as ClineProvider), + consecutiveMistakeCount: 2, + say: vi.fn().mockResolvedValue(undefined), + } + callbacks = { + askApproval: vi.fn().mockResolvedValue(true), + handleError: vi.fn().mockResolvedValue(undefined), + pushToolResult: vi.fn(), + } + }) + + afterEach(async () => { + await registry.disposeAll() + vi.restoreAllMocks() + }) + + it("forwards the task workspace instead of selecting the active editor's root and searches its manager", async () => { + const resolve = vi.spyOn(registry, "getScope") + await new CodebaseSearchTool().execute({ query: "scope lookup", path: "src/services" }, task as Task, callbacks) + + expect(resolve).toHaveBeenCalledWith(context, second.uri.fsPath) + expect(CodeIndexManager).toHaveBeenCalledWith(second.uri.fsPath, second.uri, context) + expect(getWorkspacePath).not.toHaveBeenCalled() + expect(manager.searchIndex).toHaveBeenCalledWith("scope lookup", "src/services") + expect(callbacks.pushToolResult).toHaveBeenCalledWith( + 'No relevant code snippets found for the query: "scope lookup"', + ) + expect(callbacks.handleError).not.toHaveBeenCalled() + expect(task.consecutiveMistakeCount).toBe(0) + }) + + it.each(["", " "])("forwards the fallback workspace when task cwd is %j", async (cwd) => { + task.cwd = cwd + vi.mocked(getWorkspacePath).mockReturnValue(second.uri.fsPath) + const resolve = vi.spyOn(registry, "getScope") + + await new CodebaseSearchTool().execute({ query: "fallback" }, task as Task, callbacks) + + expect(resolve).toHaveBeenCalledWith(context, second.uri.fsPath) + expect(CodeIndexManager).toHaveBeenCalledWith(second.uri.fsPath, second.uri, context) + expect(manager.searchIndex).toHaveBeenCalledWith("fallback", undefined) + expect(callbacks.handleError).not.toHaveBeenCalled() + }) + + it("reports an unavailable scope without searching or emitting a successful result", async () => { + vi.spyOn(registry, "getScope").mockReturnValue(undefined) + + await new CodebaseSearchTool().execute({ query: "missing scope" }, task as Task, callbacks) + + expect(callbacks.handleError).toHaveBeenCalledWith( + "codebase_search", + new Error("CodeIndexManager is not available."), + ) + expect(manager.searchIndex).not.toHaveBeenCalled() + expect(callbacks.pushToolResult).not.toHaveBeenCalled() + }) + + it("publishes populated results returned by the selected workspace manager", async () => { + vi.mocked(manager.searchIndex).mockResolvedValue([ + { + id: "snippet", + score: 0.9, + payload: { + filePath: "/second/src/search.ts", + startLine: 3, + endLine: 5, + codeChunk: " selected workspace code ", + }, + }, + ]) + vi.spyOn(vscode.workspace, "asRelativePath").mockReturnValue("src/search.ts") + + await new CodebaseSearchTool().execute({ query: "healthy search", path: "src" }, task as Task, callbacks) + + expect(manager.searchIndex).toHaveBeenCalledWith("healthy search", "src") + expect(task.say).toHaveBeenCalledWith( + "codebase_search_result", + JSON.stringify({ + tool: "codebaseSearch", + content: { + query: "healthy search", + results: [ + { + filePath: "src/search.ts", + score: 0.9, + startLine: 3, + endLine: 5, + codeChunk: "selected workspace code", + }, + ], + }, + }), + ) + expect(callbacks.pushToolResult).toHaveBeenCalledWith( + "Query: healthy search\nResults:\n\nFile path: src/search.ts\nScore: 0.9\nLines: 3-5\nCode Chunk: selected workspace code\n", + ) + expect(callbacks.handleError).not.toHaveBeenCalled() + }) + + it("reports a missing workspace before approval or scope resolution", async () => { + task.cwd = "" + vi.mocked(getWorkspacePath).mockReturnValue("") + const resolve = vi.spyOn(registry, "getScope") + + await new CodebaseSearchTool().execute({ query: "no workspace" }, task as Task, callbacks) + + expect(callbacks.handleError).toHaveBeenCalledWith( + "codebase_search", + new Error("Could not determine workspace path."), + ) + expect(callbacks.askApproval).not.toHaveBeenCalled() + expect(resolve).not.toHaveBeenCalled() + }) +}) diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index 86ce5d8e67..4f01b660fb 100644 --- a/src/core/webview/ClineProvider.ts +++ b/src/core/webview/ClineProvider.ts @@ -90,8 +90,8 @@ import { McpHub } from "../../services/mcp/McpHub" import { McpServerManager } from "../../services/mcp/McpServerManager" import { MarketplaceManager } from "../../services/marketplace" import { ShadowCheckpointService } from "../../services/checkpoints/ShadowCheckpointService" -import type { CodeIndexManager } from "../../services/code-index/manager" -import { CodeIndexManagerRegistry } from "../../services/code-index/code-index-manager-registry" +import type { CodeIndexWorkspaceScope } from "../../services/code-index/code-index-workspace-scope" +import { codeIndexWorkspaceScopeRegistry } from "../../services/code-index/code-index-workspace-scope-registry" import type { IndexProgressUpdate } from "../../services/code-index/interfaces/manager" import { MdmService } from "../../services/mdm/MdmService" import { SkillsManager } from "../../services/skills/SkillsManager" @@ -213,7 +213,7 @@ export class ClineProvider private static readonly delegationTransitionLocks = new Map>() private cancelledDelegationChildIds = new Set() private codeIndexStatusSubscription?: vscode.Disposable - private codeIndexManager?: CodeIndexManager + private codeIndexWorkspaceScope?: CodeIndexWorkspaceScope private _workspaceTracker?: WorkspaceTracker // workSpaceTracker read-only for access outside this class protected mcpHub?: McpHub // Change from private to protected protected skillsManager?: SkillsManager @@ -816,6 +816,9 @@ export class ClineProvider */ private clearWebviewResources() { this.rejectPendingThemeFixtureProbes(new Error("Webview was disposed before the theme fixture probe completed")) + this.codeIndexWorkspaceScope = undefined + this.codeIndexStatusSubscription?.dispose() + this.codeIndexStatusSubscription = undefined while (this.webviewDisposables.length) { const x = this.webviewDisposables.pop() if (x) { @@ -1130,8 +1133,6 @@ export class ClineProvider } else { this.log("Clearing webview resources for sidebar view") this.clearWebviewResources() - // Reset current workspace manager reference when view is disposed - this.codeIndexManager = undefined } }, null, @@ -3304,22 +3305,21 @@ export class ClineProvider } /** - * Gets the CodeIndexManager for the current active workspace - * @returns CodeIndexManager instance for the current workspace or the default one + * Gets the code-index scope for the current active workspace. + * @returns Workspace scope for the active workspace or the default one. */ - public getCurrentWorkspaceCodeIndexManager(): CodeIndexManager | undefined { - return CodeIndexManagerRegistry.getOrCreate(this.context) + public getCurrentWorkspaceCodeIndexScope(): CodeIndexWorkspaceScope | undefined { + return codeIndexWorkspaceScopeRegistry.getScope(this.context) } /** * Updates the code index status subscription to listen to the current workspace manager */ private updateCodeIndexStatusSubscription(): void { - // Get the current workspace manager - const currentManager = this.getCurrentWorkspaceCodeIndexManager() + const currentWorkspaceScope = this.getCurrentWorkspaceCodeIndexScope() - // If the manager hasn't changed, no need to update subscription - if (currentManager === this.codeIndexManager) { + // If the scope hasn't changed, no need to update subscription + if (currentWorkspaceScope === this.codeIndexWorkspaceScope) { return } @@ -3329,14 +3329,18 @@ export class ClineProvider this.codeIndexStatusSubscription = undefined } - // Update the current workspace manager reference - this.codeIndexManager = currentManager + // Update the current workspace scope reference + this.codeIndexWorkspaceScope = currentWorkspaceScope // Subscribe to the new manager's progress updates if it exists - if (currentManager) { + if (currentWorkspaceScope) { + const currentManager = currentWorkspaceScope.codeIndexManager this.codeIndexStatusSubscription = currentManager.onProgressUpdate((update: IndexProgressUpdate) => { - // Only send updates if this manager is still the current one - if (currentManager === this.getCurrentWorkspaceCodeIndexManager()) { + // Only send updates if this scope is still the current one + if ( + currentWorkspaceScope === this.codeIndexWorkspaceScope && + currentWorkspaceScope === this.getCurrentWorkspaceCodeIndexScope() + ) { // Get the full status from the manager to ensure we have all fields correctly formatted const fullStatus = currentManager.getCurrentStatus() void this.postMessageToWebview({ @@ -3346,10 +3350,6 @@ export class ClineProvider } }) - if (this.view) { - this.webviewDisposables.push(this.codeIndexStatusSubscription) - } - // Send initial status for the current workspace void this.postMessageToWebview({ type: "indexingStatusUpdate", diff --git a/src/core/webview/__tests__/ClineProvider.spec.ts b/src/core/webview/__tests__/ClineProvider.spec.ts index 97c4dd877e..b02c361064 100644 --- a/src/core/webview/__tests__/ClineProvider.spec.ts +++ b/src/core/webview/__tests__/ClineProvider.spec.ts @@ -35,6 +35,8 @@ import { webviewMessageHandler } from "../webviewMessageHandler" import { Terminal } from "../../../integrations/terminal/Terminal" import { MessageManager } from "../../message-manager" import { forceFullModelDetailsLoad, hasLoadedFullDetails } from "../../../api/providers/fetchers/lmstudio" +import { CodeIndexWorkspaceScope } from "../../../services/code-index/code-index-workspace-scope" +import { codeIndexWorkspaceScopeRegistry } from "../../../services/code-index/code-index-workspace-scope-registry" // Mock setup must come before imports. vi.mock("../../prompts/sections/custom-instructions") @@ -561,6 +563,145 @@ describe("ClineProvider", () => { }) }) + describe("workspace-scope progress subscriptions", () => { + let activeScope: CodeIndexWorkspaceScope | undefined + let changeEditor: () => void + let disposeView: () => Promise + const editorDisposable = { dispose: vi.fn() } + + function workspace(workspacePath: string) { + const scope = new CodeIndexWorkspaceScope(workspacePath, mockContext.extensionUri, mockContext) + const manager = scope.codeIndexManager + const status = { ...manager.getCurrentStatus(), message: workspacePath } + const subscriptions: { callback: (update: typeof status) => void; dispose: ReturnType }[] = [] + const subscribe = vi.fn((callback) => { + const subscription = { callback, dispose: vi.fn() } + subscriptions.push(subscription) + return subscription + }) + Object.defineProperty(manager, "onProgressUpdate", { value: subscribe }) + const getStatus = vi.spyOn(manager, "getCurrentStatus").mockReturnValue(status) + return { scope, status, subscriptions, subscribe, getStatus } + } + + beforeEach(() => { + activeScope = undefined + vi.spyOn(codeIndexWorkspaceScopeRegistry, "getScope").mockImplementation(() => activeScope) + vi.mocked(vscode.window.onDidChangeActiveTextEditor).mockImplementation((callback) => { + changeEditor = () => callback(undefined) + return editorDisposable + }) + mockWebviewView.onDidDispose = vi.fn((callback: () => Promise) => { + disposeView = callback + return { dispose: vi.fn() } + }) + }) + + afterEach(async () => { + await provider.dispose() + vi.restoreAllMocks() + }) + + it("rejects a captured stale A callback after switching to B even when workspace lookup returns A again", async () => { + const a = workspace("/a") + const b = workspace("/b") + activeScope = a.scope + await provider.resolveWebviewView(mockWebviewView) + activeScope = b.scope + changeEditor() + mockPostMessage.mockClear() + a.getStatus.mockClear() + + // The editor lookup can change before the provider receives the editor event. + activeScope = a.scope + a.subscriptions[0].callback(a.status) + + expect(mockPostMessage).not.toHaveBeenCalled() + expect(a.getStatus).not.toHaveBeenCalled() + }) + + it("reuses scope identity and publishes full current status rather than the progress payload", async () => { + const a = workspace("/a") + activeScope = a.scope + await provider.resolveWebviewView(mockWebviewView) + expect(provider.getCurrentWorkspaceCodeIndexScope()).toBe(a.scope) + expect(codeIndexWorkspaceScopeRegistry.getScope).toHaveBeenCalledWith(mockContext) + expect(mockPostMessage).toHaveBeenCalledWith({ type: "indexingStatusUpdate", values: a.status }) + mockPostMessage.mockClear() + changeEditor() + changeEditor() + expect(a.subscribe).toHaveBeenCalledOnce() + expect(a.subscriptions[0].dispose).not.toHaveBeenCalled() + expect(mockPostMessage).not.toHaveBeenCalled() + + const latest = { ...a.status, processedItems: 7, totalItems: 10 } + a.getStatus.mockReturnValue(latest) + a.subscriptions[0].callback(a.status) + expect(mockPostMessage).toHaveBeenCalledExactlyOnceWith({ type: "indexingStatusUpdate", values: latest }) + }) + + it("handles no workspace initially and detaches progress when the workspace disappears", async () => { + const a = workspace("/a") + await provider.resolveWebviewView(mockWebviewView) + changeEditor() + expect(provider.getCurrentWorkspaceCodeIndexScope()).toBeUndefined() + expect(a.subscribe).not.toHaveBeenCalled() + expect(mockPostMessage).not.toHaveBeenCalledWith(expect.objectContaining({ type: "indexingStatusUpdate" })) + + activeScope = a.scope + changeEditor() + expect(a.subscribe).toHaveBeenCalledOnce() + mockPostMessage.mockClear() + // Reject progress as soon as lookup changes, before the editor notification. + activeScope = undefined + a.subscriptions[0].callback(a.status) + expect(mockPostMessage).not.toHaveBeenCalled() + changeEditor() + changeEditor() + expect(a.subscriptions[0].dispose).toHaveBeenCalledOnce() + activeScope = a.scope + a.subscriptions[0].callback(a.status) + expect(mockPostMessage).not.toHaveBeenCalled() + await provider.dispose() + expect(a.subscriptions[0].dispose).toHaveBeenCalledOnce() + }) + + it("disposes each subscription once across A to B, sidebar disposal and reattach", async () => { + const a = workspace("/a") + const b = workspace("/b") + const disposeScope = vi.spyOn(b.scope, "dispose") + activeScope = a.scope + await provider.resolveWebviewView(mockWebviewView) + activeScope = b.scope + changeEditor() + expect(a.subscriptions[0].dispose).toHaveBeenCalledOnce() + expect(a.subscriptions[0].dispose.mock.invocationCallOrder[0]).toBeLessThan( + b.subscribe.mock.invocationCallOrder[0], + ) + + await disposeView() + expect(a.subscriptions[0].dispose).toHaveBeenCalledOnce() + expect(b.subscriptions[0].dispose).toHaveBeenCalledOnce() + expect(editorDisposable.dispose).toHaveBeenCalledOnce() + mockPostMessage.mockClear() + b.getStatus.mockClear() + b.subscriptions[0].callback(b.status) + expect(mockPostMessage).not.toHaveBeenCalled() + expect(b.getStatus).not.toHaveBeenCalled() + + await provider.resolveWebviewView(mockWebviewView) + expect(b.subscribe).toHaveBeenCalledTimes(2) + expect(b.subscriptions[0].dispose).toHaveBeenCalledOnce() + mockPostMessage.mockClear() + b.subscriptions[1].callback(b.status) + expect(mockPostMessage).toHaveBeenCalledExactlyOnceWith({ type: "indexingStatusUpdate", values: b.status }) + await provider.dispose() + expect(b.subscriptions[1].dispose).toHaveBeenCalledOnce() + expect(b.subscriptions[0].dispose).toHaveBeenCalledOnce() + expect(disposeScope).not.toHaveBeenCalled() + }) + }) + test("constructor initializes correctly", () => { expect(provider).toBeInstanceOf(ClineProvider) // Since getVisibleInstance returns the last instance where view.visible is true @@ -2957,7 +3098,7 @@ describe("webviewMessageHandler no-floating-promises coverage", () => { postMessageToWebview: vi.fn().mockResolvedValue(true), postStateToWebview: vi.fn().mockResolvedValue(undefined), getCurrentTask: vi.fn(), - getCurrentWorkspaceCodeIndexManager: vi.fn(), + getCurrentWorkspaceCodeIndexScope: vi.fn(), getMcpHub: vi.fn().mockReturnValue({ getMcpSettingsFilePath: vi.fn().mockResolvedValue("/test/mcp.json"), }), @@ -3013,7 +3154,7 @@ describe("webviewMessageHandler no-floating-promises coverage", () => { startIndexing: vi.fn().mockReturnValue(indexingPromise), }) const provider = createProvider({ - getCurrentWorkspaceCodeIndexManager: vi.fn().mockReturnValue(manager), + getCurrentWorkspaceCodeIndexScope: vi.fn().mockReturnValue({ codeIndexManager: manager }), }) await expect(webviewMessageHandler(provider, { type: "startIndexing" })).resolves.toBeUndefined() @@ -3170,13 +3311,13 @@ describe("webviewMessageHandler no-floating-promises coverage", () => { it("covers changed indexing status, secret, and missing-manager responses", async () => { const manager = createIndexManager() - const getManager = vi.fn().mockReturnValueOnce(undefined).mockReturnValue(manager) - const provider = createProvider({ getCurrentWorkspaceCodeIndexManager: getManager }) + const getScope = vi.fn().mockReturnValueOnce(undefined).mockReturnValue({ codeIndexManager: manager }) + const provider = createProvider({ getCurrentWorkspaceCodeIndexScope: getScope }) await webviewMessageHandler(provider, { type: "requestIndexingStatus" }) await webviewMessageHandler(provider, { type: "requestIndexingStatus" }) await webviewMessageHandler(provider, { type: "requestCodeIndexSecretStatus" }) - getManager.mockReturnValueOnce(undefined) + getScope.mockReturnValueOnce(undefined) await webviewMessageHandler(provider, { type: "startIndexing" }) expect(provider.postMessageToWebview).toHaveBeenCalledWith( @@ -3194,7 +3335,7 @@ describe("webviewMessageHandler no-floating-promises coverage", () => { .mockRejectedValueOnce(new Error("second failure")), }) const provider = createProvider({ - getCurrentWorkspaceCodeIndexManager: vi.fn().mockReturnValue(manager), + getCurrentWorkspaceCodeIndexScope: vi.fn().mockReturnValue({ codeIndexManager: manager }), }) await webviewMessageHandler(provider, { type: "startIndexing" }) @@ -3210,7 +3351,7 @@ describe("webviewMessageHandler no-floating-promises coverage", () => { startIndexing: vi.fn().mockRejectedValue(new Error("toggle failure")), }) const provider = createProvider({ - getCurrentWorkspaceCodeIndexManager: vi.fn().mockReturnValue(manager), + getCurrentWorkspaceCodeIndexScope: vi.fn().mockReturnValue({ codeIndexManager: manager }), }) await webviewMessageHandler(provider, { type: "stopIndexing" }) @@ -3225,7 +3366,8 @@ describe("webviewMessageHandler no-floating-promises coverage", () => { }) it("catches auto-enabled indexing failures and posts the resulting status", async () => { - const { CodeIndexManagerRegistry } = await import("../../../services/code-index/code-index-manager-registry") + const { codeIndexWorkspaceScopeRegistry } = + await import("../../../services/code-index/code-index-workspace-scope-registry") let workspaceEnabled = false const manager = createIndexManager({ setAutoEnableDefault: vi.fn().mockImplementation(async () => { @@ -3234,11 +3376,13 @@ describe("webviewMessageHandler no-floating-promises coverage", () => { startIndexing: vi.fn().mockRejectedValue(new Error("auto-enable failure")), }) Object.defineProperty(manager, "isWorkspaceEnabled", { get: () => workspaceEnabled }) - const getAllInstances = vi - .spyOn(CodeIndexManagerRegistry, "getAllInstances") - .mockReturnValue([manager] as unknown as ReturnType) + const getAllScopes = vi + .spyOn(codeIndexWorkspaceScopeRegistry, "getAllScopes") + .mockReturnValue([{ codeIndexManager: manager }] as unknown as ReturnType< + typeof codeIndexWorkspaceScopeRegistry.getAllScopes + >) const provider = createProvider({ - getCurrentWorkspaceCodeIndexManager: vi.fn().mockReturnValue(manager), + getCurrentWorkspaceCodeIndexScope: vi.fn().mockReturnValue({ codeIndexManager: manager }), }) try { @@ -3251,14 +3395,14 @@ describe("webviewMessageHandler no-floating-promises coverage", () => { expect.objectContaining({ type: "indexingStatusUpdate" }), ) } finally { - getAllInstances.mockRestore() + getAllScopes.mockRestore() } }) it("covers changed clear-index response paths", async () => { const manager = createIndexManager() - const getManager = vi.fn().mockReturnValueOnce(undefined).mockReturnValue(manager) - const provider = createProvider({ getCurrentWorkspaceCodeIndexManager: getManager }) + const getScope = vi.fn().mockReturnValueOnce(undefined).mockReturnValue({ codeIndexManager: manager }) + const provider = createProvider({ getCurrentWorkspaceCodeIndexScope: getScope }) await webviewMessageHandler(provider, { type: "clearIndexData" }) await webviewMessageHandler(provider, { type: "clearIndexData" }) diff --git a/src/core/webview/__tests__/webviewMessageHandler.auto-enable-scopes.spec.ts b/src/core/webview/__tests__/webviewMessageHandler.auto-enable-scopes.spec.ts new file mode 100644 index 0000000000..2c824307d8 --- /dev/null +++ b/src/core/webview/__tests__/webviewMessageHandler.auto-enable-scopes.spec.ts @@ -0,0 +1,138 @@ +import type { CodeIndexManager } from "../../../services/code-index/manager" +import type { CodeIndexWorkspaceScope } from "../../../services/code-index/code-index-workspace-scope" +import { codeIndexWorkspaceScopeRegistry as registry } from "../../../services/code-index/code-index-workspace-scope-registry" +import type { ContextProxy } from "../../config/ContextProxy" +import type { ClineProvider } from "../ClineProvider" +import { webviewMessageHandler } from "../webviewMessageHandler" + +vi.mock("../ClineProvider", () => ({ ClineProvider: vi.fn() })) +vi.mock("../../../services/code-index/code-index-workspace-scope-registry", () => ({ + codeIndexWorkspaceScopeRegistry: { getAllScopes: vi.fn() }, +})) + +describe("webviewMessageHandler global auto-enable across workspace scopes", () => { + let autoEnable: boolean + const contextProxy = {} as ContextProxy + + function makeScope(workspacePath: string, explicit?: boolean, configured = true) { + const manager = { + get isWorkspaceEnabled() { + return explicit ?? autoEnable + }, + isFeatureEnabled: true, + isFeatureConfigured: configured, + setAutoEnableDefault: vi.fn(async (enabled: boolean) => { + autoEnable = enabled + }), + initialize: vi.fn().mockResolvedValue({ requiresRestart: false }), + startIndexing: vi.fn().mockResolvedValue(undefined), + stopIndexing: vi.fn(), + getCurrentStatus: vi.fn().mockImplementation(() => ({ + systemStatus: "Standby", + message: workspacePath, + processedItems: 0, + totalItems: 0, + currentItemUnit: "files", + workspacePath, + workspaceEnabled: explicit ?? autoEnable, + autoEnableDefault: autoEnable, + })), + } satisfies Partial + const scope: CodeIndexWorkspaceScope = { + // Private manager infrastructure prevents structural assignment; this double implements only the consumer boundary. + codeIndexManager: manager as unknown as CodeIndexManager, + initialize: manager.initialize, + dispose: vi.fn(), + } + return { scope, manager } + } + + function makeProvider(scope?: CodeIndexWorkspaceScope) { + const provider = { + contextProxy, + getCurrentWorkspaceCodeIndexScope: vi.fn(() => scope), + postMessageToWebview: vi.fn().mockResolvedValue(undefined), + log: vi.fn(), + } + return provider + } + + async function setDefault(provider: ReturnType, bool?: boolean) { + // ClineProvider has private extension infrastructure; the handler branch only needs these four typed members. + await webviewMessageHandler(provider as unknown as ClineProvider, { type: "setAutoEnableDefault", bool }) + } + + beforeEach(() => { + vi.clearAllMocks() + autoEnable = false + }) + + it.each([true, undefined])( + "starts every newly enabled configured scope when bool=%s, isolating start rejection", + async (bool) => { + const failing = makeScope("/failing") + const healthy = makeScope("/healthy") + const optedOut = makeScope("/opted-out", false) + const alreadyEnabled = makeScope("/already-enabled", true) + const unconfigured = makeScope("/unconfigured", undefined, false) + const scopes = [failing, healthy, optedOut, alreadyEnabled, unconfigured] + vi.mocked(registry.getAllScopes).mockReturnValue(scopes.map(({ scope }) => scope)) + failing.manager.startIndexing.mockRejectedValue(new Error("first scope failed")) + const provider = makeProvider(healthy.scope) + + await setDefault(provider, bool) + + expect(healthy.manager.setAutoEnableDefault).toHaveBeenCalledWith(true) + for (const { manager } of [failing, healthy]) { + expect(manager.initialize).toHaveBeenCalledWith(contextProxy) + expect(manager.startIndexing).toHaveBeenCalledOnce() + expect(manager.stopIndexing).not.toHaveBeenCalled() + } + for (const { manager } of [optedOut, alreadyEnabled, unconfigured]) { + expect(manager.initialize).not.toHaveBeenCalled() + expect(manager.startIndexing).not.toHaveBeenCalled() + expect(manager.stopIndexing).not.toHaveBeenCalled() + } + expect(provider.log).toHaveBeenCalledWith("Indexing error: Error: first scope failed") + expect(provider.postMessageToWebview).toHaveBeenCalledWith({ + type: "indexingStatusUpdate", + values: healthy.manager.getCurrentStatus(), + }) + }, + ) + + it("stops all scopes disabled by the global default but preserves explicit workspace choices", async () => { + autoEnable = true + const first = makeScope("/first") + const second = makeScope("/second") + const optedIn = makeScope("/opted-in", true) + const optedOut = makeScope("/opted-out", false) + const scopes = [first, second, optedIn, optedOut] + vi.mocked(registry.getAllScopes).mockReturnValue(scopes.map(({ scope }) => scope)) + const provider = makeProvider(second.scope) + + await setDefault(provider, false) + + expect(second.manager.setAutoEnableDefault).toHaveBeenCalledWith(false) + for (const { manager } of [first, second]) expect(manager.stopIndexing).toHaveBeenCalledOnce() + for (const { manager } of [optedIn, optedOut]) expect(manager.stopIndexing).not.toHaveBeenCalled() + for (const { manager } of scopes) { + expect(manager.initialize).not.toHaveBeenCalled() + expect(manager.startIndexing).not.toHaveBeenCalled() + } + expect(provider.postMessageToWebview).toHaveBeenCalledWith({ + type: "indexingStatusUpdate", + values: second.manager.getCurrentStatus(), + }) + }) + + it("does not enumerate scopes or update status when there is no current workspace scope", async () => { + const provider = makeProvider() + + await setDefault(provider, true) + + expect(provider.log).toHaveBeenCalledWith("Cannot set auto-enable default: No workspace folder open") + expect(registry.getAllScopes).not.toHaveBeenCalled() + expect(provider.postMessageToWebview).not.toHaveBeenCalled() + }) +}) diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index 34a35ea3ca..049960eafe 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -62,7 +62,7 @@ import { Package } from "../../shared/package" import { type RouterName, toRouterName } from "../../shared/api" import { MessageEnhancer } from "./messageEnhancer" -import { CodeIndexManagerRegistry } from "../../services/code-index/code-index-manager-registry" +import { codeIndexWorkspaceScopeRegistry } from "../../services/code-index/code-index-workspace-scope-registry" import { checkExistKey } from "../../shared/checkExistApiConfig" import { getRouterRemovalMessage, getRouterUnavailableSignInMessage } from "../config/routerRemoval" import { experimentDefault } from "../../shared/experiments" @@ -3078,7 +3078,7 @@ export const webviewMessageHandler = async ( await provider.postStateToWebview() // Then handle validation and initialization for the current workspace - const currentCodeIndexManager = provider.getCurrentWorkspaceCodeIndexManager() + const currentCodeIndexManager = provider.getCurrentWorkspaceCodeIndexScope()?.codeIndexManager if (currentCodeIndexManager) { // If embedder provider changed, perform proactive validation if (embedderProviderChanged) { @@ -3157,7 +3157,7 @@ export const webviewMessageHandler = async ( } case "requestIndexingStatus": { - const manager = provider.getCurrentWorkspaceCodeIndexManager() + const manager = provider.getCurrentWorkspaceCodeIndexScope()?.codeIndexManager if (!manager) { // No workspace open - send error status await provider.postMessageToWebview({ @@ -3221,7 +3221,7 @@ export const webviewMessageHandler = async ( } case "startIndexing": { try { - const manager = provider.getCurrentWorkspaceCodeIndexManager() + const manager = provider.getCurrentWorkspaceCodeIndexScope()?.codeIndexManager if (!manager) { await provider.postMessageToWebview({ type: "indexingStatusUpdate", @@ -3262,7 +3262,7 @@ export const webviewMessageHandler = async ( } case "stopIndexing": { try { - const manager = provider.getCurrentWorkspaceCodeIndexManager() + const manager = provider.getCurrentWorkspaceCodeIndexScope()?.codeIndexManager if (!manager) { provider.log("Cannot stop indexing: No workspace folder open") return @@ -3279,7 +3279,7 @@ export const webviewMessageHandler = async ( } case "toggleWorkspaceIndexing": { try { - const manager = provider.getCurrentWorkspaceCodeIndexManager() + const manager = provider.getCurrentWorkspaceCodeIndexScope()?.codeIndexManager if (!manager) { provider.log("Cannot toggle workspace indexing: No workspace folder open") return @@ -3305,13 +3305,15 @@ export const webviewMessageHandler = async ( } case "setAutoEnableDefault": { try { - const manager = provider.getCurrentWorkspaceCodeIndexManager() + const manager = provider.getCurrentWorkspaceCodeIndexScope()?.codeIndexManager if (!manager) { provider.log("Cannot set auto-enable default: No workspace folder open") return } // Capture prior state for every manager before persisting the global change - const allManagers = CodeIndexManagerRegistry.getAllInstances() + const allManagers = codeIndexWorkspaceScopeRegistry + .getAllScopes() + .map((scope) => scope.codeIndexManager) const priorStates = new Map(allManagers.map((m) => [m, m.isWorkspaceEnabled])) await manager.setAutoEnableDefault(message.bool ?? true) // Apply stop/start to every affected manager @@ -3338,7 +3340,7 @@ export const webviewMessageHandler = async ( } case "clearIndexData": { try { - const manager = provider.getCurrentWorkspaceCodeIndexManager() + const manager = provider.getCurrentWorkspaceCodeIndexScope()?.codeIndexManager if (!manager) { provider.log("Cannot clear index data: No workspace folder open") await provider.postMessageToWebview({ diff --git a/src/extension.ts b/src/extension.ts index 8706de765b..76cf2b63ae 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -34,7 +34,7 @@ import { TerminalRegistry } from "./integrations/terminal/TerminalRegistry" import { openAiCodexOAuthManager } from "./integrations/openai-codex/oauth" import { kimiCodeOAuthManager } from "./integrations/kimi-code/oauth" import { McpServerManager } from "./services/mcp/McpServerManager" -import { CodeIndexManagerRegistry } from "./services/code-index/code-index-manager-registry" +import { codeIndexWorkspaceScopeRegistry } from "./services/code-index/code-index-workspace-scope-registry" import { MdmService } from "./services/mdm/MdmService" import { migrateSettings } from "./utils/migrateSettings" import { autoImportSettings } from "./utils/autoImportSettings" @@ -195,21 +195,34 @@ export async function activate(context: vscode.ExtensionContext) { }), ) - // Initialize code index managers for all workspace folders. + // The registry owns all scopes, including those created lazily after activation. + const codeIndexInitializations: Promise[] = [] + context.subscriptions.push({ + dispose: async () => { + await Promise.all(codeIndexInitializations) + try { + await codeIndexWorkspaceScopeRegistry.disposeAll() + } catch (error) { + const message = error instanceof Error ? error.message : String(error) + outputChannel.appendLine(`[CodeIndexManager] Error during workspace scope cleanup: ${message}`) + } + }, + }) + + // Initialize code index scopes for all workspace folders. if (vscode.workspace.workspaceFolders) { for (const folder of vscode.workspace.workspaceFolders) { - const manager = CodeIndexManagerRegistry.getOrCreate(context, folder.uri.fsPath) + const scope = codeIndexWorkspaceScopeRegistry.getScope(context, folder) - if (manager) { + if (scope) { // Initialize in background; do not block extension activation - void manager.initialize(contextProxy).catch((error) => { + const initialization = scope.initialize(contextProxy).catch((error) => { const message = error instanceof Error ? error.message : String(error) outputChannel.appendLine( `[CodeIndexManager] Error during background CodeIndexManager configuration/indexing for ${folder.uri.fsPath}: ${message}`, ) }) - - context.subscriptions.push(manager) + codeIndexInitializations.push(initialization) } } } @@ -408,5 +421,5 @@ export async function deactivate() { Terminal.setTerminalProfile(undefined) TerminalRegistry.cleanup() - CodeIndexManagerRegistry.disposeAll() + await codeIndexWorkspaceScopeRegistry.disposeAll() } diff --git a/src/services/code-index/__tests__/code-index-manager-registry.spec.ts b/src/services/code-index/__tests__/code-index-manager-registry.spec.ts deleted file mode 100644 index 9879ff8ea9..0000000000 --- a/src/services/code-index/__tests__/code-index-manager-registry.spec.ts +++ /dev/null @@ -1,127 +0,0 @@ -import * as vscode from "vscode" -import { makeExtensionContext, makeTextDocument, makeTextEditor, makeUri } from "../../../test-utils/vscode" -import { CodeIndexManager } from "../manager" -import { CodeIndexManagerRegistry } from "../code-index-manager-registry" - -vi.mock("vscode", () => ({ - workspace: { workspaceFolders: undefined, getWorkspaceFolder: vi.fn() }, - window: { activeTextEditor: undefined }, - Uri: { file: vi.fn() }, -})) - -vi.mock("../manager", () => ({ - CodeIndexManager: vi.fn().mockImplementation(function () { - return { dispose: vi.fn() } - }), -})) - -describe("CodeIndexManagerRegistry", () => { - let context: vscode.ExtensionContext - let first: vscode.WorkspaceFolder - let second: vscode.WorkspaceFolder - - beforeEach(() => { - vi.clearAllMocks() - context = makeExtensionContext() - first = { uri: makeUri("/first"), name: "first", index: 0 } - second = { uri: makeUri("/second"), name: "second", index: 1 } - Object.defineProperty(vscode.workspace, "workspaceFolders", { configurable: true, value: [first, second] }) - Object.defineProperty(vscode.window, "activeTextEditor", { configurable: true, value: undefined }) - vi.mocked(vscode.workspace.getWorkspaceFolder).mockReturnValue(undefined) - vi.mocked(vscode.Uri.file).mockImplementation((value) => makeUri(value)) - }) - - afterEach(() => { - CodeIndexManagerRegistry.disposeAll() - vi.restoreAllMocks() - }) - - it.each([{ folders: undefined }, { folders: [] }])("returns no manager with folders=$folders", ({ folders }) => { - Object.defineProperty(vscode.workspace, "workspaceFolders", { configurable: true, value: folders }) - expect(CodeIndexManagerRegistry.getOrCreate(context)).toBeUndefined() - expect(CodeIndexManager).not.toHaveBeenCalled() - }) - - it("uses the first workspace when there is no active editor", () => { - CodeIndexManagerRegistry.getOrCreate(context) - expect(CodeIndexManager).toHaveBeenCalledWith("/first", first.uri, context) - }) - - it("prefers the active editor's workspace", () => { - const editor = makeTextEditor({ document: makeTextDocument({ uri: makeUri("/second/file.ts") }) }) - Object.defineProperty(vscode.window, "activeTextEditor", { configurable: true, value: editor }) - vi.mocked(vscode.workspace.getWorkspaceFolder).mockReturnValue(second) - expect(CodeIndexManagerRegistry.getOrCreate(context)).toBeDefined() - expect(CodeIndexManager).toHaveBeenCalledWith("/second", second.uri, context) - }) - - it("falls back to the first workspace for an editor outside all folders", () => { - Object.defineProperty(vscode.window, "activeTextEditor", { configurable: true, value: makeTextEditor() }) - CodeIndexManagerRegistry.getOrCreate(context) - expect(CodeIndexManager).toHaveBeenCalledWith("/first", first.uri, context) - }) - - it("gives an explicit path priority over the active editor", () => { - Object.defineProperty(vscode.window, "activeTextEditor", { configurable: true, value: makeTextEditor() }) - vi.mocked(vscode.workspace.getWorkspaceFolder).mockReturnValue(first) - expect(CodeIndexManagerRegistry.getOrCreate(context, "/second")).toBeDefined() - expect(CodeIndexManager).toHaveBeenCalledWith("/second", second.uri, context) - }) - - it("preserves the actual remote workspace URI", () => { - const uri = makeUri("/remote", { scheme: "vscode-remote", authority: "ssh-remote+host" }) - Object.defineProperty(vscode.workspace, "workspaceFolders", { - configurable: true, - value: [{ uri, name: "remote", index: 0 }], - }) - CodeIndexManagerRegistry.getOrCreate(context, "/remote") - expect(CodeIndexManager).toHaveBeenCalledWith("/remote", uri, context) - expect(vi.mocked(CodeIndexManager).mock.calls[0][1]).toBe(uri) - expect(vscode.Uri.file).not.toHaveBeenCalled() - }) - - it("constructs a file URI for an explicit path without open workspaces", () => { - Object.defineProperty(vscode.workspace, "workspaceFolders", { configurable: true, value: undefined }) - const uri = makeUri("/outside folder/#name") - vi.mocked(vscode.Uri.file).mockReturnValue(uri) - CodeIndexManagerRegistry.getOrCreate(context, uri.fsPath) - expect(vscode.Uri.file).toHaveBeenCalledWith(uri.fsPath) - expect(CodeIndexManager).toHaveBeenCalledWith(uri.fsPath, uri, context) - }) - - it("constructs a file URI for an explicit path not matching any open workspace folder", () => { - // workspaceFolders contains /first and /second, but /outside/project matches neither - const uri = makeUri("/outside/project") - vi.mocked(vscode.Uri.file).mockReturnValue(uri) - CodeIndexManagerRegistry.getOrCreate(context, "/outside/project") - expect(vscode.Uri.file).toHaveBeenCalledWith("/outside/project") - expect(CodeIndexManager).toHaveBeenCalledWith("/outside/project", uri, context) - }) - - it("reuses the same path and keeps different paths isolated", () => { - const a = CodeIndexManagerRegistry.getOrCreate(context, "/first") - expect(CodeIndexManagerRegistry.getOrCreate(makeExtensionContext(), "/first")).toBe(a) - const b = CodeIndexManagerRegistry.getOrCreate(context, "/second") - expect(b).not.toBe(a) - expect(CodeIndexManager).toHaveBeenCalledTimes(2) - expect(CodeIndexManagerRegistry.getAllInstances()).toEqual([a, b]) - }) - - it("returns a snapshot that cannot mutate the cache", () => { - expect(CodeIndexManagerRegistry.getAllInstances()).toEqual([]) - const manager = CodeIndexManagerRegistry.getOrCreate(context) - CodeIndexManagerRegistry.getAllInstances().pop() - expect(CodeIndexManagerRegistry.getAllInstances()).toEqual([manager]) - }) - - it("disposes every manager, supports repeated cleanup and recreates instances", () => { - const a = CodeIndexManagerRegistry.getOrCreate(context, "/first")! - const b = CodeIndexManagerRegistry.getOrCreate(context, "/second")! - CodeIndexManagerRegistry.disposeAll() - CodeIndexManagerRegistry.disposeAll() - expect(a.dispose).toHaveBeenCalledTimes(1) - expect(b.dispose).toHaveBeenCalledTimes(1) - expect(CodeIndexManagerRegistry.getAllInstances()).toEqual([]) - expect(CodeIndexManagerRegistry.getOrCreate(context, "/first")).not.toBe(a) - }) -}) diff --git a/src/services/code-index/__tests__/code-index-workspace-scope-registry.spec.ts b/src/services/code-index/__tests__/code-index-workspace-scope-registry.spec.ts new file mode 100644 index 0000000000..fa17775d3b --- /dev/null +++ b/src/services/code-index/__tests__/code-index-workspace-scope-registry.spec.ts @@ -0,0 +1,227 @@ +import * as vscode from "vscode" +import { makeExtensionContext, makeTextDocument, makeTextEditor, makeUri } from "../../../test-utils/vscode" +import { CodeIndexManager } from "../manager" +import { codeIndexWorkspaceScopeRegistry } from "../code-index-workspace-scope-registry" + +vi.mock("vscode", () => ({ + workspace: { workspaceFolders: undefined, getWorkspaceFolder: vi.fn() }, + window: { activeTextEditor: undefined }, + Uri: { file: vi.fn() }, +})) + +vi.mock("../manager", () => ({ + CodeIndexManager: vi.fn().mockImplementation(function () { + return { initialize: vi.fn().mockResolvedValue({ requiresRestart: false }), dispose: vi.fn() } + }), +})) + +describe("CodeIndexWorkspaceScopeRegistry", () => { + let context: vscode.ExtensionContext + let first: vscode.WorkspaceFolder + let second: vscode.WorkspaceFolder + + beforeEach(() => { + vi.clearAllMocks() + context = makeExtensionContext() + first = { uri: makeUri("/first"), name: "first", index: 0 } + second = { uri: makeUri("/second"), name: "second", index: 1 } + Object.defineProperty(vscode.workspace, "workspaceFolders", { configurable: true, value: [first, second] }) + Object.defineProperty(vscode.window, "activeTextEditor", { configurable: true, value: undefined }) + vi.mocked(vscode.workspace.getWorkspaceFolder).mockReturnValue(undefined) + vi.mocked(vscode.Uri.file).mockImplementation((value) => makeUri(value)) + }) + + afterEach(async () => { + await codeIndexWorkspaceScopeRegistry.disposeAll() + vi.restoreAllMocks() + }) + + it.each([{ folders: undefined }, { folders: [] }])("returns no scope with folders=$folders", ({ folders }) => { + Object.defineProperty(vscode.workspace, "workspaceFolders", { configurable: true, value: folders }) + expect(codeIndexWorkspaceScopeRegistry.getScope(context)).toBeUndefined() + expect(CodeIndexManager).not.toHaveBeenCalled() + }) + + it("uses the first workspace when there is no active editor", () => { + codeIndexWorkspaceScopeRegistry.getScope(context) + expect(CodeIndexManager).toHaveBeenCalledWith("/first", first.uri, context) + }) + + it("prefers the active editor's workspace", () => { + const editor = makeTextEditor({ document: makeTextDocument({ uri: makeUri("/second/file.ts") }) }) + Object.defineProperty(vscode.window, "activeTextEditor", { configurable: true, value: editor }) + vi.mocked(vscode.workspace.getWorkspaceFolder).mockReturnValue(second) + codeIndexWorkspaceScopeRegistry.getScope(context) + expect(vscode.workspace.getWorkspaceFolder).toHaveBeenCalledWith(editor.document.uri) + expect(CodeIndexManager).toHaveBeenCalledWith("/second", second.uri, context) + }) + + it("falls back to the first workspace for an editor outside all folders", () => { + Object.defineProperty(vscode.window, "activeTextEditor", { configurable: true, value: makeTextEditor() }) + codeIndexWorkspaceScopeRegistry.getScope(context) + expect(CodeIndexManager).toHaveBeenCalledWith("/first", first.uri, context) + }) + + it("gives an explicit path priority over the active editor", () => { + Object.defineProperty(vscode.window, "activeTextEditor", { configurable: true, value: makeTextEditor() }) + vi.mocked(vscode.workspace.getWorkspaceFolder).mockReturnValue(first) + codeIndexWorkspaceScopeRegistry.getScope(context, "/second") + expect(CodeIndexManager).toHaveBeenCalledWith("/second", second.uri, context) + expect(vscode.workspace.getWorkspaceFolder).not.toHaveBeenCalled() + }) + + it("preserves the actual remote workspace URI", () => { + const uri = makeUri("/remote", { scheme: "vscode-remote", authority: "ssh-remote+host" }) + Object.defineProperty(vscode.workspace, "workspaceFolders", { + configurable: true, + value: [{ uri, name: "remote", index: 0 }], + }) + codeIndexWorkspaceScopeRegistry.getScope(context, "/remote") + expect(CodeIndexManager).toHaveBeenCalledWith("/remote", uri, context) + expect(vi.mocked(CodeIndexManager).mock.calls[0][1]).toBe(uri) + expect(vscode.Uri.file).not.toHaveBeenCalled() + }) + + it("keeps equal fs paths from different remote authorities isolated", () => { + const firstRemote = makeUri("/workspace", { scheme: "vscode-remote", authority: "ssh-remote+first" }) + const secondRemote = makeUri("/workspace", { scheme: "vscode-remote", authority: "ssh-remote+second" }) + vi.mocked(firstRemote.toString).mockReturnValue("vscode-remote://ssh-remote+first/workspace") + vi.mocked(secondRemote.toString).mockReturnValue("vscode-remote://ssh-remote+second/workspace") + + const firstScope = codeIndexWorkspaceScopeRegistry.getScope(context, firstRemote) + const secondScope = codeIndexWorkspaceScopeRegistry.getScope(context, secondRemote) + + expect(firstRemote.toString).toHaveBeenCalledWith(true) + expect(secondRemote.toString).toHaveBeenCalledWith(true) + expect(secondScope).not.toBe(firstScope) + expect(CodeIndexManager).toHaveBeenNthCalledWith(1, "/workspace", firstRemote, context) + expect(CodeIndexManager).toHaveBeenNthCalledWith(2, "/workspace", secondRemote, context) + }) + + it("constructs a file URI for an explicit path without open workspaces", () => { + Object.defineProperty(vscode.workspace, "workspaceFolders", { configurable: true, value: undefined }) + const uri = makeUri("/outside folder/#name") + vi.mocked(vscode.Uri.file).mockReturnValue(uri) + codeIndexWorkspaceScopeRegistry.getScope(context, uri.fsPath) + expect(vscode.Uri.file).toHaveBeenCalledWith(uri.fsPath) + expect(CodeIndexManager).toHaveBeenCalledWith(uri.fsPath, uri, context) + }) + + it("reuses the same path and keeps different paths isolated", () => { + const a = codeIndexWorkspaceScopeRegistry.getScope(context, "/first") + expect(codeIndexWorkspaceScopeRegistry.getScope(makeExtensionContext(), "/first")).toBe(a) + const b = codeIndexWorkspaceScopeRegistry.getScope(context, "/second") + expect(b).not.toBe(a) + expect(CodeIndexManager).toHaveBeenCalledTimes(2) + expect(codeIndexWorkspaceScopeRegistry.getAllScopes()).toEqual([a, b]) + }) + + it("returns a snapshot that cannot mutate the cache", () => { + expect(codeIndexWorkspaceScopeRegistry.getAllScopes()).toEqual([]) + const scope = codeIndexWorkspaceScopeRegistry.getScope(context) + codeIndexWorkspaceScopeRegistry.getAllScopes().pop() + expect(codeIndexWorkspaceScopeRegistry.getAllScopes()).toEqual([scope]) + }) + + it("disposes every scope, supports repeated cleanup and recreates scopes", async () => { + const a = codeIndexWorkspaceScopeRegistry.getScope(context, "/first")! + const b = codeIndexWorkspaceScopeRegistry.getScope(context, "/second")! + await codeIndexWorkspaceScopeRegistry.disposeAll() + await codeIndexWorkspaceScopeRegistry.disposeAll() + expect(a.codeIndexManager.dispose).toHaveBeenCalledTimes(1) + expect(b.codeIndexManager.dispose).toHaveBeenCalledTimes(1) + expect(codeIndexWorkspaceScopeRegistry.getAllScopes()).toEqual([]) + expect(codeIndexWorkspaceScopeRegistry.getScope(context, "/first")).not.toBe(a) + }) + + it("waits for lazy scope initialization before disposing its manager", async () => { + const scope = codeIndexWorkspaceScopeRegistry.getScope(context, "/lazy")! + let resolveInitialization: ((result: { requiresRestart: boolean }) => void) | undefined + vi.mocked(scope.codeIndexManager.initialize).mockImplementationOnce( + () => + new Promise((resolve) => { + resolveInitialization = resolve + }), + ) + void scope.initialize({} as never) + + const firstDisposal = codeIndexWorkspaceScopeRegistry.disposeAll() + const concurrentDisposal = codeIndexWorkspaceScopeRegistry.disposeAll() + expect(concurrentDisposal).toBe(firstDisposal) + expect(scope.codeIndexManager.dispose).not.toHaveBeenCalled() + + resolveInitialization?.({ requiresRestart: false }) + await firstDisposal + expect(scope.codeIndexManager.dispose).toHaveBeenCalledExactlyOnceWith() + }) + + it("attempts every scope and preserves all thrown values in an aggregate", async () => { + const a = codeIndexWorkspaceScopeRegistry.getScope(context, "/first")! + const b = codeIndexWorkspaceScopeRegistry.getScope(context, "/second")! + const c = codeIndexWorkspaceScopeRegistry.getScope(context, "/third")! + const error = new Error("first cleanup failed") + vi.mocked(a.codeIndexManager.dispose).mockImplementationOnce(() => { + throw error + }) + vi.mocked(b.codeIndexManager.dispose).mockImplementationOnce(() => { + throw "second cleanup failed" + }) + + let caught: unknown + try { + await codeIndexWorkspaceScopeRegistry.disposeAll() + } catch (error) { + caught = error + } + expect(caught).toBeInstanceOf(AggregateError) + if (!(caught instanceof AggregateError)) throw new Error("Expected aggregate disposal failure") + expect(caught.errors).toEqual([error, "second cleanup failed"]) + for (const scope of [a, b, c]) { + expect(scope.codeIndexManager.dispose).toHaveBeenCalledExactlyOnceWith() + } + expect(codeIndexWorkspaceScopeRegistry.getAllScopes()).toEqual([]) + await codeIndexWorkspaceScopeRegistry.disposeAll() + expect(a.codeIndexManager.dispose).toHaveBeenCalledTimes(1) + const replacement = codeIndexWorkspaceScopeRegistry.getScope(context, "/first") + expect(replacement).not.toBe(a) + expect(codeIndexWorkspaceScopeRegistry.getAllScopes()).toEqual([replacement]) + }) + + it.each([false, true])("blocks reentrant lookup and cleanup, then resets (failure=%s)", async (fails) => { + const a = codeIndexWorkspaceScopeRegistry.getScope(context, "/first")! + const b = codeIndexWorkspaceScopeRegistry.getScope(context, "/second")! + vi.mocked(a.codeIndexManager.dispose).mockImplementationOnce(() => { + expect(codeIndexWorkspaceScopeRegistry.getAllScopes()).toEqual([]) + expect(codeIndexWorkspaceScopeRegistry.getScope(context)).toBeUndefined() + expect(codeIndexWorkspaceScopeRegistry.getScope(context, "/first")).toBeUndefined() + expect(codeIndexWorkspaceScopeRegistry.getScope(context, "/new")).toBeUndefined() + void codeIndexWorkspaceScopeRegistry.disposeAll() + expect(b.codeIndexManager.dispose).not.toHaveBeenCalled() + if (fails) throw new Error("cleanup failed") + }) + + if (fails) { + await expect(codeIndexWorkspaceScopeRegistry.disposeAll()).rejects.toBeInstanceOf(AggregateError) + } else { + await codeIndexWorkspaceScopeRegistry.disposeAll() + } + expect(a.codeIndexManager.dispose).toHaveBeenCalledTimes(1) + expect(b.codeIndexManager.dispose).toHaveBeenCalledTimes(1) + expect(CodeIndexManager).toHaveBeenCalledTimes(2) + expect(codeIndexWorkspaceScopeRegistry.getAllScopes()).toEqual([]) + expect(codeIndexWorkspaceScopeRegistry.getScope(context, "/first")).not.toBe(a) + }) + + it("cleans up its own snapshot even when a caller mutates a previously returned list", async () => { + const a = codeIndexWorkspaceScopeRegistry.getScope(context, "/first")! + const b = codeIndexWorkspaceScopeRegistry.getScope(context, "/second")! + const snapshot = codeIndexWorkspaceScopeRegistry.getAllScopes() + vi.mocked(a.codeIndexManager.dispose).mockImplementationOnce(() => { + snapshot.splice(0, snapshot.length) + }) + + await codeIndexWorkspaceScopeRegistry.disposeAll() + expect(b.codeIndexManager.dispose).toHaveBeenCalledExactlyOnceWith() + expect(snapshot).toEqual([]) + }) +}) diff --git a/src/services/code-index/__tests__/code-index-workspace-scope.spec.ts b/src/services/code-index/__tests__/code-index-workspace-scope.spec.ts new file mode 100644 index 0000000000..4f929eb10e --- /dev/null +++ b/src/services/code-index/__tests__/code-index-workspace-scope.spec.ts @@ -0,0 +1,105 @@ +import { ContextProxy } from "../../../core/config/ContextProxy" +import { makeExtensionContext, makeUri } from "../../../test-utils/vscode" +import { CodeIndexManager } from "../manager" +import { CodeIndexWorkspaceScope } from "../code-index-workspace-scope" + +vi.mock("../manager", () => ({ + CodeIndexManager: vi.fn().mockImplementation(function () { + return { initialize: vi.fn().mockResolvedValue({ requiresRestart: false }), dispose: vi.fn() } + }), +})) + +describe("CodeIndexWorkspaceScope", () => { + beforeEach(() => vi.clearAllMocks()) + + it("owns, initializes and disposes its manager", async () => { + const context = makeExtensionContext() + const uri = makeUri("/workspace") + const contextProxy = {} as ContextProxy + const scope = new CodeIndexWorkspaceScope(uri.fsPath, uri, context) + + expect(CodeIndexManager).toHaveBeenCalledExactlyOnceWith(uri.fsPath, uri, context) + await expect(scope.initialize(contextProxy)).resolves.toEqual({ requiresRestart: false }) + expect(scope.codeIndexManager.initialize).toHaveBeenCalledExactlyOnceWith(contextProxy) + + await scope.dispose() + expect(scope.codeIndexManager.dispose).toHaveBeenCalledExactlyOnceWith() + }) + + it.each([true, false])("propagates requiresRestart=%s unchanged", async (requiresRestart) => { + const uri = makeUri("/workspace") + const scope = new CodeIndexWorkspaceScope(uri.fsPath, uri, makeExtensionContext()) + const result = { requiresRestart } + vi.mocked(scope.codeIndexManager.initialize).mockResolvedValueOnce(result) + + await expect(scope.initialize({} as ContextProxy)).resolves.toBe(result) + }) + + it("shares concurrent initialization and permits a later reload", async () => { + const scope = new CodeIndexWorkspaceScope("/workspace", makeUri("/workspace"), makeExtensionContext()) + const contextProxy = {} as ContextProxy + let resolveInitialization: ((result: { requiresRestart: boolean }) => void) | undefined + vi.mocked(scope.codeIndexManager.initialize).mockImplementationOnce( + () => + new Promise((resolve) => { + resolveInitialization = resolve + }), + ) + + const first = scope.initialize(contextProxy) + const concurrent = scope.initialize(contextProxy) + expect(concurrent).toBe(first) + expect(scope.codeIndexManager.initialize).toHaveBeenCalledTimes(1) + + resolveInitialization?.({ requiresRestart: false }) + await first + await scope.initialize(contextProxy) + expect(scope.codeIndexManager.initialize).toHaveBeenCalledTimes(2) + }) + + it("propagates initialization rejection without taking over consumer cleanup", async () => { + const uri = makeUri("/workspace") + const scope = new CodeIndexWorkspaceScope(uri.fsPath, uri, makeExtensionContext()) + const error = new Error("initialization failed") + vi.mocked(scope.codeIndexManager.initialize).mockRejectedValueOnce(error) + + await expect(scope.initialize({} as ContextProxy)).rejects.toBe(error) + expect(scope.codeIndexManager.dispose).not.toHaveBeenCalled() + await scope.dispose() + expect(scope.codeIndexManager.dispose).toHaveBeenCalledExactlyOnceWith() + }) + + it("waits for initialization, blocks new initialization, and disposes exactly once", async () => { + const scope = new CodeIndexWorkspaceScope("/workspace", makeUri("/workspace"), makeExtensionContext()) + const contextProxy = {} as ContextProxy + let resolveInitialization: ((result: { requiresRestart: boolean }) => void) | undefined + vi.mocked(scope.codeIndexManager.initialize).mockImplementationOnce( + () => + new Promise((resolve) => { + resolveInitialization = resolve + }), + ) + + void scope.initialize(contextProxy) + const disposal = scope.dispose() + expect(scope.codeIndexManager.dispose).not.toHaveBeenCalled() + await expect(scope.initialize(contextProxy)).rejects.toThrow("Cannot initialize a disposed") + + resolveInitialization?.({ requiresRestart: false }) + await expect(disposal).resolves.toBeUndefined() + await expect(scope.dispose()).resolves.toBeUndefined() + expect(scope.codeIndexManager.dispose).toHaveBeenCalledTimes(1) + }) + + it("propagates disposal errors to its owner", async () => { + const uri = makeUri("/workspace") + const scope = new CodeIndexWorkspaceScope(uri.fsPath, uri, makeExtensionContext()) + const error = new Error("disposal failed") + vi.mocked(scope.codeIndexManager.dispose).mockImplementationOnce(() => { + throw error + }) + + await expect(scope.dispose()).rejects.toBe(error) + expect(scope.codeIndexManager.dispose).toHaveBeenCalledExactlyOnceWith() + }) +}) diff --git a/src/services/code-index/__tests__/manager-lifecycle.spec.ts b/src/services/code-index/__tests__/manager-lifecycle.spec.ts new file mode 100644 index 0000000000..d545f6921e --- /dev/null +++ b/src/services/code-index/__tests__/manager-lifecycle.spec.ts @@ -0,0 +1,202 @@ +import type { ContextProxy } from "../../../core/config/ContextProxy" +import { makeExtensionContext, makeUri } from "../../../test-utils/vscode" +import { CodeIndexManager } from "../manager" +import { SembleProvider } from "../semble" + +const mocks = vi.hoisted(() => ({ + loadConfiguration: vi.fn<() => Promise<{ requiresRestart: boolean }>>(), + initializeCache: vi.fn<() => Promise>(), + initializeProvider: vi.fn<() => Promise>(), + startIndexing: vi.fn<() => Promise>(), + stopIndexing: vi.fn(), + disposeProvider: vi.fn(), + disposeState: vi.fn(), + setSystemState: vi.fn(), +})) + +vi.mock("../config-manager", () => ({ + CodeIndexConfigManager: vi.fn().mockImplementation(function () { + return { + loadConfiguration: mocks.loadConfiguration, + isFeatureEnabled: true, + isFeatureConfigured: true, + currentEmbedderProvider: "semble", + } + }), +})) +vi.mock("../cache-manager", () => ({ + CacheManager: vi.fn().mockImplementation(function () { + return { initialize: mocks.initializeCache } + }), +})) +vi.mock("../state-manager", () => ({ + CodeIndexStateManager: vi.fn().mockImplementation(function () { + return { dispose: mocks.disposeState, setSystemState: mocks.setSystemState } + }), +})) +vi.mock("../semble", () => ({ + SembleProvider: vi.fn().mockImplementation(function () { + return { + initialize: mocks.initializeProvider, + startIndexing: mocks.startIndexing, + stopIndexing: mocks.stopIndexing, + dispose: mocks.disposeProvider, + } + }), +})) +vi.mock("../service-factory") +vi.mock("../search-service") +vi.mock("../orchestrator") +vi.mock("../../../core/ignore/RooIgnoreController") + +function deferred() { + let resolve!: (value: T | PromiseLike) => void + const promise = new Promise((resolvePromise) => { + resolve = resolvePromise + }) + return { promise, resolve } +} + +describe("CodeIndexManager consumer-owned lifecycle", () => { + let manager: CodeIndexManager + // Configuration is mocked, so this dependency is only passed through. + const contextProxy = {} as ContextProxy + + beforeEach(() => { + vi.clearAllMocks() + mocks.loadConfiguration.mockReset() + mocks.initializeCache.mockReset() + mocks.initializeProvider.mockReset() + mocks.loadConfiguration.mockResolvedValue({ requiresRestart: false }) + mocks.initializeCache.mockResolvedValue(undefined) + mocks.initializeProvider.mockResolvedValue(undefined) + mocks.startIndexing.mockResolvedValue(undefined) + const uri = makeUri("/workspace") + manager = new CodeIndexManager(uri.fsPath, uri, makeExtensionContext()) + vi.spyOn(manager, "isWorkspaceEnabled", "get").mockReturnValue(true) + }) + + it("starts resources before the consumer performs its single disposal", async () => { + await expect(manager.initialize(contextProxy)).resolves.toEqual({ requiresRestart: false }) + expect(mocks.initializeCache).toHaveBeenCalledOnce() + expect(mocks.initializeProvider).toHaveBeenCalledOnce() + expect(mocks.startIndexing).toHaveBeenCalledOnce() + manager.dispose() + expect(mocks.stopIndexing).toHaveBeenCalledOnce() + expect(mocks.disposeProvider).toHaveBeenCalledOnce() + expect(mocks.disposeState).toHaveBeenCalledOnce() + }) + + it("supports intentional sequential configuration reload without recreating unchanged services", async () => { + await manager.initialize(contextProxy) + await expect(manager.initialize(contextProxy)).resolves.toEqual({ requiresRestart: false }) + expect(mocks.loadConfiguration).toHaveBeenCalledTimes(2) + expect(SembleProvider).toHaveBeenCalledOnce() + expect(mocks.startIndexing).toHaveBeenCalledOnce() + manager.dispose() + }) + + it("recreates services for a sequential restart request", async () => { + await manager.initialize(contextProxy) + mocks.loadConfiguration.mockResolvedValueOnce({ requiresRestart: true }) + await expect(manager.initialize(contextProxy)).resolves.toEqual({ requiresRestart: true }) + expect(SembleProvider).toHaveBeenCalledTimes(2) + expect(mocks.disposeProvider).toHaveBeenCalledOnce() + expect(mocks.startIndexing).toHaveBeenCalledTimes(2) + manager.dispose() + }) + + it("allows initialization after explicit error recovery", async () => { + const error = new Error("configuration unavailable") + mocks.loadConfiguration.mockRejectedValueOnce(error) + await expect(manager.initialize(contextProxy)).rejects.toBe(error) + await manager.recoverFromError() + await expect(manager.initialize(contextProxy)).resolves.toEqual({ requiresRestart: false }) + expect(mocks.startIndexing).toHaveBeenCalledOnce() + manager.dispose() + }) + + it.each(["configuration", "cache", "provider"] as const)( + "propagates %s initialization rejection", + async (stage) => { + const error = new Error(`${stage} unavailable`) + const operation = { + configuration: mocks.loadConfiguration, + cache: mocks.initializeCache, + provider: mocks.initializeProvider, + }[stage] + operation.mockRejectedValueOnce(error) + await expect(manager.initialize(contextProxy)).rejects.toBe(error) + expect(mocks.startIndexing).not.toHaveBeenCalled() + manager.dispose() + }, + ) + + // Characterize unsupported ordering, not desired safety guarantees. Consumers must + // await initialization before disposal and must not initialize concurrently. + it.each(["configuration", "cache"] as const)( + "documents resources starting after disposal during %s initialization", + async (stage) => { + const entered = deferred() + const release = deferred() + if (stage === "configuration") { + mocks.loadConfiguration.mockImplementationOnce(async () => { + entered.resolve() + await release.promise + return { requiresRestart: false } + }) + } else { + mocks.initializeCache.mockImplementationOnce(() => { + entered.resolve() + return release.promise + }) + } + const initialization = manager.initialize(contextProxy) + await entered.promise + manager.dispose() + release.resolve() + await initialization + expect(mocks.disposeState).toHaveBeenCalledOnce() + expect(mocks.startIndexing).toHaveBeenCalledOnce() + expect(mocks.disposeState.mock.invocationCallOrder[0]).toBeLessThan( + mocks.startIndexing.mock.invocationCallOrder[0], + ) + expect(mocks.disposeProvider).not.toHaveBeenCalled() + }, + ) + + it("documents initialization reporting success after disposal clears a pending provider", async () => { + const entered = deferred() + const release = deferred() + mocks.initializeProvider.mockImplementationOnce(() => { + entered.resolve() + return release.promise + }) + const initialization = manager.initialize(contextProxy) + await entered.promise + manager.dispose() + release.resolve() + await expect(initialization).resolves.toEqual({ requiresRestart: false }) + expect(manager.isInitialized).toBe(false) + expect(mocks.disposeProvider).toHaveBeenCalledOnce() + expect(mocks.disposeState).toHaveBeenCalledOnce() + expect(mocks.startIndexing).not.toHaveBeenCalled() + }) + + it("documents concurrent initialization starting indexing before shared cache readiness", async () => { + const entered = deferred() + const release = deferred() + mocks.initializeCache.mockImplementationOnce(() => { + entered.resolve() + return release.promise + }) + const first = manager.initialize(contextProxy) + await entered.promise + await manager.initialize(contextProxy) + const startedBeforeCacheReady = mocks.startIndexing.mock.calls.length + release.resolve() + await first + manager.dispose() + expect(startedBeforeCacheReady).toBe(1) + }) +}) diff --git a/src/services/code-index/__tests__/manager.spec.ts b/src/services/code-index/__tests__/manager.spec.ts index 9faf06627e..a5af3a09fa 100644 --- a/src/services/code-index/__tests__/manager.spec.ts +++ b/src/services/code-index/__tests__/manager.spec.ts @@ -1,5 +1,5 @@ import { CodeIndexManager } from "../manager" -import { CodeIndexManagerRegistry } from "../code-index-manager-registry" +import { codeIndexWorkspaceScopeRegistry } from "../code-index-workspace-scope-registry" import { CodeIndexServiceFactory } from "../service-factory" import type { MockedClass } from "vitest" import * as path from "path" @@ -125,9 +125,9 @@ describe("CodeIndexManager - handleSettingsChange regression", () => { const testGlobalStoragePath = path.join(path.sep, "test", "global-storage") const testLogPath = path.join(path.sep, "test", "log") - beforeEach(() => { + beforeEach(async () => { // Clear all instances before each test - CodeIndexManagerRegistry.disposeAll() + await codeIndexWorkspaceScopeRegistry.disposeAll() const workspaceStateStore: Record = {} const globalStateStore: Record = {} @@ -161,11 +161,11 @@ describe("CodeIndexManager - handleSettingsChange regression", () => { languageModelAccessInformation: {} as any, } - manager = CodeIndexManagerRegistry.getOrCreate(mockContext)! + manager = codeIndexWorkspaceScopeRegistry.getScope(mockContext)!.codeIndexManager }) - afterEach(() => { - CodeIndexManagerRegistry.disposeAll() + afterEach(async () => { + await codeIndexWorkspaceScopeRegistry.disposeAll() }) describe("handleSettingsChange", () => { @@ -734,7 +734,7 @@ describe("CodeIndexManager - handleSettingsChange regression", () => { }) it("should store enablement per folder URI, not per window", async () => { - CodeIndexManagerRegistry.disposeAll() + await codeIndexWorkspaceScopeRegistry.disposeAll() const vscode = await import("vscode") @@ -765,8 +765,8 @@ describe("CodeIndexManager - handleSettingsChange regression", () => { { uri: folderBUri, name: "folderB", index: 1 }, ] - const managerA = CodeIndexManagerRegistry.getOrCreate(sharedContext, folderAPath)! - const managerB = CodeIndexManagerRegistry.getOrCreate(sharedContext, folderBPath)! + const managerA = codeIndexWorkspaceScopeRegistry.getScope(sharedContext, folderAPath)!.codeIndexManager + const managerB = codeIndexWorkspaceScopeRegistry.getScope(sharedContext, folderBPath)!.codeIndexManager // Both start disabled (autoEnableDefault is false via globalState mock) expect(managerA.isWorkspaceEnabled).toBe(false) @@ -785,7 +785,7 @@ describe("CodeIndexManager - handleSettingsChange regression", () => { expect(managerA.isWorkspaceEnabled).toBe(false) expect(managerB.isWorkspaceEnabled).toBe(true) - CodeIndexManagerRegistry.disposeAll() + await codeIndexWorkspaceScopeRegistry.disposeAll() }) }) diff --git a/src/services/code-index/code-index-manager-registry.ts b/src/services/code-index/code-index-manager-registry.ts deleted file mode 100644 index 635ec62647..0000000000 --- a/src/services/code-index/code-index-manager-registry.ts +++ /dev/null @@ -1,53 +0,0 @@ -import * as vscode from "vscode" -import { CodeIndexManager } from "./manager" - -/** Resolves workspaces and owns their cached CodeIndexManager instances. */ -export class CodeIndexManagerRegistry { - private static instances = new Map() - - public static getOrCreate(context: vscode.ExtensionContext, workspacePath?: string): CodeIndexManager | undefined { - const folder = this.resolveWorkspaceFolder(workspacePath) - const resolvedPath = workspacePath || folder?.uri.fsPath - if (!resolvedPath) { - return undefined - } - - const existing = this.instances.get(resolvedPath) - if (existing) { - return existing - } - - // Preserve real workspace URIs, including remote schemes and authorities. - const folderUri = folder?.uri ?? vscode.Uri.file(resolvedPath) - const manager = new CodeIndexManager(resolvedPath, folderUri, context) - this.instances.set(resolvedPath, manager) - return manager - } - - public static getAllInstances(): CodeIndexManager[] { - return Array.from(this.instances.values()) - } - - public static disposeAll(): void { - for (const instance of this.instances.values()) { - instance.dispose() - } - this.instances.clear() - } - - private static resolveWorkspaceFolder(workspacePath?: string): vscode.WorkspaceFolder | undefined { - if (workspacePath) { - return vscode.workspace.workspaceFolders?.find((folder) => folder.uri.fsPath === workspacePath) - } - - const activeEditor = vscode.window.activeTextEditor - if (activeEditor) { - const folder = vscode.workspace.getWorkspaceFolder(activeEditor.document.uri) - if (folder) { - return folder - } - } - - return vscode.workspace.workspaceFolders?.[0] - } -} diff --git a/src/services/code-index/code-index-workspace-scope-registry.ts b/src/services/code-index/code-index-workspace-scope-registry.ts new file mode 100644 index 0000000000..f04015c667 --- /dev/null +++ b/src/services/code-index/code-index-workspace-scope-registry.ts @@ -0,0 +1,96 @@ +import * as vscode from "vscode" + +import { CodeIndexWorkspaceScope } from "./code-index-workspace-scope" + +/** Resolves workspaces and owns their cached code-index scopes. */ +export class CodeIndexWorkspaceScopeRegistry { + public static readonly instance = new CodeIndexWorkspaceScopeRegistry() + + private readonly scopes = new Map() + private disposal?: Promise + + private constructor() {} + + public getScope( + context: vscode.ExtensionContext, + workspace?: string | vscode.Uri | vscode.WorkspaceFolder, + ): CodeIndexWorkspaceScope | undefined { + if (this.disposal) { + return undefined + } + const folder = this.resolveWorkspaceFolder(typeof workspace === "string" ? workspace : undefined) + const folderUri = + typeof workspace === "string" + ? (folder?.uri ?? vscode.Uri.file(workspace)) + : workspace === undefined + ? folder?.uri + : "uri" in workspace + ? workspace.uri + : workspace + const resolvedPath = typeof workspace === "string" ? workspace : folderUri?.fsPath + if (!resolvedPath || !folderUri) { + return undefined + } + + const scopeKey = folderUri.toString(true) + const existing = this.scopes.get(scopeKey) + if (existing) { + return existing + } + + const scope = new CodeIndexWorkspaceScope(resolvedPath, folderUri, context) + this.scopes.set(scopeKey, scope) + return scope + } + + public getAllScopes(): CodeIndexWorkspaceScope[] { + return Array.from(this.scopes.values()) + } + + public disposeAll(): Promise { + if (this.disposal) { + return this.disposal + } + + const scopes = this.getAllScopes() + this.scopes.clear() + const disposal = (async () => { + const errors: unknown[] = [] + for (const scope of scopes) { + try { + await scope.dispose() + } catch (error) { + errors.push(error) + } + } + if (errors.length > 0) { + throw new AggregateError(errors, "Failed to dispose code index workspace scopes") + } + })().finally(() => { + if (this.disposal === disposal) { + this.disposal = undefined + } + }) + this.disposal = disposal + return disposal + } + + private resolveWorkspaceFolder(workspacePath?: string): vscode.WorkspaceFolder | undefined { + if (workspacePath) { + return vscode.workspace.workspaceFolders?.find((folder) => folder.uri.fsPath === workspacePath) + } + + const activeEditor = vscode.window.activeTextEditor + if (activeEditor) { + const folder = vscode.workspace.getWorkspaceFolder(activeEditor.document.uri) + if (folder) { + return folder + } + } + + return vscode.workspace.workspaceFolders?.[0] + } +} + +/** Shared workspace scope registry used by the extension runtime. */ +export const codeIndexWorkspaceScopeRegistry = CodeIndexWorkspaceScopeRegistry.instance diff --git a/src/services/code-index/code-index-workspace-scope.ts b/src/services/code-index/code-index-workspace-scope.ts new file mode 100644 index 0000000000..bfccb2f3eb --- /dev/null +++ b/src/services/code-index/code-index-workspace-scope.ts @@ -0,0 +1,53 @@ +import * as vscode from "vscode" + +import { ContextProxy } from "../../core/config/ContextProxy" +import { CodeIndexManager } from "./manager" + +/** + * Owns code-index services whose lifetime is bound to one workspace. + * The consumer owns initialization ordering and single disposal. Registry-owned + * scopes must be disposed through the registry, not independently by borrowers. + */ +export class CodeIndexWorkspaceScope implements vscode.Disposable { + public readonly codeIndexManager: CodeIndexManager + private initialization?: Promise<{ requiresRestart: boolean }> + private disposal?: Promise + + public constructor(workspacePath: string, folderUri: vscode.Uri, context: vscode.ExtensionContext) { + this.codeIndexManager = new CodeIndexManager(workspacePath, folderUri, context) + } + + public initialize(contextProxy: ContextProxy): Promise<{ requiresRestart: boolean }> { + if (this.disposal) { + return Promise.reject(new Error("Cannot initialize a disposed code index workspace scope")) + } + if (this.initialization) { + return this.initialization + } + + const initialization = this.codeIndexManager.initialize(contextProxy).finally(() => { + if (this.initialization === initialization) { + this.initialization = undefined + } + }) + this.initialization = initialization + return initialization + } + + public dispose(): Promise { + if (this.disposal) { + return this.disposal + } + + const initialization = this.initialization + this.disposal = (async () => { + try { + await initialization + } catch { + // Initialization failures do not release ownership; the manager still needs disposal. + } + this.codeIndexManager.dispose() + })() + return this.disposal + } +} diff --git a/src/services/code-index/manager.ts b/src/services/code-index/manager.ts index fd3e6b0553..31e5fbbfe7 100644 --- a/src/services/code-index/manager.ts +++ b/src/services/code-index/manager.ts @@ -34,7 +34,7 @@ export class CodeIndexManager { private readonly _folderUri: vscode.Uri private readonly context: vscode.ExtensionContext - /** @internal — construct only via {@link CodeIndexManagerRegistry} */ + /** @internal — construct only via CodeIndexWorkspaceScope */ public constructor(workspacePath: string, folderUri: vscode.Uri, context: vscode.ExtensionContext) { this.workspacePath = workspacePath this._folderUri = folderUri