Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -36,9 +36,10 @@ vi.mock("vscode", () => ({

// Mock modelCache getModels/flushModels used by the handler
const getModelsMock = vi.fn()
const flushModelsMock = vi.fn()
vi.mock("../../../api/providers/fetchers/modelCache", () => ({
getModels: (...args: any[]) => getModelsMock(...args),
flushModels: vi.fn(),
flushModels: (...args: any[]) => flushModelsMock(...args),
}))

describe("webviewMessageHandler - requestRouterModels provider filter", () => {
Expand Down Expand Up @@ -164,4 +165,60 @@ describe("webviewMessageHandler - requestRouterModels provider filter", () => {
const providersCalled = getModelsMock.mock.calls.map((c: any[]) => c[0]?.provider)
expect(providersCalled).toEqual(["openrouter"])
})

it("flushes cache when LiteLLM credentials are provided in message values", async () => {
// Provide LiteLLM credentials via message.values (simulating Refresh Models button)
await webviewMessageHandler(
mockProvider as any,
{
type: "requestRouterModels",
values: {
litellmApiKey: "test-api-key",
litellmBaseUrl: "http://localhost:4000",
},
} as any,
)

// flushModels should have been called for litellm with refresh=true
expect(flushModelsMock).toHaveBeenCalledWith("litellm", true)

// getModels should have been called with the provided credentials
const litellmCalls = getModelsMock.mock.calls.filter((c: any[]) => c[0]?.provider === "litellm")
expect(litellmCalls.length).toBe(1)
expect(litellmCalls[0][0]).toEqual({
provider: "litellm",
apiKey: "test-api-key",
baseUrl: "http://localhost:4000",
})
})

it("does not flush cache when using stored LiteLLM credentials", async () => {
// Provide stored credentials via apiConfiguration
mockProvider.getState.mockResolvedValue({
apiConfiguration: {
litellmApiKey: "stored-api-key",
litellmBaseUrl: "http://stored:4000",
},
})

await webviewMessageHandler(
mockProvider as any,
{
type: "requestRouterModels",
} as any,
)

// flushModels should NOT have been called for litellm
const litellmFlushCalls = flushModelsMock.mock.calls.filter((c: any[]) => c[0] === "litellm")
expect(litellmFlushCalls.length).toBe(0)

// getModels should still have been called with stored credentials
const litellmCalls = getModelsMock.mock.calls.filter((c: any[]) => c[0]?.provider === "litellm")
expect(litellmCalls.length).toBe(1)
expect(litellmCalls[0][0]).toEqual({
provider: "litellm",
apiKey: "stored-api-key",
baseUrl: "http://stored:4000",
})
})
})
6 changes: 6 additions & 0 deletions src/core/webview/webviewMessageHandler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -882,6 +882,12 @@ export const webviewMessageHandler = async (
const litellmBaseUrl = apiConfiguration.litellmBaseUrl || message?.values?.litellmBaseUrl

if (litellmApiKey && litellmBaseUrl) {
// If explicit credentials are provided in message.values (from Refresh Models button),
// flush the cache first to ensure we fetch fresh data with the new credentials
if (message?.values?.litellmApiKey || message?.values?.litellmBaseUrl) {
await flushModels("litellm", true)
}

candidates.push({
key: "litellm",
options: { provider: "litellm", apiKey: litellmApiKey, baseUrl: litellmBaseUrl },
Expand Down
Loading