diff --git a/src/api/providers/__tests__/native-ollama.spec.ts b/src/api/providers/__tests__/native-ollama.spec.ts index 8fcbf4a0a1..26136d309c 100644 --- a/src/api/providers/__tests__/native-ollama.spec.ts +++ b/src/api/providers/__tests__/native-ollama.spec.ts @@ -912,6 +912,50 @@ describe("NativeOllamaHandler", () => { }) }) + describe("ensureModelFetched", () => { + it("makes the detected context window available before a request", async () => { + mockGetOllamaModels.mockResolvedValueOnce({ + llama2: { + contextWindow: 1_000_000, + maxTokens: 4096, + supportsImages: false, + supportsPromptCache: true, + }, + }) + + expect(handler.getModel().info.contextWindow).toBe(128_000) + + await handler.ensureModelFetched() + + expect(handler.getModel().info.contextWindow).toBe(1_000_000) + expect(mockGetOllamaModels).toHaveBeenCalledTimes(1) + }) + + it("skips subsequent fetches after models are populated", async () => { + await handler.ensureModelFetched() + mockGetOllamaModels.mockClear() + + await handler.fetchModel() + + expect(mockGetOllamaModels).not.toHaveBeenCalled() + }) + + it("deduplicates concurrent metadata fetches", async () => { + await Promise.all([handler.ensureModelFetched(), handler.ensureModelFetched()]) + + expect(mockGetOllamaModels).toHaveBeenCalledTimes(1) + }) + + it("allows a later fetch after a rejected request", async () => { + mockGetOllamaModels.mockRejectedValueOnce(new Error("network down")) + + await expect(handler.ensureModelFetched()).rejects.toThrow("network down") + await handler.ensureModelFetched() + + expect(mockGetOllamaModels).toHaveBeenCalledTimes(2) + }) + }) + describe("tool calling", () => { it("should include tools when tools are provided", async () => { // Model metadata should not gate tool inclusion; metadata.tools controls it. diff --git a/src/api/providers/base-provider.ts b/src/api/providers/base-provider.ts index 89366fb619..621e85fa17 100644 --- a/src/api/providers/base-provider.ts +++ b/src/api/providers/base-provider.ts @@ -1,6 +1,6 @@ import { Anthropic } from "@anthropic-ai/sdk" -import type { ModelInfo } from "@roo-code/types" +import type { ModelInfo, ModelRecord } from "@roo-code/types" import type { ApiHandler, ApiHandlerCreateMessageMetadata } from "../index" import { ApiStream } from "../transform/stream" @@ -13,6 +13,7 @@ import { getApiRequestTimeout } from "./utils/timeout-config" */ export abstract class BaseProvider implements ApiHandler { protected readonly timeoutMs: number = getApiRequestTimeout() + private modelFetchPromise?: Promise<{ id: string; info: ModelInfo }> abstract createMessage( systemPrompt: string, @@ -22,6 +23,34 @@ export abstract class BaseProvider implements ApiHandler { abstract getModel(): { id: string; info: ModelInfo } + /** + * Single-flights model metadata requests and clears the in-flight promise + * after either success or failure so a later call can retry. + */ + protected fetchModelWithSingleFlight( + models: ModelRecord, + fetchModels: () => Promise, + setModels: (models: ModelRecord) => void, + getModel: () => { id: string; info: ModelInfo }, + ): Promise<{ id: string; info: ModelInfo }> { + if (Object.keys(models).length > 0) { + return Promise.resolve(getModel()) + } + + if (!this.modelFetchPromise) { + this.modelFetchPromise = fetchModels() + .then((fetchedModels) => { + setModels(fetchedModels) + return getModel() + }) + .finally(() => { + this.modelFetchPromise = undefined + }) + } + + return this.modelFetchPromise + } + /** * Converts an array of tools to be compatible with OpenAI's strict mode. * Filters for function tools, applies schema conversion to their parameters, diff --git a/src/api/providers/native-ollama.ts b/src/api/providers/native-ollama.ts index eb380d1eb2..6f80e16f14 100644 --- a/src/api/providers/native-ollama.ts +++ b/src/api/providers/native-ollama.ts @@ -538,8 +538,18 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio } async fetchModel() { - this.models = await getOllamaModels(this.options.ollamaBaseUrl, this.options.ollamaApiKey) - return this.getModel() + return this.fetchModelWithSingleFlight( + this.models, + () => getOllamaModels(this.options.ollamaBaseUrl, this.options.ollamaApiKey), + (models) => { + this.models = models + }, + () => this.getModel(), + ) + } + + async ensureModelFetched(): Promise { + await this.fetchModel() } override getModel(): { id: string; info: ModelInfo } { diff --git a/src/api/providers/router-provider.ts b/src/api/providers/router-provider.ts index cbdd49e58b..7affa129b1 100644 --- a/src/api/providers/router-provider.ts +++ b/src/api/providers/router-provider.ts @@ -56,29 +56,20 @@ export abstract class RouterProvider extends BaseProvider { }) } - private modelFetchPromise?: Promise<{ id: string; info: ModelInfo }> - public async fetchModel() { - if (Object.keys(this.models).length > 0) { - return this.getModel() - } - - if (!this.modelFetchPromise) { - this.modelFetchPromise = getModels({ - provider: this.name, - apiKey: this.client.apiKey, - baseUrl: this.client.baseURL, - }) - .then((models) => { - this.models = models - return this.getModel() - }) - .finally(() => { - this.modelFetchPromise = undefined - }) - } - - return this.modelFetchPromise + return this.fetchModelWithSingleFlight( + this.models, + () => + getModels({ + provider: this.name, + apiKey: this.client.apiKey, + baseUrl: this.client.baseURL, + }), + (models) => { + this.models = models + }, + () => this.getModel(), + ) } async ensureModelFetched(): Promise { diff --git a/src/core/task/__tests__/Task.spec.ts b/src/core/task/__tests__/Task.spec.ts index 60bc2f3192..e8efffaf1c 100644 --- a/src/core/task/__tests__/Task.spec.ts +++ b/src/core/task/__tests__/Task.spec.ts @@ -26,6 +26,7 @@ import { ContextProxy } from "../../config/ContextProxy" import { processUserContentMentions } from "../../mentions/processUserContentMentions" import { MultiSearchReplaceDiffStrategy } from "../../diff/strategies/multi-search-replace" import type { ApiMessage } from "../../task-persistence" +import { manageContext } from "../../context-management" type TaskTestAccess = { getSystemPrompt: () => Promise @@ -237,6 +238,14 @@ vi.mock("../../condense", async (importOriginal) => { } }) +vi.mock("../../context-management", async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + manageContext: vi.fn(actual.manageContext), + } +}) + vi.mock("../../prompts/system", async (importOriginal) => { const actual = await importOriginal() return { @@ -3000,6 +3009,67 @@ describe("Cline", () => { }) }) + describe("context management model metadata", () => { + it("passes the fetched context window to context management", async () => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + startTask: false, + }) + + let modelInfo: ModelInfo = { + supportsImages: false, + supportsPromptCache: true, + contextWindow: 128_000, + maxTokens: 4096, + } + vi.spyOn(task.api, "getModel").mockImplementation(() => ({ + id: mockApiConfig.apiModelId!, + info: modelInfo, + })) + const ensureModelFetched = vi.fn().mockImplementation(async () => { + modelInfo = { ...modelInfo, contextWindow: 1_000_000 } + }) + Object.assign(task.api, { ensureModelFetched }) + vi.spyOn(getTaskTestAccess(task), "getSystemPrompt").mockResolvedValue("mock system prompt") + vi.spyOn(task, "getTokenUsage").mockReturnValue({ + totalCost: 0, + totalTokensIn: 0, + totalTokensOut: 0, + contextTokens: 50_000, + }) + vi.spyOn(task.api, "createMessage").mockReturnValue({ + async *[Symbol.asyncIterator]() { + yield { type: "text", text: "ok" } + }, + async next() { + return { done: true, value: undefined } + }, + async return() { + return { done: true, value: undefined } + }, + async throw(error: unknown) { + throw error + }, + async [Symbol.asyncDispose]() {}, + } as AsyncGenerator) + task.apiConversationHistory = [ + { + role: "user" as const, + content: [{ type: "text" as const, text: "test message" }], + ts: Date.now(), + }, + ] + + vi.mocked(manageContext).mockClear() + await task.attemptApiRequest(0).next() + + expect(ensureModelFetched).toHaveBeenCalledTimes(1) + expect(manageContext).toHaveBeenCalledWith(expect.objectContaining({ contextWindow: 1_000_000 })) + }) + }) + describe("start()", () => { it("should be a no-op if the task was already started in the constructor", () => { const task = new Task({