From 872902754a83c6b724a894df924e01382fe74930 Mon Sep 17 00:00:00 2001 From: Chris Estreich Date: Fri, 23 May 2025 02:32:21 -0700 Subject: [PATCH] Fix max tokens in the task header (#3893) --- .changeset/ten-bags-hang.md | 5 ++ src/api/transform/model-params.ts | 12 ++-- src/shared/__tests__/api.test.ts | 61 +++++++++++++++---- src/shared/api.ts | 12 ++++ .../__tests__/ContextWindowProgress.test.tsx | 1 + webview-ui/src/components/chat/TaskHeader.tsx | 12 +++- 6 files changed, 85 insertions(+), 18 deletions(-) create mode 100644 .changeset/ten-bags-hang.md diff --git a/.changeset/ten-bags-hang.md b/.changeset/ten-bags-hang.md new file mode 100644 index 0000000000..838e208a1c --- /dev/null +++ b/.changeset/ten-bags-hang.md @@ -0,0 +1,5 @@ +--- +"roo-cline": patch +--- + +Fix max tokens in task header diff --git a/src/api/transform/model-params.ts b/src/api/transform/model-params.ts index 1a081eb53b..5e99e770bf 100644 --- a/src/api/transform/model-params.ts +++ b/src/api/transform/model-params.ts @@ -87,16 +87,20 @@ export function getModelParams({ reasoningEffort = customReasoningEffort ?? model.reasoningEffort } + // TODO: We should consolidate this logic to compute `maxTokens` with + // `getModelMaxOutputTokens` in order to maintain a single source of truth. + + const isAnthropic = format === "anthropic" || (format === "openrouter" && modelId.startsWith("anthropic/")) + // For "Hybrid" reasoning models, we should discard the model's actual - // `maxTokens` value if we're not using reasoning. - if (model.supportsReasoningBudget && !reasoningBudget) { + // `maxTokens` value if we're not using reasoning. We do this for Anthropic + // models only for now. Should we do this for Gemini too? + if (model.supportsReasoningBudget && !reasoningBudget && isAnthropic) { maxTokens = ANTHROPIC_DEFAULT_MAX_TOKENS } // For Anthropic models we should always make sure a `maxTokens` value is // set. - const isAnthropic = format === "anthropic" || (format === "openrouter" && modelId.startsWith("anthropic/")) - if (!maxTokens && isAnthropic) { maxTokens = ANTHROPIC_DEFAULT_MAX_TOKENS } diff --git a/src/shared/__tests__/api.test.ts b/src/shared/__tests__/api.test.ts index 875faa8f2d..19dbe1dbcc 100644 --- a/src/shared/__tests__/api.test.ts +++ b/src/shared/__tests__/api.test.ts @@ -7,8 +7,11 @@ import { shouldUseReasoningBudget, shouldUseReasoningEffort, } from "../api" +import { ANTHROPIC_DEFAULT_MAX_TOKENS } from "../../api/providers/constants" describe("getMaxTokensForModel", () => { + const modelId = "test" + /** * Testing the specific fix in commit cc79178f: * For thinking models, use apiConfig.modelMaxTokens if available, @@ -27,7 +30,7 @@ describe("getMaxTokensForModel", () => { modelMaxTokens: 4000, } - expect(getModelMaxOutputTokens({ model, settings })).toBe(4000) + expect(getModelMaxOutputTokens({ modelId, model, settings })).toBe(4000) }) it("should return 16_384 for thinking models when modelMaxTokens not provided", () => { @@ -40,7 +43,7 @@ describe("getMaxTokensForModel", () => { const settings = {} - expect(getModelMaxOutputTokens({ model, settings })).toBe(16_384) + expect(getModelMaxOutputTokens({ modelId, model, settings })).toBe(16_384) }) it("should return 16_384 for thinking models when apiConfig is undefined", () => { @@ -51,7 +54,7 @@ describe("getMaxTokensForModel", () => { maxTokens: 8000, } - expect(getModelMaxOutputTokens({ model, settings: undefined })).toBe(16_384) + expect(getModelMaxOutputTokens({ modelId, model, settings: undefined })).toBe(16_384) }) it("should return modelInfo.maxTokens for non-thinking models", () => { @@ -65,7 +68,7 @@ describe("getMaxTokensForModel", () => { modelMaxTokens: 4000, } - expect(getModelMaxOutputTokens({ model, settings })).toBe(8000) + expect(getModelMaxOutputTokens({ modelId, model, settings })).toBe(8000) }) it("should return undefined for non-thinking models with undefined maxTokens", () => { @@ -78,7 +81,7 @@ describe("getMaxTokensForModel", () => { modelMaxTokens: 4000, } - expect(getModelMaxOutputTokens({ model, settings })).toBeUndefined() + expect(getModelMaxOutputTokens({ modelId, model, settings })).toBeUndefined() }) test("should return maxTokens from modelInfo when thinking is false", () => { @@ -92,7 +95,7 @@ describe("getMaxTokensForModel", () => { modelMaxTokens: 4096, } - const result = getModelMaxOutputTokens({ model, settings }) + const result = getModelMaxOutputTokens({ modelId, model, settings }) expect(result).toBe(2048) }) @@ -108,7 +111,7 @@ describe("getMaxTokensForModel", () => { modelMaxTokens: 4096, } - const result = getModelMaxOutputTokens({ model, settings }) + const result = getModelMaxOutputTokens({ modelId, model, settings }) expect(result).toBe(4096) }) @@ -122,7 +125,7 @@ describe("getMaxTokensForModel", () => { const settings: ProviderSettings = {} - const result = getModelMaxOutputTokens({ model, settings: undefined }) + const result = getModelMaxOutputTokens({ modelId, model, settings: undefined }) expect(result).toBe(16_384) }) @@ -133,7 +136,7 @@ describe("getMaxTokensForModel", () => { maxTokens: 2048, } - expect(getModelMaxOutputTokens({ model: modelInfoOnly, settings: undefined })).toBe(2048) + expect(getModelMaxOutputTokens({ modelId, model: modelInfoOnly, settings: undefined })).toBe(2048) }) test("should handle missing properties gracefully", () => { @@ -147,7 +150,7 @@ describe("getMaxTokensForModel", () => { modelMaxTokens: 4096, } - expect(getModelMaxOutputTokens({ model: modelInfoWithoutMaxTokens, settings })).toBe(4096) + expect(getModelMaxOutputTokens({ modelId, model: modelInfoWithoutMaxTokens, settings })).toBe(4096) const modelInfoWithoutThinking: ModelInfo = { contextWindow: 200_000, @@ -155,7 +158,43 @@ describe("getMaxTokensForModel", () => { maxTokens: 2048, } - expect(getModelMaxOutputTokens({ model: modelInfoWithoutThinking, settings: undefined })).toBe(2048) + expect(getModelMaxOutputTokens({ modelId, model: modelInfoWithoutThinking, settings: undefined })).toBe(2048) + }) + + test("should return ANTHROPIC_DEFAULT_MAX_TOKENS for Anthropic models that support reasoning budget but aren't using it", () => { + // Test case for models that support reasoning budget but enableReasoningEffort is false + const anthropicModelId = "claude-sonnet-4-20250514" + const model: ModelInfo = { + contextWindow: 200_000, + supportsPromptCache: true, + supportsReasoningBudget: true, + maxTokens: 64_000, // This should be ignored + } + + const settings: ProviderSettings = { + enableReasoningEffort: false, // Not using reasoning + } + + const result = getModelMaxOutputTokens({ modelId: anthropicModelId, model, settings }) + expect(result).toBe(ANTHROPIC_DEFAULT_MAX_TOKENS) // Should be 8192, not 64_000 + }) + + test("should return model.maxTokens for non-Anthropic models that support reasoning budget but aren't using it", () => { + // Test case for non-Anthropic models - should still use model.maxTokens + const geminiModelId = "gemini-2.5-flash-preview-04-17" + const model: ModelInfo = { + contextWindow: 1_048_576, + supportsPromptCache: false, + supportsReasoningBudget: true, + maxTokens: 65_535, + } + + const settings: ProviderSettings = { + enableReasoningEffort: false, // Not using reasoning + } + + const result = getModelMaxOutputTokens({ modelId: geminiModelId, model, settings }) + expect(result).toBe(65_535) // Should use model.maxTokens, not ANTHROPIC_DEFAULT_MAX_TOKENS }) }) diff --git a/src/shared/api.ts b/src/shared/api.ts index 713949c433..92cc406854 100644 --- a/src/shared/api.ts +++ b/src/shared/api.ts @@ -1,3 +1,4 @@ +import { ANTHROPIC_DEFAULT_MAX_TOKENS } from "../api/providers/constants" import { ModelInfo, ProviderName, ProviderSettings } from "../schemas" export type { ModelInfo, ProviderName, ProviderSettings } @@ -1936,9 +1937,11 @@ export const DEFAULT_HYBRID_REASONING_MODEL_MAX_TOKENS = 16_384 export const DEFAULT_HYBRID_REASONING_MODEL_THINKING_TOKENS = 8_192 export const getModelMaxOutputTokens = ({ + modelId, model, settings, }: { + modelId: string model: ModelInfo settings?: ProviderSettings }): number | undefined => { @@ -1946,5 +1949,14 @@ export const getModelMaxOutputTokens = ({ return settings?.modelMaxTokens || DEFAULT_HYBRID_REASONING_MODEL_MAX_TOKENS } + const isAnthropicModel = modelId.includes("claude") + + // For "Hybrid" reasoning models, we should discard the model's actual + // `maxTokens` value if we're not using reasoning. We do this for Anthropic + // models only for now. Should we do this for Gemini too? + if (model.supportsReasoningBudget && isAnthropicModel) { + return ANTHROPIC_DEFAULT_MAX_TOKENS + } + return model.maxTokens ?? undefined } diff --git a/webview-ui/src/__tests__/ContextWindowProgress.test.tsx b/webview-ui/src/__tests__/ContextWindowProgress.test.tsx index 15c6f4be20..a4db28cf6a 100644 --- a/webview-ui/src/__tests__/ContextWindowProgress.test.tsx +++ b/webview-ui/src/__tests__/ContextWindowProgress.test.tsx @@ -38,6 +38,7 @@ jest.mock("@src/components/chat/TaskHeader", () => { // Mock useSelectedModel hook jest.mock("@src/components/ui/hooks/useSelectedModel", () => ({ useSelectedModel: jest.fn(() => ({ + id: "test", info: { contextWindow: 4000 }, })), })) diff --git a/webview-ui/src/components/chat/TaskHeader.tsx b/webview-ui/src/components/chat/TaskHeader.tsx index 736057ae65..038a903d04 100644 --- a/webview-ui/src/components/chat/TaskHeader.tsx +++ b/webview-ui/src/components/chat/TaskHeader.tsx @@ -48,7 +48,7 @@ const TaskHeader = ({ }: TaskHeaderProps) => { const { t } = useTranslation() const { apiConfiguration, currentTaskItem } = useExtensionState() - const { info: model } = useSelectedModel(apiConfiguration) + const { id: modelId, info: model } = useSelectedModel(apiConfiguration) const [isTaskExpanded, setIsTaskExpanded] = useState(false) const textContainerRef = useRef(null) @@ -101,7 +101,9 @@ const TaskHeader = ({ contextWindow={contextWindow} contextTokens={contextTokens || 0} maxTokens={ - model ? getModelMaxOutputTokens({ model, settings: apiConfiguration }) : undefined + model + ? getModelMaxOutputTokens({ modelId, model, settings: apiConfiguration }) + : undefined } /> {!!totalCost && ${totalCost.toFixed(2)}} @@ -140,7 +142,11 @@ const TaskHeader = ({ contextTokens={contextTokens || 0} maxTokens={ model - ? getModelMaxOutputTokens({ model, settings: apiConfiguration }) + ? getModelMaxOutputTokens({ + modelId, + model, + settings: apiConfiguration, + }) : undefined } />