Skip to content

Commit a20d047

Browse files
authored
fix: flush LiteLLM cache when credentials change on refresh (#9536)
1 parent 3ac5bec commit a20d047

File tree

2 files changed

+64
-1
lines changed

2 files changed

+64
-1
lines changed

src/core/webview/__tests__/webviewMessageHandler.routerModels.spec.ts

Lines changed: 58 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,10 @@ vi.mock("vscode", () => ({
3636

3737
// Mock modelCache getModels/flushModels used by the handler
3838
const getModelsMock = vi.fn()
39+
const flushModelsMock = vi.fn()
3940
vi.mock("../../../api/providers/fetchers/modelCache", () => ({
4041
getModels: (...args: any[]) => getModelsMock(...args),
41-
flushModels: vi.fn(),
42+
flushModels: (...args: any[]) => flushModelsMock(...args),
4243
}))
4344

4445
describe("webviewMessageHandler - requestRouterModels provider filter", () => {
@@ -164,4 +165,60 @@ describe("webviewMessageHandler - requestRouterModels provider filter", () => {
164165
const providersCalled = getModelsMock.mock.calls.map((c: any[]) => c[0]?.provider)
165166
expect(providersCalled).toEqual(["openrouter"])
166167
})
168+
169+
it("flushes cache when LiteLLM credentials are provided in message values", async () => {
170+
// Provide LiteLLM credentials via message.values (simulating Refresh Models button)
171+
await webviewMessageHandler(
172+
mockProvider as any,
173+
{
174+
type: "requestRouterModels",
175+
values: {
176+
litellmApiKey: "test-api-key",
177+
litellmBaseUrl: "http://localhost:4000",
178+
},
179+
} as any,
180+
)
181+
182+
// flushModels should have been called for litellm with refresh=true
183+
expect(flushModelsMock).toHaveBeenCalledWith("litellm", true)
184+
185+
// getModels should have been called with the provided credentials
186+
const litellmCalls = getModelsMock.mock.calls.filter((c: any[]) => c[0]?.provider === "litellm")
187+
expect(litellmCalls.length).toBe(1)
188+
expect(litellmCalls[0][0]).toEqual({
189+
provider: "litellm",
190+
apiKey: "test-api-key",
191+
baseUrl: "http://localhost:4000",
192+
})
193+
})
194+
195+
it("does not flush cache when using stored LiteLLM credentials", async () => {
196+
// Provide stored credentials via apiConfiguration
197+
mockProvider.getState.mockResolvedValue({
198+
apiConfiguration: {
199+
litellmApiKey: "stored-api-key",
200+
litellmBaseUrl: "http://stored:4000",
201+
},
202+
})
203+
204+
await webviewMessageHandler(
205+
mockProvider as any,
206+
{
207+
type: "requestRouterModels",
208+
} as any,
209+
)
210+
211+
// flushModels should NOT have been called for litellm
212+
const litellmFlushCalls = flushModelsMock.mock.calls.filter((c: any[]) => c[0] === "litellm")
213+
expect(litellmFlushCalls.length).toBe(0)
214+
215+
// getModels should still have been called with stored credentials
216+
const litellmCalls = getModelsMock.mock.calls.filter((c: any[]) => c[0]?.provider === "litellm")
217+
expect(litellmCalls.length).toBe(1)
218+
expect(litellmCalls[0][0]).toEqual({
219+
provider: "litellm",
220+
apiKey: "stored-api-key",
221+
baseUrl: "http://stored:4000",
222+
})
223+
})
167224
})

src/core/webview/webviewMessageHandler.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -882,6 +882,12 @@ export const webviewMessageHandler = async (
882882
const litellmBaseUrl = apiConfiguration.litellmBaseUrl || message?.values?.litellmBaseUrl
883883

884884
if (litellmApiKey && litellmBaseUrl) {
885+
// If explicit credentials are provided in message.values (from Refresh Models button),
886+
// flush the cache first to ensure we fetch fresh data with the new credentials
887+
if (message?.values?.litellmApiKey || message?.values?.litellmBaseUrl) {
888+
await flushModels("litellm", true)
889+
}
890+
885891
candidates.push({
886892
key: "litellm",
887893
options: { provider: "litellm", apiKey: litellmApiKey, baseUrl: litellmBaseUrl },

0 commit comments

Comments
 (0)