diff --git a/packages/types/src/codebase-index.ts b/packages/types/src/codebase-index.ts index 61009ba3011..932253070f6 100644 --- a/packages/types/src/codebase-index.ts +++ b/packages/types/src/codebase-index.ts @@ -31,6 +31,7 @@ export const codebaseIndexConfigSchema = z.object({ "vercel-ai-gateway", "bedrock", "openrouter", + "roo", ]) .optional(), codebaseIndexEmbedderBaseUrl: z.string().optional(), @@ -67,6 +68,7 @@ export const codebaseIndexModelsSchema = z.object({ "vercel-ai-gateway": z.record(z.string(), z.object({ dimension: z.number() })).optional(), openrouter: z.record(z.string(), z.object({ dimension: z.number() })).optional(), bedrock: z.record(z.string(), z.object({ dimension: z.number() })).optional(), + roo: z.record(z.string(), z.object({ dimension: z.number() })).optional(), }) export type CodebaseIndexModels = z.infer diff --git a/packages/types/src/embedding.ts b/packages/types/src/embedding.ts index 1c5a92e1acb..31159da0010 100644 --- a/packages/types/src/embedding.ts +++ b/packages/types/src/embedding.ts @@ -6,7 +6,8 @@ export type EmbedderProvider = | "mistral" | "vercel-ai-gateway" | "bedrock" - | "openrouter" // Add other providers as needed. + | "openrouter" + | "roo" // Add other providers as needed. export interface EmbeddingModelProfile { dimension: number diff --git a/packages/types/src/vscode-extension-host.ts b/packages/types/src/vscode-extension-host.ts index 51c7fa49d5e..4adf02b8ccc 100644 --- a/packages/types/src/vscode-extension-host.ts +++ b/packages/types/src/vscode-extension-host.ts @@ -687,6 +687,7 @@ export interface WebviewMessage { | "vercel-ai-gateway" | "bedrock" | "openrouter" + | "roo" codebaseIndexEmbedderBaseUrl?: string codebaseIndexEmbedderModelId: string codebaseIndexEmbedderModelDimension?: number // Generic dimension for all providers diff --git a/src/i18n/locales/en/embeddings.json b/src/i18n/locales/en/embeddings.json index 5819e45c1a8..94eb8b2e981 100644 --- a/src/i18n/locales/en/embeddings.json +++ b/src/i18n/locales/en/embeddings.json @@ -54,6 +54,7 @@ "geminiConfigMissing": "Gemini configuration missing for embedder creation", "mistralConfigMissing": "Mistral configuration missing for embedder creation", "openRouterConfigMissing": "OpenRouter configuration missing for embedder creation", + "rooConfigMissing": "Roo Code Cloud authentication required for embedder creation. Please sign in to Roo Code Cloud.", "vercelAiGatewayConfigMissing": "Vercel AI Gateway configuration missing for embedder creation", "bedrockConfigMissing": "Amazon Bedrock configuration missing for embedder creation", "invalidEmbedderType": "Invalid embedder type configured: {{embedderProvider}}", diff --git a/src/services/code-index/__tests__/config-manager.spec.ts b/src/services/code-index/__tests__/config-manager.spec.ts index 27815c0bef0..e6e3c2c8129 100644 --- a/src/services/code-index/__tests__/config-manager.spec.ts +++ b/src/services/code-index/__tests__/config-manager.spec.ts @@ -6,6 +6,25 @@ import { PreviousConfigSnapshot } from "../interfaces/config" // Mock ContextProxy vi.mock("../../../core/config/ContextProxy") +// Mock CloudService - use vi.hoisted so variables are available when vi.mock runs +const { mockIsAuthenticated, mockCloudHasInstance } = vi.hoisted(() => ({ + mockIsAuthenticated: vi.fn().mockReturnValue(false), + mockCloudHasInstance: vi.fn().mockReturnValue(false), +})) +vi.mock("@roo-code/cloud", () => ({ + CloudService: { + hasInstance: mockCloudHasInstance, + get instance() { + return { + isAuthenticated: mockIsAuthenticated, + authService: { + getSessionToken: vi.fn().mockReturnValue("test-session-token"), + }, + } + }, + }, +})) + // Mock embeddingModels module vi.mock("../../../shared/embeddingModels") @@ -1684,6 +1703,57 @@ describe("CodeIndexConfigManager", () => { expect(configManager.isConfigured()).toBe(false) }) + it("should return true when Roo provider is authenticated and Qdrant configured", () => { + mockCloudHasInstance.mockReturnValue(true) + mockIsAuthenticated.mockReturnValue(true) + + mockContextProxy.getGlobalState.mockReturnValue({ + codebaseIndexEnabled: true, + codebaseIndexEmbedderProvider: "roo", + codebaseIndexQdrantUrl: "http://localhost:6333", + }) + mockContextProxy.getSecret.mockReturnValue(undefined) + + configManager = new CodeIndexConfigManager(mockContextProxy) + expect(configManager.isConfigured()).toBe(true) + + // Cleanup + mockCloudHasInstance.mockReturnValue(false) + mockIsAuthenticated.mockReturnValue(false) + }) + + it("should return false when Roo provider is not authenticated", () => { + mockCloudHasInstance.mockReturnValue(true) + mockIsAuthenticated.mockReturnValue(false) + + mockContextProxy.getGlobalState.mockReturnValue({ + codebaseIndexEnabled: true, + codebaseIndexEmbedderProvider: "roo", + codebaseIndexQdrantUrl: "http://localhost:6333", + }) + mockContextProxy.getSecret.mockReturnValue(undefined) + + configManager = new CodeIndexConfigManager(mockContextProxy) + expect(configManager.isConfigured()).toBe(false) + + // Cleanup + mockCloudHasInstance.mockReturnValue(false) + }) + + it("should return false when Roo provider has no CloudService instance", () => { + mockCloudHasInstance.mockReturnValue(false) + + mockContextProxy.getGlobalState.mockReturnValue({ + codebaseIndexEnabled: true, + codebaseIndexEmbedderProvider: "roo", + codebaseIndexQdrantUrl: "http://localhost:6333", + }) + mockContextProxy.getSecret.mockReturnValue(undefined) + + configManager = new CodeIndexConfigManager(mockContextProxy) + expect(configManager.isConfigured()).toBe(false) + }) + describe("currentModelDimension", () => { beforeEach(() => { vi.clearAllMocks() diff --git a/src/services/code-index/__tests__/service-factory.spec.ts b/src/services/code-index/__tests__/service-factory.spec.ts index 3e943ebd82e..1fd83a76d1c 100644 --- a/src/services/code-index/__tests__/service-factory.spec.ts +++ b/src/services/code-index/__tests__/service-factory.spec.ts @@ -5,12 +5,14 @@ import { CodeIndexOllamaEmbedder } from "../embedders/ollama" import { OpenAICompatibleEmbedder } from "../embedders/openai-compatible" import { GeminiEmbedder } from "../embedders/gemini" import { QdrantVectorStore } from "../vector-store/qdrant-client" +import { RooEmbedder } from "../embedders/roo" // Mock the embedders and vector store vitest.mock("../embedders/openai") vitest.mock("../embedders/ollama") vitest.mock("../embedders/openai-compatible") vitest.mock("../embedders/gemini") +vitest.mock("../embedders/roo") vitest.mock("../vector-store/qdrant-client") // Mock the embedding models module @@ -33,6 +35,7 @@ const MockedCodeIndexOllamaEmbedder = CodeIndexOllamaEmbedder as MockedClass const MockedGeminiEmbedder = GeminiEmbedder as MockedClass const MockedQdrantVectorStore = QdrantVectorStore as MockedClass +const MockedRooEmbedder = RooEmbedder as MockedClass // Import the mocked functions import { getDefaultModelId, getModelDimension } from "../../../shared/embeddingModels" @@ -345,6 +348,21 @@ describe("CodeIndexServiceFactory", () => { expect(() => factory.createEmbedder()).toThrow("serviceFactory.geminiConfigMissing") }) + it("should create RooEmbedder when using Roo provider", () => { + // Arrange + const testConfig = { + embedderProvider: "roo", + modelId: "text-embedding-3-small", + } + mockConfigManager.getConfig.mockReturnValue(testConfig as any) + + // Act + factory.createEmbedder() + + // Assert + expect(MockedRooEmbedder).toHaveBeenCalledWith("text-embedding-3-small") + }) + it("should throw error for invalid embedder provider", () => { // Arrange const testConfig = { diff --git a/src/services/code-index/config-manager.ts b/src/services/code-index/config-manager.ts index e7f239e621f..6d098d58d9d 100644 --- a/src/services/code-index/config-manager.ts +++ b/src/services/code-index/config-manager.ts @@ -1,3 +1,5 @@ +import { CloudService } from "@roo-code/cloud" + import { ApiHandlerOptions } from "../../shared/api" import { ContextProxy } from "../../core/config/ContextProxy" import { EmbedderProvider } from "./interfaces/manager" @@ -120,6 +122,8 @@ export class CodeIndexConfigManager { this.embedderProvider = "bedrock" } else if (codebaseIndexEmbedderProvider === "openrouter") { this.embedderProvider = "openrouter" + } else if (codebaseIndexEmbedderProvider === "roo") { + this.embedderProvider = "roo" } else { this.embedderProvider = "openai" } @@ -272,6 +276,12 @@ export class CodeIndexConfigManager { const qdrantUrl = this.qdrantUrl const isConfigured = !!(apiKey && qdrantUrl) return isConfigured + } else if (this.embedderProvider === "roo") { + // Roo Code Router uses CloudService session token for auth. + // Use isAuthenticated() for a stable auth check (not raw session token parsing). + const qdrantUrl = this.qdrantUrl + const isAuthenticated = CloudService.hasInstance() && CloudService.instance.isAuthenticated() + return !!(isAuthenticated && qdrantUrl) } return false // Should not happen if embedderProvider is always set correctly } diff --git a/src/services/code-index/embedders/__tests__/roo.spec.ts b/src/services/code-index/embedders/__tests__/roo.spec.ts new file mode 100644 index 00000000000..0db598682ac --- /dev/null +++ b/src/services/code-index/embedders/__tests__/roo.spec.ts @@ -0,0 +1,282 @@ +import type { MockedClass, MockedFunction } from "vitest" +import { describe, it, expect, beforeEach, vi, afterEach } from "vitest" +import { OpenAI } from "openai" +import { RooEmbedder } from "../roo" +import { getDefaultModelId } from "../../../../shared/embeddingModels" + +// Mock the OpenAI SDK +vi.mock("openai") + +// Mock CloudService - use vi.hoisted so variables are available when vi.mock runs +const { mockGetSessionToken, mockCloudIsAuthenticated, mockCloudHasInstance } = vi.hoisted(() => ({ + mockGetSessionToken: vi.fn().mockReturnValue("test-session-token"), + mockCloudIsAuthenticated: vi.fn().mockReturnValue(true), + mockCloudHasInstance: vi.fn().mockReturnValue(true), +})) +vi.mock("@roo-code/cloud", () => ({ + CloudService: { + hasInstance: mockCloudHasInstance, + get instance() { + return { + authService: { + getSessionToken: mockGetSessionToken, + }, + isAuthenticated: mockCloudIsAuthenticated, + } + }, + }, +})) + +// Mock TelemetryService +vi.mock("@roo-code/telemetry", () => ({ + TelemetryService: { + instance: { + captureEvent: vi.fn(), + }, + }, + TelemetryEventName: {}, +})) + +// Mock i18n +vi.mock("../../../../i18n", () => ({ + t: (key: string, params?: Record) => { + const translations: Record = { + "embeddings:validation.apiKeyRequired": "validation.apiKeyRequired", + "embeddings:authenticationFailed": + "Failed to create embeddings: Authentication failed. Please check your API key.", + "embeddings:failedWithStatus": `Failed to create embeddings after ${params?.attempts} attempts: HTTP ${params?.statusCode} - ${params?.errorMessage}`, + "embeddings:failedWithError": `Failed to create embeddings after ${params?.attempts} attempts: ${params?.errorMessage}`, + "embeddings:failedMaxAttempts": `Failed to create embeddings after ${params?.attempts} attempts`, + "embeddings:textExceedsTokenLimit": `Text at index ${params?.index} exceeds maximum token limit (${params?.itemTokens} > ${params?.maxTokens}). Skipping.`, + "embeddings:rateLimitRetry": `Rate limit hit, retrying in ${params?.delayMs}ms (attempt ${params?.attempt}/${params?.maxRetries})`, + } + return translations[key] || key + }, +})) + +const MockedOpenAI = OpenAI as MockedClass + +describe("RooEmbedder", () => { + let mockEmbeddingsCreate: MockedFunction + let mockOpenAIInstance: any + + beforeEach(() => { + vi.clearAllMocks() + vi.spyOn(console, "warn").mockImplementation(() => {}) + vi.spyOn(console, "error").mockImplementation(() => {}) + + // Setup mock OpenAI instance + mockEmbeddingsCreate = vi.fn() + mockOpenAIInstance = { + embeddings: { + create: mockEmbeddingsCreate, + }, + apiKey: "test-session-token", + } + + MockedOpenAI.mockImplementation(() => mockOpenAIInstance) + }) + + afterEach(() => { + vi.restoreAllMocks() + }) + + describe("constructor", () => { + it("should create an instance using CloudService session token", () => { + const embedder = new RooEmbedder() + expect(embedder).toBeInstanceOf(RooEmbedder) + }) + + it("should use default model when none specified", () => { + const embedder = new RooEmbedder() + const expectedDefault = getDefaultModelId("roo") + expect(expectedDefault).toBe("text-embedding-3-small") + expect(embedder.embedderInfo.name).toBe("roo") + }) + + it("should use custom model when specified", () => { + const customModel = "text-embedding-3-large" + const embedder = new RooEmbedder(customModel) + expect(embedder.embedderInfo.name).toBe("roo") + }) + + it("should initialize OpenAI client with correct headers and base URL ending in /v1", () => { + new RooEmbedder() + + // Verify client was created with correct headers + const callArgs = MockedOpenAI.mock.calls[0][0] as any + expect(callArgs.defaultHeaders).toEqual({ + "HTTP-Referer": "https://github.com/RooCodeInc/Roo-Code", + "X-Title": "Roo Code", + }) + + // Verify the baseURL ends with /v1 + expect(callArgs.baseURL).toMatch(/\/v1$/) + }) + }) + + describe("createEmbeddings", () => { + it("should create embeddings for a batch of texts", async () => { + // Create a proper base64-encoded float32 array + const float32Array = new Float32Array([0.1, 0.2, 0.3]) + const buffer = Buffer.from(float32Array.buffer) + const base64Embedding = buffer.toString("base64") + + mockEmbeddingsCreate.mockResolvedValueOnce({ + data: [{ embedding: base64Embedding }, { embedding: base64Embedding }], + usage: { prompt_tokens: 10, total_tokens: 15 }, + }) + + const embedder = new RooEmbedder() + const result = await embedder.createEmbeddings(["text1", "text2"]) + + expect(result.embeddings).toHaveLength(2) + expect(result.embeddings[0]).toHaveLength(3) // 3 floats in our test array + expect(result.usage?.promptTokens).toBe(10) + expect(result.usage?.totalTokens).toBe(15) + }) + + it("should request base64 encoding format", async () => { + const float32Array = new Float32Array([0.1]) + const buffer = Buffer.from(float32Array.buffer) + const base64Embedding = buffer.toString("base64") + + mockEmbeddingsCreate.mockResolvedValueOnce({ + data: [{ embedding: base64Embedding }], + usage: { prompt_tokens: 5, total_tokens: 5 }, + }) + + const embedder = new RooEmbedder() + await embedder.createEmbeddings(["test"]) + + expect(mockEmbeddingsCreate).toHaveBeenCalledWith( + expect.objectContaining({ + encoding_format: "base64", + }), + ) + }) + + it("should refresh session token before making request", async () => { + const float32Array = new Float32Array([0.1]) + const buffer = Buffer.from(float32Array.buffer) + const base64Embedding = buffer.toString("base64") + + mockEmbeddingsCreate.mockResolvedValueOnce({ + data: [{ embedding: base64Embedding }], + usage: { prompt_tokens: 5, total_tokens: 5 }, + }) + + const embedder = new RooEmbedder() + await embedder.createEmbeddings(["test"]) + + // The apiKey should have been refreshed via the setter + // (we can't directly verify the setter was called since it's + // a mock object, but the call to createEmbeddings succeeds) + expect(mockEmbeddingsCreate).toHaveBeenCalled() + }) + + it("should handle numeric array embeddings (non-base64)", async () => { + mockEmbeddingsCreate.mockResolvedValueOnce({ + data: [{ embedding: [0.1, 0.2, 0.3] }], + usage: { prompt_tokens: 5, total_tokens: 5 }, + }) + + const embedder = new RooEmbedder() + const result = await embedder.createEmbeddings(["test"]) + + expect(result.embeddings).toHaveLength(1) + expect(result.embeddings[0]).toEqual([0.1, 0.2, 0.3]) + }) + + it("should skip texts exceeding token limit", async () => { + const longText = "a".repeat(400000) // Exceeds MAX_ITEM_TOKENS + + const float32Array = new Float32Array([0.1]) + const buffer = Buffer.from(float32Array.buffer) + const base64Embedding = buffer.toString("base64") + + mockEmbeddingsCreate.mockResolvedValueOnce({ + data: [{ embedding: base64Embedding }], + usage: { prompt_tokens: 5, total_tokens: 5 }, + }) + + const embedder = new RooEmbedder() + const result = await embedder.createEmbeddings([longText, "short text"]) + + // Only the short text should have been embedded + expect(result.embeddings).toHaveLength(1) + }) + }) + + describe("validateConfiguration", () => { + it("should return valid when API responds correctly", async () => { + const float32Array = new Float32Array([0.1]) + const buffer = Buffer.from(float32Array.buffer) + const base64Embedding = buffer.toString("base64") + + mockEmbeddingsCreate.mockResolvedValueOnce({ + data: [{ embedding: base64Embedding }], + usage: { prompt_tokens: 1, total_tokens: 1 }, + }) + + const embedder = new RooEmbedder() + const result = await embedder.validateConfiguration() + + expect(result.valid).toBe(true) + }) + + it("should return invalid when API returns empty response", async () => { + mockEmbeddingsCreate.mockResolvedValueOnce({ + data: [], + usage: { prompt_tokens: 0, total_tokens: 0 }, + }) + + const embedder = new RooEmbedder() + const result = await embedder.validateConfiguration() + + expect(result.valid).toBe(false) + }) + + it("should handle API errors during validation", async () => { + const apiError = new Error("API connection failed") + ;(apiError as any).status = 500 + + mockEmbeddingsCreate.mockRejectedValueOnce(apiError) + + const embedder = new RooEmbedder() + const result = await embedder.validateConfiguration() + + expect(result.valid).toBe(false) + }) + }) + + describe("embedderInfo", () => { + it("should return roo as the embedder name", () => { + const embedder = new RooEmbedder() + expect(embedder.embedderInfo).toEqual({ name: "roo" }) + }) + }) + + describe("rate limiting", () => { + it("should retry on 429 errors with exponential backoff", async () => { + const rateLimitError = new Error("Rate limit exceeded") + ;(rateLimitError as any).status = 429 + + const float32Array = new Float32Array([0.1]) + const buffer = Buffer.from(float32Array.buffer) + const base64Embedding = buffer.toString("base64") + + // First call fails with 429, second succeeds + mockEmbeddingsCreate.mockRejectedValueOnce(rateLimitError).mockResolvedValueOnce({ + data: [{ embedding: base64Embedding }], + usage: { prompt_tokens: 5, total_tokens: 5 }, + }) + + const embedder = new RooEmbedder() + const result = await embedder.createEmbeddings(["test"]) + + expect(result.embeddings).toHaveLength(1) + expect(mockEmbeddingsCreate).toHaveBeenCalledTimes(2) + }) + }) +}) diff --git a/src/services/code-index/embedders/roo.ts b/src/services/code-index/embedders/roo.ts new file mode 100644 index 00000000000..e638ecda469 --- /dev/null +++ b/src/services/code-index/embedders/roo.ts @@ -0,0 +1,393 @@ +import { OpenAI } from "openai" +import { CloudService } from "@roo-code/cloud" +import { IEmbedder, EmbeddingResponse, EmbedderInfo } from "../interfaces/embedder" +import { + MAX_BATCH_TOKENS, + MAX_ITEM_TOKENS, + MAX_BATCH_RETRIES as MAX_RETRIES, + INITIAL_RETRY_DELAY_MS as INITIAL_DELAY_MS, +} from "../constants" +import { getDefaultModelId, getModelQueryPrefix } from "../../../shared/embeddingModels" +import { t } from "../../../i18n" +import { withValidationErrorHandling, HttpError, formatEmbeddingError } from "../shared/validation-helpers" +import { TelemetryEventName } from "@roo-code/types" +import { TelemetryService } from "@roo-code/telemetry" +import { Mutex } from "async-mutex" +import { handleOpenAIError } from "../../../api/providers/utils/openai-error-handler" + +interface EmbeddingItem { + embedding: string | number[] + [key: string]: any +} + +interface RooEmbeddingResponse { + data: EmbeddingItem[] + usage?: { + prompt_tokens?: number + total_tokens?: number + } +} + +/** + * Returns the current session token from CloudService, or "unauthenticated" if unavailable. + */ +function getSessionToken(): string { + const token = CloudService.hasInstance() ? CloudService.instance.authService?.getSessionToken() : undefined + return token ?? "unauthenticated" +} + +/** + * Returns the base URL for the Roo Code Router proxy, with /v1 suffix. + */ +function getRooBaseUrl(): string { + let baseURL = process.env.ROO_CODE_PROVIDER_URL ?? "https://api.roocode.com/proxy" + if (!baseURL.endsWith("/v1")) { + baseURL = `${baseURL}/v1` + } + return baseURL +} + +/** + * Roo Code Router implementation of the embedder interface with batching and rate limiting. + * Uses CloudService session token for authentication against the Roo Code proxy, + * which forwards embedding requests to the underlying model provider (e.g. OpenAI). + * No third-party API key is required -- users authenticate via Roo Code Cloud. + */ +export class RooEmbedder implements IEmbedder { + private embeddingsClient: OpenAI + private readonly defaultModelId: string + private readonly maxItemTokens: number + + // Global rate limiting state shared across all instances + private static globalRateLimitState = { + isRateLimited: false, + rateLimitResetTime: 0, + consecutiveRateLimitErrors: 0, + lastRateLimitError: 0, + mutex: new Mutex(), + } + + /** + * Creates a new Roo Code Router embedder. + * Authentication is handled via CloudService session token. + * @param modelId Optional model identifier (defaults to "text-embedding-3-small") + * @param maxItemTokens Optional maximum tokens per item (defaults to MAX_ITEM_TOKENS) + */ + constructor(modelId?: string, maxItemTokens?: number) { + const sessionToken = getSessionToken() + const baseURL = getRooBaseUrl() + + try { + this.embeddingsClient = new OpenAI({ + baseURL, + apiKey: sessionToken, + defaultHeaders: { + "HTTP-Referer": "https://github.com/RooCodeInc/Roo-Code", + "X-Title": "Roo Code", + }, + }) + } catch (error) { + throw handleOpenAIError(error, "Roo Code Cloud") + } + + this.defaultModelId = modelId || getDefaultModelId("roo") + this.maxItemTokens = maxItemTokens || MAX_ITEM_TOKENS + } + + /** + * Creates embeddings for the given texts with batching and rate limiting. + * Refreshes the session token before each top-level call to ensure freshness. + * @param texts Array of text strings to embed + * @param model Optional model identifier + * @returns Promise resolving to embedding response + */ + async createEmbeddings(texts: string[], model?: string): Promise { + // Refresh API key to use the latest session token + this.embeddingsClient.apiKey = getSessionToken() + + const modelToUse = model || this.defaultModelId + + // Apply model-specific query prefix if required + const queryPrefix = getModelQueryPrefix("roo", modelToUse) + const processedTexts = queryPrefix + ? texts.map((text, index) => { + if (text.startsWith(queryPrefix)) { + return text + } + const prefixedText = `${queryPrefix}${text}` + const estimatedTokens = Math.ceil(prefixedText.length / 4) + if (estimatedTokens > MAX_ITEM_TOKENS) { + console.warn( + t("embeddings:textWithPrefixExceedsTokenLimit", { + index, + estimatedTokens, + maxTokens: MAX_ITEM_TOKENS, + }), + ) + return text + } + return prefixedText + }) + : texts + + const allEmbeddings: number[][] = [] + const usage = { promptTokens: 0, totalTokens: 0 } + const remainingTexts = [...processedTexts] + + while (remainingTexts.length > 0) { + const currentBatch: string[] = [] + let currentBatchTokens = 0 + const processedIndices: number[] = [] + + for (let i = 0; i < remainingTexts.length; i++) { + const text = remainingTexts[i] + const itemTokens = Math.ceil(text.length / 4) + + if (itemTokens > this.maxItemTokens) { + console.warn( + t("embeddings:textExceedsTokenLimit", { + index: i, + itemTokens, + maxTokens: this.maxItemTokens, + }), + ) + processedIndices.push(i) + continue + } + + if (currentBatchTokens + itemTokens <= MAX_BATCH_TOKENS) { + currentBatch.push(text) + currentBatchTokens += itemTokens + processedIndices.push(i) + } else { + break + } + } + + // Remove processed items from remainingTexts (in reverse order to maintain correct indices) + for (let i = processedIndices.length - 1; i >= 0; i--) { + remainingTexts.splice(processedIndices[i], 1) + } + + if (currentBatch.length > 0) { + const batchResult = await this._embedBatchWithRetries(currentBatch, modelToUse) + allEmbeddings.push(...batchResult.embeddings) + usage.promptTokens += batchResult.usage.promptTokens + usage.totalTokens += batchResult.usage.totalTokens + } + } + + return { embeddings: allEmbeddings, usage } + } + + /** + * Helper method to handle batch embedding with retries and exponential backoff. + * Refreshes the session token on each retry attempt. + */ + private async _embedBatchWithRetries( + batchTexts: string[], + model: string, + ): Promise<{ embeddings: number[][]; usage: { promptTokens: number; totalTokens: number } }> { + for (let attempts = 0; attempts < MAX_RETRIES; attempts++) { + await this.waitForGlobalRateLimit() + + // Refresh token on each attempt + this.embeddingsClient.apiKey = getSessionToken() + + try { + const requestParams: any = { + input: batchTexts, + model: model, + // Request base64 encoding to bypass the OpenAI package's parser + // which truncates embedding dimensions to 256 for numeric arrays. + encoding_format: "base64", + } + + const response = (await this.embeddingsClient.embeddings.create(requestParams)) as RooEmbeddingResponse + + // Convert base64 embeddings to float32 arrays + const processedEmbeddings = response.data.map((item: EmbeddingItem) => { + if (typeof item.embedding === "string") { + const buffer = Buffer.from(item.embedding, "base64") + const float32Array = new Float32Array(buffer.buffer, buffer.byteOffset, buffer.byteLength / 4) + return { + ...item, + embedding: Array.from(float32Array), + } + } + return item + }) + + response.data = processedEmbeddings + + const embeddings = response.data.map((item) => item.embedding as number[]) + + return { + embeddings, + usage: { + promptTokens: response.usage?.prompt_tokens || 0, + totalTokens: response.usage?.total_tokens || 0, + }, + } + } catch (error) { + TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { + error: error instanceof Error ? error.message : String(error), + stack: error instanceof Error ? error.stack : undefined, + location: "RooEmbedder:_embedBatchWithRetries", + attempt: attempts + 1, + }) + + const hasMoreAttempts = attempts < MAX_RETRIES - 1 + + const httpError = error as HttpError + if (httpError?.status === 429) { + await this.updateGlobalRateLimitState(httpError) + + if (hasMoreAttempts) { + const baseDelay = INITIAL_DELAY_MS * Math.pow(2, attempts) + const globalDelay = await this.getGlobalRateLimitDelay() + const delayMs = Math.max(baseDelay, globalDelay) + + console.warn( + t("embeddings:rateLimitRetry", { + delayMs, + attempt: attempts + 1, + maxRetries: MAX_RETRIES, + }), + ) + await new Promise((resolve) => setTimeout(resolve, delayMs)) + continue + } + } + + console.error(`Roo embedder error (attempt ${attempts + 1}/${MAX_RETRIES}):`, error) + throw formatEmbeddingError(error, MAX_RETRIES) + } + } + + throw new Error(t("embeddings:failedMaxAttempts", { attempts: MAX_RETRIES })) + } + + /** + * Validates the Roo embedder configuration by testing API connectivity. + */ + async validateConfiguration(): Promise<{ valid: boolean; error?: string }> { + return withValidationErrorHandling(async () => { + // Refresh token before validation + this.embeddingsClient.apiKey = getSessionToken() + + try { + const testTexts = ["test"] + const modelToUse = this.defaultModelId + + const requestParams: any = { + input: testTexts, + model: modelToUse, + encoding_format: "base64", + } + + const response = (await this.embeddingsClient.embeddings.create(requestParams)) as RooEmbeddingResponse + + if (!response?.data || response.data.length === 0) { + return { + valid: false, + error: "embeddings:validation.invalidResponse", + } + } + + return { valid: true } + } catch (error) { + TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { + error: error instanceof Error ? error.message : String(error), + stack: error instanceof Error ? error.stack : undefined, + location: "RooEmbedder:validateConfiguration", + }) + throw error + } + }, "roo") + } + + /** + * Returns information about this embedder. + */ + get embedderInfo(): EmbedderInfo { + return { + name: "roo", + } + } + + /** + * Waits if there's an active global rate limit. + */ + private async waitForGlobalRateLimit(): Promise { + const release = await RooEmbedder.globalRateLimitState.mutex.acquire() + let mutexReleased = false + + try { + const state = RooEmbedder.globalRateLimitState + + if (state.isRateLimited && state.rateLimitResetTime > Date.now()) { + const waitTime = state.rateLimitResetTime - Date.now() + release() + mutexReleased = true + await new Promise((resolve) => setTimeout(resolve, waitTime)) + return + } + + if (state.isRateLimited && state.rateLimitResetTime <= Date.now()) { + state.isRateLimited = false + state.consecutiveRateLimitErrors = 0 + } + } finally { + if (!mutexReleased) { + release() + } + } + } + + /** + * Updates global rate limit state when a 429 error occurs. + */ + private async updateGlobalRateLimitState(error: HttpError): Promise { + const release = await RooEmbedder.globalRateLimitState.mutex.acquire() + try { + const state = RooEmbedder.globalRateLimitState + const now = Date.now() + + if (now - state.lastRateLimitError < 60000) { + state.consecutiveRateLimitErrors++ + } else { + state.consecutiveRateLimitErrors = 1 + } + + state.lastRateLimitError = now + + const baseDelay = 5000 + const maxDelay = 300000 + const exponentialDelay = Math.min(baseDelay * Math.pow(2, state.consecutiveRateLimitErrors - 1), maxDelay) + + state.isRateLimited = true + state.rateLimitResetTime = now + exponentialDelay + } finally { + release() + } + } + + /** + * Gets the current global rate limit delay. + */ + private async getGlobalRateLimitDelay(): Promise { + const release = await RooEmbedder.globalRateLimitState.mutex.acquire() + try { + const state = RooEmbedder.globalRateLimitState + + if (state.isRateLimited && state.rateLimitResetTime > Date.now()) { + return state.rateLimitResetTime - Date.now() + } + + return 0 + } finally { + release() + } + } +} diff --git a/src/services/code-index/interfaces/embedder.ts b/src/services/code-index/interfaces/embedder.ts index aa8515d6986..73fcacb6ad7 100644 --- a/src/services/code-index/interfaces/embedder.ts +++ b/src/services/code-index/interfaces/embedder.ts @@ -37,6 +37,7 @@ export type AvailableEmbedders = | "vercel-ai-gateway" | "bedrock" | "openrouter" + | "roo" export interface EmbedderInfo { name: AvailableEmbedders diff --git a/src/services/code-index/interfaces/manager.ts b/src/services/code-index/interfaces/manager.ts index 28ff5523277..ae70c0a3a49 100644 --- a/src/services/code-index/interfaces/manager.ts +++ b/src/services/code-index/interfaces/manager.ts @@ -79,6 +79,7 @@ export type EmbedderProvider = | "vercel-ai-gateway" | "bedrock" | "openrouter" + | "roo" export interface IndexProgressUpdate { systemStatus: IndexingState diff --git a/src/services/code-index/service-factory.ts b/src/services/code-index/service-factory.ts index d23eff4810b..6213ec9f1ea 100644 --- a/src/services/code-index/service-factory.ts +++ b/src/services/code-index/service-factory.ts @@ -20,6 +20,7 @@ import { MistralEmbedder } from "./embedders/mistral" import { VercelAiGatewayEmbedder } from "./embedders/vercel-ai-gateway" import { BedrockEmbedder } from "./embedders/bedrock" import { OpenRouterEmbedder } from "./embedders/openrouter" +import { RooEmbedder } from "./embedders/roo" import { QdrantVectorStore } from "./vector-store/qdrant-client" import { codeParser, DirectoryScanner, FileWatcher } from "./processors" import { ICodeParser, IEmbedder, IFileWatcher, IVectorStore } from "./interfaces" @@ -103,6 +104,9 @@ export class CodeIndexServiceFactory { undefined, // maxItemTokens config.openRouterOptions.specificProvider, ) + } else if (provider === "roo") { + // Roo Code Router uses CloudService session token -- no API key needed + return new RooEmbedder(config.modelId) } throw new Error( diff --git a/src/shared/embeddingModels.ts b/src/shared/embeddingModels.ts index 0b59c5b4b28..78d73e23167 100644 --- a/src/shared/embeddingModels.ts +++ b/src/shared/embeddingModels.ts @@ -85,6 +85,11 @@ export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = { "qwen/qwen3-embedding-4b": { dimension: 2560, scoreThreshold: 0.4 }, "qwen/qwen3-embedding-8b": { dimension: 4096, scoreThreshold: 0.4 }, }, + roo: { + // Roo Code Router proxies OpenAI embedding models + "text-embedding-3-small": { dimension: 1536, scoreThreshold: 0.4 }, + "text-embedding-3-large": { dimension: 3072, scoreThreshold: 0.4 }, + }, } /** @@ -183,6 +188,9 @@ export function getDefaultModelId(provider: EmbedderProvider): string { case "openrouter": return "openai/text-embedding-3-large" + case "roo": + return "text-embedding-3-small" + default: // Fallback for unknown providers console.warn(`Unknown provider for default model ID: ${provider}. Falling back to OpenAI default.`) diff --git a/webview-ui/src/components/chat/CodeIndexPopover.tsx b/webview-ui/src/components/chat/CodeIndexPopover.tsx index 4fcf6406e3b..63daffa7f92 100644 --- a/webview-ui/src/components/chat/CodeIndexPopover.tsx +++ b/webview-ui/src/components/chat/CodeIndexPopover.tsx @@ -176,6 +176,14 @@ const createValidationSchema = (provider: EmbedderProvider, t: any) => { .min(1, t("settings:codeIndex.validation.modelSelectionRequired")), }) + case "roo": + // Roo Code Cloud uses session token auth -- no API key needed + return baseSchema.extend({ + codebaseIndexEmbedderModelId: z + .string() + .min(1, t("settings:codeIndex.validation.modelSelectionRequired")), + }) + default: return baseSchema } @@ -187,7 +195,8 @@ export const CodeIndexPopover: React.FC = ({ }) => { const SECRET_PLACEHOLDER = "••••••••••••••••" const { t } = useAppTranslation() - const { codebaseIndexConfig, codebaseIndexModels, cwd, apiConfiguration } = useExtensionState() + const { codebaseIndexConfig, codebaseIndexModels, cwd, apiConfiguration, cloudIsAuthenticated } = + useExtensionState() const [open, setOpen] = useState(false) const [isAdvancedSettingsOpen, setIsAdvancedSettingsOpen] = useState(false) const [isSetupSettingsOpen, setIsSetupSettingsOpen] = useState(false) @@ -761,6 +770,9 @@ export const CodeIndexPopover: React.FC = ({ {t("settings:codeIndex.openRouterProvider")} + + {t("settings:codeIndex.rooCodeCloudProvider")} + @@ -1430,6 +1442,62 @@ export const CodeIndexPopover: React.FC = ({ )} + {currentSettings.codebaseIndexEmbedderProvider === "roo" && ( + <> + {/* Auth status for Roo Code Cloud */} +
+ {cloudIsAuthenticated ? ( +

+ {t("settings:codeIndex.rooAuthenticated")} +

+ ) : ( +

+ {t("settings:codeIndex.rooAuthenticationRequired")} +

+ )} +
+ +
+ + + updateSetting("codebaseIndexEmbedderModelId", e.target.value) + } + className={cn("w-full", { + "border-red-500": formErrors.codebaseIndexEmbedderModelId, + })}> + + {t("settings:codeIndex.selectModel")} + + {getAvailableModels().map((modelId) => { + const model = + codebaseIndexModels?.[ + currentSettings.codebaseIndexEmbedderProvider as keyof typeof codebaseIndexModels + ]?.[modelId] + return ( + + {modelId}{" "} + {model + ? t("settings:codeIndex.modelDimensions", { + dimension: model.dimension, + }) + : ""} + + ) + })} + + {formErrors.codebaseIndexEmbedderModelId && ( +

+ {formErrors.codebaseIndexEmbedderModelId} +

+ )} +
+ + )} + {/* Qdrant Settings */}