Switched GetModelsOptions to a discriminated union based on api provider. Adjusted related code and tests to match

Author: slytechnical

src/api/providers/__tests__/litellm.test.ts

@@ -77,11 +77,11 @@ describe("LiteLLMHandler", () => {
 		it("returns correct model info when modelId is provided and found in getModels", async () => {
 			const handler = new LiteLLMHandler(defaultMockOptions)
 			const result = await handler.fetchModel()
-			expect(mockGetModels).toHaveBeenCalledWith(
-				"litellm",
-				defaultMockOptions.litellmApiKey,
-				defaultMockOptions.litellmBaseUrl,
-			)
+			expect(mockGetModels).toHaveBeenCalledWith({
+				provider: "litellm",
+				apiKey: defaultMockOptions.litellmApiKey,
+				baseUrl: defaultMockOptions.litellmBaseUrl,
+			})
 			expect(result).toEqual({ id: defaultMockOptions.litellmModelId, info: mockModelInfo })
 		})
 

src/api/providers/fetchers/modelCache.ts

@@ -5,7 +5,7 @@ import NodeCache from "node-cache"
 
 import { ContextProxy } from "../../../core/config/ContextProxy"
 import { getCacheDirectoryPath } from "../../../shared/storagePathManager"
-import { RouterName, ModelRecord } from "../../../shared/api"
+import { RouterName, ModelRecord, GetModelsOptions } from "../../../shared/api"
 import { fileExistsAtPath } from "../../../utils/fs"
 
 import { getOpenRouterModels } from "./openrouter"
@@ -30,18 +30,6 @@ async function readModels(router: RouterName): Promise<ModelRecord | undefined>
 	return exists ? JSON.parse(await fs.readFile(filePath, "utf8")) : undefined
 }
 
-/**
- * Options for fetching models from different routers.
- * This is a discriminated union type where the router property determines
- * which other properties are required.
- */
-export type GetModelsOptions =
-	| { router: "openrouter" }
-	| { router: "glama" }
-	| { router: "requesty"; apiKey?: string }
-	| { router: "unbound"; apiKey?: string }
-	| { router: "litellm"; apiKey: string; baseUrl: string }
-
 /**
  * Get models from the cache or fetch them from the provider and cache them.
  * There are two caches:
@@ -52,14 +40,14 @@ export type GetModelsOptions =
  * @returns The models from the cache or the fetched models.
  */
 export const getModels = async (options: GetModelsOptions): Promise<ModelRecord> => {
-	const { router } = options
-	let models = memoryCache.get<ModelRecord>(router)
+	const { provider } = options
+	let models = memoryCache.get<ModelRecord>(provider)
 	if (models) {
 		return models
 	}
 
 	try {
-		switch (router) {
+		switch (provider) {
 			case "openrouter":
 				models = await getOpenRouterModels()
 				break
@@ -80,26 +68,26 @@ export const getModels = async (options: GetModelsOptions): Promise<ModelRecord>
 				break
 			default:
 				// Ensures router is exhaustively checked if RouterName is a strict union
-				const exhaustiveCheck: never = router
+				const exhaustiveCheck: never = provider
 				throw new Error(`Unknown router: ${exhaustiveCheck}`)
 		}
 
 		// Cache the fetched models (even if empty, to signify a successful fetch with no models)
-		memoryCache.set(router, models)
-		await writeModels(router, models).catch((err) =>
-			console.error(`[getModels] Error writing ${router} models to file cache:`, err),
+		memoryCache.set(provider, models)
+		await writeModels(provider, models).catch((err) =>
+			console.error(`[getModels] Error writing ${provider} models to file cache:`, err),
 		)
 
 		try {
-			models = await readModels(router)
+			models = await readModels(provider)
 			// console.log(`[getModels] read ${router} models from file cache`)
 		} catch (error) {
-			console.error(`[getModels] error reading ${router} models from file cache`, error)
+			console.error(`[getModels] error reading ${provider} models from file cache`, error)
 		}
 		return models || {}
 	} catch (error) {
 		// Log the error and re-throw it so the caller can handle it (e.g., show a UI message).
-		console.error(`[getModels] Failed to fetch models for ${router}:`, error)
+		console.error(`[getModels] Failed to fetch models in modelCache for ${provider}:`, error)
 
 		throw error // Re-throw the original error to be handled by the caller.
 	}

src/api/providers/litellm.ts

@@ -19,7 +19,7 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa
 			options,
 			name: "litellm",
 			baseURL: `${options.litellmBaseUrl || "http://localhost:4000"}`,
-			apiKey: options.litellmApiKey || "dummy-key",
+			apiKey: options.litellmApiKey || "sk-1234",
 			modelId: options.litellmModelId,
 			defaultModelId: litellmDefaultModelId,
 			defaultModelInfo: litellmDefaultModelInfo,

src/api/providers/openrouter.ts

@@ -171,7 +171,7 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH
 
 	public async fetchModel() {
 		const [models, endpoints] = await Promise.all([
-			getModels("openrouter"),
+			getModels({ provider: "openrouter" }),
 			getModelEndpoints({
 				router: "openrouter",
 				modelId: this.options.openRouterModelId,

src/api/providers/requesty.ts

@@ -45,7 +45,7 @@ export class RequestyHandler extends BaseProvider implements SingleCompletionHan
 	}
 
 	public async fetchModel() {
-		this.models = await getModels("requesty")
+		this.models = await getModels({ provider: "requesty", apiKey: this.options.requestyApiKey })
 		return this.getModel()
 	}
 

src/api/providers/router-provider.ts

@@ -1,6 +1,6 @@
 import OpenAI from "openai"
 
-import { ApiHandlerOptions, RouterName, ModelRecord, ModelInfo } from "../../shared/api"
+import { ApiHandlerOptions, RouterName, ModelRecord, ModelInfo, GetModelsOptions } from "../../shared/api"
 import { BaseProvider } from "./base-provider"
 import { getModels } from "./fetchers/modelCache"
 
@@ -51,7 +51,31 @@ export abstract class RouterProvider extends BaseProvider {
 	}
 
 	public async fetchModel() {
-		this.models = await getModels(this.name, this.apiKey, this.baseURL)
+		// Create the appropriate options based on router type
+		let options: GetModelsOptions
+
+		switch (this.name) {
+			case "openrouter":
+				options = { provider: "openrouter" }
+				break
+			case "glama":
+				options = { provider: "glama" }
+				break
+			case "requesty":
+				options = { provider: "requesty", apiKey: this.apiKey }
+				break
+			case "unbound":
+				options = { provider: "unbound", apiKey: this.apiKey }
+				break
+			case "litellm":
+				options = { provider: "litellm", apiKey: this.apiKey, baseUrl: this.baseURL }
+				break
+			default:
+				const exhaustiveCheck: never = this.name
+				throw new Error(`Unknown provider: ${exhaustiveCheck}`)
+		}
+
+		this.models = await getModels(options)
 		return this.getModel()
 	}
 

src/core/webview/__tests__/webviewMessageHandler.test.ts

@@ -33,10 +33,11 @@ describe("webviewMessageHandler", () => {
 	describe("requestRouterModels", () => {
 		test("handles all successful model fetches correctly", async () => {
 			// Mock all getModels calls to succeed with different data
-			;(getModels as jest.Mock).mockImplementation((router) => {
+			;(getModels as jest.Mock).mockImplementation((options) => {
+				const provider = options.provider
 				return Promise.resolve({
-					[`${router}-model-1`]: { name: `${router} Model 1` },
-					[`${router}-model-2`]: { name: `${router} Model 2` },
+					[`${provider}-model-1`]: { name: `${provider} Model 1` },
+					[`${provider}-model-2`]: { name: `${provider} Model 2` },
 				})
 			})
 
@@ -75,14 +76,15 @@ describe("webviewMessageHandler", () => {
 
 		test("handles some failed model fetches correctly", async () => {
 			// Mock some getModels calls to succeed and others to fail
-			;(getModels as jest.Mock).mockImplementation((router) => {
-				if (router === "openrouter" || router === "litellm") {
+			;(getModels as jest.Mock).mockImplementation((options) => {
+				const provider = options.provider
+				if (provider === "openrouter" || provider === "litellm") {
 					return Promise.resolve({
-						[`${router}-model-1`]: { name: `${router} Model 1` },
+						[`${provider}-model-1`]: { name: `${provider} Model 1` },
 					})
 				}
-				// For other routers, throw an error
-				return Promise.reject(new Error(`Failed to fetch ${router} models`))
+				// For other providers, throw an error
+				return Promise.reject(new Error(`Failed to fetch ${provider} models`))
 			})
 
 			// Call the handler
@@ -129,4 +131,46 @@ describe("webviewMessageHandler", () => {
 			})
 		})
 	})
+
+	describe("requestProviderModels", () => {
+		test("when getModels succeeds, it posts a providerModelsResponse with models", async () => {
+			const mockLiteLLMModels = { "litellm-model-1": { name: "LiteLLM Model 1" } }
+			;(getModels as jest.Mock).mockResolvedValueOnce(mockLiteLLMModels)
+
+			await webviewMessageHandler(mockProvider as any, {
+				type: "requestProviderModels",
+				payload: { provider: "litellm", apiKey: "test-key", baseUrl: "test-url" },
+			})
+
+			expect(mockProvider.postMessageToWebview).toHaveBeenCalledWith({
+				type: "providerModelsResponse",
+				payload: {
+					provider: "litellm",
+					models: mockLiteLLMModels,
+					error: undefined, // Explicitly check error is undefined on success
+				},
+			})
+			expect(getModels).toHaveBeenCalledWith({ provider: "litellm", apiKey: "test-key", baseUrl: "test-url" })
+		})
+
+		test("when getModels fails, it posts a providerModelsResponse with an error and empty models", async () => {
+			const errorMessage = "Failed to fetch LiteLLM models: No response from server."
+			;(getModels as jest.Mock).mockRejectedValueOnce(new Error(errorMessage))
+
+			await webviewMessageHandler(mockProvider as any, {
+				type: "requestProviderModels",
+				payload: { provider: "litellm", apiKey: "test-key", baseUrl: "test-url" },
+			})
+
+			expect(mockProvider.postMessageToWebview).toHaveBeenCalledWith({
+				type: "providerModelsResponse",
+				payload: {
+					provider: "litellm",
+					models: {},
+					error: errorMessage,
+				},
+			})
+			expect(getModels).toHaveBeenCalledWith({ provider: "litellm", apiKey: "test-key", baseUrl: "test-url" })
+		})
+	})
 })