mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-08-28 05:27:24 +00:00
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:
parent
cd02cd8a03
commit
e710d970b7
2 changed files with 185 additions and 14 deletions
132
src/core/webview/__tests__/webviewMessageHandler.test.ts
Normal file
132
src/core/webview/__tests__/webviewMessageHandler.test.ts
Normal 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: {},
|
||||
},
|
||||
})
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
@ -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":
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue