From f3102f8134fcea80b49db77070f02ff259de5b35 Mon Sep 17 00:00:00 2001 From: Roo Code Date: Mon, 20 Oct 2025 15:46:20 +0000 Subject: [PATCH] feat: add API key support for Ollama Embedder Provider - Add ollamaApiKey parameter to CodeIndexOllamaEmbedder constructor - Include Authorization header in all API requests when API key is provided - Update CodeIndexConfigManager to read and store Ollama API key from secrets - Add UI field for Ollama API key input in CodeIndexPopover component - Update type definitions to include codebaseIndexOllamaApiKey secret - Add comprehensive tests for API key authentication - Maintain backward compatibility for local Ollama instances without auth Closes RooCodeInc/Roo-Code#8737 --- packages/types/src/codebase-index.ts | 1 + packages/types/src/global-settings.ts | 1 + src/core/webview/webviewMessageHandler.ts | 8 + src/services/code-index/config-manager.ts | 7 +- .../embedders/__tests__/ollama.spec.ts | 268 ++++++++++++++++++ src/services/code-index/embedders/ollama.ts | 38 ++- src/services/code-index/interfaces/config.ts | 1 + src/shared/WebviewMessage.ts | 1 + .../src/components/chat/CodeIndexPopover.tsx | 32 ++- 9 files changed, 346 insertions(+), 11 deletions(-) diff --git a/packages/types/src/codebase-index.ts b/packages/types/src/codebase-index.ts index be7778f53875..ff708e07ff2a 100644 --- a/packages/types/src/codebase-index.ts +++ b/packages/types/src/codebase-index.ts @@ -68,6 +68,7 @@ export const codebaseIndexProviderSchema = z.object({ codebaseIndexGeminiApiKey: z.string().optional(), codebaseIndexMistralApiKey: z.string().optional(), codebaseIndexVercelAiGatewayApiKey: z.string().optional(), + codebaseIndexOllamaApiKey: z.string().optional(), }) export type CodebaseIndexProvider = z.infer diff --git a/packages/types/src/global-settings.ts b/packages/types/src/global-settings.ts index a56a00fc355a..ffa91cc9d719 100644 --- a/packages/types/src/global-settings.ts +++ b/packages/types/src/global-settings.ts @@ -199,6 +199,7 @@ export const SECRET_STATE_KEYS = [ "codebaseIndexGeminiApiKey", "codebaseIndexMistralApiKey", "codebaseIndexVercelAiGatewayApiKey", + "codebaseIndexOllamaApiKey", "huggingFaceApiKey", "sambaNovaApiKey", "zaiApiKey", diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index af5f9925c353..26f37d80c150 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -2493,6 +2493,12 @@ export const webviewMessageHandler = async ( settings.codebaseIndexVercelAiGatewayApiKey, ) } + if (settings.codebaseIndexOllamaApiKey !== undefined) { + await provider.contextProxy.storeSecret( + "codebaseIndexOllamaApiKey", + settings.codebaseIndexOllamaApiKey, + ) + } // Send success response first - settings are saved regardless of validation await provider.postMessageToWebview({ @@ -2630,6 +2636,7 @@ export const webviewMessageHandler = async ( const hasVercelAiGatewayApiKey = !!(await provider.context.secrets.get( "codebaseIndexVercelAiGatewayApiKey", )) + const hasOllamaApiKey = !!(await provider.context.secrets.get("codebaseIndexOllamaApiKey")) provider.postMessageToWebview({ type: "codeIndexSecretStatus", @@ -2640,6 +2647,7 @@ export const webviewMessageHandler = async ( hasGeminiApiKey, hasMistralApiKey, hasVercelAiGatewayApiKey, + hasOllamaApiKey, }, }) break diff --git a/src/services/code-index/config-manager.ts b/src/services/code-index/config-manager.ts index 2c0e8bb5c9e4..ca1a3c301ab6 100644 --- a/src/services/code-index/config-manager.ts +++ b/src/services/code-index/config-manager.ts @@ -71,6 +71,7 @@ export class CodeIndexConfigManager { const geminiApiKey = this.contextProxy?.getSecret("codebaseIndexGeminiApiKey") ?? "" const mistralApiKey = this.contextProxy?.getSecret("codebaseIndexMistralApiKey") ?? "" const vercelAiGatewayApiKey = this.contextProxy?.getSecret("codebaseIndexVercelAiGatewayApiKey") ?? "" + const ollamaApiKey = this.contextProxy?.getSecret("codebaseIndexOllamaApiKey") ?? "" // Update instance variables with configuration this.codebaseIndexEnabled = codebaseIndexEnabled ?? true @@ -116,6 +117,7 @@ export class CodeIndexConfigManager { this.ollamaOptions = { ollamaBaseUrl: codebaseIndexEmbedderBaseUrl, + ollamaApiKey: ollamaApiKey || undefined, } this.openAiCompatibleOptions = @@ -162,6 +164,7 @@ export class CodeIndexConfigManager { modelDimension: this.modelDimension, openAiKey: this.openAiOptions?.openAiNativeApiKey ?? "", ollamaBaseUrl: this.ollamaOptions?.ollamaBaseUrl ?? "", + ollamaApiKey: this.ollamaOptions?.ollamaApiKey ?? "", openAiCompatibleBaseUrl: this.openAiCompatibleOptions?.baseUrl ?? "", openAiCompatibleApiKey: this.openAiCompatibleOptions?.apiKey ?? "", geminiApiKey: this.geminiOptions?.apiKey ?? "", @@ -263,6 +266,7 @@ export class CodeIndexConfigManager { const prevProvider = prev?.embedderProvider ?? "openai" const prevOpenAiKey = prev?.openAiKey ?? "" const prevOllamaBaseUrl = prev?.ollamaBaseUrl ?? "" + const prevOllamaApiKey = prev?.ollamaApiKey ?? "" const prevOpenAiCompatibleBaseUrl = prev?.openAiCompatibleBaseUrl ?? "" const prevOpenAiCompatibleApiKey = prev?.openAiCompatibleApiKey ?? "" const prevModelDimension = prev?.modelDimension @@ -301,6 +305,7 @@ export class CodeIndexConfigManager { // Authentication changes (API keys) const currentOpenAiKey = this.openAiOptions?.openAiNativeApiKey ?? "" const currentOllamaBaseUrl = this.ollamaOptions?.ollamaBaseUrl ?? "" + const currentOllamaApiKey = this.ollamaOptions?.ollamaApiKey ?? "" const currentOpenAiCompatibleBaseUrl = this.openAiCompatibleOptions?.baseUrl ?? "" const currentOpenAiCompatibleApiKey = this.openAiCompatibleOptions?.apiKey ?? "" const currentModelDimension = this.modelDimension @@ -314,7 +319,7 @@ export class CodeIndexConfigManager { return true } - if (prevOllamaBaseUrl !== currentOllamaBaseUrl) { + if (prevOllamaBaseUrl !== currentOllamaBaseUrl || prevOllamaApiKey !== currentOllamaApiKey) { return true } diff --git a/src/services/code-index/embedders/__tests__/ollama.spec.ts b/src/services/code-index/embedders/__tests__/ollama.spec.ts index 650140beacf0..8a5f3d9d4cb3 100644 --- a/src/services/code-index/embedders/__tests__/ollama.spec.ts +++ b/src/services/code-index/embedders/__tests__/ollama.spec.ts @@ -81,6 +81,15 @@ describe("CodeIndexOllamaEmbedder", () => { expect(embedderWithDefaults.embedderInfo.name).toBe("ollama") }) + it("should initialize with API key when provided", () => { + const embedderWithApiKey = new CodeIndexOllamaEmbedder({ + ollamaModelId: "nomic-embed-text", + ollamaBaseUrl: "http://localhost:11434", + ollamaApiKey: "test-api-key-123", + }) + expect(embedderWithApiKey.embedderInfo.name).toBe("ollama") + }) + it("should normalize URLs with trailing slashes", async () => { // Create embedder with URL that has a trailing slash const embedderWithTrailingSlash = new CodeIndexOllamaEmbedder({ @@ -166,6 +175,128 @@ describe("CodeIndexOllamaEmbedder", () => { }) }) + describe("API Key Authentication", () => { + it("should include Authorization header when API key is provided", async () => { + const embedderWithApiKey = new CodeIndexOllamaEmbedder({ + ollamaModelId: "nomic-embed-text", + ollamaBaseUrl: "http://localhost:11434", + ollamaApiKey: "test-api-key-123", + }) + + // Mock successful /api/tags call + mockFetch.mockImplementationOnce(() => + Promise.resolve({ + ok: true, + status: 200, + json: () => + Promise.resolve({ + models: [{ name: "nomic-embed-text" }], + }), + } as Response), + ) + + // Mock successful /api/embed test call + mockFetch.mockImplementationOnce(() => + Promise.resolve({ + ok: true, + status: 200, + json: () => + Promise.resolve({ + embeddings: [[0.1, 0.2, 0.3]], + }), + } as Response), + ) + + await embedderWithApiKey.validateConfiguration() + + // Check that Authorization header is included in both calls + expect(mockFetch).toHaveBeenCalledTimes(2) + + // First call to /api/tags + expect(mockFetch.mock.calls[0][1]?.headers).toEqual({ + "Content-Type": "application/json", + Authorization: "Bearer test-api-key-123", + }) + + // Second call to /api/embed + expect(mockFetch.mock.calls[1][1]?.headers).toEqual({ + "Content-Type": "application/json", + Authorization: "Bearer test-api-key-123", + }) + }) + + it("should not include Authorization header when API key is not provided", async () => { + // Mock successful /api/tags call + mockFetch.mockImplementationOnce(() => + Promise.resolve({ + ok: true, + status: 200, + json: () => + Promise.resolve({ + models: [{ name: "nomic-embed-text" }], + }), + } as Response), + ) + + // Mock successful /api/embed test call + mockFetch.mockImplementationOnce(() => + Promise.resolve({ + ok: true, + status: 200, + json: () => + Promise.resolve({ + embeddings: [[0.1, 0.2, 0.3]], + }), + } as Response), + ) + + await embedder.validateConfiguration() + + // Check that Authorization header is NOT included + expect(mockFetch).toHaveBeenCalledTimes(2) + + // First call to /api/tags + expect(mockFetch.mock.calls[0][1]?.headers).toEqual({ + "Content-Type": "application/json", + }) + + // Second call to /api/embed + expect(mockFetch.mock.calls[1][1]?.headers).toEqual({ + "Content-Type": "application/json", + }) + }) + + it("should handle authentication errors with API key", async () => { + const embedderWithApiKey = new CodeIndexOllamaEmbedder({ + ollamaModelId: "nomic-embed-text", + ollamaBaseUrl: "http://localhost:11434", + ollamaApiKey: "invalid-api-key", + }) + + // Mock 401 Unauthorized response + mockFetch.mockImplementationOnce(() => + Promise.resolve({ + ok: false, + status: 401, + } as Response), + ) + + const result = await embedderWithApiKey.validateConfiguration() + + expect(result.valid).toBe(false) + expect(result.error).toBe("embeddings:ollama.serviceUnavailable") + expect(mockFetch).toHaveBeenCalledWith( + "http://localhost:11434/api/tags", + expect.objectContaining({ + headers: { + "Content-Type": "application/json", + Authorization: "Bearer invalid-api-key", + }, + }), + ) + }) + }) + describe("validateConfiguration", () => { it("should validate successfully when service is available and model exists", async () => { // Mock successful /api/tags call @@ -323,5 +454,142 @@ describe("CodeIndexOllamaEmbedder", () => { expect(result.valid).toBe(false) expect(result.error).toBe("Network timeout") }) + + describe("createEmbeddings", () => { + it("should create embeddings successfully without API key", async () => { + const texts = ["Hello world", "Test embedding"] + + // Mock successful /api/embed call + mockFetch.mockImplementationOnce(() => + Promise.resolve({ + ok: true, + status: 200, + json: () => + Promise.resolve({ + embeddings: [ + [0.1, 0.2, 0.3], + [0.4, 0.5, 0.6], + ], + }), + } as Response), + ) + + const result = await embedder.createEmbeddings(texts) + + expect(result).toEqual({ + embeddings: [ + [0.1, 0.2, 0.3], + [0.4, 0.5, 0.6], + ], + }) + + expect(mockFetch).toHaveBeenCalledWith( + "http://localhost:11434/api/embed", + expect.objectContaining({ + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ + model: "nomic-embed-text", + input: texts, + }), + }), + ) + }) + + it("should create embeddings with API key in Authorization header", async () => { + const embedderWithApiKey = new CodeIndexOllamaEmbedder({ + ollamaModelId: "nomic-embed-text", + ollamaBaseUrl: "http://localhost:11434", + ollamaApiKey: "test-api-key-123", + }) + + const texts = ["Hello world", "Test embedding"] + + // Mock successful /api/embed call + mockFetch.mockImplementationOnce(() => + Promise.resolve({ + ok: true, + status: 200, + json: () => + Promise.resolve({ + embeddings: [ + [0.1, 0.2, 0.3], + [0.4, 0.5, 0.6], + ], + }), + } as Response), + ) + + const result = await embedderWithApiKey.createEmbeddings(texts) + + expect(result).toEqual({ + embeddings: [ + [0.1, 0.2, 0.3], + [0.4, 0.5, 0.6], + ], + }) + + // Verify Authorization header is included + expect(mockFetch).toHaveBeenCalledWith( + "http://localhost:11434/api/embed", + expect.objectContaining({ + method: "POST", + headers: { + "Content-Type": "application/json", + Authorization: "Bearer test-api-key-123", + }, + body: JSON.stringify({ + model: "nomic-embed-text", + input: texts, + }), + }), + ) + }) + + it("should handle authentication error when creating embeddings", async () => { + const embedderWithApiKey = new CodeIndexOllamaEmbedder({ + ollamaModelId: "nomic-embed-text", + ollamaBaseUrl: "http://localhost:11434", + ollamaApiKey: "invalid-api-key", + }) + + const texts = ["Hello world"] + + // Mock 401 Unauthorized response + mockFetch.mockImplementationOnce(() => + Promise.resolve({ + ok: false, + status: 401, + statusText: "Unauthorized", + } as Response), + ) + + await expect(embedderWithApiKey.createEmbeddings(texts)).rejects.toThrow( + "embeddings:ollama.embeddingFailed", + ) + + // Verify request included the API key + expect(mockFetch).toHaveBeenCalledWith( + "http://localhost:11434/api/embed", + expect.objectContaining({ + headers: { + "Content-Type": "application/json", + Authorization: "Bearer invalid-api-key", + }, + }), + ) + }) + + it("should handle network errors when creating embeddings", async () => { + const texts = ["Hello world"] + + // Mock network error + mockFetch.mockRejectedValueOnce(new Error("Network error")) + + await expect(embedder.createEmbeddings(texts)).rejects.toThrow("embeddings:ollama.embeddingFailed") + }) + }) }) }) diff --git a/src/services/code-index/embedders/ollama.ts b/src/services/code-index/embedders/ollama.ts index 9688a15ff041..d2ab66522f55 100644 --- a/src/services/code-index/embedders/ollama.ts +++ b/src/services/code-index/embedders/ollama.ts @@ -17,6 +17,7 @@ const OLLAMA_VALIDATION_TIMEOUT_MS = 30000 // 30 seconds for validation requests export class CodeIndexOllamaEmbedder implements IEmbedder { private readonly baseUrl: string private readonly defaultModelId: string + private readonly apiKey?: string constructor(options: ApiHandlerOptions) { // Ensure ollamaBaseUrl and ollamaModelId exist on ApiHandlerOptions or add defaults @@ -27,6 +28,7 @@ export class CodeIndexOllamaEmbedder implements IEmbedder { this.baseUrl = baseUrl this.defaultModelId = options.ollamaModelId || "nomic-embed-text:latest" + this.apiKey = options.ollamaApiKey } /** @@ -72,11 +74,17 @@ export class CodeIndexOllamaEmbedder implements IEmbedder { const controller = new AbortController() const timeoutId = setTimeout(() => controller.abort(), OLLAMA_EMBEDDING_TIMEOUT_MS) + // Build headers with optional API key + const headers: Record = { + "Content-Type": "application/json", + } + if (this.apiKey) { + headers["Authorization"] = `Bearer ${this.apiKey}` + } + const response = await fetch(url, { method: "POST", - headers: { - "Content-Type": "application/json", - }, + headers, body: JSON.stringify({ model: modelToUse, input: processedTexts, // Using 'input' as requested @@ -151,11 +159,17 @@ export class CodeIndexOllamaEmbedder implements IEmbedder { const controller = new AbortController() const timeoutId = setTimeout(() => controller.abort(), OLLAMA_VALIDATION_TIMEOUT_MS) + // Build headers with optional API key + const headers: Record = { + "Content-Type": "application/json", + } + if (this.apiKey) { + headers["Authorization"] = `Bearer ${this.apiKey}` + } + const modelsResponse = await fetch(modelsUrl, { method: "GET", - headers: { - "Content-Type": "application/json", - }, + headers, signal: controller.signal, }) clearTimeout(timeoutId) @@ -208,11 +222,17 @@ export class CodeIndexOllamaEmbedder implements IEmbedder { const testController = new AbortController() const testTimeoutId = setTimeout(() => testController.abort(), OLLAMA_VALIDATION_TIMEOUT_MS) + // Build headers with optional API key for test request + const testHeaders: Record = { + "Content-Type": "application/json", + } + if (this.apiKey) { + testHeaders["Authorization"] = `Bearer ${this.apiKey}` + } + const testResponse = await fetch(testUrl, { method: "POST", - headers: { - "Content-Type": "application/json", - }, + headers: testHeaders, body: JSON.stringify({ model: this.defaultModelId, input: ["test"], diff --git a/src/services/code-index/interfaces/config.ts b/src/services/code-index/interfaces/config.ts index f168e268691a..d27705b25e39 100644 --- a/src/services/code-index/interfaces/config.ts +++ b/src/services/code-index/interfaces/config.ts @@ -32,6 +32,7 @@ export type PreviousConfigSnapshot = { modelDimension?: number // Generic dimension property openAiKey?: string ollamaBaseUrl?: string + ollamaApiKey?: string openAiCompatibleBaseUrl?: string openAiCompatibleApiKey?: string geminiApiKey?: string diff --git a/src/shared/WebviewMessage.ts b/src/shared/WebviewMessage.ts index d43a2fce0434..622b000238a1 100644 --- a/src/shared/WebviewMessage.ts +++ b/src/shared/WebviewMessage.ts @@ -302,6 +302,7 @@ export interface WebviewMessage { codebaseIndexGeminiApiKey?: string codebaseIndexMistralApiKey?: string codebaseIndexVercelAiGatewayApiKey?: string + codebaseIndexOllamaApiKey?: string } } diff --git a/webview-ui/src/components/chat/CodeIndexPopover.tsx b/webview-ui/src/components/chat/CodeIndexPopover.tsx index 45bf4224a12f..0a1ca49bc1fe 100644 --- a/webview-ui/src/components/chat/CodeIndexPopover.tsx +++ b/webview-ui/src/components/chat/CodeIndexPopover.tsx @@ -73,6 +73,7 @@ interface LocalCodeIndexSettings { codebaseIndexGeminiApiKey?: string codebaseIndexMistralApiKey?: string codebaseIndexVercelAiGatewayApiKey?: string + codebaseIndexOllamaApiKey?: string } // Validation schema for codebase index settings @@ -101,6 +102,7 @@ const createValidationSchema = (provider: EmbedderProvider, t: any) => { .string() .min(1, t("settings:codeIndex.validation.ollamaBaseUrlRequired")) .url(t("settings:codeIndex.validation.invalidOllamaUrl")), + codebaseIndexOllamaApiKey: z.string().optional(), // API key is optional for Ollama codebaseIndexEmbedderModelId: z.string().min(1, t("settings:codeIndex.validation.modelIdRequired")), codebaseIndexEmbedderModelDimension: z .number() @@ -194,6 +196,7 @@ export const CodeIndexPopover: React.FC = ({ codebaseIndexGeminiApiKey: "", codebaseIndexMistralApiKey: "", codebaseIndexVercelAiGatewayApiKey: "", + codebaseIndexOllamaApiKey: "", }) // Initial settings state - stores the settings when popover opens @@ -229,6 +232,7 @@ export const CodeIndexPopover: React.FC = ({ codebaseIndexGeminiApiKey: "", codebaseIndexMistralApiKey: "", codebaseIndexVercelAiGatewayApiKey: "", + codebaseIndexOllamaApiKey: "", } setInitialSettings(settings) setCurrentSettings(settings) @@ -345,6 +349,9 @@ export const CodeIndexPopover: React.FC = ({ ? SECRET_PLACEHOLDER : "" } + if (!prev.codebaseIndexOllamaApiKey || prev.codebaseIndexOllamaApiKey === SECRET_PLACEHOLDER) { + updated.codebaseIndexOllamaApiKey = secretStatus.hasOllamaApiKey ? SECRET_PLACEHOLDER : "" + } return updated } @@ -418,7 +425,8 @@ export const CodeIndexPopover: React.FC = ({ key === "codebaseIndexOpenAiCompatibleApiKey" || key === "codebaseIndexGeminiApiKey" || key === "codebaseIndexMistralApiKey" || - key === "codebaseIndexVercelAiGatewayApiKey" + key === "codebaseIndexVercelAiGatewayApiKey" || + key === "codebaseIndexOllamaApiKey" ) { dataToValidate[key] = "placeholder-valid" } @@ -772,6 +780,28 @@ export const CodeIndexPopover: React.FC = ({ )} +
+ + + updateSetting("codebaseIndexOllamaApiKey", e.target.value) + } + placeholder={t("settings:codeIndex.ollamaApiKeyPlaceholder")} + className={cn("w-full", { + "border-red-500": formErrors.codebaseIndexOllamaApiKey, + })} + /> + {formErrors.codebaseIndexOllamaApiKey && ( +

+ {formErrors.codebaseIndexOllamaApiKey} +

+ )} +
+