fix: URI-keyed enablement, throw AbortError in back-pressure, stopWatcher on early-return

Author: James Mtendamema

src/__mocks__/vscode.js

@@ -10,8 +10,8 @@ const mockDisposable = {
 }
 
 const mockUri = {
-	file: (path) => ({ fsPath: path, path, scheme: "file" }),
-	parse: (path) => ({ fsPath: path, path, scheme: "file" }),
+	file: (path) => ({ fsPath: path, path, scheme: "file", toString: () => `file://${path}` }),
+	parse: (path) => ({ fsPath: path, path, scheme: "file", toString: () => path }),
 }
 
 const mockRange = class {

src/services/code-index/__tests__/manager.spec.ts

@@ -3,18 +3,45 @@ import { CodeIndexServiceFactory } from "../service-factory"
 import type { MockedClass } from "vitest"
 import * as path from "path"
 
+// Helper: create a mock vscode.Uri from an fsPath
+function mockUri(fsPath: string, scheme = "file") {
+	return {
+		fsPath,
+		scheme,
+		authority: "",
+		path: fsPath,
+		toString: (skipEncoding?: boolean) => `${scheme}://${fsPath}`,
+	}
+}
+
 // Mock vscode module
 vi.mock("vscode", () => {
 	const testPath = require("path")
 	const testWorkspacePath = testPath.join(testPath.sep, "test", "workspace")
 	return {
+		Uri: {
+			file: (p: string) => ({
+				fsPath: p,
+				scheme: "file",
+				authority: "",
+				path: p,
+				toString: (_skipEncoding?: boolean) => `file://${p}`,
+			}),
+			joinPath: vi.fn((...args: any[]) => ({ fsPath: args.join("/") })),
+		},
 		window: {
 			activeTextEditor: null,
 		},
 		workspace: {
 			workspaceFolders: [
 				{
-					uri: { fsPath: testWorkspacePath },
+					uri: {
+						fsPath: testWorkspacePath,
+						scheme: "file",
+						authority: "",
+						path: testWorkspacePath,
+						toString: (_skipEncoding?: boolean) => `file://${testWorkspacePath}`,
+					},
 					name: "test",
 					index: 0,
 				},
@@ -25,8 +52,9 @@ vi.mock("vscode", () => {
 				onDidDelete: vi.fn().mockReturnValue({ dispose: vi.fn() }),
 				dispose: vi.fn(),
 			}),
+			getWorkspaceFolder: vi.fn(),
 		},
-		RelativePattern: vi.fn().mockImplementation((base, pattern) => ({ base, pattern })),
+		RelativePattern: vi.fn().mockImplementation((base: any, pattern: any) => ({ base, pattern })),
 	}
 })
 
@@ -672,18 +700,59 @@ describe("CodeIndexManager - handleSettingsChange regression", () => {
 			expect(manager.isWorkspaceEnabled).toBe(false)
 		})
 
-		it("should store enablement per folder, not per window", async () => {
-			await manager.setAutoEnableDefault(false)
+		it("should store enablement per folder URI, not per window", async () => {
+			CodeIndexManager.disposeAll()
+
+			const vscode = await import("vscode")
+
+			const folderAPath = path.join(path.sep, "test", "folderA")
+			const folderBPath = path.join(path.sep, "test", "folderB")
+			const folderAUri = mockUri(folderAPath)
+			const folderBUri = mockUri(folderBPath)
+
+			// Both folders share the same workspaceState (same window)
+			const sharedStore: Record<string, any> = {}
+			const sharedContext = {
+				...mockContext,
+				workspaceState: {
+					get: vi.fn((key: string, defaultValue?: any) => sharedStore[key] ?? defaultValue),
+					update: vi.fn(async (key: string, value: any) => {
+						sharedStore[key] = value
+					}),
+				} as any,
+				globalState: {
+					get: vi.fn((_key: string, _defaultValue?: any) => false),
+					update: vi.fn(),
+				} as any,
+			}
 
-			// Enable indexing for the current manager's folder
-			await manager.setWorkspaceEnabled(true)
-			expect(manager.isWorkspaceEnabled).toBe(true)
+			// Patch workspaceFolders to include both folders
+			;(vscode.workspace as any).workspaceFolders = [
+				{ uri: folderAUri, name: "folderA", index: 0 },
+				{ uri: folderBUri, name: "folderB", index: 1 },
+			]
+
+			const managerA = CodeIndexManager.getInstance(sharedContext as any, folderAPath)!
+			const managerB = CodeIndexManager.getInstance(sharedContext as any, folderBPath)!
+
+			// Both start disabled (autoEnableDefault is false via globalState mock)
+			expect(managerA.isWorkspaceEnabled).toBe(false)
+			expect(managerB.isWorkspaceEnabled).toBe(false)
+
+			// Enable A only
+			await managerA.setWorkspaceEnabled(true)
+
+			expect(managerA.isWorkspaceEnabled).toBe(true)
+			expect(managerB.isWorkspaceEnabled).toBe(false)
+
+			// Enable B, disable A
+			await managerB.setWorkspaceEnabled(true)
+			await managerA.setWorkspaceEnabled(false)
 
-			// Create a second manager for a different folder path
-			const otherManager = CodeIndexManager.getInstance(mockContext as any, "/other/workspace")!
+			expect(managerA.isWorkspaceEnabled).toBe(false)
+			expect(managerB.isWorkspaceEnabled).toBe(true)
 
-			// The other folder should NOT be enabled — keys are per-folder
-			expect(otherManager.isWorkspaceEnabled).toBe(false)
+			CodeIndexManager.disposeAll()
 		})
 	})
 

src/services/code-index/manager.ts

@@ -32,26 +32,31 @@ export class CodeIndexManager {
 	private _isRecoveringFromError = false
 
 	public static getInstance(context: vscode.ExtensionContext, workspacePath?: string): CodeIndexManager | undefined {
-		// If workspacePath is not provided, try to get it from the active editor or first workspace folder
-		if (!workspacePath) {
+		// Resolve the workspace folder to get both fsPath and the real URI
+		let folder: vscode.WorkspaceFolder | undefined
+
+		if (workspacePath) {
+			folder = vscode.workspace.workspaceFolders?.find((f) => f.uri.fsPath === workspacePath)
+		} else {
 			const activeEditor = vscode.window.activeTextEditor
 			if (activeEditor) {
-				const workspaceFolder = vscode.workspace.getWorkspaceFolder(activeEditor.document.uri)
-				workspacePath = workspaceFolder?.uri.fsPath
+				folder = vscode.workspace.getWorkspaceFolder(activeEditor.document.uri)
 			}
-
-			if (!workspacePath) {
+			if (!folder) {
 				const workspaceFolders = vscode.workspace.workspaceFolders
 				if (!workspaceFolders || workspaceFolders.length === 0) {
 					return undefined
 				}
-				// Use the first workspace folder as fallback
-				workspacePath = workspaceFolders[0].uri.fsPath
+				folder = workspaceFolders[0]
 			}
+			workspacePath = folder.uri.fsPath
 		}
 
 		if (!CodeIndexManager.instances.has(workspacePath)) {
-			CodeIndexManager.instances.set(workspacePath, new CodeIndexManager(workspacePath, context))
+			// folder may be undefined when workspacePath was provided but doesn't match
+			// any workspace folder (e.g. cwd passed from a tool). Fall back to file:// URI.
+			const folderUri = folder?.uri ?? vscode.Uri.file(workspacePath)
+			CodeIndexManager.instances.set(workspacePath, new CodeIndexManager(workspacePath, folderUri, context))
 		}
 		return CodeIndexManager.instances.get(workspacePath)!
 	}
@@ -64,28 +69,35 @@ export class CodeIndexManager {
 	}
 
 	private readonly workspacePath: string
+	private readonly _folderUri: vscode.Uri
 	private readonly context: vscode.ExtensionContext
 
 	// Private constructor for singleton pattern
-	private constructor(workspacePath: string, context: vscode.ExtensionContext) {
+	private constructor(workspacePath: string, folderUri: vscode.Uri, context: vscode.ExtensionContext) {
 		this.workspacePath = workspacePath
+		this._folderUri = folderUri
 		this.context = context
 		this._stateManager = new CodeIndexStateManager()
 	}
 
 	// --- Public API ---
 
+	/**
+	 * Returns the workspaceState key for per-folder indexing enablement,
+	 * keyed by the real workspace folder URI so local/remote schemes cannot collide.
+	 */
+	private _workspaceEnabledKey(): string {
+		return "codeIndexWorkspaceEnabled:" + this._folderUri.toString(true)
+	}
+
 	public get isWorkspaceEnabled(): boolean {
-		const explicit = this.context.workspaceState.get<boolean | undefined>(
-			`codeIndexWorkspaceEnabled:${this.workspacePath}`,
-			undefined,
-		)
+		const explicit = this.context.workspaceState.get<boolean | undefined>(this._workspaceEnabledKey(), undefined)
 		if (explicit !== undefined) return explicit
 		return this.autoEnableDefault
 	}
 
 	public async setWorkspaceEnabled(enabled: boolean): Promise<void> {
-		await this.context.workspaceState.update(`codeIndexWorkspaceEnabled:${this.workspacePath}`, enabled)
+		await this.context.workspaceState.update(this._workspaceEnabledKey(), enabled)
 	}
 
 	public get autoEnableDefault(): boolean {

src/services/code-index/orchestrator.ts

@@ -186,6 +186,7 @@ export class CodeIndexOrchestrator {
 
 				if (signal.aborted) {
 					await this.cacheManager.flush()
+					this.stopWatcher()
 					this.stateManager.setSystemState("Standby", t("embeddings:orchestrator.indexingStopped"))
 					return
 				}
@@ -246,6 +247,7 @@ export class CodeIndexOrchestrator {
 
 				if (signal.aborted) {
 					await this.cacheManager.flush()
+					this.stopWatcher()
 					this.stateManager.setSystemState("Standby", t("embeddings:orchestrator.indexingStopped"))
 					return
 				}

src/services/code-index/processors/scanner.ts

@@ -185,7 +185,9 @@ export class DirectoryScanner implements IDirectoryScanner {
 									if (currentBatchBlocks.length >= this.batchSegmentThreshold) {
 										// Wait if we've reached the maximum pending batches
 										while (pendingBatchCount >= MAX_PENDING_BATCHES) {
-											if (signal?.aborted) break
+											if (signal?.aborted) {
+												throw new DOMException("Indexing aborted", "AbortError")
+											}
 											await Promise.race(activeBatchPromises)
 										}