feat: Add dynamic 1M context window switching for Claude Sonnet models

- Implement dynamic context window switching based on actual context size
- Add shouldUse1MContext helper function for centralized decision logic
- Update Anthropic and Bedrock providers to support dynamic switching
- Pass context tokens from Task.ts through to API providers
- Update UI labels to clarify dynamic behavior
- Add comprehensive test coverage

When the anthropicBeta1MContext or awsBedrock1MContext settings are enabled,
the system now automatically switches between 200K and 1M context windows
based on actual usage (default threshold: 190K tokens).

Fixes #9250
This commit is contained in:
Roo Code 2025-11-13 23:43:30 +00:00
parent ad1e9a82f9
commit c133994163
7 changed files with 293 additions and 25 deletions

View file

@ -277,16 +277,43 @@ describe("AnthropicHandler", () => {
expect(model.info.supportsReasoningBudget).toBe(true)
})
it("should enable 1M context for Claude 4.5 Sonnet when beta flag is set", () => {
it("should use 200K context for Claude 4.5 Sonnet when beta flag is set but context is low", () => {
const handler = new AnthropicHandler({
apiKey: "test-api-key",
apiModelId: "claude-sonnet-4-5",
anthropicBeta1MContext: true,
})
// Without contextTokens in metadata, it defaults to 0 which is below threshold
const model = handler.getModel()
expect(model.info.contextWindow).toBe(200000)
expect(model.info.inputPrice).toBe(3.0)
expect(model.info.outputPrice).toBe(15.0)
})
it("should enable 1M context for Claude 4.5 Sonnet when beta flag is set and context is high", () => {
const handler = new AnthropicHandler({
apiKey: "test-api-key",
apiModelId: "claude-sonnet-4-5",
anthropicBeta1MContext: true,
})
// Simulate high context by passing contextTokens > 190K threshold
const model = handler.getModel(195000)
expect(model.info.contextWindow).toBe(1000000)
expect(model.info.inputPrice).toBe(6.0)
expect(model.info.outputPrice).toBe(22.5)
})
it("should always use 200K context when beta flag is not set regardless of context size", () => {
const handler = new AnthropicHandler({
apiKey: "test-api-key",
apiModelId: "claude-sonnet-4-5",
anthropicBeta1MContext: false,
})
// Even with high context, should stay at 200K when flag is false
const model = handler.getModel(195000)
expect(model.info.contextWindow).toBe(200000)
expect(model.info.inputPrice).toBe(3.0)
expect(model.info.outputPrice).toBe(15.0)
})
})
})

View file

