Skip to content
Open
Show file tree
Hide file tree
Changes from all 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
44 changes: 44 additions & 0 deletions src/api/providers/__tests__/native-ollama.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -912,6 +912,50 @@ describe("NativeOllamaHandler", () => {
})
})

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.

This unit test proves ensureModelFetched populates contextWindow, but no test asserts the fetched value actually reaches context management — Task.spec stubs getModel to a fixed 200k, so nothing guards against the 128k fallback being used instead of the fetched window, which is exactly the symptom this PR fixes. Worth an integration assertion that the fetched value reaches modelInfo.contextWindow?

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.
Expand Down
31 changes: 30 additions & 1 deletion src/api/providers/base-provider.ts
Original file line number Diff line number Diff line change
@@ -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"
Expand All @@ -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,
Expand All @@ -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<ModelRecord>,
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,
Expand Down
14 changes: 12 additions & 2 deletions src/api/providers/native-ollama.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<void> {
await this.fetchModel()
}

override getModel(): { id: string; info: ModelInfo } {
Expand Down
35 changes: 13 additions & 22 deletions src/api/providers/router-provider.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<void> {
Expand Down
70 changes: 70 additions & 0 deletions src/core/task/__tests__/Task.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<string>
Expand Down Expand Up @@ -237,6 +238,14 @@ vi.mock("../../condense", async (importOriginal) => {
}
})

vi.mock("../../context-management", async (importOriginal) => {
const actual = await importOriginal<typeof import("../../context-management")>()
return {
...actual,
manageContext: vi.fn(actual.manageContext),
}
})

vi.mock("../../prompts/system", async (importOriginal) => {
const actual = await importOriginal<typeof import("../../prompts/system")>()
return {
Expand Down Expand Up @@ -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<ApiStreamChunk>)
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({
Expand Down
Loading