Refactor model fetching in webviewMessageHandler to handle failures independently

- Introduced a helper function to safely fetch models, ensuring that one failure does not affect others.
- Replaced Promise.all with Promise.allSettled to manage multiple model fetch requests concurrently while capturing errors.
This commit is contained in:
slytechnical 2025-05-19 11:56:29 -05:00
parent cd02cd8a03
commit e710d970b7
2 changed files with 185 additions and 14 deletions

View file

@ -0,0 +1,132 @@
// npx jest src/core/webview/__tests__/webviewMessageHandler.test.ts
import { webviewMessageHandler } from "../webviewMessageHandler"
import { getModels } from "../../../api/providers/fetchers/modelCache"
// Mock dependencies
jest.mock("../../../api/providers/fetchers/modelCache", () => ({
getModels: jest.fn(),
flushModels: jest.fn().mockResolvedValue(undefined),
}))
describe("webviewMessageHandler", () => {
// Set up provider mock with essential methods needed by the handler
const mockProvider = {
postMessageToWebview: jest.fn(),
getState: jest.fn().mockResolvedValue({
apiConfiguration: {
openRouterApiKey: "mock-openrouter-key",
requestyApiKey: "mock-requesty-key",
glamaApiKey: "mock-glama-key",
unboundApiKey: "mock-unbound-key",
litellmApiKey: "mock-litellm-key",
litellmBaseUrl: "https://mock-litellm-url",
},
}),
log: jest.fn(),
}
beforeEach(() => {
jest.clearAllMocks()
})
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) => {
return Promise.resolve({
[`${router}-model-1`]: { name: `${router} Model 1` },
[`${router}-model-2`]: { name: `${router} Model 2` },
})
})
// Call the handler
await webviewMessageHandler(mockProvider as any, {
type: "requestRouterModels",
})
// Verify the provider posted the correct message with all models
expect(mockProvider.postMessageToWebview).toHaveBeenCalledWith({
type: "routerModels",
routerModels: {
openrouter: {
"openrouter-model-1": { name: "openrouter Model 1" },
"openrouter-model-2": { name: "openrouter Model 2" },
},
requesty: {
"requesty-model-1": { name: "requesty Model 1" },
"requesty-model-2": { name: "requesty Model 2" },
},
glama: {
"glama-model-1": { name: "glama Model 1" },
"glama-model-2": { name: "glama Model 2" },
},
unbound: {
"unbound-model-1": { name: "unbound Model 1" },
"unbound-model-2": { name: "unbound Model 2" },
},
litellm: {
"litellm-model-1": { name: "litellm Model 1" },
"litellm-model-2": { name: "litellm Model 2" },
},
},
})
})
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") {
return Promise.resolve({
[`${router}-model-1`]: { name: `${router} Model 1` },
})
}
// For other routers, throw an error
return Promise.reject(new Error(`Failed to fetch ${router} models`))
})
// Call the handler
await webviewMessageHandler(mockProvider as any, {
type: "requestRouterModels",
})
// Verify the provider posted the correct message with only successful models
expect(mockProvider.postMessageToWebview).toHaveBeenCalledWith({
type: "routerModels",
routerModels: {
openrouter: {
"openrouter-model-1": { name: "openrouter Model 1" },
},
requesty: {},
glama: {},
unbound: {},
litellm: {
"litellm-model-1": { name: "litellm Model 1" },
},
},
})
})
test("handles all failed model fetches correctly", async () => {
// Mock all getModels calls to fail
;(getModels as jest.Mock).mockRejectedValue(new Error("API Error"))
// Call the handler
await webviewMessageHandler(mockProvider as any, {
type: "requestRouterModels",
})
// Verify the provider posted the correct message with empty objects for each router
expect(mockProvider.postMessageToWebview).toHaveBeenCalledWith({
type: "routerModels",
routerModels: {
openrouter: {},
requesty: {},
glama: {},
unbound: {},
litellm: {},
},
})
})
})
})

View file

@ -318,23 +318,62 @@ export const webviewMessageHandler = async (provider: ClineProvider, message: We
case "requestRouterModels":
const { apiConfiguration } = await provider.getState()
const [openRouterModels, requestyModels, glamaModels, unboundModels, litellmModels] = await Promise.all([
getModels("openrouter", apiConfiguration.openRouterApiKey),
getModels("requesty", apiConfiguration.requestyApiKey),
getModels("glama", apiConfiguration.glamaApiKey),
getModels("unbound", apiConfiguration.unboundApiKey),
getModels("litellm", apiConfiguration.litellmApiKey, apiConfiguration.litellmBaseUrl),
])
// Handle each model fetch independently to avoid one failure affecting others
const routerModels = {
openrouter: {},
requesty: {},
glama: {},
unbound: {},
litellm: {},
}
// Helper function to safely fetch models
const safeGetModels = async (router: RouterName, apiKey?: string, baseUrl?: string) => {
try {
return await getModels(router, apiKey, baseUrl)
} catch (error) {
console.error(`Failed to fetch models for ${router}:`, error)
return {} // Return empty object on failure
}
}
// Fetch all models in parallel but handle failures independently
const results = await Promise.allSettled(
[
{ key: "openrouter", promise: safeGetModels("openrouter", apiConfiguration.openRouterApiKey) },
{ key: "requesty", promise: safeGetModels("requesty", apiConfiguration.requestyApiKey) },
{ key: "glama", promise: safeGetModels("glama", apiConfiguration.glamaApiKey) },
{ key: "unbound", promise: safeGetModels("unbound", apiConfiguration.unboundApiKey) },
{
key: "litellm",
promise: safeGetModels(
"litellm",
apiConfiguration.litellmApiKey,
apiConfiguration.litellmBaseUrl,
),
},
].map(async ({ key, promise }) => {
try {
const models = await promise
return { key, models }
} catch (error) {
console.error(`Error in router models fetch for ${key}:`, error)
return { key, models: {} }
}
}),
)
// Process results and assign to routerModels
results.forEach((result) => {
if (result.status === "fulfilled") {
const key = result.value.key as keyof typeof routerModels
routerModels[key] = result.value.models
}
})
provider.postMessageToWebview({
type: "routerModels",
routerModels: {
openrouter: openRouterModels,
requesty: requestyModels,
glama: glamaModels,
unbound: unboundModels,
litellm: litellmModels,
},
routerModels,
})
break
case "requestOpenAiModels":