Skip to content
This repository was archived by the owner on May 15, 2026. It is now read-only.

Commit 91fb55d

Browse files
hannesrudolphhannesrudolph
committed
fix(vertex): implement 1M context beta functionality
- Add tier pricing for Claude Sonnet 4 models with 1M context support - Implement getModel() to return betas array when vertex1MContext enabled - Update createMessage() to pass anthropic-beta header via requestOptions - Add comprehensive tests for 1M context beta feature - Update existing tests to handle the new requestOptions parameter The vertex1MContext checkbox now properly enables the 1M token context window for supported Vertex AI Claude models (claude-sonnet-4@20250514, claude-sonnet-4-5@20250929) by passing the 'context-1m-2025-08-07' beta header and updating the model pricing accordingly.
1 parent 020cd61 commit 91fb55d

3 files changed

Lines changed: 268 additions & 43 deletions

File tree

‎packages/types/src/providers/vertex.ts‎

Lines changed: 30 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -275,29 +275,49 @@ export const vertexModels = {
275275
},
276276
"claude-sonnet-4@20250514": {
277277
maxTokens: 8192,
278-
contextWindow: 200_000,
278+
contextWindow: 200_000, // Default 200K, extendable to 1M with beta flag 'context-1m-2025-08-07'
279279
supportsImages: true,
280280
supportsPromptCache: true,
281281
supportsNativeTools: true,
282282
defaultToolProtocol: "native",
283-
inputPrice: 3.0,
284-
outputPrice: 15.0,
285-
cacheWritesPrice: 3.75,
286-
cacheReadsPrice: 0.3,
283+
inputPrice: 3.0, // $3 per million input tokens (≤200K context)
284+
outputPrice: 15.0, // $15 per million output tokens (≤200K context)
285+
cacheWritesPrice: 3.75, // $3.75 per million tokens
286+
cacheReadsPrice: 0.3, // $0.30 per million tokens
287287
supportsReasoningBudget: true,
288+
// Tiered pricing for extended context (requires beta flag 'context-1m-2025-08-07')
289+
tiers: [
290+
{
291+
contextWindow: 1_000_000, // 1M tokens with beta flag
292+
inputPrice: 6.0, // $6 per million input tokens (>200K context)
293+
outputPrice: 22.5, // $22.50 per million output tokens (>200K context)
294+
cacheWritesPrice: 7.5, // $7.50 per million tokens (>200K context)
295+
cacheReadsPrice: 0.6, // $0.60 per million tokens (>200K context)
296+
},
297+
],
288298
},
289299
"claude-sonnet-4-5@20250929": {
290300
maxTokens: 8192,
291-
contextWindow: 200_000,
301+
contextWindow: 200_000, // Default 200K, extendable to 1M with beta flag 'context-1m-2025-08-07'
292302
supportsImages: true,
293303
supportsPromptCache: true,
294304
supportsNativeTools: true,
295305
defaultToolProtocol: "native",
296-
inputPrice: 3.0,
297-
outputPrice: 15.0,
298-
cacheWritesPrice: 3.75,
299-
cacheReadsPrice: 0.3,
306+
inputPrice: 3.0, // $3 per million input tokens (≤200K context)
307+
outputPrice: 15.0, // $15 per million output tokens (≤200K context)
308+
cacheWritesPrice: 3.75, // $3.75 per million tokens
309+
cacheReadsPrice: 0.3, // $0.30 per million tokens
300310
supportsReasoningBudget: true,
311+
// Tiered pricing for extended context (requires beta flag 'context-1m-2025-08-07')
312+
tiers: [
313+
{
314+
contextWindow: 1_000_000, // 1M tokens with beta flag
315+
inputPrice: 6.0, // $6 per million input tokens (>200K context)
316+
outputPrice: 22.5, // $22.50 per million output tokens (>200K context)
317+
cacheWritesPrice: 7.5, // $7.50 per million tokens (>200K context)
318+
cacheReadsPrice: 0.6, // $0.60 per million tokens (>200K context)
319+
},
320+
],
301321
},
302322
"claude-haiku-4-5@20251001": {
303323
maxTokens: 8192,

‎src/api/providers/__tests__/anthropic-vertex.spec.ts‎

Lines changed: 195 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@
33
import { Anthropic } from "@anthropic-ai/sdk"
44
import { AnthropicVertex } from "@anthropic-ai/vertex-sdk"
55

6+
import { VERTEX_1M_CONTEXT_MODEL_IDS } from "@roo-code/types"
7+
68
import { ApiStreamChunk } from "../../transform/stream"
79

810
import { AnthropicVertexHandler } from "../anthropic-vertex"
@@ -159,35 +161,39 @@ describe("VertexHandler", () => {
159161
outputTokens: 5,
160162
})
161163

162-
expect(mockCreate).toHaveBeenCalledWith({
163-
model: "claude-3-5-sonnet-v2@20241022",
164-
max_tokens: 8192,
165-
temperature: 0,
166-
system: [
167-
{
168-
type: "text",
169-
text: "You are a helpful assistant",
170-
cache_control: { type: "ephemeral" },
171-
},
172-
],
173-
messages: [
174-
{
175-
role: "user",
176-
content: [
177-
{
178-
type: "text",
179-
text: "Hello",
180-
cache_control: { type: "ephemeral" },
181-
},
182-
],
183-
},
184-
{
185-
role: "assistant",
186-
content: "Hi there!",
187-
},
188-
],
189-
stream: true,
190-
})
164+
expect(mockCreate).toHaveBeenCalledWith(
165+
{
166+
model: "claude-3-5-sonnet-v2@20241022",
167+
max_tokens: 8192,
168+
temperature: 0,
169+
thinking: undefined,
170+
system: [
171+
{
172+
type: "text",
173+
text: "You are a helpful assistant",
174+
cache_control: { type: "ephemeral" },
175+
},
176+
],
177+
messages: [
178+
{
179+
role: "user",
180+
content: [
181+
{
182+
type: "text",
183+
text: "Hello",
184+
cache_control: { type: "ephemeral" },
185+
},
186+
],
187+
},
188+
{
189+
role: "assistant",
190+
content: "Hi there!",
191+
},
192+
],
193+
stream: true,
194+
},
195+
undefined,
196+
)
191197
})
192198

