diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index ec7c07999f..58a2ca4d31 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -63,6 +63,7 @@ import { BrowserSession } from "../../services/browser/BrowserSession" import { McpHub } from "../../services/mcp/McpHub" import { McpServerManager } from "../../services/mcp/McpServerManager" import { RepoPerTaskCheckpointService } from "../../services/checkpoints" +import { CodeIndexManager } from "../../services/code-index/manager" // integrations import { DiffViewProvider } from "../../integrations/editor/DiffViewProvider" @@ -1684,6 +1685,78 @@ export class Task extends EventEmitter implements TaskLike { // Kicks off the checkpoints initialization process in the background. getCheckpointService(this) + // Check if code indexing is in progress and gate semantic search if needed + const context = this.providerRef.deref()?.context + if (context) { + const codeIndexManager = CodeIndexManager.getInstance(context) + if (codeIndexManager && codeIndexManager.isFeatureEnabled && codeIndexManager.isFeatureConfigured) { + const status = codeIndexManager.getCurrentStatus() + + // If indexing is in progress, offer to pause or skip + if (status.systemStatus === "Indexing") { + const progressPercentage = + status.totalItems > 0 ? Math.round((status.processedItems / status.totalItems) * 100) : 0 + + const { response } = await this.ask( + "followup", + `šŸ”„ Code indexing is currently ${progressPercentage}% complete (${status.processedItems}/${status.totalItems} ${status.currentItemUnit || "items"}). ` + + `\n\nSemantic search will provide better results once indexing completes. ` + + `\n\nWould you like to wait for indexing to finish, or continue without semantic search?` + + `\n\nā€¢ **Wait**: Pause until indexing completes for best search results` + + `\nā€¢ **Skip**: Continue now using file-based search tools`, + ) + + if (response === "messageResponse" && this.askResponseText?.toLowerCase().includes("wait")) { + // User chose to wait - show progress updates + await this.say("text", "ā³ Waiting for code indexing to complete...") + + // Poll for indexing completion + while (!this.abort) { + const currentStatus = codeIndexManager.getCurrentStatus() + if (currentStatus.systemStatus === "Indexed") { + await this.say("text", "āœ… Code indexing complete! Semantic search is now available.") + break + } else if (currentStatus.systemStatus === "Error") { + await this.say( + "error", + "Code indexing encountered an error. Proceeding without semantic search.", + ) + break + } else if (currentStatus.systemStatus === "Indexing") { + const currentProgress = + currentStatus.totalItems > 0 + ? Math.round((currentStatus.processedItems / currentStatus.totalItems) * 100) + : 0 + await this.say( + "text", + `šŸ“Š Indexing progress: ${currentProgress}% (${currentStatus.processedItems}/${currentStatus.totalItems} ${currentStatus.currentItemUnit || "items"})`, + undefined, + true, // partial update + ) + await delay(2000) // Check every 2 seconds + } else { + // Standby or other state - stop waiting + break + } + } + } else { + // User chose to skip - continue without semantic search + await this.say( + "text", + "ā­ļø Continuing without semantic search. File-based search tools will be used instead.", + ) + } + } else if (status.systemStatus === "Standby" && !codeIndexManager.isInitialized) { + // Index hasn't started yet + await this.say( + "text", + "ā„¹ļø Code index has not been built yet. Semantic search will not be available for this task. " + + "File-based search tools will be used instead.", + ) + } + } + } + let nextUserContent = userContent let includeFileDetails = true diff --git a/src/core/tools/__tests__/codebaseSearchTool.spec.ts b/src/core/tools/__tests__/codebaseSearchTool.spec.ts new file mode 100644 index 0000000000..87acd69d40 --- /dev/null +++ b/src/core/tools/__tests__/codebaseSearchTool.spec.ts @@ -0,0 +1,422 @@ +import { describe, it, expect, vi, beforeEach } from "vitest" +import { codebaseSearchTool } from "../codebaseSearchTool" +import { CodeIndexManager } from "../../../services/code-index/manager" +import { Task } from "../../task/Task" +import { formatResponse } from "../../prompts/responses" +import * as vscode from "vscode" + +// Mock dependencies +vi.mock("../../../services/code-index/manager") +vi.mock("../../prompts/responses") +vi.mock("vscode", () => ({ + workspace: { + asRelativePath: vi.fn((path: string) => path.replace("/test/", "")), + }, +})) + +describe("codebaseSearchTool", () => { + let mockTask: any + let mockAskApproval: any + let mockHandleError: any + let mockPushToolResult: any + let mockRemoveClosingTag: any + let mockCodeIndexManager: any + + beforeEach(() => { + vi.clearAllMocks() + + // Setup mock task + mockTask = { + cwd: "/test/workspace", + consecutiveMistakeCount: 0, + providerRef: { + deref: vi.fn().mockReturnValue({ + context: {}, + }), + }, + ask: vi.fn(), + say: vi.fn(), + sayAndCreateMissingParamError: vi.fn(), + } + + // Setup mock functions + mockAskApproval = vi.fn().mockResolvedValue(true) + mockHandleError = vi.fn() + mockPushToolResult = vi.fn() + mockRemoveClosingTag = vi.fn((tag, value) => value) + + // Setup mock CodeIndexManager + mockCodeIndexManager = { + isFeatureEnabled: true, + isFeatureConfigured: true, + isInitialized: true, + getCurrentStatus: vi.fn(), + searchIndex: vi.fn(), + } + + vi.mocked(CodeIndexManager.getInstance).mockReturnValue(mockCodeIndexManager) + vi.mocked(formatResponse.toolDenied).mockReturnValue("Tool denied") + vi.mocked(formatResponse.missingToolParameterError).mockReturnValue("Missing parameter") + }) + + describe("indexing status checks", () => { + it("should throw error when indexing is in progress", async () => { + // Arrange + mockCodeIndexManager.getCurrentStatus.mockReturnValue({ + systemStatus: "Indexing", + processedItems: 50, + totalItems: 100, + currentItemUnit: "blocks", + }) + + const block = { + type: "tool_use" as const, + name: "codebase_search" as const, + params: { + query: "test query", + }, + partial: false, + } + + // Act & Assert + await codebaseSearchTool( + mockTask, + block, + mockAskApproval, + mockHandleError, + mockPushToolResult, + mockRemoveClosingTag, + ) + + expect(mockHandleError).toHaveBeenCalledWith( + "codebase_search", + expect.objectContaining({ + message: expect.stringContaining("Code indexing is currently in progress (50% complete)"), + }), + ) + }) + + it("should throw error when index is in standby state", async () => { + // Arrange + mockCodeIndexManager.getCurrentStatus.mockReturnValue({ + systemStatus: "Standby", + processedItems: 0, + totalItems: 0, + }) + + const block = { + type: "tool_use" as const, + name: "codebase_search" as const, + params: { + query: "test query", + }, + partial: false, + } + + // Act & Assert + await codebaseSearchTool( + mockTask, + block, + mockAskApproval, + mockHandleError, + mockPushToolResult, + mockRemoveClosingTag, + ) + + expect(mockHandleError).toHaveBeenCalledWith( + "codebase_search", + expect.objectContaining({ + message: expect.stringContaining("Code index is not ready (status: Standby)"), + }), + ) + }) + + it("should throw error when index is in error state", async () => { + // Arrange + mockCodeIndexManager.getCurrentStatus.mockReturnValue({ + systemStatus: "Error", + processedItems: 0, + totalItems: 0, + message: "Index failed", + }) + + const block = { + type: "tool_use" as const, + name: "codebase_search" as const, + params: { + query: "test query", + }, + partial: false, + } + + // Act & Assert + await codebaseSearchTool( + mockTask, + block, + mockAskApproval, + mockHandleError, + mockPushToolResult, + mockRemoveClosingTag, + ) + + expect(mockHandleError).toHaveBeenCalledWith( + "codebase_search", + expect.objectContaining({ + message: expect.stringContaining("Code index is not ready (status: Error)"), + }), + ) + }) + + it("should proceed with search when index is ready", async () => { + // Arrange + mockCodeIndexManager.getCurrentStatus.mockReturnValue({ + systemStatus: "Indexed", + processedItems: 100, + totalItems: 100, + }) + + mockCodeIndexManager.searchIndex.mockResolvedValue([ + { + score: 0.95, + payload: { + filePath: "/test/file.ts", + startLine: 10, + endLine: 20, + codeChunk: "test code", + }, + }, + ]) + + const block = { + type: "tool_use" as const, + name: "codebase_search" as const, + params: { + query: "test query", + }, + partial: false, + } + + // Mock say method to capture the result + mockTask.say = vi.fn() + + // Act + await codebaseSearchTool( + mockTask, + block, + mockAskApproval, + mockHandleError, + mockPushToolResult, + mockRemoveClosingTag, + ) + + // Assert + expect(mockCodeIndexManager.searchIndex).toHaveBeenCalledWith("test query", undefined) + expect(mockPushToolResult).toHaveBeenCalled() + expect(mockHandleError).not.toHaveBeenCalled() + + // Verify the result was pushed with correct format + const pushCall = mockPushToolResult.mock.calls[0][0] + expect(pushCall).toContain("Query: test query") + expect(pushCall).toContain("file.ts") + expect(pushCall).toContain("Score: 0.95") + expect(pushCall).toContain("Lines: 10-20") + expect(pushCall).toContain("test code") + }) + + it("should calculate progress percentage correctly", async () => { + // Arrange + mockCodeIndexManager.getCurrentStatus.mockReturnValue({ + systemStatus: "Indexing", + processedItems: 75, + totalItems: 150, + currentItemUnit: "files", + }) + + const block = { + type: "tool_use" as const, + name: "codebase_search" as const, + params: { + query: "test query", + }, + partial: false, + } + + // Act + await codebaseSearchTool( + mockTask, + block, + mockAskApproval, + mockHandleError, + mockPushToolResult, + mockRemoveClosingTag, + ) + + // Assert + expect(mockHandleError).toHaveBeenCalledWith( + "codebase_search", + expect.objectContaining({ + message: expect.stringContaining("50% complete"), + }), + ) + }) + + it("should handle zero total items when calculating progress", async () => { + // Arrange + mockCodeIndexManager.getCurrentStatus.mockReturnValue({ + systemStatus: "Indexing", + processedItems: 0, + totalItems: 0, + currentItemUnit: "blocks", + }) + + const block = { + type: "tool_use" as const, + name: "codebase_search" as const, + params: { + query: "test query", + }, + partial: false, + } + + // Act + await codebaseSearchTool( + mockTask, + block, + mockAskApproval, + mockHandleError, + mockPushToolResult, + mockRemoveClosingTag, + ) + + // Assert + expect(mockHandleError).toHaveBeenCalledWith( + "codebase_search", + expect.objectContaining({ + message: expect.stringContaining("0% complete"), + }), + ) + }) + }) + + describe("existing functionality", () => { + beforeEach(() => { + // Set index to ready state for existing functionality tests + mockCodeIndexManager.getCurrentStatus.mockReturnValue({ + systemStatus: "Indexed", + processedItems: 100, + totalItems: 100, + }) + }) + + it("should handle missing query parameter", async () => { + // Arrange + const block = { + type: "tool_use" as const, + name: "codebase_search" as const, + params: {}, + partial: false, + } + + mockTask.sayAndCreateMissingParamError.mockResolvedValue("error message") + + // Act + await codebaseSearchTool( + mockTask, + block, + mockAskApproval, + mockHandleError, + mockPushToolResult, + mockRemoveClosingTag, + ) + + // Assert + expect(mockTask.sayAndCreateMissingParamError).toHaveBeenCalledWith("codebase_search", "query") + expect(mockPushToolResult).toHaveBeenCalledWith("error message") + }) + + it("should handle user denial", async () => { + // Arrange + mockAskApproval.mockResolvedValue(false) + + const block = { + type: "tool_use" as const, + name: "codebase_search" as const, + params: { + query: "test query", + }, + partial: false, + } + + // Act + await codebaseSearchTool( + mockTask, + block, + mockAskApproval, + mockHandleError, + mockPushToolResult, + mockRemoveClosingTag, + ) + + // Assert + expect(mockPushToolResult).toHaveBeenCalledWith("Tool denied") + expect(mockCodeIndexManager.searchIndex).not.toHaveBeenCalled() + }) + + it("should handle empty search results", async () => { + // Arrange + mockCodeIndexManager.searchIndex.mockResolvedValue([]) + + const block = { + type: "tool_use" as const, + name: "codebase_search" as const, + params: { + query: "test query", + }, + partial: false, + } + + // Act + await codebaseSearchTool( + mockTask, + block, + mockAskApproval, + mockHandleError, + mockPushToolResult, + mockRemoveClosingTag, + ) + + // Assert + expect(mockPushToolResult).toHaveBeenCalledWith( + 'No relevant code snippets found for the query: "test query"', + ) + }) + + it("should handle search with directory prefix", async () => { + // Arrange + mockCodeIndexManager.searchIndex.mockResolvedValue([]) + + const block = { + type: "tool_use" as const, + name: "codebase_search" as const, + params: { + query: "test query", + path: "src/components", + }, + partial: false, + } + + // Act + await codebaseSearchTool( + mockTask, + block, + mockAskApproval, + mockHandleError, + mockPushToolResult, + mockRemoveClosingTag, + ) + + // Assert + expect(mockCodeIndexManager.searchIndex).toHaveBeenCalledWith("test query", "src/components") + }) + }) +}) diff --git a/src/core/tools/codebaseSearchTool.ts b/src/core/tools/codebaseSearchTool.ts index 700d1b7c7c..a7cd90d790 100644 --- a/src/core/tools/codebaseSearchTool.ts +++ b/src/core/tools/codebaseSearchTool.ts @@ -17,7 +17,7 @@ export async function codebaseSearchTool( removeClosingTag: RemoveClosingTag, ) { const toolName = "codebase_search" - const workspacePath = (cline.cwd && cline.cwd.trim() !== '') ? cline.cwd : getWorkspacePath() + const workspacePath = cline.cwd && cline.cwd.trim() !== "" ? cline.cwd : getWorkspacePath() if (!workspacePath) { // This case should ideally not happen if Cline is initialized correctly @@ -82,6 +82,26 @@ export async function codebaseSearchTool( throw new Error("Code Indexing is not configured (Missing OpenAI Key or Qdrant URL).") } + // Check if indexing is in progress + const status = manager.getCurrentStatus() + if (status.systemStatus === "Indexing") { + const progressPercentage = + status.totalItems > 0 ? Math.round((status.processedItems / status.totalItems) * 100) : 0 + throw new Error( + `Code indexing is currently in progress (${progressPercentage}% complete). ` + + `Please wait for indexing to complete before using semantic search, or use file-based search tools instead.`, + ) + } + + // Check if index is ready + if (status.systemStatus === "Standby" || status.systemStatus === "Error") { + throw new Error( + `Code index is not ready (status: ${status.systemStatus}). ` + + `The index needs to be built before semantic search can be used. ` + + `Please ensure indexing is enabled and configured, then wait for the initial index to build.`, + ) + } + const searchResults: VectorStoreSearchResult[] = await manager.searchIndex(query, directoryPrefix) // 3. Format and push results