Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 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
7 changes: 7 additions & 0 deletions src/api/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,13 @@ export interface ApiHandler {

getModel(): { id: string; info: ModelInfo }

/**
* Ensures model metadata has been fetched from the remote API so that getModel()
* returns accurate info (context window, pricing, etc.) instead of hardcoded defaults.
* Only router providers that discover models over the network implement this.
*/
ensureModelFetched?(): Promise<void>

/**
* Optional context window for context-management / auto-condense when it must differ from
* getModel().info.contextWindow. Only VS Code LM overrides it (static `maxInputTokens` vs its
Expand Down
57 changes: 57 additions & 0 deletions src/api/providers/__tests__/zoo-gateway.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -635,4 +635,61 @@ describe("ZooGatewayHandler", () => {
)
})
})

describe("ensureModelFetched", () => {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is the rejection path covered anywhere? If getModels rejects once and the finally cleanup in router-provider.ts (L71-74) ever regresses, the stored rejected promise would poison every subsequent call. A reject-then-resolve test here would lock that in.

it("fetches models when instance models are empty", async () => {
const handler = new ZooGatewayHandler(mockOptions)
const { getModels } = await import("../fetchers/modelCache")

expect(handler.getModel().info.contextWindow).toBe(200000)

await handler.ensureModelFetched()

expect(getModels).toHaveBeenCalled()
})

it("skips the fetch when models are already populated", async () => {
const handler = new ZooGatewayHandler(mockOptions)
const { getModels } = await import("../fetchers/modelCache")

await handler.ensureModelFetched()
vitest.mocked(getModels).mockClear()

await handler.ensureModelFetched()
expect(getModels).not.toHaveBeenCalled()
})

it("deduplicates concurrent calls into a single fetch", async () => {
const handler = new ZooGatewayHandler(mockOptions)
const { getModels } = await import("../fetchers/modelCache")
vitest.mocked(getModels).mockClear()

await Promise.all([handler.ensureModelFetched(), handler.ensureModelFetched()])

expect(getModels).toHaveBeenCalledTimes(1)
})

it("makes getModel return the fetched context window instead of the default", async () => {
const { getModels } = await import("../fetchers/modelCache")
vitest.mocked(getModels).mockResolvedValueOnce({
"google/gemini-2.5-pro": {
maxTokens: 65536,
contextWindow: 1048576,
supportsImages: true,
supportsPromptCache: false,
},
})

const handler = new ZooGatewayHandler({
...mockOptions,
zooGatewayModelId: "google/gemini-2.5-pro",
})

expect(handler.getModel().info.contextWindow).toBe(200000)

await handler.ensureModelFetched()

expect(handler.getModel().info.contextWindow).toBe(1048576)
})
})
})
15 changes: 15 additions & 0 deletions src/api/providers/router-provider.ts
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,21 @@ export abstract class RouterProvider extends BaseProvider {
return this.getModel()
}

private modelFetchPromise?: Promise<void>

async ensureModelFetched(): Promise<void> {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For auth-scoped providers, does the first request now fetch the model list twice — once here and again unconditionally in createMessage (zoo-gateway.ts L184)? getModels bypasses both caches for auth-scoped providers (modelCache.ts L265), so the second call is another network round trip. Worth making fetchModel() single-flight (or short-circuiting when this.models is populated) so both callers share one fetch?

if (Object.keys(this.models).length === 0) {
const fetchPromise = (this.modelFetchPromise ??= this.fetchModel().then(() => undefined))
try {
await fetchPromise
} finally {
if (this.modelFetchPromise === fetchPromise) {
this.modelFetchPromise = undefined
}
}
}
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.

override getModel(): { id: string; info: ModelInfo } {
// Use `||` (not `??`) so an empty-string modelId also falls back to the default,
// guaranteeing a non-empty id rather than forwarding "" to the API as an invalid
Expand Down
4 changes: 4 additions & 0 deletions src/core/task/Task.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2762,6 +2762,8 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {

await this.diffViewProvider.reset()

await this.api.ensureModelFetched?.()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The litellm/deepseek/moonshot fetchers re-throw, so a failure here would reject into the catch at L3743, which return trues and ends the task silently — the comment there assumes only attemptApiRequest can throw. Pre-change, the same failure surfaced via the createMessage retry path. Would catching and logging here be safer, letting getModel() fall back to defaults?


// Cache model info once per API request to avoid repeated calls during streaming
// This is especially important for tools and background usage collection
this.cachedStreamingModel = this.api.getModel()
Expand Down Expand Up @@ -3842,6 +3844,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
const { profileThresholds = {}, mode, apiConfiguration } = state ?? {}

const { contextTokens } = this.getTokenUsage()
await this.api.ensureModelFetched?.()
const modelInfo = this.api.getModel().info

const maxTokens = getModelMaxOutputTokens({
Expand Down Expand Up @@ -4042,6 +4045,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
const { contextTokens } = this.getTokenUsage()

if (contextTokens) {
await this.api.ensureModelFetched?.()
const modelInfo = this.api.getModel().info
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated

const maxTokens = getModelMaxOutputTokens({
Expand Down
Loading