193199
it("should handle multiple content blocks with line breaks for Claude", async () => {
@@ -401,6 +407,7 @@ describe("VertexHandler", () => {
401407
}),
402408
],
403409
}),
410+
undefined,
404411
)
405412
})
406413

@@ -858,6 +865,162 @@ describe("VertexHandler", () => {
858865
expect(result.reasoningBudget).toBeUndefined()
859866
expect(result.temperature).toBe(0)
860867
})
868+
869+
it("should enable 1M context for Claude Sonnet 4 when beta flag is set", () => {
870+
const handler = new AnthropicVertexHandler({
871+
apiModelId: VERTEX_1M_CONTEXT_MODEL_IDS[0],
872+
vertexProjectId: "test-project",
873+
vertexRegion: "us-central1",
874+
vertex1MContext: true,
875+
})
876+
877+
const model = handler.getModel()
878+
expect(model.info.contextWindow).toBe(1_000_000)
879+
expect(model.info.inputPrice).toBe(6.0)
880+
expect(model.info.outputPrice).toBe(22.5)
881+
expect(model.betas).toContain("context-1m-2025-08-07")
882+
})
883+
884+
it("should enable 1M context for Claude Sonnet 4.5 when beta flag is set", () => {
885+
const handler = new AnthropicVertexHandler({
886+
apiModelId: VERTEX_1M_CONTEXT_MODEL_IDS[1],
887+
vertexProjectId: "test-project",
888+
vertexRegion: "us-central1",
889+
vertex1MContext: true,
890+
})
891+
892+
const model = handler.getModel()
893+
expect(model.info.contextWindow).toBe(1_000_000)
894+
expect(model.info.inputPrice).toBe(6.0)
895+
expect(model.info.outputPrice).toBe(22.5)
896+
expect(model.betas).toContain("context-1m-2025-08-07")
897+
})
898+
899+
it("should not enable 1M context when flag is disabled", () => {
900+
const handler = new AnthropicVertexHandler({
901+
apiModelId: VERTEX_1M_CONTEXT_MODEL_IDS[0],
902+
vertexProjectId: "test-project",
903+
vertexRegion: "us-central1",
904+
vertex1MContext: false,
905+
})
906+
907+
const model = handler.getModel()
908+
expect(model.info.contextWindow).toBe(200_000)
909+
expect(model.info.inputPrice).toBe(3.0)
910+
expect(model.info.outputPrice).toBe(15.0)
911+
expect(model.betas).toBeUndefined()
912+
})
913+
914+
it("should not enable 1M context for non-supported models even with flag", () => {
915+
const handler = new AnthropicVertexHandler({
916+
apiModelId: "claude-3-5-sonnet-v2@20241022",
917+
vertexProjectId: "test-project",
918+
vertexRegion: "us-central1",
919+
vertex1MContext: true,
920+
})
921+
922+
const model = handler.getModel()
923+
expect(model.info.contextWindow).toBe(200_000)
924+
expect(model.betas).toBeUndefined()
925+
})
926+
})
927+
928+
describe("1M context beta header", () => {
929+
const mockMessages: Anthropic.Messages.MessageParam[] = [
930+
{
931+
role: "user",
932+
content: "Hello",
933+
},
934+
]
935+
936+
const systemPrompt = "You are a helpful assistant"
937+
938+
it("should include anthropic-beta header when 1M context is enabled", async () => {
939+
const handler = new AnthropicVertexHandler({
940+
apiModelId: VERTEX_1M_CONTEXT_MODEL_IDS[0],
941+
vertexProjectId: "test-project",
942+
vertexRegion: "us-central1",
943+
vertex1MContext: true,
944+
})
945+
946+
const mockStream = [
947+
{
948+
type: "message_start",
949+
message: {
950+
usage: {
951+
input_tokens: 10,
952+
output_tokens: 0,
953+
},
954+
},
955+
},
956+
]
957+
958+
const asyncIterator = {
959+
async *[Symbol.asyncIterator]() {
960+
for (const chunk of mockStream) {
961+
yield chunk
962+
}
963+
},
964+
}
965+
966+
const mockCreate = vitest.fn().mockResolvedValue(asyncIterator)
967+
;(handler["client"].messages as any).create = mockCreate
968+
969+
const stream = handler.createMessage(systemPrompt, mockMessages)
970+
971+
for await (const _chunk of stream) {
972+
// Just consume
973+
}
974+
975+
// Verify the API was called with the beta header
976+
expect(mockCreate).toHaveBeenCalledWith(
977+
expect.anything(),
978+
expect.objectContaining({
979+
headers: { "anthropic-beta": "context-1m-2025-08-07" },
980+
}),
981+
)
982+
})
983+
984+
it("should not include anthropic-beta header when 1M context is disabled", async () => {
985+
const handler = new AnthropicVertexHandler({
986+
apiModelId: VERTEX_1M_CONTEXT_MODEL_IDS[0],
987+
vertexProjectId: "test-project",
988+
vertexRegion: "us-central1",
989+
vertex1MContext: false,
990+
})
991+
992+
const mockStream = [
993+
{
994+
type: "message_start",
995+
message: {
996+
usage: {
997+
input_tokens: 10,
998+
output_tokens: 0,
999+
},
1000+
},
1001+
},
1002+
]
1003+
1004+
const asyncIterator = {
1005+
async *[Symbol.asyncIterator]() {
1006+
for (const chunk of mockStream) {
1007+
yield chunk
1008+
}
1009+
},
1010+
}
1011+
1012+
const mockCreate = vitest.fn().mockResolvedValue(asyncIterator)
1013+
;(handler["client"].messages as any).create = mockCreate
1014+
1015+
const stream = handler.createMessage(systemPrompt, mockMessages)
1016+
1017+
for await (const _chunk of stream) {
1018+
// Just consume
1019+
}
1020+
1021+
// Verify the API was called without the beta header
1022+
expect(mockCreate).toHaveBeenCalledWith(expect.anything(), undefined)
1023+
})
8611024
})
8621025

8631026
describe("thinking model configuration", () => {
@@ -946,6 +1109,7 @@ describe("VertexHandler", () => {
9461109
thinking: { type: "enabled", budget_tokens: 4096 },
9471110
temperature: 1.0, // Thinking requires temperature 1.0
9481111
}),
1112+
undefined,
9491113
)
9501114
})
9511115
})
@@ -1032,6 +1196,7 @@ describe("VertexHandler", () => {
10321196
]),
10331197
tool_choice: { type: "auto", disable_parallel_tool_use: true },
10341198
}),
1199+
undefined,
10351200
)
10361201
})
10371202

@@ -1080,6 +1245,7 @@ describe("VertexHandler", () => {
10801245
expect.not.objectContaining({
10811246
tools: expect.anything(),
10821247
}),
1248+
undefined,
10831249
)
10841250
})
10851251

0 commit comments

Comments
 (0)