From 2233a14a87587cd22ab1e69091ec88ed37033404 Mon Sep 17 00:00:00 2001 From: SIGTERM Date: Fri, 18 Sep 2026 16:30:45 +0200 Subject: [PATCH] fix(vertex): preserve Claude prompt caching and usage metadata (#13220) * fix(vertex): preserve Claude prompt caching * fix(vertex): normalize unsupported cache TTLs * docs(changelog): note Vertex prompt caching fix --- .../13220-vertex-claude-prompt-caching.md | 1 + open-sse/executors/vertex.ts | 52 +++++++- open-sse/handlers/usageExtractor.ts | 2 + open-sse/translator/helpers/claudeHelper.ts | 29 ++++- .../translator/request/openai-to-claude.ts | 9 +- open-sse/utils/cacheControlPolicy.ts | 8 +- tests/unit/cache-control-policy.test.ts | 19 +++ .../claude-cache-control-passthrough.test.ts | 57 +++++++++ tests/unit/executor-vertex-extended.test.ts | 111 +++++++++++++++++- tests/unit/usage-extractor.test.ts | 23 ++++ 10 files changed, 300 insertions(+), 11 deletions(-) create mode 100644 changelog.d/fixes/13220-vertex-claude-prompt-caching.md diff --git a/changelog.d/fixes/13220-vertex-claude-prompt-caching.md b/changelog.d/fixes/13220-vertex-claude-prompt-caching.md new file mode 100644 index 0000000000..c43d672e2b --- /dev/null +++ b/changelog.d/fixes/13220-vertex-claude-prompt-caching.md @@ -0,0 +1 @@ +- **fix(vertex):** preserve Claude prompt-cache breakpoints for Vertex and Vertex Partner, use the documented five-minute ephemeral TTL by default, and forward cache usage metadata through streaming responses ([#13220](https://github.com/diegosouzapw/OmniRoute/pull/13220)) — fixes #13219 diff --git a/open-sse/executors/vertex.ts b/open-sse/executors/vertex.ts index 80fc0c7879..c244a59988 100644 --- a/open-sse/executors/vertex.ts +++ b/open-sse/executors/vertex.ts @@ -218,6 +218,45 @@ function buildProjectScopedVertexUrl( return `https://aiplatform.googleapis.com/v1/projects/${project}/locations/${region}/publishers/google/models/${canonicalModel}:${operation}${querySeparator}${opaqueApiKey ? `key=${opaqueApiKey}` : ""}`; } +// Vertex does not support Anthropic's optional one-hour prompt-cache TTL on these +// legacy Claude models. Keep the breakpoint, but omit ttl so Vertex uses its +// documented five-minute ephemeral cache instead of rejecting the request. +const VERTEX_ONE_HOUR_TTL_UNSUPPORTED = new Set([ + "claude-3-7-sonnet", + "claude-3-5-sonnet-v2", + "claude-3-5-sonnet", + "claude-3-opus", +]); + +function downgradeUnsupportedVertexClaudeTtl(body: Record, model: string): void { + const normalizedModel = model.toLowerCase().split("@", 1)[0]; + if (!VERTEX_ONE_HOUR_TTL_UNSUPPORTED.has(normalizedModel)) return; + + const normalizeBlock = (block: unknown) => { + if (!block || typeof block !== "object" || Array.isArray(block)) return; + const record = block as Record; + const cacheControl = record.cache_control; + if (!cacheControl || typeof cacheControl !== "object" || Array.isArray(cacheControl)) return; + const control = cacheControl as Record; + if (control.type === "ephemeral" && control.ttl === "1h") delete control.ttl; + }; + + const system = body.system; + if (Array.isArray(system)) system.forEach(normalizeBlock); + + const messages = body.messages; + if (Array.isArray(messages)) { + for (const message of messages) { + if (!message || typeof message !== "object" || Array.isArray(message)) continue; + const content = (message as Record).content; + if (Array.isArray(content)) content.forEach(normalizeBlock); + } + } + + const tools = body.tools; + if (Array.isArray(tools)) tools.forEach(normalizeBlock); +} + // Defensive normalizer: target-format resolution for manually-added custom Claude models under // "vertex"/"vertex-partner" was observed sending a Gemini-shaped body (contents/parts) to the // Anthropic rawPredict endpoint instead of the configured "claude" format, causing a hard @@ -260,6 +299,16 @@ function synthesizeClaudeSse(response: Record): string { const stopReason = typeof response.stop_reason === "string" ? response.stop_reason : "end_turn"; const stopSequence = (response.stop_sequence as string | null | undefined) ?? null; const content = Array.isArray(response.content) ? response.content : []; + const inputUsage: Record = { + input_tokens: usage.input_tokens || 0, + output_tokens: 0, + }; + if (typeof usage.cache_creation_input_tokens === "number") { + inputUsage.cache_creation_input_tokens = usage.cache_creation_input_tokens; + } + if (typeof usage.cache_read_input_tokens === "number") { + inputUsage.cache_read_input_tokens = usage.cache_read_input_tokens; + } const events: Array<{ event: string; data: Record }> = []; @@ -275,7 +324,7 @@ function synthesizeClaudeSse(response: Record): string { model, stop_reason: null, stop_sequence: null, - usage: { input_tokens: usage.input_tokens || 0, output_tokens: 0 }, + usage: inputUsage, }, }, }); @@ -408,6 +457,7 @@ export class VertexExecutor extends BaseExecutor { // "model: Extra inputs are not permitted" if the translated request body still carries // one (the openai→claude request translator copies the client's model field over). delete body.model; + downgradeUnsupportedVertexClaudeTtl(body, model); } const result = await super.execute(input); diff --git a/open-sse/handlers/usageExtractor.ts b/open-sse/handlers/usageExtractor.ts index e0076112a9..2a7ccc036a 100644 --- a/open-sse/handlers/usageExtractor.ts +++ b/open-sse/handlers/usageExtractor.ts @@ -10,6 +10,8 @@ export function extractUsageFromResponse(responseBody, provider) { const isClaudeProvider = providerId === "claude" || providerId === "anthropic" || + providerId === "vertex" || + providerId === "vertex-partner" || providerId.startsWith("anthropic-compatible"); // OpenAI format (has prompt_tokens / completion_tokens) diff --git a/open-sse/translator/helpers/claudeHelper.ts b/open-sse/translator/helpers/claudeHelper.ts index a3878b7256..a3a18d88fd 100644 --- a/open-sse/translator/helpers/claudeHelper.ts +++ b/open-sse/translator/helpers/claudeHelper.ts @@ -301,6 +301,20 @@ function markMessageCacheControl(msg: ClaudeMessage, ttl?: string): boolean { return true; } +/** + * Build the cache marker OmniRoute adds when the client did not provide one. + * Vertex defaults to five minutes when ttl is omitted; unlike 1h, that mode is + * supported by every cache-capable Claude model on Vertex and has cheaper writes. + */ +export function createDefaultClaudeCacheControl(provider?: string | null): { + type: string; + ttl?: string; +} { + return provider === "vertex" || provider === "vertex-partner" + ? { type: "ephemeral" } + : { type: "ephemeral", ttl: "1h" }; +} + /** True when the body carries at least one cache_control marker anywhere * (system blocks, message content blocks, or tools). Used to decide whether * preserve-mode has anything to preserve. */ @@ -361,10 +375,15 @@ export function prepareClaudeRequest( preserveCacheControl = false; } - // 1. System: remove all cache_control, add only to last block with ttl 1h + // 1. System: remove all cache_control, add only to the last block with the provider TTL // In passthrough mode, preserve existing cache_control markers + const isVertexClaudeProvider = provider === "vertex" || provider === "vertex-partner"; const supportsPromptCaching = - provider === "claude" || provider?.startsWith?.("anthropic-compatible-"); + provider === "claude" || + isVertexClaudeProvider || + provider?.startsWith?.("anthropic-compatible-"); + // Vertex's documented default is a five-minute ephemeral cache. Omitting ttl is + // both cheaper and compatible with Claude models that reject the optional 1h TTL. const isKimiCoding = provider === "kimi-coding" || provider === "kimi-coding-apikey"; // Non-Anthropic Claude-shape providers (kimi-coding, glmt, zai, …) cannot @@ -389,7 +408,7 @@ export function prepareClaudeRequest( body.system = systemBlocks.map((block, i) => { const { cache_control, ...rest } = block; if (i === systemBlocks.length - 1 && supportsPromptCaching) { - return { ...rest, cache_control: { type: "ephemeral", ttl: "1h" } }; + return { ...rest, cache_control: createDefaultClaudeCacheControl(provider) }; } return rest; }); @@ -721,7 +740,7 @@ export function prepareClaudeRequest( } } - // 3. Tools: remove all cache_control, add only to last non-deferred tool with ttl 1h + // 3. Tools: remove all cache_control, add only to the last non-deferred tool // Tools with defer_loading=true cannot have cache_control (API rejects it) // In passthrough mode, preserve existing cache_control markers if (body.tools && Array.isArray(body.tools) && !preserveCacheControl) { @@ -732,7 +751,7 @@ export function prepareClaudeRequest( if (supportsPromptCaching) { for (let i = body.tools.length - 1; i >= 0; i--) { if (!body.tools[i].defer_loading) { - body.tools[i].cache_control = { type: "ephemeral", ttl: "1h" }; + body.tools[i].cache_control = createDefaultClaudeCacheControl(provider); break; } } diff --git a/open-sse/translator/request/openai-to-claude.ts b/open-sse/translator/request/openai-to-claude.ts index 2a636f5ab6..dd5d13f5c5 100644 --- a/open-sse/translator/request/openai-to-claude.ts +++ b/open-sse/translator/request/openai-to-claude.ts @@ -5,7 +5,10 @@ import { supportsClaudeMaxEffort, supportsXHighEffort } from "../../config/provi import { adjustMaxTokens } from "../helpers/maxTokensHelper.ts"; import { sanitizeToolId } from "../helpers/schemaCoercion.ts"; import { safeParseJSON } from "../helpers/jsonUtil.ts"; -import { applyKimiCodingThinking } from "../helpers/claudeHelper.ts"; +import { + applyKimiCodingThinking, + createDefaultClaudeCacheControl, +} from "../helpers/claudeHelper.ts"; import { DEFAULT_THINKING_CLAUDE_SIGNATURE } from "../../config/defaultThinkingSignature.ts"; import { getDefaultThinkingBudget, @@ -454,7 +457,7 @@ export function openaiToClaudeRequest(model, body, stream, credentials = null) { // rejects cache_control on defer_loading tools. for (let i = result.tools.length - 1; i >= 0; i--) { if (!result.tools[i].defer_loading) { - result.tools[i].cache_control = { type: "ephemeral", ttl: "1h" }; + result.tools[i].cache_control = createDefaultClaudeCacheControl(routedProvider); break; } } @@ -491,7 +494,7 @@ export function openaiToClaudeRequest(model, body, stream, credentials = null) { const systemBlock = { type: "text", text: systemText, - cache_control: { type: "ephemeral", ttl: "1h" }, + cache_control: createDefaultClaudeCacheControl(routedProvider), }; // Merge with existing body.system if present if (Array.isArray(body.system)) { diff --git a/open-sse/utils/cacheControlPolicy.ts b/open-sse/utils/cacheControlPolicy.ts index 7e1abe6c54..76d4f74024 100644 --- a/open-sse/utils/cacheControlPolicy.ts +++ b/open-sse/utils/cacheControlPolicy.ts @@ -216,7 +216,13 @@ export function providerSupportsCaching( return connectionCacheOverride.supportsPromptCaching; } if (!provider) return false; - if (CACHING_PROVIDERS.has(provider.toLowerCase())) return true; + const providerId = provider.toLowerCase(); + // Vertex is a mixed-format provider. Only its Anthropic Claude path accepts + // cache_control; Gemini and OpenAI-format partner models use other mechanisms. + if (providerId === "vertex" || providerId === "vertex-partner") { + return targetFormat?.toLowerCase() === "claude"; + } + if (CACHING_PROVIDERS.has(providerId)) return true; // All Claude-protocol providers support prompt caching if (targetFormat === "claude") return true; return false; diff --git a/tests/unit/cache-control-policy.test.ts b/tests/unit/cache-control-policy.test.ts index 934c39a8c7..6dae448edc 100644 --- a/tests/unit/cache-control-policy.test.ts +++ b/tests/unit/cache-control-policy.test.ts @@ -45,6 +45,11 @@ describe("Cache Control Policy", () => { assert.equal(providerSupportsCaching("openai"), true); assert.equal(providerSupportsCaching("codex"), true); assert.equal(providerSupportsCaching("azure"), true); + // Vertex is mixed-format: only Claude partner models use cache_control. + assert.equal(providerSupportsCaching("vertex", "claude"), true); + assert.equal(providerSupportsCaching("vertex-partner", "claude"), true); + assert.equal(providerSupportsCaching("vertex", "gemini"), false); + assert.equal(providerSupportsCaching("vertex-partner", "gemini"), false); }); test("rejects non-caching providers", () => { @@ -108,6 +113,20 @@ describe("Cache Control Policy", () => { ); }); + test("preserves Claude Code cache markers for both Vertex provider IDs", () => { + for (const targetProvider of ["vertex", "vertex-partner"]) { + assert.equal( + shouldPreserveCacheControl({ + userAgent: "claude-code/0.1.0", + isCombo: false, + targetProvider, + targetFormat: "claude", + }), + true + ); + } + }); + test("preserves for combo with priority strategy + Claude client + caching provider", () => { assert.equal( shouldPreserveCacheControl({ diff --git a/tests/unit/claude-cache-control-passthrough.test.ts b/tests/unit/claude-cache-control-passthrough.test.ts index 3b217175ac..31deabfe5f 100644 --- a/tests/unit/claude-cache-control-passthrough.test.ts +++ b/tests/unit/claude-cache-control-passthrough.test.ts @@ -1,6 +1,8 @@ import { describe, test } from "node:test"; import assert from "node:assert/strict"; import { prepareClaudeRequest } from "../../open-sse/translator/helpers/claudeHelper.ts"; +import { FORMATS } from "../../open-sse/translator/formats.ts"; +import { translateRequest } from "../../open-sse/translator/index.ts"; describe("Claude cache_control passthrough", () => { test("preserveCacheControl=true preserves cache_control in system blocks", () => { @@ -181,4 +183,59 @@ describe("Claude cache_control passthrough", () => { assert.equal(result.messages[2].content[0].cache_control, undefined); assert.deepEqual(result.tools[0].cache_control, { type: "ephemeral", ttl: "5m" }); }); + + for (const provider of ["vertex", "vertex-partner"]) { + test(`${provider} supports prompt caching with the cost-sensitive default TTL`, () => { + const body = { + system: [{ type: "text", text: "Stable system prefix" }], + messages: [{ role: "user", content: [{ type: "text", text: "Dynamic question" }] }], + tools: [ + { + name: "lookup", + description: "Stable tool definition", + input_schema: { type: "object" }, + }, + ], + }; + + const result = prepareClaudeRequest(body, provider, false, "claude-sonnet-4-6"); + + // Omitting ttl selects Anthropic's 5-minute default. Vertex does not support + // ttl:1h on every Claude model, and one-hour writes cost more. + assert.deepEqual(result.system[0].cache_control, { type: "ephemeral" }); + assert.deepEqual(result.tools[0].cache_control, { type: "ephemeral" }); + }); + + test(`${provider} uses the five-minute default through the full translation path`, () => { + const result = translateRequest( + FORMATS.OPENAI, + FORMATS.CLAUDE, + "claude-3-7-sonnet", + { + messages: [ + { role: "system", content: "Stable system prefix" }, + { role: "user", content: "Dynamic question" }, + ], + tools: [ + { + type: "function", + function: { + name: "lookup", + description: "Stable tool definition", + parameters: { type: "object" }, + }, + }, + ], + }, + true, + null, + provider, + null, + { preserveCacheControl: true } + ); + + assert.deepEqual(result.system[0].cache_control, { type: "ephemeral" }); + assert.deepEqual(result.tools[0].cache_control, { type: "ephemeral" }); + }); + } }); diff --git a/tests/unit/executor-vertex-extended.test.ts b/tests/unit/executor-vertex-extended.test.ts index 4efc1cac99..48074323a6 100644 --- a/tests/unit/executor-vertex-extended.test.ts +++ b/tests/unit/executor-vertex-extended.test.ts @@ -401,6 +401,104 @@ test("VertexExecutor.execute strips the client's model field and injects anthrop } }); +test("VertexExecutor downgrades unsupported Claude 1h cache TTLs without changing supported or 5m TTLs", async () => { + const executor = new VertexExecutor(); + const originalFetch = globalThis.fetch; + type CapturedBody = { + system: Array<{ cache_control?: Record }>; + messages: Array<{ content: Array<{ cache_control?: Record }> }>; + tools: Array<{ cache_control?: Record }>; + }; + const sentBodies: CapturedBody[] = []; + + globalThis.fetch = async (_url, options) => { + sentBodies.push(JSON.parse(String(options?.body || "{}")) as CapturedBody); + return new Response( + JSON.stringify({ + id: "msg_cache_ttl", + type: "message", + role: "assistant", + model: "claude-test", + content: [{ type: "text", text: "ok" }], + stop_reason: "end_turn", + usage: { input_tokens: 1, output_tokens: 1 }, + }), + { status: 200, headers: { "Content-Type": "application/json" } } + ); + }; + + const credentials = { + apiKey: createServiceAccountJson({ projectId: "proj-claude-cache" }), + accessToken: "ya29.claude-cache", + }; + const body = { + system: [{ type: "text", text: "stable", cache_control: { type: "ephemeral", ttl: "1h" } }], + messages: [ + { + role: "user", + content: [ + { type: "text", text: "question", cache_control: { type: "ephemeral", ttl: "1h" } }, + { type: "text", text: "five-minute", cache_control: { type: "ephemeral", ttl: "5m" } }, + { type: "text", text: "no ttl", cache_control: { type: "ephemeral" } }, + ], + }, + ], + tools: [ + { + name: "lookup", + input_schema: { type: "object" }, + cache_control: { type: "ephemeral", ttl: "1h" }, + }, + ], + }; + + try { + for (const model of [ + "claude-3-7-sonnet", + "claude-3-5-sonnet-v2@20241022", + "claude-3-5-sonnet", + "claude-3-opus@20240229", + ]) { + await executor.execute({ + model, + body: structuredClone(body), + stream: false, + credentials: { ...credentials }, + }); + } + + await executor.execute({ + model: "claude-sonnet-4-6", + body: structuredClone(body), + stream: false, + credentials: { ...credentials }, + }); + + const unsupportedBodies = sentBodies.slice(0, 4); + for (const sent of unsupportedBodies) { + assert.deepEqual(sent.system[0].cache_control, { type: "ephemeral" }); + assert.deepEqual(sent.messages[0].content[0].cache_control, { type: "ephemeral" }); + assert.deepEqual(sent.messages[0].content[1].cache_control, { + type: "ephemeral", + ttl: "5m", + }); + assert.deepEqual(sent.messages[0].content[2].cache_control, { type: "ephemeral" }); + assert.deepEqual(sent.tools[0].cache_control, { type: "ephemeral" }); + } + + assert.deepEqual(sentBodies[4].system[0].cache_control, { + type: "ephemeral", + ttl: "1h", + }); + assert.deepEqual(sentBodies[4].tools[0].cache_control, { + type: "ephemeral", + ttl: "1h", + }); + } finally { + globalThis.fetch = originalFetch; + } +}); + test("VertexExecutor.execute synthesizes a genuine Anthropic-format SSE stream when rawPredict returns a complete JSON body for a streaming request", async () => { const executor = new VertexExecutor(); const originalFetch = globalThis.fetch; @@ -420,7 +518,12 @@ test("VertexExecutor.execute synthesizes a genuine Anthropic-format SSE stream w content: [{ type: "text", text: "hello" }], stop_reason: "end_turn", stop_sequence: null, - usage: { input_tokens: 5, output_tokens: 2 }, + usage: { + input_tokens: 5, + output_tokens: 2, + cache_creation_input_tokens: 1_024, + cache_read_input_tokens: 4_096, + }, }), { status: 200, headers: { "Content-Type": "application/json" } } ); @@ -458,6 +561,12 @@ test("VertexExecutor.execute synthesizes a genuine Anthropic-format SSE stream w "message_delta", "message_stop", ]); + assert.deepEqual(dataLines[0].message.usage, { + input_tokens: 5, + output_tokens: 0, + cache_creation_input_tokens: 1_024, + cache_read_input_tokens: 4_096, + }); } finally { globalThis.fetch = originalFetch; } diff --git a/tests/unit/usage-extractor.test.ts b/tests/unit/usage-extractor.test.ts index af328e52b6..58c94cf224 100644 --- a/tests/unit/usage-extractor.test.ts +++ b/tests/unit/usage-extractor.test.ts @@ -192,6 +192,29 @@ test("extractUsageFromResponse totals Claude prompt tokens with cache read and c }); }); +for (const provider of ["vertex", "vertex-partner"]) { + test(`extractUsageFromResponse totals Claude cache tokens for ${provider}`, () => { + const usage = extractUsageFromResponse( + { + usage: { + input_tokens: 10, + output_tokens: 7, + cache_read_input_tokens: 4_000, + cache_creation_input_tokens: 1_000, + }, + }, + provider + ); + + assert.deepEqual(usage, { + prompt_tokens: 5_010, + completion_tokens: 7, + cache_read_input_tokens: 4_000, + cache_creation_input_tokens: 1_000, + }); + }); +} + test("extractUsageFromResponse surfaces Claude thinking tokens without inflating completion", () => { const usage = extractUsageFromResponse( {