@ -11,6 +11,7 @@ import {
} from "@roo-code/types"
import type { ApiHandlerOptions } from "../../shared/api"
import { shouldUse1MContext } from "../../shared/api"
import { ApiStream } from "../transform/stream"
import { getModelParams } from "../transform/model-params"
@ -39,16 +40,23 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa
async *createMessage(
systemPrompt: string,
messages: Anthropic.Messages.MessageParam[],
metadata?: ApiHandlerCreateMessageMetadata,
metadata?: ApiHandlerCreateMessageMetadata & { contextTokens?: number },
): ApiStream {
let stream: AnthropicStream<Anthropic.Messages.RawMessageStreamEvent>
const cacheControl: CacheControlEphemeral = { type: "ephemeral" }
let { id: modelId, betas = [], maxTokens, temperature, reasoning: thinking } = this.getModel()
// Add 1M context beta flag if enabled for Claude Sonnet 4 and 4.5
// Pass context tokens to getModel for dynamic context window selection
const contextTokens = metadata?.contextTokens ?? 0
let { id: modelId, betas = [], maxTokens, temperature, reasoning: thinking } = this.getModel(contextTokens)
// Add 1M context beta flag if dynamic mode is enabled and context size warrants it
if (
(modelId === "claude-sonnet-4-20250514" || modelId === "claude-sonnet-4-5") &&
this.options.anthropicBeta1MContext
shouldUse1MContext({
baseModel: modelId,
dynamicEnabled: this.options.anthropicBeta1MContext ?? false,
contextTokens,
})
) {
betas.push("context-1m-2025-08-07")
}
@ -247,13 +255,20 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa
}
}
getModel() {
getModel(contextTokens?: number) {
const modelId = this.options.apiModelId
let id = modelId && modelId in anthropicModels ? (modelId as AnthropicModelId) : anthropicDefaultModelId
let info: ModelInfo = anthropicModels[id]
// If 1M context beta is enabled for Claude Sonnet 4 or 4.5, update the model info
if ((id === "claude-sonnet-4-20250514" || id === "claude-sonnet-4-5") && this.options.anthropicBeta1MContext) {
// If dynamic context switching is enabled and context warrants 1M, update the model info
if (
(id === "claude-sonnet-4-20250514" || id === "claude-sonnet-4-5") &&
shouldUse1MContext({
baseModel: id,
dynamicEnabled: this.options.anthropicBeta1MContext ?? false,
contextTokens: contextTokens ?? 0,
})
) {
// Use the tier pricing for 1M context
const tier = info.tiers?.[0]
if (tier) {

View file

@ -33,7 +33,7 @@ import { MultiPointStrategy } from "../transform/cache-strategy/multi-point-stra
import { ModelInfo as CacheModelInfo } from "../transform/cache-strategy/types"
import { convertToBedrockConverseMessages as sharedConverter } from "../transform/bedrock-converse-format"
import { getModelParams } from "../transform/model-params"
import { shouldUseReasoningBudget } from "../../shared/api"
import { shouldUseReasoningBudget, shouldUse1MContext } from "../../shared/api"
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
/************************************************************************************
@ -322,9 +322,12 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
maxTokens?: number
maxThinkingTokens?: number
}
contextTokens?: number
},
): ApiStream {
const modelConfig = this.getModel()
// Pass context tokens to getModel for dynamic context window selection
const contextTokens = metadata?.contextTokens ?? 0
const modelConfig = this.getModel(contextTokens)
const usePromptCache = Boolean(this.options.awsUsePromptCache && this.supportsAwsPromptCache(modelConfig))
const conversationId =
@ -376,11 +379,22 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
temperature: modelConfig.temperature ?? (this.options.modelTemperature as number),
}
// Check if 1M context is enabled for Claude Sonnet 4
// Check if dynamic 1M context should be used based on actual context size
// Use parseBaseModelId to handle cross-region inference prefixes
const baseModelId = this.parseBaseModelId(modelConfig.id)
const is1MContextEnabled =
BEDROCK_1M_CONTEXT_MODEL_IDS.includes(baseModelId as any) && this.options.awsBedrock1MContext
BEDROCK_1M_CONTEXT_MODEL_IDS.includes(baseModelId as any) &&
(baseModelId === "anthropic.claude-sonnet-4-20250514" || baseModelId === "anthropic.claude-sonnet-4-5"
? shouldUse1MContext({
baseModel:
baseModelId === "anthropic.claude-sonnet-4-20250514"
? "claude-sonnet-4-20250514"
: "claude-sonnet-4-5",
dynamicEnabled: this.options.awsBedrock1MContext ?? false,
contextTokens,
})
: // For non-Sonnet 4.x models, use static flag behavior
this.options.awsBedrock1MContext)
// Add anthropic_beta for 1M context to additionalModelRequestFields
if (is1MContextEnabled) {
@ -936,7 +950,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
return model
}
override getModel(): {
override getModel(contextTokens?: number): {
id: BedrockModelId | string
info: ModelInfo
maxTokens?: number
@ -987,14 +1001,34 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
}
}
// Check if 1M context is enabled for Claude Sonnet 4 / 4.5
// Check if dynamic 1M context should be used based on actual context size
// Use parseBaseModelId to handle cross-region inference prefixes
const baseModelId = this.parseBaseModelId(modelConfig.id)
if (BEDROCK_1M_CONTEXT_MODEL_IDS.includes(baseModelId as any) && this.options.awsBedrock1MContext) {
// Update context window to 1M tokens when 1M context beta is enabled
modelConfig.info = {
...modelConfig.info,
contextWindow: 1_000_000,
if (BEDROCK_1M_CONTEXT_MODEL_IDS.includes(baseModelId as any)) {
// For Claude Sonnet 4.x models, use dynamic switching
if (baseModelId === "anthropic.claude-sonnet-4-20250514" || baseModelId === "anthropic.claude-sonnet-4-5") {
if (
shouldUse1MContext({
baseModel:
baseModelId === "anthropic.claude-sonnet-4-20250514"
? "claude-sonnet-4-20250514"
: "claude-sonnet-4-5",
dynamicEnabled: this.options.awsBedrock1MContext ?? false,
contextTokens: contextTokens ?? 0,
})
) {
// Update context window to 1M tokens when dynamic switching warrants it
modelConfig.info = {
...modelConfig.info,
contextWindow: 1_000_000,
}
}
} else if (this.options.awsBedrock1MContext) {
// For non-Sonnet 4.x models, use static flag behavior (backward compatibility)
modelConfig.info = {
...modelConfig.info,
contextWindow: 1_000_000,
}
}
}

View file

@ -2847,6 +2847,10 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
const systemPrompt = await this.getSystemPrompt()
const { contextTokens } = this.getTokenUsage()
// Store the projected context tokens for dynamic 1M context switching
// This will be passed to the API handler to decide whether to use 1M or 200K context
let projectedContextTokens = 0
if (contextTokens) {
const modelInfo = this.api.getModel().info
@ -2895,6 +2899,11 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
contextCondense,
)
}
// Store the projected context tokens from the truncation result
if ("prevContextTokens" in truncateResult) {
projectedContextTokens = truncateResult.prevContextTokens
}
}
// Properly type cleaned conversation history to include either standard Anthropic messages
@ -2953,11 +2962,13 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
allTools = [...nativeTools, ...mcpTools]
}
const metadata: ApiHandlerCreateMessageMetadata = {
const metadata: ApiHandlerCreateMessageMetadata & { contextTokens?: number } = {
mode: mode,
taskId: this.taskId,
// Include tools and tool protocol when using native protocol and model supports it
...(shouldIncludeTools ? { tools: allTools, tool_choice: "auto", toolProtocol } : {}),
// Include projected context tokens for dynamic 1M context switching
...(projectedContextTokens > 0 ? { contextTokens: projectedContextTokens } : {}),
}
// The provider accepts reasoning items alongside standard messages; cast to the expected parameter type.

View file

@ -5,7 +5,7 @@ import {
ANTHROPIC_DEFAULT_MAX_TOKENS,
} from "@roo-code/types"
import { getModelMaxOutputTokens, shouldUseReasoningBudget, shouldUseReasoningEffort } from "../api"
import { getModelMaxOutputTokens, shouldUseReasoningBudget, shouldUseReasoningEffort, shouldUse1MContext } from "../api"
describe("getModelMaxOutputTokens", () => {
const mockModel: ModelInfo = {
@ -474,3 +474,153 @@ describe("shouldUseReasoningEffort", () => {
expect(shouldUseReasoningEffort({ model, settings: settingsHigh })).toBe(true)
})
})
describe("shouldUse1MContext", () => {
test("should return false when dynamicEnabled is false", () => {
expect(
shouldUse1MContext({ baseModel: "claude-sonnet-4-20250514", dynamicEnabled: false, contextTokens: 50000 }),
).toBe(false)
expect(
shouldUse1MContext({ baseModel: "claude-sonnet-4-5", dynamicEnabled: false, contextTokens: 250000 }),
).toBe(false)
})
test("should return false for non-Sonnet 4 models", () => {
// Non-Sonnet models
expect(
shouldUse1MContext({ baseModel: "claude-3-opus-20240229", dynamicEnabled: true, contextTokens: 250000 }),
).toBe(false)
expect(
shouldUse1MContext({ baseModel: "claude-3-5-haiku-20241022", dynamicEnabled: true, contextTokens: 250000 }),
).toBe(false)
// Sonnet models that aren't 4.x series
expect(
shouldUse1MContext({
baseModel: "claude-3-5-sonnet-20241022",
dynamicEnabled: true,
contextTokens: 250000,
}),
).toBe(false)
expect(
shouldUse1MContext({ baseModel: "claude-3-sonnet-20240229", dynamicEnabled: true, contextTokens: 250000 }),
).toBe(false)
})
test("should return true for Sonnet 4 models above default threshold", () => {
// Default threshold is 190,000 tokens
expect(
shouldUse1MContext({ baseModel: "claude-sonnet-4-20250514", dynamicEnabled: true, contextTokens: 190001 }),
).toBe(true)
expect(
shouldUse1MContext({ baseModel: "claude-sonnet-4-5", dynamicEnabled: true, contextTokens: 200000 }),
).toBe(true)
expect(
shouldUse1MContext({ baseModel: "claude-sonnet-4-20250514", dynamicEnabled: true, contextTokens: 250000 }),
).toBe(true)
expect(
shouldUse1MContext({ baseModel: "claude-sonnet-4-5", dynamicEnabled: true, contextTokens: 190001 }),
).toBe(true)
})
test("should return false for Sonnet 4 models below default threshold", () => {
// Default threshold is 190,000 tokens
expect(
shouldUse1MContext({ baseModel: "claude-sonnet-4-20250514", dynamicEnabled: true, contextTokens: 189999 }),
).toBe(false)
expect(
shouldUse1MContext({ baseModel: "claude-sonnet-4-5", dynamicEnabled: true, contextTokens: 100000 }),
).toBe(false)
expect(
shouldUse1MContext({ baseModel: "claude-sonnet-4-20250514", dynamicEnabled: true, contextTokens: 50000 }),
).toBe(false)
expect(shouldUse1MContext({ baseModel: "claude-sonnet-4-5", dynamicEnabled: true, contextTokens: 0 })).toBe(
false,
)
})
test("should respect custom threshold parameter", () => {
const customThreshold = 150000
// Above custom threshold
expect(
shouldUse1MContext({
baseModel: "claude-sonnet-4-20250514",
dynamicEnabled: true,
contextTokens: 150001,
threshold: customThreshold,
}),
).toBe(true)
expect(
shouldUse1MContext({
baseModel: "claude-sonnet-4-5",
dynamicEnabled: true,
contextTokens: 160000,
threshold: customThreshold,
}),
).toBe(true)
// Below custom threshold
expect(
shouldUse1MContext({
baseModel: "claude-sonnet-4-20250514",
dynamicEnabled: true,
contextTokens: 149999,
threshold: customThreshold,
}),
).toBe(false)
expect(
shouldUse1MContext({
baseModel: "claude-sonnet-4-5",
dynamicEnabled: true,
contextTokens: 140000,
threshold: customThreshold,
}),
).toBe(false)
})
test("should handle edge cases at exact threshold", () => {
// At exact default threshold
expect(
shouldUse1MContext({ baseModel: "claude-sonnet-4-20250514", dynamicEnabled: true, contextTokens: 190000 }),
).toBe(false)
// At exact custom threshold
expect(
shouldUse1MContext({
baseModel: "claude-sonnet-4-5",
dynamicEnabled: true,
contextTokens: 100000,
threshold: 100000,
}),
).toBe(false)
})
test("should handle undefined and null contextTokens gracefully", () => {
// Should treat undefined/0/null as below threshold
expect(
shouldUse1MContext({
baseModel: "claude-sonnet-4-20250514",
dynamicEnabled: true,
contextTokens: undefined as any,
}),
).toBe(false)
expect(
shouldUse1MContext({ baseModel: "claude-sonnet-4-5", dynamicEnabled: true, contextTokens: null as any }),
).toBe(false)
expect(
shouldUse1MContext({ baseModel: "claude-sonnet-4-20250514", dynamicEnabled: true, contextTokens: NaN }),
).toBe(false)
})
test("should not crash with invalid model IDs", () => {
expect(shouldUse1MContext({ baseModel: "", dynamicEnabled: true, contextTokens: 250000 })).toBe(false)
expect(shouldUse1MContext({ baseModel: " ", dynamicEnabled: true, contextTokens: 250000 })).toBe(false)
expect(shouldUse1MContext({ baseModel: "invalid-model", dynamicEnabled: true, contextTokens: 250000 })).toBe(
false,
)
expect(
shouldUse1MContext({ baseModel: "claude-3-6-opus-20250219", dynamicEnabled: true, contextTokens: 250000 }),
).toBe(false) // opus not sonnet
})
})

View file

@ -46,6 +46,37 @@ export type ModelRecord = Record<string, ModelInfo>
export type RouterModels = Record<RouterName, ModelRecord>
// Context Window Selection
/**
* Determines whether to use 1M context window based on actual context size.
* When dynamic context switching is enabled, this function decides whether to use
* the standard 200K context window or switch to the 1M context window based on
* the projected token count.
*
* @param options Configuration for context window selection
* @returns true if 1M context should be used, false for standard 200K
*/
export function shouldUse1MContext(options: {
baseModel: string
dynamicEnabled: boolean
contextTokens: number
threshold?: number
}): boolean {
if (!options.dynamicEnabled) return false
// Check if this is a Claude Sonnet 4.x model that supports 1M context
// The Sonnet 4 models are:
// - claude-sonnet-4-20250514
// - claude-sonnet-4-5
const isSonnet4x = options.baseModel === "claude-sonnet-4-20250514" || options.baseModel === "claude-sonnet-4-5"
if (!isSonnet4x) return false
// Use 1M when context exceeds threshold (default 190K - leaves 10K buffer before 200K limit)
const threshold = options.threshold ?? 190_000
return options.contextTokens > threshold
}
// Reasoning
export const shouldUseReasoningBudget = ({

View file

@ -280,10 +280,10 @@
"anthropicApiKey": "Anthropic API Key",
"getAnthropicApiKey": "Get Anthropic API Key",
"anthropicUseAuthToken": "Pass Anthropic API Key as Authorization header instead of X-Api-Key",
"anthropic1MContextBetaLabel": "Enable 1M context window (Beta)",
"anthropic1MContextBetaDescription": "Extends context window to 1 million tokens for Claude Sonnet 4",
"awsBedrock1MContextBetaLabel": "Enable 1M context window (Beta)",
"awsBedrock1MContextBetaDescription": "Extends context window to 1 million tokens for Claude Sonnet 4",
"anthropic1MContextBetaLabel": "Enable Dynamic 1M context window (Beta)",
"anthropic1MContextBetaDescription": "Automatically switches between 200K and 1M context windows for Claude Sonnet 4 based on actual context size",
"awsBedrock1MContextBetaLabel": "Enable Dynamic 1M context window (Beta)",
"awsBedrock1MContextBetaDescription": "Automatically switches between 200K and 1M context windows for Claude Sonnet 4 based on actual context size",
"cerebrasApiKey": "Cerebras API Key",
"getCerebrasApiKey": "Get Cerebras API Key",
"chutesApiKey": "Chutes API Key",