diff --git a/src/extension.ts b/src/extension.ts index beb69b30b5..93db0785cd 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -30,6 +30,7 @@ import { MdmService } from "./services/mdm/MdmService" import { migrateSettings } from "./utils/migrateSettings" import { autoImportSettings } from "./utils/autoImportSettings" import { API } from "./extension/api" +import { AiCompletionProvider } from "./integrations/completion/AiCompletionProvider" import { handleUri, @@ -192,6 +193,45 @@ export async function activate(context: vscode.ExtensionContext) { registerCodeActions(context) registerTerminalActions(context) + // Register AI completion provider + const aiCompletionEnabled = vscode.workspace + .getConfiguration(Package.name) + .get("aiTabCompletion.enabled", false) + if (aiCompletionEnabled) { + const aiCompletionProvider = new AiCompletionProvider(outputChannel) + context.subscriptions.push( + vscode.languages.registerInlineCompletionItemProvider({ pattern: "**/*" }, aiCompletionProvider), + aiCompletionProvider, + ) + outputChannel.appendLine("AI Tab Completion: Provider registered") + } + + // Watch for configuration changes to enable/disable AI completion + context.subscriptions.push( + vscode.workspace.onDidChangeConfiguration((e) => { + if (e.affectsConfiguration(`${Package.name}.aiTabCompletion.enabled`)) { + const nowEnabled = vscode.workspace + .getConfiguration(Package.name) + .get("aiTabCompletion.enabled", false) + if (nowEnabled) { + outputChannel.appendLine( + "AI Tab Completion: Configuration changed - reloading extension recommended to enable", + ) + vscode.window + .showInformationMessage( + "AI Tab Completion has been enabled. Please reload the window for changes to take effect.", + "Reload Window", + ) + .then((selection) => { + if (selection === "Reload Window") { + vscode.commands.executeCommand("workbench.action.reloadWindow") + } + }) + } + } + }), + ) + // Allows other extensions to activate once Roo is ready. vscode.commands.executeCommand(`${Package.name}.activationCompleted`) diff --git a/src/integrations/completion/AiCompletionProvider.ts b/src/integrations/completion/AiCompletionProvider.ts new file mode 100644 index 0000000000..3f35f26d91 --- /dev/null +++ b/src/integrations/completion/AiCompletionProvider.ts @@ -0,0 +1,268 @@ +import * as vscode from "vscode" +import { Anthropic } from "@anthropic-ai/sdk" +import debounce from "lodash.debounce" + +import { buildApiHandler, ApiHandler } from "../../api" +import { ProviderSettings } from "@roo-code/types" +import { Package } from "../../shared/package" + +export class AiCompletionProvider implements vscode.InlineCompletionItemProvider { + private apiHandler: ApiHandler | null = null + private outputChannel: vscode.OutputChannel + private debouncedProvideCompletions: ReturnType + private lastCompletionRequestId = 0 + private activeCompletionRequest: AbortController | null = null + + constructor(outputChannel: vscode.OutputChannel) { + this.outputChannel = outputChannel + + // Initialize debounced completion function + const debounceDelay = vscode.workspace + .getConfiguration(Package.name) + .get("aiTabCompletion.debounceDelay", 300) + this.debouncedProvideCompletions = debounce(this.provideCompletionsInternal.bind(this), debounceDelay) + + // Update configuration when settings change + vscode.workspace.onDidChangeConfiguration((e) => { + if (e.affectsConfiguration(`${Package.name}.aiTabCompletion`)) { + this.updateConfiguration() + } + }) + + this.updateConfiguration() + } + + private updateConfiguration() { + const config = vscode.workspace.getConfiguration(Package.name) + const enabled = config.get("aiTabCompletion.enabled", false) + + if (!enabled) { + this.apiHandler = null + return + } + + const provider = config.get("aiTabCompletion.provider", "anthropic") + const model = config.get("aiTabCompletion.model", "claude-3-haiku-20240307") + + // Build provider settings based on configuration + const providerSettings: ProviderSettings = { + apiProvider: provider as any, + apiModelId: model, + } + + // Add API keys based on provider + switch (provider) { + case "anthropic": { + const anthropicKey = config.get("anthropicApiKey") + if (anthropicKey) { + providerSettings.apiKey = anthropicKey + } + break + } + case "openai": { + const openaiKey = config.get("openaiApiKey") + if (openaiKey) { + providerSettings.openAiApiKey = openaiKey + } + break + } + case "openrouter": { + const openrouterKey = config.get("openRouterApiKey") + if (openrouterKey) { + providerSettings.openRouterApiKey = openrouterKey + } + break + } + // Add other providers as needed + } + + try { + this.apiHandler = buildApiHandler(providerSettings) + this.outputChannel.appendLine(`AI Tab Completion: Initialized with provider ${provider} and model ${model}`) + } catch (error) { + this.outputChannel.appendLine(`AI Tab Completion: Failed to initialize - ${error}`) + this.apiHandler = null + } + + // Update debounce delay + const newDebounceDelay = config.get("aiTabCompletion.debounceDelay", 300) + this.debouncedProvideCompletions = debounce(this.provideCompletionsInternal.bind(this), newDebounceDelay) + } + + async provideInlineCompletionItems( + document: vscode.TextDocument, + position: vscode.Position, + context: vscode.InlineCompletionContext, + token: vscode.CancellationToken, + ): Promise { + if (!this.apiHandler) { + return undefined + } + + // Cancel any existing completion request + if (this.activeCompletionRequest) { + this.activeCompletionRequest.abort() + } + + // Create new abort controller for this request + this.activeCompletionRequest = new AbortController() + const requestId = ++this.lastCompletionRequestId + + // Use debounced function + return new Promise((resolve) => { + this.debouncedProvideCompletions(document, position, context, token, requestId, resolve) + }) + } + + private async provideCompletionsInternal( + document: vscode.TextDocument, + position: vscode.Position, + context: vscode.InlineCompletionContext, + token: vscode.CancellationToken, + requestId: number, + resolve: (value: vscode.InlineCompletionItem[] | undefined) => void, + ) { + // Check if this request is still the latest + if (requestId !== this.lastCompletionRequestId) { + resolve(undefined) + return + } + + try { + const config = vscode.workspace.getConfiguration(Package.name) + const maxTokens = config.get("aiTabCompletion.maxTokens", 150) + const temperature = config.get("aiTabCompletion.temperature", 0.2) + + // Get context around cursor + const linePrefix = document.lineAt(position.line).text.substring(0, position.character) + const lineSuffix = document.lineAt(position.line).text.substring(position.character) + + // Get preceding lines for context (up to 50 lines) + const precedingLines: string[] = [] + for (let i = Math.max(0, position.line - 50); i < position.line; i++) { + precedingLines.push(document.lineAt(i).text) + } + + // Get following lines for context (up to 10 lines) + const followingLines: string[] = [] + for (let i = position.line + 1; i < Math.min(document.lineCount, position.line + 10); i++) { + followingLines.push(document.lineAt(i).text) + } + + // Build prompt for completion + const prompt = this.buildCompletionPrompt( + document.languageId, + precedingLines.join("\n"), + linePrefix, + lineSuffix, + followingLines.join("\n"), + ) + + // Create messages for API + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user", + content: prompt, + }, + ] + + // Use streaming for faster response + const stream = this.apiHandler!.createMessage( + "You are a code completion assistant. Complete the code at the cursor position. Only provide the completion text, no explanations.", + messages, + { + taskId: `completion-${requestId}`, + mode: "completion", + }, + ) + + let completion = "" + for await (const chunk of stream) { + if (token.isCancellationRequested || requestId !== this.lastCompletionRequestId) { + break + } + + if (chunk.type === "text") { + completion += chunk.text + } + } + + // Clean up the completion + completion = this.cleanCompletion(completion, linePrefix, lineSuffix) + + if (completion && !token.isCancellationRequested && requestId === this.lastCompletionRequestId) { + const item = new vscode.InlineCompletionItem(completion, new vscode.Range(position, position)) + resolve([item]) + } else { + resolve(undefined) + } + } catch (error) { + this.outputChannel.appendLine(`AI Tab Completion Error: ${error}`) + resolve(undefined) + } finally { + if (requestId === this.lastCompletionRequestId) { + this.activeCompletionRequest = null + } + } + } + + private buildCompletionPrompt( + languageId: string, + precedingContext: string, + linePrefix: string, + lineSuffix: string, + followingContext: string, + ): string { + return `Language: ${languageId} + +Context before cursor: +${precedingContext} + +Current line before cursor: ${linePrefix} +Current line after cursor: ${lineSuffix} + +Context after cursor: +${followingContext} + +Complete the code at the cursor position. Provide only the code to insert, nothing else. The completion should fit naturally between the prefix and suffix.` + } + + private cleanCompletion(completion: string, linePrefix: string, lineSuffix: string): string { + // Remove any markdown code blocks + completion = completion.replace(/```[\w]*\n?/g, "").replace(/```$/g, "") + + // Trim whitespace + completion = completion.trim() + + // Remove duplicate prefix if AI included it + if (completion.startsWith(linePrefix.trimEnd())) { + completion = completion.substring(linePrefix.trimEnd().length) + } + + // Remove duplicate suffix if AI included it + if (lineSuffix && completion.endsWith(lineSuffix.trimStart())) { + completion = completion.substring(0, completion.length - lineSuffix.trimStart().length) + } + + // Handle proper spacing + if ( + linePrefix && + !linePrefix.endsWith(" ") && + completion && + !completion.startsWith(" ") && + /\w$/.test(linePrefix) + ) { + // Add space if needed between words + completion = " " + completion + } + + return completion + } + + dispose() { + this.debouncedProvideCompletions.cancel() + if (this.activeCompletionRequest) { + this.activeCompletionRequest.abort() + } + } +} diff --git a/src/integrations/completion/__tests__/AiCompletionProvider.spec.ts b/src/integrations/completion/__tests__/AiCompletionProvider.spec.ts new file mode 100644 index 0000000000..2d8d73b2ae --- /dev/null +++ b/src/integrations/completion/__tests__/AiCompletionProvider.spec.ts @@ -0,0 +1,329 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest" +import * as vscode from "vscode" +import { AiCompletionProvider } from "../AiCompletionProvider" +import { buildApiHandler } from "../../../api" + +// Mock vscode +vi.mock("vscode", () => ({ + workspace: { + getConfiguration: vi.fn(), + onDidChangeConfiguration: vi.fn(() => ({ dispose: vi.fn() })), + }, + languages: { + registerInlineCompletionItemProvider: vi.fn(), + }, + window: { + showInformationMessage: vi.fn(), + }, + commands: { + executeCommand: vi.fn(), + }, + Position: vi.fn((line, character) => ({ line, character })), + Range: vi.fn((start, end) => ({ start, end })), + InlineCompletionItem: vi.fn((text, range) => ({ text, range })), + CancellationTokenSource: vi.fn(() => ({ + token: { isCancellationRequested: false }, + cancel: vi.fn(), + dispose: vi.fn(), + })), +})) + +// Mock API handler +vi.mock("../../../api", () => ({ + buildApiHandler: vi.fn(), +})) + +// Mock lodash.debounce +vi.mock("lodash.debounce", () => ({ + default: (fn: any) => { + const debounced = (...args: any[]) => fn(...args) + debounced.cancel = vi.fn() + return debounced + }, +})) + +describe("AiCompletionProvider", () => { + let provider: AiCompletionProvider + let mockOutputChannel: any + let mockConfig: any + let mockApiHandler: any + + beforeEach(() => { + // Reset mocks + vi.clearAllMocks() + + // Mock output channel + mockOutputChannel = { + appendLine: vi.fn(), + } + + // Mock configuration + mockConfig = { + get: vi.fn((key: string, defaultValue?: any) => { + const configMap: Record = { + "aiTabCompletion.enabled": true, + "aiTabCompletion.provider": "anthropic", + "aiTabCompletion.model": "claude-3-haiku-20240307", + "aiTabCompletion.debounceDelay": 300, + "aiTabCompletion.maxTokens": 150, + "aiTabCompletion.temperature": 0.2, + anthropicApiKey: "test-api-key", + } + return configMap[key] ?? defaultValue + }), + } + + // Mock API handler + mockApiHandler = { + createMessage: vi.fn(() => { + // Return an async generator + return (async function* () { + yield { type: "text", text: "test completion" } + })() + }), + } + + vi.mocked(vscode.workspace.getConfiguration).mockReturnValue(mockConfig as any) + vi.mocked(buildApiHandler).mockReturnValue(mockApiHandler) + + provider = new AiCompletionProvider(mockOutputChannel) + }) + + afterEach(() => { + provider.dispose() + }) + + describe("initialization", () => { + it("should initialize with enabled configuration", () => { + expect(buildApiHandler).toHaveBeenCalledWith({ + apiProvider: "anthropic", + apiModelId: "claude-3-haiku-20240307", + apiKey: "test-api-key", + }) + expect(mockOutputChannel.appendLine).toHaveBeenCalledWith( + "AI Tab Completion: Initialized with provider anthropic and model claude-3-haiku-20240307", + ) + }) + + it("should not initialize API handler when disabled", () => { + vi.clearAllMocks() + mockConfig.get.mockImplementation((key: string, defaultValue?: any) => { + if (key === "aiTabCompletion.enabled") return false + return defaultValue + }) + + new AiCompletionProvider(mockOutputChannel) + + expect(buildApiHandler).not.toHaveBeenCalled() + }) + + it("should handle initialization errors gracefully", () => { + vi.clearAllMocks() + vi.mocked(buildApiHandler).mockImplementation(() => { + throw new Error("API initialization failed") + }) + + new AiCompletionProvider(mockOutputChannel) + + expect(mockOutputChannel.appendLine).toHaveBeenCalledWith( + "AI Tab Completion: Failed to initialize - Error: API initialization failed", + ) + }) + }) + + describe("provideInlineCompletionItems", () => { + it("should return undefined when API handler is not initialized", async () => { + mockConfig.get.mockImplementation((key: string, defaultValue?: any) => { + if (key === "aiTabCompletion.enabled") return false + return defaultValue + }) + const disabledProvider = new AiCompletionProvider(mockOutputChannel) + + const mockDocument = { + languageId: "typescript", + lineAt: vi.fn(() => ({ text: "const x = " })), + lineCount: 10, + } + const mockPosition = new (vscode as any).Position(0, 10) + const mockContext = {} + const mockToken = { isCancellationRequested: false } + + const result = await disabledProvider.provideInlineCompletionItems( + mockDocument as any, + mockPosition as any, + mockContext as any, + mockToken as any, + ) + + expect(result).toBeUndefined() + }) + + it("should provide completion items", async () => { + const mockDocument = { + languageId: "typescript", + lineAt: vi.fn((line: number) => { + if (line === 5) return { text: "const x = " } + return { text: "// some code" } + }), + lineCount: 10, + } + const mockPosition = new (vscode as any).Position(5, 10) + const mockContext = {} + const mockToken = { isCancellationRequested: false } + + const result = await provider.provideInlineCompletionItems( + mockDocument as any, + mockPosition as any, + mockContext as any, + mockToken as any, + ) + + expect(result).toBeDefined() + expect(result).toHaveLength(1) + expect(result![0]).toHaveProperty("text", "test completion") + }) + + it("should handle cancellation", async () => { + const mockDocument = { + languageId: "typescript", + lineAt: vi.fn(() => ({ text: "const x = " })), + lineCount: 10, + } + const mockPosition = new (vscode as any).Position(0, 10) + const mockContext = {} + const mockToken = { isCancellationRequested: true } + + mockApiHandler.createMessage.mockReturnValue( + (async function* () { + yield { type: "text", text: "test" } + })(), + ) + + const result = await provider.provideInlineCompletionItems( + mockDocument as any, + mockPosition as any, + mockContext as any, + mockToken as any, + ) + + expect(result).toBeUndefined() + }) + + it("should handle API errors gracefully", async () => { + mockApiHandler.createMessage.mockImplementation(() => { + throw new Error("API error") + }) + + const mockDocument = { + languageId: "typescript", + lineAt: vi.fn(() => ({ text: "const x = " })), + lineCount: 10, + } + const mockPosition = new (vscode as any).Position(0, 10) + const mockContext = {} + const mockToken = { isCancellationRequested: false } + + const result = await provider.provideInlineCompletionItems( + mockDocument as any, + mockPosition as any, + mockContext as any, + mockToken as any, + ) + + expect(result).toBeUndefined() + expect(mockOutputChannel.appendLine).toHaveBeenCalledWith("AI Tab Completion Error: Error: API error") + }) + }) + + describe("completion cleaning", () => { + it("should remove markdown code blocks from completion", async () => { + mockApiHandler.createMessage.mockReturnValue( + (async function* () { + yield { type: "text", text: "```typescript\nconst y = 42\n```" } + })(), + ) + + const mockDocument = { + languageId: "typescript", + lineAt: vi.fn((line: number) => { + if (line === 0) return { text: "const x = " } + return { text: "" } + }), + lineCount: 1, + } + const mockPosition = new (vscode as any).Position(0, 10) + const mockContext = {} + const mockToken = { isCancellationRequested: false } + + const result = await provider.provideInlineCompletionItems( + mockDocument as any, + mockPosition as any, + mockContext as any, + mockToken as any, + ) + + expect(result).toBeDefined() + expect(result![0]).toHaveProperty("text", "const y = 42") + }) + + it("should handle proper spacing between words", async () => { + mockApiHandler.createMessage.mockReturnValue( + (async function* () { + yield { type: "text", text: "world" } + })(), + ) + + const mockDocument = { + languageId: "typescript", + lineAt: vi.fn(() => ({ text: "hello" })), + lineCount: 1, + } + const mockPosition = new (vscode as any).Position(0, 5) + const mockContext = {} + const mockToken = { isCancellationRequested: false } + + const result = await provider.provideInlineCompletionItems( + mockDocument as any, + mockPosition as any, + mockContext as any, + mockToken as any, + ) + + expect(result).toBeDefined() + expect(result![0]).toHaveProperty("text", " world") + }) + }) + + describe("configuration updates", () => { + it("should update configuration when settings change", () => { + const onDidChangeConfiguration = vi.mocked(vscode.workspace.onDidChangeConfiguration) + const changeHandler = onDidChangeConfiguration.mock.calls[0][0] + + // Clear previous calls + vi.clearAllMocks() + + // Update configuration + mockConfig.get.mockImplementation((key: string, defaultValue?: any) => { + const configMap: Record = { + "aiTabCompletion.enabled": true, + "aiTabCompletion.provider": "openai", + "aiTabCompletion.model": "gpt-4", + "aiTabCompletion.debounceDelay": 500, + openaiApiKey: "new-api-key", + } + return configMap[key] ?? defaultValue + }) + + // Trigger configuration change + changeHandler({ + affectsConfiguration: (section: string) => section.includes("aiTabCompletion"), + } as any) + + expect(buildApiHandler).toHaveBeenCalledWith({ + apiProvider: "openai", + apiModelId: "gpt-4", + openAiApiKey: "new-api-key", + }) + }) + }) +}) diff --git a/src/package.json b/src/package.json index aa2110dfd5..3f27184a7f 100644 --- a/src/package.json +++ b/src/package.json @@ -391,6 +391,51 @@ "type": "boolean", "default": true, "description": "%settings.useAgentRules.description%" + }, + "roo-cline.aiTabCompletion.enabled": { + "type": "boolean", + "default": false, + "description": "%settings.aiTabCompletion.enabled.description%" + }, + "roo-cline.aiTabCompletion.provider": { + "type": "string", + "enum": [ + "anthropic", + "openai", + "openrouter", + "bedrock", + "ollama", + "lmstudio", + "vscode-lm" + ], + "default": "anthropic", + "description": "%settings.aiTabCompletion.provider.description%" + }, + "roo-cline.aiTabCompletion.model": { + "type": "string", + "default": "claude-3-haiku-20240307", + "description": "%settings.aiTabCompletion.model.description%" + }, + "roo-cline.aiTabCompletion.debounceDelay": { + "type": "number", + "default": 300, + "minimum": 100, + "maximum": 1000, + "description": "%settings.aiTabCompletion.debounceDelay.description%" + }, + "roo-cline.aiTabCompletion.maxTokens": { + "type": "number", + "default": 150, + "minimum": 50, + "maximum": 500, + "description": "%settings.aiTabCompletion.maxTokens.description%" + }, + "roo-cline.aiTabCompletion.temperature": { + "type": "number", + "default": 0.2, + "minimum": 0, + "maximum": 1, + "description": "%settings.aiTabCompletion.temperature.description%" } } } diff --git a/src/package.nls.json b/src/package.nls.json index 36ef72a823..2f3c08c2d0 100644 --- a/src/package.nls.json +++ b/src/package.nls.json @@ -37,5 +37,11 @@ "settings.customStoragePath.description": "Custom storage path. Leave empty to use the default location. Supports absolute paths (e.g. 'D:\\RooCodeStorage')", "settings.enableCodeActions.description": "Enable Roo Code quick fixes", "settings.autoImportSettingsPath.description": "Path to a RooCode configuration file to automatically import on extension startup. Supports absolute paths and paths relative to the home directory (e.g. '~/Documents/roo-code-settings.json'). Leave empty to disable auto-import.", - "settings.useAgentRules.description": "Enable loading of AGENTS.md files for agent-specific rules (see https://agent-rules.org/)" + "settings.useAgentRules.description": "Enable loading of AGENTS.md files for agent-specific rules (see https://agent-rules.org/)", + "settings.aiTabCompletion.enabled.description": "Enable AI-powered tab completion for code suggestions", + "settings.aiTabCompletion.provider.description": "AI provider to use for tab completions", + "settings.aiTabCompletion.model.description": "AI model to use for tab completions (e.g., 'claude-3-haiku-20240307' for Anthropic)", + "settings.aiTabCompletion.debounceDelay.description": "Delay in milliseconds before triggering completion after typing stops (100-1000ms)", + "settings.aiTabCompletion.maxTokens.description": "Maximum number of tokens to generate for each completion (50-500)", + "settings.aiTabCompletion.temperature.description": "Temperature for AI completions - lower values are more deterministic (0-1)" }