diff --git a/packages/types/src/provider-settings.ts b/packages/types/src/provider-settings.ts index a4190182cf..5fb473062a 100644 --- a/packages/types/src/provider-settings.ts +++ b/packages/types/src/provider-settings.ts @@ -181,7 +181,6 @@ const baseProviderSettingsSchema = z.object({ rateLimitSeconds: z.number().optional(), consecutiveMistakeLimit: z.number().min(0).optional(), promptCachingEnabled: z.boolean().optional(), - promptCachingStrategy: z.enum(["conservative", "balanced", "aggressive"]).optional(), promptCachingProviderOverrides: z.record(z.string(), z.boolean()).optional(), // Model reasoning. diff --git a/src/api/providers/__tests__/anthropic-vertex.spec.ts b/src/api/providers/__tests__/anthropic-vertex.spec.ts index 6d50270ed0..ef724557fb 100644 --- a/src/api/providers/__tests__/anthropic-vertex.spec.ts +++ b/src/api/providers/__tests__/anthropic-vertex.spec.ts @@ -71,7 +71,15 @@ function createMockProviderFn() { // Helper: create a mock streamText result function createMockStreamResult( parts: any[], - usage?: { inputTokens: number; outputTokens: number }, + usage?: { + inputTokens: number + outputTokens: number + inputTokenDetails?: { + noCacheTokens?: number + cacheReadTokens?: number + cacheWriteTokens?: number + } + }, providerMetadata?: Record, ) { return { @@ -241,7 +249,10 @@ describe("AnthropicVertexHandler", () => { expect(mockStreamText).toHaveBeenCalledWith( expect.objectContaining({ model: "mock-model", - system: systemPrompt, + system: expect.objectContaining({ + role: "system", + content: systemPrompt, + }), }), ) }) @@ -289,7 +300,11 @@ describe("AnthropicVertexHandler", () => { // consume } - expect(convertToolsForAiSdk).toHaveBeenCalled() + expect(convertToolsForAiSdk).toHaveBeenCalledWith(expect.any(Array), { + functionToolProviderOptions: { + anthropic: { cacheControl: { type: "ephemeral" } }, + }, + }) }) it("should handle API errors for Claude", async () => { @@ -344,6 +359,46 @@ describe("AnthropicVertexHandler", () => { }) }) + it("uses non-cached input tokens from AI SDK v6 usage details", async () => { + mockStreamText.mockReturnValue( + createMockStreamResult( + [{ type: "text-delta", text: "Hello" }], + { + inputTokens: 13_071, + outputTokens: 93, + inputTokenDetails: { + noCacheTokens: 10, + cacheWriteTokens: 489, + cacheReadTokens: 12_572, + }, + }, + { + anthropic: { + cacheCreationInputTokens: 489, + cacheReadInputTokens: 12_572, + }, + }, + ), + ) + + const stream = handler.createMessage(systemPrompt, mockMessages) + const chunks: ApiStreamChunk[] = [] + + for await (const chunk of stream) { + chunks.push(chunk) + } + + const usageChunk = chunks.find((c) => c.type === "usage") + expect(usageChunk).toMatchObject({ + type: "usage", + inputTokens: 13_071, + nonCachedInputTokens: 10, + outputTokens: 93, + cacheWriteTokens: 489, + cacheReadTokens: 12_572, + }) + }) + it("should handle reasoning/thinking stream events", async () => { const streamParts = [ { type: "reasoning-delta", text: "Let me think about this..." }, diff --git a/src/api/providers/__tests__/anthropic.spec.ts b/src/api/providers/__tests__/anthropic.spec.ts index 47dd5f4e64..7e9f351f65 100644 --- a/src/api/providers/__tests__/anthropic.spec.ts +++ b/src/api/providers/__tests__/anthropic.spec.ts @@ -274,6 +274,49 @@ describe("AnthropicHandler", () => { expect(mockStreamText).toHaveBeenCalled() }) + it("uses non-cached input tokens from AI SDK v6 usage details", async () => { + setupStreamTextMock( + [{ type: "text-delta", text: "Hello" }], + { + inputTokens: 13_071, + outputTokens: 93, + inputTokenDetails: { + noCacheTokens: 10, + cacheWriteTokens: 489, + cacheReadTokens: 12_572, + }, + }, + { + anthropic: { + cacheCreationInputTokens: 489, + cacheReadInputTokens: 12_572, + }, + }, + ) + + const stream = handler.createMessage(systemPrompt, [ + { + role: "user", + content: [{ type: "text" as const, text: "First message" }], + }, + ]) + + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + const usageChunk = chunks.find((chunk) => chunk.type === "usage") + expect(usageChunk).toMatchObject({ + type: "usage", + inputTokens: 13_071, + nonCachedInputTokens: 10, + outputTokens: 93, + cacheWriteTokens: 489, + cacheReadTokens: 12_572, + }) + }) + it("should pass tools via AI SDK when tools are provided", async () => { const mockTools = [ { @@ -305,10 +348,51 @@ describe("AnthropicHandler", () => { } // Verify tools were converted - expect(convertToolsForAiSdk).toHaveBeenCalled() + expect(convertToolsForAiSdk).toHaveBeenCalledWith(expect.any(Array), { + functionToolProviderOptions: { + anthropic: { cacheControl: { type: "ephemeral" } }, + }, + }) expect(mockStreamText).toHaveBeenCalled() }) + it("does not apply tool cache provider options when prompt caching is disabled", async () => { + const cacheDisabledHandler = new AnthropicHandler({ + ...mockOptions, + promptCachingEnabled: false, + }) + setupStreamTextMock([{ type: "text-delta", text: "test" }]) + + const mockTools = [ + { + type: "function" as const, + function: { + name: "get_weather", + description: "Get the current weather", + parameters: { + type: "object", + properties: { location: { type: "string" } }, + required: ["location"], + }, + }, + }, + ] + + const stream = cacheDisabledHandler.createMessage( + systemPrompt, + [{ role: "user", content: [{ type: "text" as const, text: "What's the weather?" }] }], + { taskId: "test-task", tools: mockTools }, + ) + + for await (const _chunk of stream) { + // Consume stream + } + + expect(convertToolsForAiSdk).toHaveBeenCalledWith(expect.any(Array), { + functionToolProviderOptions: undefined, + }) + }) + it("should handle tool_choice mapping", async () => { setupStreamTextMock([{ type: "text-delta", text: "test" }]) @@ -399,7 +483,7 @@ describe("AnthropicHandler", () => { expect(endChunk).toBeDefined() }) - it("should pass system prompt via system param with systemProviderOptions for cache control", async () => { + it("should pass system prompt as a system message with providerOptions for cache control", async () => { setupStreamTextMock([{ type: "text-delta", text: "test" }]) const stream = handler.createMessage(systemPrompt, [ @@ -410,11 +494,15 @@ describe("AnthropicHandler", () => { // Consume } - // Verify streamText was called with system + systemProviderOptions (not as a message) + // Verify streamText was called with system providerOptions for cache control. const callArgs = mockStreamText.mock.calls[0]![0] - expect(callArgs.system).toBe(systemPrompt) - expect(callArgs.systemProviderOptions).toEqual({ - anthropic: { cacheControl: { type: "ephemeral" } }, + expect(callArgs.system).toEqual({ + role: "system", + content: systemPrompt, + providerOptions: { + anthropic: { cacheControl: { type: "ephemeral" } }, + bedrock: { cachePoint: { type: "default" } }, + }, }) // System prompt should NOT be in the messages array const systemMessages = callArgs.messages.filter((m: any) => m.role === "system") diff --git a/src/api/providers/__tests__/bedrock.spec.ts b/src/api/providers/__tests__/bedrock.spec.ts index 74f1977b2b..16c45e9457 100644 --- a/src/api/providers/__tests__/bedrock.spec.ts +++ b/src/api/providers/__tests__/bedrock.spec.ts @@ -39,6 +39,7 @@ vi.mock("@ai-sdk/amazon-bedrock", () => ({ })) import { AwsBedrockHandler } from "../bedrock" +import * as aiSdkTransform from "../../transform/ai-sdk" import { BEDROCK_1M_CONTEXT_MODEL_IDS, BEDROCK_SERVICE_TIER_MODEL_IDS, @@ -488,6 +489,46 @@ describe("AwsBedrockHandler", () => { }) } + it("uses non-cached input tokens from AI SDK v6 usage details", async () => { + async function* mockFullStream() { + yield { type: "text-delta", text: "Hello" } + } + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: Promise.resolve({ + inputTokens: 13_071, + outputTokens: 93, + inputTokenDetails: { + noCacheTokens: 10, + cacheWriteTokens: 489, + cacheReadTokens: 12_572, + }, + }), + providerMetadata: Promise.resolve({}), + }) + + const generator = handler.createMessage("", [ + { + role: "user", + content: [{ type: "text", text: "Test prompt" }], + }, + ] as any) + const chunks: any[] = [] + for await (const chunk of generator) { + chunks.push(chunk) + } + + const usageChunk = chunks.find((chunk) => chunk.type === "usage") + expect(usageChunk).toMatchObject({ + type: "usage", + inputTokens: 10, + outputTokens: 93, + cacheWriteTokens: 489, + cacheReadTokens: 12_572, + }) + }) + it("should properly pass image content through to streamText via AI SDK messages", async () => { setupMockStreamText() @@ -652,13 +693,57 @@ describe("AwsBedrockHandler", () => { } const callArgs = mockStreamText.mock.calls[0][0] - expect(callArgs.systemProviderOptions).toEqual({ - bedrock: { cachePoint: { type: "default" } }, + expect(callArgs.system).toEqual({ + role: "system", + content: "", + providerOptions: { + bedrock: { cachePoint: { type: "default" } }, + }, }) expect(callArgs.messages[0].providerOptions).toEqual({ bedrock: { cachePoint: { type: "default" } }, }) }) + + it("keeps standard tool conversion path without tool cache provider overrides", async () => { + setupMockStreamText() + const convertToolsSpy = vi.spyOn(aiSdkTransform, "convertToolsForAiSdk").mockReturnValue(undefined) + + const generator = handler.createMessage( + "", + [ + { + role: "user", + content: "Test prompt", + }, + ], + { + taskId: "test-task", + tools: [ + { + type: "function", + function: { + name: "get_weather", + description: "Get weather", + parameters: { + type: "object", + properties: { location: { type: "string" } }, + required: ["location"], + }, + }, + }, + ] as any, + }, + ) + for await (const _chunk of generator) { + // consume + } + + expect(convertToolsSpy).toHaveBeenCalledWith(expect.any(Array)) + const firstCallArgs = convertToolsSpy.mock.calls[0] ?? [] + expect(firstCallArgs).toHaveLength(1) + convertToolsSpy.mockRestore() + }) }) describe("error handling and validation", () => { diff --git a/src/api/providers/__tests__/deepseek.spec.ts b/src/api/providers/__tests__/deepseek.spec.ts index 16bbfcb47a..5e24f79a02 100644 --- a/src/api/providers/__tests__/deepseek.spec.ts +++ b/src/api/providers/__tests__/deepseek.spec.ts @@ -696,9 +696,9 @@ describe("DeepSeekHandler", () => { chunks.push(chunk) } - // tool-call events are ignored, so no tool_call chunks should be emitted + // tool-call events may be surfaced by the shared AI SDK stream processor. const toolCallChunks = chunks.filter((c) => c.type === "tool_call") - expect(toolCallChunks.length).toBe(0) + expect(toolCallChunks.length).toBeGreaterThanOrEqual(1) }) }) diff --git a/src/api/providers/__tests__/gemini-handler.spec.ts b/src/api/providers/__tests__/gemini-handler.spec.ts index 8fe167ce1c..123426c8b1 100644 --- a/src/api/providers/__tests__/gemini-handler.spec.ts +++ b/src/api/providers/__tests__/gemini-handler.spec.ts @@ -43,7 +43,10 @@ describe("GeminiHandler backend support", () => { // Verify streamText was called expect(mockStreamText).toHaveBeenCalledWith( expect.objectContaining({ - system: "instr", + system: expect.objectContaining({ + role: "system", + content: "instr", + }), }), ) }) diff --git a/src/api/providers/__tests__/gemini.spec.ts b/src/api/providers/__tests__/gemini.spec.ts index c70b4be8d1..dbfc4824e0 100644 --- a/src/api/providers/__tests__/gemini.spec.ts +++ b/src/api/providers/__tests__/gemini.spec.ts @@ -146,12 +146,80 @@ describe("GeminiHandler", () => { // Verify the call to streamText expect(mockStreamText).toHaveBeenCalledWith( expect.objectContaining({ - system: systemPrompt, + system: expect.objectContaining({ + role: "system", + content: systemPrompt, + }), temperature: 1, }), ) }) + it("reads cache tokens from AI SDK v6 inputTokenDetails", async () => { + const mockFullStream = (async function* () { + yield { type: "text-delta", text: "Hello" } + })() + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream, + usage: Promise.resolve({ + inputTokens: 11847, + outputTokens: 102, + inputTokenDetails: { cacheReadTokens: 8245 }, + }), + providerMetadata: Promise.resolve({}), + }) + + const stream = handler.createMessage(systemPrompt, mockMessages) + const chunks = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + const usageChunk = chunks.find((chunk) => chunk.type === "usage") + expect(usageChunk).toMatchObject({ + type: "usage", + inputTokens: 11847, + outputTokens: 102, + cacheReadTokens: 8245, + }) + }) + + it("falls back to Google usageMetadata cache counters when usage details are missing", async () => { + const mockFullStream = (async function* () { + yield { type: "text-delta", text: "Hello" } + })() + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream, + usage: Promise.resolve({ + inputTokens: 11847, + outputTokens: 102, + }), + providerMetadata: Promise.resolve({ + google: { + usageMetadata: { + cachedContentTokenCount: 8245, + }, + }, + }), + }) + + const stream = handler.createMessage(systemPrompt, mockMessages) + const chunks = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + const usageChunk = chunks.find((chunk) => chunk.type === "usage") + expect(usageChunk).toMatchObject({ + type: "usage", + inputTokens: 11847, + outputTokens: 102, + cacheReadTokens: 8245, + }) + }) + it("should yield informative message when stream produces no text content", async () => { // Stream with only reasoning (no text-delta) simulates thinking-only response const mockFullStream = (async function* () { diff --git a/src/api/providers/__tests__/lite-llm.spec.ts b/src/api/providers/__tests__/lite-llm.spec.ts index ae582d946a..d4e431cca6 100644 --- a/src/api/providers/__tests__/lite-llm.spec.ts +++ b/src/api/providers/__tests__/lite-llm.spec.ts @@ -267,7 +267,12 @@ describe("LiteLLMHandler", () => { expect(mockStreamText).toHaveBeenCalledTimes(1) const callArgs = mockStreamText.mock.calls[0][0] - expect(callArgs.system).toBe(systemPrompt) + expect(callArgs.system).toEqual( + expect.objectContaining({ + role: "system", + content: systemPrompt, + }), + ) expect(callArgs.model).toBeDefined() }) diff --git a/src/api/providers/__tests__/minimax.spec.ts b/src/api/providers/__tests__/minimax.spec.ts index 60e7483ef8..11f8b1a9ec 100644 --- a/src/api/providers/__tests__/minimax.spec.ts +++ b/src/api/providers/__tests__/minimax.spec.ts @@ -8,6 +8,7 @@ import { minimaxDefaultModelId } from "@roo-code/types" import type { ApiHandlerOptions } from "../../../shared/api" import type { ApiStream, ApiStreamChunk } from "../../transform/stream" import { MiniMaxHandler } from "../minimax" +import * as aiSdkTransform from "../../transform/ai-sdk" const { mockStreamText, @@ -65,7 +66,15 @@ function createHandler(options: HandlerOptions = {}) { function createMockStream( chunks: Array>, - usage: { inputTokens?: number; outputTokens?: number } = { inputTokens: 10, outputTokens: 5 }, + usage: { + inputTokens?: number + outputTokens?: number + inputTokenDetails?: { + noCacheTokens?: number + cacheReadTokens?: number + cacheWriteTokens?: number + } + } = { inputTokens: 10, outputTokens: 5 }, providerMetadata: Record> = { anthropic: { cacheReadInputTokens: 0, @@ -238,7 +247,14 @@ describe("MiniMaxHandler", () => { expect(mockStreamText).toHaveBeenCalledWith( expect.objectContaining({ model: "mock-model-instance", - system: systemPrompt, + system: { + role: "system", + content: systemPrompt, + providerOptions: { + anthropic: { cacheControl: { type: "ephemeral" } }, + bedrock: { cachePoint: { type: "default" } }, + }, + }, temperature: 1, messages: expect.any(Array), }), @@ -323,6 +339,42 @@ describe("MiniMaxHandler", () => { expect(typeof usageChunk?.totalCost).toBe("number") }) + it("uses non-cached input tokens from AI SDK v6 usage details", async () => { + mockStreamText.mockReturnValue( + createMockStream( + [{ type: "text-delta", text: "Done" }], + { + inputTokens: 13_071, + outputTokens: 93, + inputTokenDetails: { + noCacheTokens: 10, + cacheWriteTokens: 489, + cacheReadTokens: 12_572, + }, + }, + { + anthropic: { + cacheCreationInputTokens: 489, + cacheReadInputTokens: 12_572, + }, + }, + ), + ) + + const handler = createHandler() + const chunks = await collectChunks(handler.createMessage(systemPrompt, messages)) + const usageChunk = chunks.find((chunk) => chunk.type === "usage") + + expect(usageChunk).toMatchObject({ + type: "usage", + inputTokens: 13_071, + nonCachedInputTokens: 10, + outputTokens: 93, + cacheWriteTokens: 489, + cacheReadTokens: 12_572, + }) + }) + it("calls mergeEnvironmentDetailsForMiniMax before conversion", async () => { const mergedMessages: RooMessage[] = [ { @@ -347,12 +399,48 @@ describe("MiniMaxHandler", () => { anthropic: { cacheControl: { type: "ephemeral" }, }, + bedrock: { + cachePoint: { type: "default" }, + }, }, }), ]), ) }) + it("passes anthropic tool cache provider options to function-tool conversion when caching is enabled", async () => { + const convertToolsSpy = vi.spyOn(aiSdkTransform, "convertToolsForAiSdk").mockReturnValue(undefined) + mockStreamText.mockReturnValue(createMockStream([{ type: "text-delta", text: "OK" }])) + + const handler = createHandler() + await collectChunks( + handler.createMessage(systemPrompt, messages, { + taskId: "test-task", + tools: [ + { + type: "function", + function: { + name: "get_weather", + description: "Get weather", + parameters: { + type: "object", + properties: { location: { type: "string" } }, + required: ["location"], + }, + }, + }, + ] as any, + }), + ) + + expect(convertToolsSpy).toHaveBeenCalledWith(expect.any(Array), { + functionToolProviderOptions: { + anthropic: { cacheControl: { type: "ephemeral" } }, + }, + }) + convertToolsSpy.mockRestore() + }) + it("handles errors via handleAiSdkError", async () => { mockStreamText.mockImplementation(() => { throw new Error("API Error") diff --git a/src/api/providers/__tests__/mistral.spec.ts b/src/api/providers/__tests__/mistral.spec.ts index 5f25690056..1970f56739 100644 --- a/src/api/providers/__tests__/mistral.spec.ts +++ b/src/api/providers/__tests__/mistral.spec.ts @@ -449,9 +449,9 @@ describe("MistralHandler", () => { chunks.push(chunk) } - // tool-call events are ignored, so no tool_call chunks should be emitted + // tool-call events may be surfaced by the shared AI SDK stream processor. const toolCallChunks = chunks.filter((c) => c.type === "tool_call") - expect(toolCallChunks.length).toBe(0) + expect(toolCallChunks.length).toBeGreaterThanOrEqual(1) }) }) diff --git a/src/api/providers/__tests__/moonshot.spec.ts b/src/api/providers/__tests__/moonshot.spec.ts index 340a405ef3..359010c6b1 100644 --- a/src/api/providers/__tests__/moonshot.spec.ts +++ b/src/api/providers/__tests__/moonshot.spec.ts @@ -219,7 +219,7 @@ describe("MoonshotHandler", () => { const usageChunks = chunks.filter((chunk) => chunk.type === "usage") expect(usageChunks.length).toBeGreaterThan(0) - expect(usageChunks[0].cacheWriteTokens).toBe(0) + expect(usageChunks[0].cacheWriteTokens).toBeUndefined() expect(usageChunks[0].cacheReadTokens).toBe(2) }) }) @@ -266,7 +266,7 @@ describe("MoonshotHandler", () => { expect(result.type).toBe("usage") expect(result.inputTokens).toBe(100) expect(result.outputTokens).toBe(50) - expect(result.cacheWriteTokens).toBe(0) + expect(result.cacheWriteTokens).toBeUndefined() expect(result.cacheReadTokens).toBe(20) }) @@ -291,7 +291,7 @@ describe("MoonshotHandler", () => { expect(result.type).toBe("usage") expect(result.inputTokens).toBe(100) expect(result.outputTokens).toBe(50) - expect(result.cacheWriteTokens).toBe(0) + expect(result.cacheWriteTokens).toBeUndefined() expect(result.cacheReadTokens).toBeUndefined() }) }) @@ -468,9 +468,9 @@ describe("MoonshotHandler", () => { chunks.push(chunk) } - // tool-call events are ignored, so no tool_call chunks should be emitted + // tool-call events may be surfaced by the shared AI SDK stream processor. const toolCallChunks = chunks.filter((c) => c.type === "tool_call") - expect(toolCallChunks.length).toBe(0) + expect(toolCallChunks.length).toBeGreaterThanOrEqual(1) }) }) }) diff --git a/src/api/providers/__tests__/native-ollama.spec.ts b/src/api/providers/__tests__/native-ollama.spec.ts index df6f7cb841..31297c7847 100644 --- a/src/api/providers/__tests__/native-ollama.spec.ts +++ b/src/api/providers/__tests__/native-ollama.spec.ts @@ -84,7 +84,7 @@ describe("NativeOllamaHandler", () => { expect(results).toHaveLength(3) expect(results[0]).toEqual({ type: "text", text: "Hello" }) expect(results[1]).toEqual({ type: "text", text: " world" }) - expect(results[2]).toEqual({ type: "usage", inputTokens: 10, outputTokens: 2 }) + expect(results[2]).toMatchObject({ type: "usage", inputTokens: 10, outputTokens: 2 }) }) it("should not include providerOptions by default (no num_ctx)", async () => { diff --git a/src/api/providers/__tests__/openai-codex-native-tool-calls.spec.ts b/src/api/providers/__tests__/openai-codex-native-tool-calls.spec.ts index 39fcc77b46..ee86d475dc 100644 --- a/src/api/providers/__tests__/openai-codex-native-tool-calls.spec.ts +++ b/src/api/providers/__tests__/openai-codex-native-tool-calls.spec.ts @@ -192,4 +192,98 @@ describe("OpenAiCodexHandler native tool calls", () => { totalCost: 0, }) }) + + it("extracts cached and reasoning tokens from OpenAI responses usage shapes", async () => { + vi.spyOn(openAiCodexOAuthManager, "getAccessToken").mockResolvedValue("test-token") + vi.spyOn(openAiCodexOAuthManager, "getAccountId").mockResolvedValue("acct_test") + + async function* mockFullStream() { + yield { type: "text-delta", text: "response" } + } + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: Promise.resolve({ + inputTokens: 120, + outputTokens: 30, + raw: { + input_tokens_details: { cached_tokens: 40 }, + output_tokens_details: { reasoning_tokens: 12 }, + }, + }), + providerMetadata: Promise.resolve({ + openai: { + responseId: "resp_usage_shape", + cachedPromptTokens: 40, + reasoningTokens: 12, + }, + }), + content: Promise.resolve([]), + }) + + const stream = handler.createMessage("system", [{ role: "user", content: "hello" } as any], { + taskId: "t", + tools: [], + }) + + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + const usageChunks = chunks.filter((c) => c.type === "usage") + expect(usageChunks.length).toBe(1) + expect(usageChunks[0]).toMatchObject({ + type: "usage", + inputTokens: 120, + outputTokens: 30, + cacheReadTokens: 40, + reasoningTokens: 12, + totalCost: 0, + }) + }) + + it("extracts cached and reasoning tokens from AI SDK detail fields", async () => { + vi.spyOn(openAiCodexOAuthManager, "getAccessToken").mockResolvedValue("test-token") + vi.spyOn(openAiCodexOAuthManager, "getAccountId").mockResolvedValue("acct_test") + + async function* mockFullStream() { + yield { type: "text-delta", text: "response" } + } + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: Promise.resolve({ + inputTokens: 90, + outputTokens: 20, + inputTokenDetails: { cacheReadTokens: 25 }, + outputTokenDetails: { reasoningTokens: 8 }, + }), + providerMetadata: Promise.resolve({ + openai: { responseId: "resp_usage_details" }, + }), + content: Promise.resolve([]), + }) + + const stream = handler.createMessage("system", [{ role: "user", content: "hello" } as any], { + taskId: "t", + tools: [], + }) + + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + const usageChunks = chunks.filter((c) => c.type === "usage") + expect(usageChunks.length).toBe(1) + expect(usageChunks[0]).toMatchObject({ + type: "usage", + inputTokens: 90, + outputTokens: 20, + cacheReadTokens: 25, + reasoningTokens: 8, + totalCost: 0, + }) + }) }) diff --git a/src/api/providers/__tests__/openai-native-usage.spec.ts b/src/api/providers/__tests__/openai-native-usage.spec.ts index 0fbb9614be..589579f959 100644 --- a/src/api/providers/__tests__/openai-native-usage.spec.ts +++ b/src/api/providers/__tests__/openai-native-usage.spec.ts @@ -100,6 +100,35 @@ describe("OpenAiNativeHandler - usage metrics", () => { }) describe("cache metrics", () => { + it("should handle cached input tokens from AI SDK v6 inputTokenDetails", async () => { + async function* mockFullStream() { + yield { type: "text-delta", text: "Test" } + } + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: Promise.resolve({ + inputTokens: 100, + outputTokens: 50, + inputTokenDetails: { + cacheReadTokens: 30, + }, + }), + providerMetadata: Promise.resolve({}), + content: Promise.resolve([]), + }) + + const stream = handler.createMessage(systemPrompt, messages) + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + const usageChunks = chunks.filter((c) => c.type === "usage") + expect(usageChunks).toHaveLength(1) + expect(usageChunks[0].cacheReadTokens).toBe(30) + }) + it("should handle cached input tokens from usage details", async () => { async function* mockFullStream() { yield { type: "text-delta", text: "Test" } diff --git a/src/api/providers/__tests__/openrouter.spec.ts b/src/api/providers/__tests__/openrouter.spec.ts index 763d0ef606..f39470c691 100644 --- a/src/api/providers/__tests__/openrouter.spec.ts +++ b/src/api/providers/__tests__/openrouter.spec.ts @@ -505,7 +505,7 @@ describe("OpenRouterHandler", () => { expect(chunks[3]).toEqual({ type: "tool_call_end", id: "call_1" }) }) - it("ignores tool-call events (handled by tool-input-start/delta/end)", async () => { + it("handles tool-call events without breaking usage emission", async () => { const handler = new OpenRouterHandler(mockOptions) const mockFullStream = (async function* () { @@ -530,10 +530,8 @@ describe("OpenRouterHandler", () => { chunks.push(chunk) } - // tool-call is intentionally ignored by processAiSdkStreamPart, - // only usage chunk should be present - expect(chunks).toHaveLength(1) - expect(chunks[0]).toMatchObject({ type: "usage" }) + const usageChunks = chunks.filter((chunk) => chunk.type === "usage") + expect(usageChunks).toHaveLength(1) }) it("handles API errors gracefully", async () => { @@ -672,12 +670,10 @@ describe("OpenRouterHandler", () => { }), ) - // Verify that providerOptions does NOT contain extended_thinking - expect(mockStreamText).toHaveBeenCalledWith( - expect.objectContaining({ - providerOptions: undefined, - }), - ) + // Provider options are omitted unless provider routing/patches are present. + const request = mockStreamText.mock.calls.at(-1)?.[0] + expect(request).toBeDefined() + expect(request).not.toHaveProperty("providerOptions") }) it("does not pass reasoning via extraBody when reasoning is disabled", async () => { @@ -710,12 +706,10 @@ describe("OpenRouterHandler", () => { }), ) - // Verify that providerOptions is undefined when no provider routing - expect(mockStreamText).toHaveBeenCalledWith( - expect.objectContaining({ - providerOptions: undefined, - }), - ) + // Provider options are omitted unless provider routing/patches are present. + const request = mockStreamText.mock.calls.at(-1)?.[0] + expect(request).toBeDefined() + expect(request).not.toHaveProperty("providerOptions") }) }) @@ -984,7 +978,12 @@ describe("OpenRouterHandler", () => { // System prompt should be passed directly via streamText const streamTextCall = mockStreamText.mock.calls[0][0] - expect(streamTextCall.system).toBe("test") + expect(streamTextCall.system).toEqual( + expect.objectContaining({ + role: "system", + content: "test", + }), + ) // Messages should be the converted AI SDK messages (no system-role message injected) const systemMsgs = streamTextCall.messages.filter((m: any) => m.role === "system") @@ -1018,7 +1017,12 @@ describe("OpenRouterHandler", () => { // System prompt should be passed directly via streamText const streamTextCall = mockStreamText.mock.calls[0][0] - expect(streamTextCall.system).toBe("test") + expect(streamTextCall.system).toEqual( + expect.objectContaining({ + role: "system", + content: "test", + }), + ) // No system-role message should be injected const systemMsgs = streamTextCall.messages.filter((m: any) => m.role === "system") diff --git a/src/api/providers/__tests__/requesty.spec.ts b/src/api/providers/__tests__/requesty.spec.ts index 0967eefdd7..ff4c507f04 100644 --- a/src/api/providers/__tests__/requesty.spec.ts +++ b/src/api/providers/__tests__/requesty.spec.ts @@ -173,6 +173,7 @@ describe("RequestyHandler", () => { expect(chunks[1]).toEqual({ type: "usage", inputTokens: 10, + nonCachedInputTokens: 3, outputTokens: 20, cacheWriteTokens: 5, cacheReadTokens: 2, @@ -200,7 +201,10 @@ describe("RequestyHandler", () => { expect(mockStreamText).toHaveBeenCalledWith( expect.objectContaining({ - system: "test system prompt", + system: expect.objectContaining({ + role: "system", + content: "test system prompt", + }), temperature: 0, maxOutputTokens: 8192, }), @@ -468,7 +472,7 @@ describe("RequestyHandler", () => { expect(result.type).toBe("usage") expect(result.inputTokens).toBe(100) expect(result.outputTokens).toBe(50) - expect(result.cacheWriteTokens).toBe(0) + expect(result.cacheWriteTokens).toBeUndefined() expect(result.cacheReadTokens).toBe(15) expect(result.reasoningTokens).toBe(25) }) @@ -492,9 +496,9 @@ describe("RequestyHandler", () => { expect(result.type).toBe("usage") expect(result.inputTokens).toBe(100) expect(result.outputTokens).toBe(50) - expect(result.cacheWriteTokens).toBe(0) - expect(result.cacheReadTokens).toBe(0) - expect(result.totalCost).toBe(0) + expect(result.cacheWriteTokens).toBeUndefined() + expect(result.cacheReadTokens).toBeUndefined() + expect(result.totalCost).toBeUndefined() }) }) diff --git a/src/api/providers/__tests__/roo.spec.ts b/src/api/providers/__tests__/roo.spec.ts index 1be5d11656..0a58496081 100644 --- a/src/api/providers/__tests__/roo.spec.ts +++ b/src/api/providers/__tests__/roo.spec.ts @@ -10,10 +10,12 @@ import type { RooMessage } from "../../../core/task-persistence/rooMessage" const mockStreamText = vitest.fn() const mockGenerateText = vitest.fn() const mockCreateOpenAICompatible = vitest.fn() +const mockCreateGateway = vitest.fn() vitest.mock("ai", () => ({ streamText: (...args: unknown[]) => mockStreamText(...args), generateText: (...args: unknown[]) => mockGenerateText(...args), + createGateway: (...args: unknown[]) => mockCreateGateway(...args), tool: vitest.fn((t) => t), jsonSchema: vitest.fn((s) => s), })) @@ -96,6 +98,8 @@ vitest.mock("../../providers/fetchers/modelCache", () => ({ import { RooHandler } from "../roo" import { CloudService } from "@roo-code/cloud" +const mockGatewayProvider = vitest.fn((modelId: string) => ({ modelId, provider: "roo-gateway" })) + /** * Helper to create a mock stream result for streamText. */ @@ -105,6 +109,15 @@ function createMockStreamResult(options?: { toolCallParts?: Array<{ type: string; id?: string; toolName?: string; delta?: string }> inputTokens?: number outputTokens?: number + usage?: { + inputTokens?: number + outputTokens?: number + inputTokenDetails?: { + noCacheTokens?: number + cacheReadTokens?: number + cacheWriteTokens?: number + } + } providerMetadata?: Record }) { const { @@ -130,7 +143,7 @@ function createMockStreamResult(options?: { return { fullStream, - usage: Promise.resolve({ inputTokens, outputTokens }), + usage: Promise.resolve(options?.usage ?? { inputTokens, outputTokens }), providerMetadata: Promise.resolve(providerMetadata), } } @@ -156,6 +169,8 @@ describe("RooHandler", () => { mockStreamText.mockClear() mockGenerateText.mockClear() mockCreateOpenAICompatible.mockClear() + mockCreateGateway.mockClear() + mockCreateGateway.mockReturnValue(mockGatewayProvider) vitest.clearAllMocks() }) @@ -313,7 +328,10 @@ describe("RooHandler", () => { // Verify streamText was called with system prompt and converted messages expect(mockStreamText).toHaveBeenCalledWith( expect.objectContaining({ - system: systemPrompt, + system: expect.objectContaining({ + role: "system", + content: systemPrompt, + }), messages: expect.any(Array), }), ) @@ -353,6 +371,32 @@ describe("RooHandler", () => { }), ) }) + + it("uses ai-sdk gateway provider when ROO_CODE_ROUTER_USE_GATEWAY_SDK is enabled", async () => { + process.env.ROO_CODE_ROUTER_USE_GATEWAY_SDK = "true" + + mockStreamText.mockReturnValue(createMockStreamResult()) + + const gatewayHandler = new RooHandler(mockOptions) + const stream = gatewayHandler.createMessage(systemPrompt, messages, { taskId: "gw-task-1" }) + for await (const _chunk of stream) { + // consume + } + + expect(mockCreateGateway).toHaveBeenCalledWith( + expect.objectContaining({ + apiKey: "test-session-token", + baseURL: "https://api.roocode.com/proxy/v3/ai", + headers: expect.objectContaining({ + "X-Roo-App-Version": expect.any(String), + "X-Roo-Task-ID": "gw-task-1", + }), + }), + ) + expect(mockCreateOpenAICompatible).not.toHaveBeenCalled() + + delete process.env.ROO_CODE_ROUTER_USE_GATEWAY_SDK + }) }) describe("completePrompt", () => { @@ -768,6 +812,124 @@ describe("RooHandler", () => { expect(usageChunk.cacheWriteTokens).toBe(20) expect(usageChunk.cacheReadTokens).toBe(30) }) + + it("should read anthropic/gateway usage metadata when roo metadata is absent", async () => { + const anthropicHandler = new RooHandler({ + apiModelId: "anthropic/claude-haiku-4.5", + }) + + mockStreamText.mockReturnValue( + createMockStreamResult({ + usage: { + inputTokens: 12_582, + outputTokens: 100, + }, + providerMetadata: { + anthropic: { + usage: { + cache_creation_input_tokens: 12_572, + cache_read_input_tokens: 0, + }, + }, + gateway: { + cost: "0.081125", + }, + }, + }), + ) + + const stream = anthropicHandler.createMessage(systemPrompt, messages) + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + const usageChunk = chunks.find((c) => c.type === "usage") + expect(usageChunk).toBeDefined() + expect(usageChunk.inputTokens).toBe(12_582) + expect(usageChunk.nonCachedInputTokens).toBe(10) + expect(usageChunk.outputTokens).toBe(100) + expect(usageChunk.cacheWriteTokens).toBe(12_572) + expect(usageChunk.cacheReadTokens).toBe(0) + expect(usageChunk.totalCost).toBe(0.081125) + }) + + it("should fall back to gateway cache metadata when anthropic/roo cache fields are absent", async () => { + const anthropicHandler = new RooHandler({ + apiModelId: "anthropic/claude-haiku-4.5", + }) + + mockStreamText.mockReturnValue( + createMockStreamResult({ + usage: { + inputTokens: 12_592, + outputTokens: 100, + }, + providerMetadata: { + gateway: { + cache_creation_input_tokens: 459, + cached_tokens: 12_572, + cost: "0.01157975", + }, + }, + }), + ) + + const stream = anthropicHandler.createMessage(systemPrompt, messages) + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + const usageChunk = chunks.find((c) => c.type === "usage") + expect(usageChunk).toBeDefined() + expect(usageChunk.inputTokens).toBe(12_592) + expect(usageChunk.nonCachedInputTokens).toBe(0) + expect(usageChunk.outputTokens).toBe(100) + expect(usageChunk.cacheWriteTokens).toBe(459) + expect(usageChunk.cacheReadTokens).toBe(12_572) + expect(usageChunk.totalCost).toBe(0.01157975) + }) + + it("uses non-cached input tokens for anthropic protocol models", async () => { + const anthropicHandler = new RooHandler({ + apiModelId: "anthropic/claude-haiku-4.5", + }) + + mockStreamText.mockReturnValue( + createMockStreamResult({ + usage: { + inputTokens: 13_071, + outputTokens: 93, + inputTokenDetails: { + noCacheTokens: 10, + cacheWriteTokens: 489, + cacheReadTokens: 12_572, + }, + }, + providerMetadata: { + roo: { + cache_creation_input_tokens: 489, + cache_read_input_tokens: 12_572, + }, + }, + }), + ) + + const stream = anthropicHandler.createMessage(systemPrompt, messages) + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + const usageChunk = chunks.find((c) => c.type === "usage") + expect(usageChunk).toBeDefined() + expect(usageChunk.inputTokens).toBe(13_071) + expect(usageChunk.nonCachedInputTokens).toBe(10) + expect(usageChunk.outputTokens).toBe(93) + expect(usageChunk.cacheWriteTokens).toBe(489) + expect(usageChunk.cacheReadTokens).toBe(12_572) + }) }) describe("isAiSdkProvider", () => { diff --git a/src/api/providers/__tests__/usage-metrics.spec.ts b/src/api/providers/__tests__/usage-metrics.spec.ts new file mode 100644 index 0000000000..f17aa7e3c8 --- /dev/null +++ b/src/api/providers/__tests__/usage-metrics.spec.ts @@ -0,0 +1,156 @@ +import { normalizeProviderUsage } from "../utils/normalize-provider-usage" + +const baseModelInfo = { + maxTokens: 200_000, + contextWindow: 200_000, + supportsPromptCache: true, + inputPrice: 1, + outputPrice: 2, + cacheWritesPrice: 1.25, + cacheReadsPrice: 0.1, +} as any + +describe("normalizeProviderUsage", () => { + it("normalizes Anthropic-style usage with both total and non-cached semantics", () => { + const normalized = normalizeProviderUsage({ + provider: "anthropic", + apiProtocol: "anthropic", + usage: { + inputTokens: 13_011, + outputTokens: 90, + }, + providerMetadata: { + anthropic: { + usage: { + input_tokens: 3, + cache_creation_input_tokens: 458, + cache_read_input_tokens: 12_550, + }, + }, + }, + modelInfo: baseModelInfo, + }) + + expect(normalized.chunk.inputTokens).toBe(13_011) + expect(normalized.chunk.nonCachedInputTokens).toBe(3) + expect(normalized.chunk.cacheWriteTokens).toBe(458) + expect(normalized.chunk.cacheReadTokens).toBe(12_550) + expect(normalized.chunk.outputTokens).toBe(90) + }) + + it("applies precedence policy: providerMetadata > usage > raw", () => { + const normalized = normalizeProviderUsage({ + provider: "openai", + apiProtocol: "openai", + usage: { + inputTokens: 1000, + outputTokens: 100, + inputTokenDetails: { + cacheReadTokens: 400, + cacheWriteTokens: 120, + }, + raw: { + input_tokens_details: { + cached_tokens: 700, + }, + cache_creation_input_tokens: 350, + }, + }, + providerMetadata: { + openai: { + cachedPromptTokens: 500, + }, + }, + modelInfo: baseModelInfo, + }) + + expect(normalized.chunk.cacheReadTokens).toBe(500) + expect(normalized.chunk.cacheWriteTokens).toBe(120) + }) + + it("derives non-cached tokens from total-cache when needed", () => { + const normalized = normalizeProviderUsage({ + provider: "openai", + apiProtocol: "openai", + usage: { + inputTokens: 100, + outputTokens: 50, + inputTokenDetails: { + cacheReadTokens: 20, + cacheWriteTokens: 10, + }, + }, + modelInfo: baseModelInfo, + }) + + expect(normalized.chunk.inputTokens).toBe(100) + expect(normalized.chunk.nonCachedInputTokens).toBe(70) + }) + + it("supports array-based metadata extraction (Gemini cacheTokensDetails)", () => { + const normalized = normalizeProviderUsage({ + provider: "gemini", + apiProtocol: "openai", + usage: { + inputTokens: 500, + outputTokens: 40, + }, + providerMetadata: { + google: { + usageMetadata: { + cacheTokensDetails: [{ tokenCount: 20 }, { tokenCount: 30 }], + }, + }, + }, + modelInfo: baseModelInfo, + }) + + expect(normalized.chunk.cacheReadTokens).toBe(50) + }) + + it("supports dynamic protocol for vercel-ai-gateway profile", () => { + const normalized = normalizeProviderUsage({ + provider: "vercel-ai-gateway", + apiProtocol: "anthropic", + usage: { + inputTokens: 13_011, + outputTokens: 90, + }, + providerMetadata: { + anthropic: { + usage: { + input_tokens: 3, + cache_creation_input_tokens: 458, + cache_read_input_tokens: 12_550, + }, + }, + }, + modelInfo: baseModelInfo, + }) + + expect(normalized.chunk.nonCachedInputTokens).toBe(3) + expect(normalized.chunk.cacheWriteTokens).toBe(458) + expect(normalized.chunk.cacheReadTokens).toBe(12_550) + }) + + it("coerces string values and omits zero caches by default", () => { + const normalized = normalizeProviderUsage({ + provider: "openai", + apiProtocol: "openai", + usage: { + inputTokens: "42", + outputTokens: "7", + inputTokenDetails: { + cacheReadTokens: "0", + cacheWriteTokens: "0", + }, + } as any, + modelInfo: baseModelInfo, + }) + + expect(normalized.chunk.inputTokens).toBe(42) + expect(normalized.chunk.outputTokens).toBe(7) + expect(normalized.chunk.cacheReadTokens).toBeUndefined() + expect(normalized.chunk.cacheWriteTokens).toBeUndefined() + }) +}) diff --git a/src/api/providers/__tests__/vercel-ai-gateway.spec.ts b/src/api/providers/__tests__/vercel-ai-gateway.spec.ts index f482c7cf2d..8e6657e1ca 100644 --- a/src/api/providers/__tests__/vercel-ai-gateway.spec.ts +++ b/src/api/providers/__tests__/vercel-ai-gateway.spec.ts @@ -135,7 +135,16 @@ describe("VercelAiGatewayHandler", () => { describe("createMessage", () => { function createMockStreamResult(options?: { - usage?: { inputTokens: number; outputTokens: number; details?: Record } + usage?: { + inputTokens: number + outputTokens: number + inputTokenDetails?: { + noCacheTokens?: number + cacheReadTokens?: number + cacheWriteTokens?: number + } + details?: Record + } providerMetadata?: Record> fullStream?: AsyncGenerator }) { @@ -187,10 +196,12 @@ describe("VercelAiGatewayHandler", () => { expect(chunks[1]).toEqual({ type: "usage", inputTokens: 10, + nonCachedInputTokens: 5, outputTokens: 5, cacheWriteTokens: 2, cacheReadTokens: 3, totalCost: 0.005, + reasoningTokens: undefined, }) }) @@ -277,10 +288,55 @@ describe("VercelAiGatewayHandler", () => { expect(usageChunk).toEqual({ type: "usage", inputTokens: 10, + nonCachedInputTokens: 5, outputTokens: 5, cacheWriteTokens: 2, cacheReadTokens: 3, totalCost: 0.005, + reasoningTokens: undefined, + }) + }) + + it("uses non-cached input tokens for anthropic protocol models", async () => { + mockStreamText.mockReturnValue( + createMockStreamResult({ + usage: { + inputTokens: 13_071, + outputTokens: 93, + inputTokenDetails: { + noCacheTokens: 10, + cacheWriteTokens: 489, + cacheReadTokens: 12_572, + }, + }, + providerMetadata: { + gateway: { + cost: 0.005, + }, + }, + }), + ) + + const handler = new VercelAiGatewayHandler(mockOptions) + const systemPrompt = "You are a helpful assistant." + const messages: RooMessage[] = [{ role: "user", content: "Hello" }] + + const stream = handler.createMessage(systemPrompt, messages) + const chunks = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + const usageChunk = chunks.find((chunk) => chunk.type === "usage") + expect(usageChunk).toEqual({ + type: "usage", + inputTokens: 13_071, + nonCachedInputTokens: 10, + outputTokens: 93, + cacheWriteTokens: 489, + cacheReadTokens: 12_572, + totalCost: 0.005, + reasoningTokens: undefined, }) }) @@ -424,7 +480,10 @@ describe("VercelAiGatewayHandler", () => { expect(mockStreamText).toHaveBeenCalledWith( expect.objectContaining({ - system: "test prompt", + system: expect.objectContaining({ + role: "system", + content: "test prompt", + }), }), ) }) diff --git a/src/api/providers/__tests__/vertex.spec.ts b/src/api/providers/__tests__/vertex.spec.ts index 7ae0e5fdff..cf6819302f 100644 --- a/src/api/providers/__tests__/vertex.spec.ts +++ b/src/api/providers/__tests__/vertex.spec.ts @@ -191,11 +191,49 @@ describe("VertexHandler", () => { expect(mockStreamText).toHaveBeenCalledWith( expect.objectContaining({ - system: systemPrompt, + system: expect.objectContaining({ + role: "system", + content: systemPrompt, + }), temperature: 1, }), ) }) + + it("maps cache read tokens from Google usageMetadata fallback", async () => { + const mockFullStream = (async function* () { + yield { type: "text-delta", text: "Hello" } + })() + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream, + usage: Promise.resolve({ + inputTokens: 11847, + outputTokens: 102, + }), + providerMetadata: Promise.resolve({ + google: { + usageMetadata: { + cachedContentTokenCount: 8245, + }, + }, + }), + }) + + const stream = handler.createMessage(systemPrompt, mockMessages) + const chunks: ApiStreamChunk[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + const usageChunk = chunks.find((chunk) => chunk.type === "usage") + expect(usageChunk).toMatchObject({ + type: "usage", + inputTokens: 11847, + outputTokens: 102, + cacheReadTokens: 8245, + }) + }) }) describe("completePrompt", () => { diff --git a/src/api/providers/__tests__/zai.spec.ts b/src/api/providers/__tests__/zai.spec.ts index 789915d6ef..3e53770c30 100644 --- a/src/api/providers/__tests__/zai.spec.ts +++ b/src/api/providers/__tests__/zai.spec.ts @@ -342,7 +342,10 @@ describe("ZAiHandler", () => { expect(mockStreamText).toHaveBeenCalledWith( expect.objectContaining({ - system: systemPrompt, + system: expect.objectContaining({ + role: "system", + content: systemPrompt, + }), temperature: expect.any(Number), }), ) diff --git a/src/api/providers/anthropic-vertex.ts b/src/api/providers/anthropic-vertex.ts index 34c4c57624..0225ac0fbf 100644 --- a/src/api/providers/anthropic-vertex.ts +++ b/src/api/providers/anthropic-vertex.ts @@ -26,7 +26,7 @@ import { yieldResponseMessage, } from "../transform/ai-sdk" import { applyPromptCacheToMessages } from "../transform/prompt-cache" -import { calculateApiCostAnthropic } from "../../shared/cost" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import { DEFAULT_HEADERS } from "./constants" import { BaseProvider } from "./base-provider" @@ -93,9 +93,22 @@ export class AnthropicVertexHandler extends BaseProvider implements SingleComple // Convert messages to AI SDK format const aiSdkMessages = messages as ModelMessage[] + const promptCache = applyPromptCacheToMessages({ + adapter: "anthropic-vertex", + overrideKey: "vertex", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: modelConfig.info.supportsPromptCache, + promptCacheRetention: modelConfig.info.promptCacheRetention, + }, + settings: this.options, + }) + // Convert tools to AI SDK format const openAiTools = this.convertToolsForOpenAI(metadata?.tools) - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + const aiSdkTools = convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined // Build Anthropic provider options const anthropicProviderOptions: Record = {} @@ -119,25 +132,13 @@ export class AnthropicVertexHandler extends BaseProvider implements SingleComple anthropicProviderOptions.disableParallelToolUse = true } - const promptCache = applyPromptCacheToMessages({ - adapter: "anthropic", - overrideKey: "vertex", - messages: aiSdkMessages, - modelInfo: { - supportsPromptCache: modelConfig.info.supportsPromptCache, - promptCacheRetention: modelConfig.info.promptCacheRetention, - }, - settings: this.options, - }) - // Build streamText request // Cast providerOptions to any to bypass strict JSONObject typing — the AI SDK accepts the correct runtime values const requestOptions: Parameters[0] = { model: this.provider(modelConfig.id), - system: systemPrompt, - ...(promptCache.systemProviderOptions - ? ({ systemProviderOptions: promptCache.systemProviderOptions } as Record) - : {}), + system: promptCache.systemProviderOptions + ? ({ role: "system", content: systemPrompt, providerOptions: promptCache.systemProviderOptions } as any) + : systemPrompt, messages: aiSdkMessages, temperature: modelConfig.temperature, maxOutputTokens: modelConfig.maxTokens ?? ANTHROPIC_DEFAULT_MAX_TOKENS, @@ -189,36 +190,27 @@ export class AnthropicVertexHandler extends BaseProvider implements SingleComple * Process usage metrics from the AI SDK response, including Anthropic's cache metrics. */ private processUsageMetrics( - usage: { inputTokens?: number; outputTokens?: number }, + usage: { + inputTokens?: number + outputTokens?: number + inputTokenDetails?: { + noCacheTokens?: number + cacheReadTokens?: number + cacheWriteTokens?: number + } + cachedInputTokens?: number + }, info: ModelInfo, providerMetadata?: Record>, ): ApiStreamUsageChunk { - const inputTokens = usage.inputTokens ?? 0 - const outputTokens = usage.outputTokens ?? 0 - - // Extract cache metrics from Anthropic's providerMetadata - const anthropicMeta = providerMetadata?.anthropic as - | { cacheCreationInputTokens?: number; cacheReadInputTokens?: number } - | undefined - const cacheWriteTokens = anthropicMeta?.cacheCreationInputTokens ?? 0 - const cacheReadTokens = anthropicMeta?.cacheReadInputTokens ?? 0 - - const { totalCost } = calculateApiCostAnthropic( - info, - inputTokens, - outputTokens, - cacheWriteTokens, - cacheReadTokens, - ) - - return { - type: "usage", - inputTokens, - outputTokens, - cacheWriteTokens: cacheWriteTokens > 0 ? cacheWriteTokens : undefined, - cacheReadTokens: cacheReadTokens > 0 ? cacheReadTokens : undefined, - totalCost, - } + const { chunk } = normalizeProviderUsage({ + provider: "anthropic-vertex", + apiProtocol: "anthropic", + usage: usage as any, + providerMetadata: providerMetadata as Record | undefined, + modelInfo: info, + }) + return chunk } getModel() { diff --git a/src/api/providers/anthropic.ts b/src/api/providers/anthropic.ts index 03584f4115..0446b59dd8 100644 --- a/src/api/providers/anthropic.ts +++ b/src/api/providers/anthropic.ts @@ -24,7 +24,7 @@ import { yieldResponseMessage, } from "../transform/ai-sdk" import { applyPromptCacheToMessages } from "../transform/prompt-cache" -import { calculateApiCostAnthropic } from "../../shared/cost" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import { DEFAULT_HEADERS } from "./constants" import { BaseProvider } from "./base-provider" @@ -79,9 +79,22 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa // Convert messages to AI SDK format const aiSdkMessages = messages as ModelMessage[] + const promptCache = applyPromptCacheToMessages({ + adapter: "anthropic", + overrideKey: "anthropic", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: modelConfig.info.supportsPromptCache, + promptCacheRetention: modelConfig.info.promptCacheRetention, + }, + settings: this.options, + }) + // Convert tools to AI SDK format const openAiTools = this.convertToolsForOpenAI(metadata?.tools) - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + const aiSdkTools = convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined // Build Anthropic provider options const anthropicProviderOptions: Record = {} @@ -105,25 +118,13 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa anthropicProviderOptions.disableParallelToolUse = true } - const promptCache = applyPromptCacheToMessages({ - adapter: "anthropic", - overrideKey: "anthropic", - messages: aiSdkMessages, - modelInfo: { - supportsPromptCache: modelConfig.info.supportsPromptCache, - promptCacheRetention: modelConfig.info.promptCacheRetention, - }, - settings: this.options, - }) - // Build streamText request // Cast providerOptions to any to bypass strict JSONObject typing — the AI SDK accepts the correct runtime values const requestOptions: Parameters[0] = { model: this.provider(modelConfig.id), - system: systemPrompt, - ...(promptCache.systemProviderOptions - ? ({ systemProviderOptions: promptCache.systemProviderOptions } as Record) - : {}), + system: promptCache.systemProviderOptions + ? ({ role: "system", content: systemPrompt, providerOptions: promptCache.systemProviderOptions } as any) + : systemPrompt, messages: aiSdkMessages, temperature: modelConfig.temperature, maxOutputTokens: modelConfig.maxTokens ?? ANTHROPIC_DEFAULT_MAX_TOKENS, @@ -175,36 +176,27 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa * Process usage metrics from the AI SDK response, including Anthropic's cache metrics. */ private processUsageMetrics( - usage: { inputTokens?: number; outputTokens?: number }, + usage: { + inputTokens?: number + outputTokens?: number + inputTokenDetails?: { + noCacheTokens?: number + cacheReadTokens?: number + cacheWriteTokens?: number + } + cachedInputTokens?: number + }, info: ModelInfo, providerMetadata?: Record>, ): ApiStreamUsageChunk { - const inputTokens = usage.inputTokens ?? 0 - const outputTokens = usage.outputTokens ?? 0 - - // Extract cache metrics from Anthropic's providerMetadata - const anthropicMeta = providerMetadata?.anthropic as - | { cacheCreationInputTokens?: number; cacheReadInputTokens?: number } - | undefined - const cacheWriteTokens = anthropicMeta?.cacheCreationInputTokens ?? 0 - const cacheReadTokens = anthropicMeta?.cacheReadInputTokens ?? 0 - - const { totalCost } = calculateApiCostAnthropic( - info, - inputTokens, - outputTokens, - cacheWriteTokens, - cacheReadTokens, - ) - - return { - type: "usage", - inputTokens, - outputTokens, - cacheWriteTokens: cacheWriteTokens > 0 ? cacheWriteTokens : undefined, - cacheReadTokens: cacheReadTokens > 0 ? cacheReadTokens : undefined, - totalCost, - } + const { chunk } = normalizeProviderUsage({ + provider: "anthropic", + apiProtocol: "anthropic", + usage: usage as any, + providerMetadata: providerMetadata as Record | undefined, + modelInfo: info, + }) + return chunk } getModel() { diff --git a/src/api/providers/azure.ts b/src/api/providers/azure.ts index baad07558c..bec8c2ec27 100644 --- a/src/api/providers/azure.ts +++ b/src/api/providers/azure.ts @@ -13,8 +13,10 @@ import { mapToolChoice, handleAiSdkError, } from "../transform/ai-sdk" +import { applyPromptCacheToMessages, mergeProviderOptions } from "../transform/prompt-cache" import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" import { getModelParams } from "../transform/model-params" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import { DEFAULT_HEADERS } from "./constants" import { BaseProvider } from "./base-provider" @@ -100,20 +102,14 @@ export class AzureHandler extends BaseProvider implements SingleCompletionHandle } }, ): ApiStreamUsageChunk { - // Extract cache metrics from Azure's providerMetadata if available - const cacheReadTokens = providerMetadata?.azure?.promptCacheHitTokens ?? usage.details?.cachedInputTokens - // Azure uses OpenAI-compatible caching which does not report cache write tokens separately; - // promptCacheMissTokens represents tokens NOT found in cache (processed from scratch), not tokens written to cache. - const cacheWriteTokens = undefined - - return { - type: "usage", - inputTokens: usage.inputTokens || 0, - outputTokens: usage.outputTokens || 0, - cacheReadTokens, - cacheWriteTokens, - reasoningTokens: usage.details?.reasoningTokens, - } + const { chunk } = normalizeProviderUsage({ + provider: "azure", + apiProtocol: "openai", + usage: usage as any, + providerMetadata: providerMetadata as Record | undefined, + modelInfo: this.getModel().info, + }) + return chunk } /** @@ -135,25 +131,43 @@ export class AzureHandler extends BaseProvider implements SingleCompletionHandle messages: RooMessage[], metadata?: ApiHandlerCreateMessageMetadata, ): ApiStream { - const { temperature } = this.getModel() + const { temperature, info } = this.getModel() const languageModel = this.getLanguageModel() // Convert messages to AI SDK format const aiSdkMessages = messages as ModelMessage[] + const promptCache = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: "azure", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: info.supportsPromptCache, + promptCacheRetention: info.promptCacheRetention, + }, + settings: this.options, + }) + // Convert tools to OpenAI format first, then to AI SDK format const openAiTools = this.convertToolsForOpenAI(metadata?.tools) - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + const aiSdkTools = convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined + + const providerOptions = mergeProviderOptions(undefined, promptCache.providerOptionsPatch) // Build the request options const requestOptions: Parameters[0] = { model: languageModel, - system: systemPrompt, + system: promptCache.systemProviderOptions + ? ({ role: "system", content: systemPrompt, providerOptions: promptCache.systemProviderOptions } as any) + : systemPrompt, messages: aiSdkMessages, temperature: this.options.modelTemperature ?? temperature ?? AZURE_DEFAULT_TEMPERATURE, maxOutputTokens: this.getMaxOutputTokens(), tools: aiSdkTools, toolChoice: mapToolChoice(metadata?.tool_choice), + ...(providerOptions ? ({ providerOptions } as Record) : {}), } // Use streamText for streaming responses diff --git a/src/api/providers/baseten.ts b/src/api/providers/baseten.ts index 4057b4714d..3a23a73c44 100644 --- a/src/api/providers/baseten.ts +++ b/src/api/providers/baseten.ts @@ -13,8 +13,10 @@ import { mapToolChoice, handleAiSdkError, } from "../transform/ai-sdk" +import { applyPromptCacheToMessages, mergeProviderOptions } from "../transform/prompt-cache" import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" import { getModelParams } from "../transform/model-params" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import { DEFAULT_HEADERS } from "./constants" import { BaseProvider } from "./base-provider" @@ -66,20 +68,26 @@ export class BasetenHandler extends BaseProvider implements SingleCompletionHand /** * Process usage metrics from the AI SDK response. */ - protected processUsageMetrics(usage: { - inputTokens?: number - outputTokens?: number - details?: { - cachedInputTokens?: number - reasoningTokens?: number - } - }): ApiStreamUsageChunk { - return { - type: "usage", - inputTokens: usage.inputTokens || 0, - outputTokens: usage.outputTokens || 0, - reasoningTokens: usage.details?.reasoningTokens, - } + protected processUsageMetrics( + usage: { + inputTokens?: number + outputTokens?: number + details?: { + cachedInputTokens?: number + reasoningTokens?: number + } + raw?: Record + }, + providerMetadata?: Record, + ): ApiStreamUsageChunk { + const { chunk } = normalizeProviderUsage({ + provider: "baseten", + apiProtocol: "openai", + usage: usage as any, + providerMetadata, + modelInfo: this.getModel().info, + }) + return chunk } /** @@ -98,22 +106,40 @@ export class BasetenHandler extends BaseProvider implements SingleCompletionHand messages: RooMessage[], metadata?: ApiHandlerCreateMessageMetadata, ): ApiStream { - const { temperature } = this.getModel() + const { temperature, info } = this.getModel() const languageModel = this.getLanguageModel() const aiSdkMessages = messages as ModelMessage[] + const promptCache = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: "baseten", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: info.supportsPromptCache, + promptCacheRetention: info.promptCacheRetention, + }, + settings: this.options, + }) + const openAiTools = this.convertToolsForOpenAI(metadata?.tools) - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + const aiSdkTools = convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined + + const providerOptions = mergeProviderOptions(undefined, promptCache.providerOptionsPatch) const requestOptions: Parameters[0] = { model: languageModel, - system: systemPrompt, + system: promptCache.systemProviderOptions + ? ({ role: "system", content: systemPrompt, providerOptions: promptCache.systemProviderOptions } as any) + : systemPrompt, messages: aiSdkMessages, temperature: this.options.modelTemperature ?? temperature ?? BASETEN_DEFAULT_TEMPERATURE, maxOutputTokens: this.getMaxOutputTokens(), tools: aiSdkTools, toolChoice: mapToolChoice(metadata?.tool_choice), + ...(providerOptions ? ({ providerOptions } as Record) : {}), } const result = streamText(requestOptions) @@ -121,8 +147,8 @@ export class BasetenHandler extends BaseProvider implements SingleCompletionHand try { const processUsage = this.processUsageMetrics.bind(this) yield* consumeAiSdkStream(result, async function* () { - const usage = await result.usage - yield processUsage(usage) + const [usage, providerMetadata] = await Promise.all([result.usage, result.providerMetadata]) + yield processUsage(usage as any, providerMetadata as Record | undefined) }) } catch (error) { throw handleAiSdkError(error, "Baseten") diff --git a/src/api/providers/bedrock.ts b/src/api/providers/bedrock.ts index 076449256f..fe2390cbbf 100644 --- a/src/api/providers/bedrock.ts +++ b/src/api/providers/bedrock.ts @@ -36,6 +36,7 @@ import { applyPromptCacheToMessages } from "../transform/prompt-cache" import { shouldUseReasoningBudget } from "../../shared/api" import { BaseProvider } from "./base-provider" import { DEFAULT_HEADERS } from "./constants" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import { logger } from "../../utils/logging" import { Package } from "../../shared/package" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" @@ -266,10 +267,9 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH // Cast providerOptions to any to bypass strict JSONObject typing — the AI SDK accepts the correct runtime values const requestOptions: Parameters[0] = { model: this.provider(modelConfig.id), - system: systemPrompt, - ...(promptCache.systemProviderOptions - ? ({ systemProviderOptions: promptCache.systemProviderOptions } as Record) - : {}), + system: promptCache.systemProviderOptions + ? ({ role: "system", content: systemPrompt, providerOptions: promptCache.systemProviderOptions } as any) + : systemPrompt, messages: aiSdkMessages, temperature: modelConfig.temperature ?? (this.options.modelTemperature as number), maxOutputTokens: modelConfig.maxTokens || (modelConfig.info.maxTokens as number), @@ -332,17 +332,29 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH * Process usage metrics from the AI SDK response. */ private processUsageMetrics( - usage: { inputTokens?: number; outputTokens?: number }, + usage: { + inputTokens?: number + outputTokens?: number + inputTokenDetails?: { + noCacheTokens?: number + cacheReadTokens?: number + cacheWriteTokens?: number + } + cachedInputTokens?: number + reasoningTokens?: number + outputTokenDetails?: { + reasoningTokens?: number + } + }, info: ModelInfo, providerMetadata?: Record>, ): ApiStreamUsageChunk { - const inputTokens = usage.inputTokens ?? 0 + const inputTokens = usage.inputTokenDetails?.noCacheTokens ?? usage.inputTokens ?? 0 const outputTokens = usage.outputTokens ?? 0 // The AI SDK exposes reasoningTokens as a top-level field on usage, and also // under outputTokenDetails.reasoningTokens — there is no .details property. - const reasoningTokens = - (usage as any).reasoningTokens ?? (usage as any).outputTokenDetails?.reasoningTokens ?? 0 + const reasoningTokens = usage.reasoningTokens ?? usage.outputTokenDetails?.reasoningTokens ?? 0 // Extract cache metrics primarily from usage (AI SDK standard locations), // falling back to providerMetadata.bedrock.usage for provider-specific fields. @@ -350,12 +362,11 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH | { cacheReadInputTokens?: number; cacheWriteInputTokens?: number } | undefined const cacheReadTokens = - (usage as any).inputTokenDetails?.cacheReadTokens ?? - (usage as any).cachedInputTokens ?? + usage.inputTokenDetails?.cacheReadTokens ?? + usage.cachedInputTokens ?? bedrockUsage?.cacheReadInputTokens ?? 0 - const cacheWriteTokens = - (usage as any).inputTokenDetails?.cacheWriteTokens ?? bedrockUsage?.cacheWriteInputTokens ?? 0 + const cacheWriteTokens = usage.inputTokenDetails?.cacheWriteTokens ?? bedrockUsage?.cacheWriteInputTokens ?? 0 // For prompt routers, the AI SDK surfaces the invoked model ID in // providerMetadata.bedrock.trace.promptRouter.invokedModelId. @@ -383,19 +394,27 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH } } - return { - type: "usage", - inputTokens, - outputTokens, - cacheReadTokens: cacheReadTokens > 0 ? cacheReadTokens : undefined, - cacheWriteTokens: cacheWriteTokens > 0 ? cacheWriteTokens : undefined, - reasoningTokens: reasoningTokens > 0 ? reasoningTokens : undefined, - totalCost: this.calculateCost({ + const normalized = normalizeProviderUsage({ + provider: "bedrock", + apiProtocol: "anthropic", + usage: { + ...(usage as any), inputTokens, outputTokens, - cacheWriteTokens, - cacheReadTokens, reasoningTokens, + } as any, + providerMetadata: providerMetadata as Record | undefined, + modelInfo: costInfo, + }) + + return { + ...normalized.chunk, + totalCost: this.calculateCost({ + inputTokens: normalized.canonical.inputTokensNonCached ?? inputTokens, + outputTokens: normalized.canonical.outputTokens, + cacheWriteTokens: normalized.canonical.cacheWriteTokens, + cacheReadTokens: normalized.canonical.cacheReadTokens, + reasoningTokens: normalized.canonical.reasoningTokens ?? 0, info: costInfo, }), } diff --git a/src/api/providers/deepseek.ts b/src/api/providers/deepseek.ts index 6dc3016356..1b0232b987 100644 --- a/src/api/providers/deepseek.ts +++ b/src/api/providers/deepseek.ts @@ -13,8 +13,10 @@ import { mapToolChoice, handleAiSdkError, } from "../transform/ai-sdk" +import { applyPromptCacheToMessages, mergeProviderOptions } from "../transform/prompt-cache" import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" import { getModelParams } from "../transform/model-params" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import { DEFAULT_HEADERS } from "./constants" import { BaseProvider } from "./base-provider" @@ -82,18 +84,14 @@ export class DeepSeekHandler extends BaseProvider implements SingleCompletionHan } }, ): ApiStreamUsageChunk { - // Extract cache metrics from DeepSeek's providerMetadata - const cacheReadTokens = providerMetadata?.deepseek?.promptCacheHitTokens ?? usage.details?.cachedInputTokens - const cacheWriteTokens = providerMetadata?.deepseek?.promptCacheMissTokens - - return { - type: "usage", - inputTokens: usage.inputTokens || 0, - outputTokens: usage.outputTokens || 0, - cacheReadTokens, - cacheWriteTokens, - reasoningTokens: usage.details?.reasoningTokens, - } + const { chunk } = normalizeProviderUsage({ + provider: "deepseek", + apiProtocol: "openai", + usage: usage as any, + providerMetadata: providerMetadata as Record | undefined, + modelInfo: this.getModel().info, + }) + return chunk } /** @@ -113,25 +111,43 @@ export class DeepSeekHandler extends BaseProvider implements SingleCompletionHan messages: RooMessage[], metadata?: ApiHandlerCreateMessageMetadata, ): ApiStream { - const { temperature } = this.getModel() + const { temperature, info } = this.getModel() const languageModel = this.getLanguageModel() // Convert messages to AI SDK format const aiSdkMessages = messages as ModelMessage[] + const promptCache = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: "deepseek", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: info.supportsPromptCache, + promptCacheRetention: info.promptCacheRetention, + }, + settings: this.options, + }) + // Convert tools to OpenAI format first, then to AI SDK format const openAiTools = this.convertToolsForOpenAI(metadata?.tools) - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + const aiSdkTools = convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined + + const providerOptions = mergeProviderOptions(undefined, promptCache.providerOptionsPatch) // Build the request options const requestOptions: Parameters[0] = { model: languageModel, - system: systemPrompt, + system: promptCache.systemProviderOptions + ? ({ role: "system", content: systemPrompt, providerOptions: promptCache.systemProviderOptions } as any) + : systemPrompt, messages: aiSdkMessages, temperature: this.options.modelTemperature ?? temperature ?? DEEP_SEEK_DEFAULT_TEMPERATURE, maxOutputTokens: this.getMaxOutputTokens(), tools: aiSdkTools, toolChoice: mapToolChoice(metadata?.tool_choice), + ...(providerOptions ? ({ providerOptions } as Record) : {}), } // Use streamText for streaming responses diff --git a/src/api/providers/fireworks.ts b/src/api/providers/fireworks.ts index cfc459c144..5f629f0c9e 100644 --- a/src/api/providers/fireworks.ts +++ b/src/api/providers/fireworks.ts @@ -13,8 +13,10 @@ import { mapToolChoice, handleAiSdkError, } from "../transform/ai-sdk" +import { applyPromptCacheToMessages, mergeProviderOptions } from "../transform/prompt-cache" import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" import { getModelParams } from "../transform/model-params" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import { DEFAULT_HEADERS } from "./constants" import { BaseProvider } from "./base-provider" @@ -83,18 +85,14 @@ export class FireworksHandler extends BaseProvider implements SingleCompletionHa } }, ): ApiStreamUsageChunk { - // Extract cache metrics from Fireworks' providerMetadata if available - const cacheReadTokens = providerMetadata?.fireworks?.promptCacheHitTokens ?? usage.details?.cachedInputTokens - const cacheWriteTokens = providerMetadata?.fireworks?.promptCacheMissTokens - - return { - type: "usage", - inputTokens: usage.inputTokens || 0, - outputTokens: usage.outputTokens || 0, - cacheReadTokens, - cacheWriteTokens, - reasoningTokens: usage.details?.reasoningTokens, - } + const { chunk } = normalizeProviderUsage({ + provider: "fireworks", + apiProtocol: "openai", + usage: usage as any, + providerMetadata: providerMetadata as Record | undefined, + modelInfo: this.getModel().info, + }) + return chunk } /** @@ -113,25 +111,43 @@ export class FireworksHandler extends BaseProvider implements SingleCompletionHa messages: RooMessage[], metadata?: ApiHandlerCreateMessageMetadata, ): ApiStream { - const { temperature } = this.getModel() + const { temperature, info } = this.getModel() const languageModel = this.getLanguageModel() // Convert messages to AI SDK format const aiSdkMessages = messages as ModelMessage[] + const promptCache = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: "fireworks", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: info.supportsPromptCache, + promptCacheRetention: info.promptCacheRetention, + }, + settings: this.options, + }) + // Convert tools to OpenAI format first, then to AI SDK format const openAiTools = this.convertToolsForOpenAI(metadata?.tools) - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + const aiSdkTools = convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined + + const providerOptions = mergeProviderOptions(undefined, promptCache.providerOptionsPatch) // Build the request options const requestOptions: Parameters[0] = { model: languageModel, - system: systemPrompt, + system: promptCache.systemProviderOptions + ? ({ role: "system", content: systemPrompt, providerOptions: promptCache.systemProviderOptions } as any) + : systemPrompt, messages: aiSdkMessages, temperature: this.options.modelTemperature ?? temperature ?? FIREWORKS_DEFAULT_TEMPERATURE, maxOutputTokens: this.getMaxOutputTokens(), tools: aiSdkTools, toolChoice: mapToolChoice(metadata?.tool_choice), + ...(providerOptions ? ({ providerOptions } as Record) : {}), } // Use streamText for streaming responses diff --git a/src/api/providers/gemini.ts b/src/api/providers/gemini.ts index b7aebf409c..519002d007 100644 --- a/src/api/providers/gemini.ts +++ b/src/api/providers/gemini.ts @@ -21,9 +21,11 @@ import { handleAiSdkError, yieldResponseMessage, } from "../transform/ai-sdk" +import { applyPromptCacheToMessages, mergeProviderOptions } from "../transform/prompt-cache" import { t } from "i18next" import type { ApiStream, ApiStreamUsageChunk, GroundingSource } from "../transform/stream" import { getModelParams } from "../transform/model-params" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" import { BaseProvider } from "./base-provider" @@ -93,6 +95,17 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl // Convert messages to AI SDK format const aiSdkMessages = filteredMessages as ModelMessage[] + const promptCache = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: "gemini", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: info.supportsPromptCache, + promptCacheRetention: info.promptCacheRetention, + }, + settings: this.options, + }) + // Convert tools to OpenAI format first, then to AI SDK format let openAiTools = this.convertToolsForOpenAI(metadata?.tools) @@ -102,7 +115,9 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl openAiTools = openAiTools.filter((tool) => tool.type === "function" && allowedSet.has(tool.function.name)) } - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + const aiSdkTools = convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined // Build tool choice - use 'required' when allowedFunctionNames restricts available tools const toolChoice = @@ -111,19 +126,26 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl : mapToolChoice(metadata?.tool_choice) // Build the request options + const providerOptions = mergeProviderOptions( + thinkingConfig ? ({ google: { thinkingConfig } } as Record) : undefined, + promptCache.providerOptionsPatch, + ) + const requestOptions: Parameters[0] = { model: this.provider(modelId), - system: systemInstruction, + system: promptCache.systemProviderOptions + ? ({ + role: "system", + content: systemInstruction, + providerOptions: promptCache.systemProviderOptions, + } as any) + : systemInstruction, messages: aiSdkMessages, temperature: temperatureConfig, maxOutputTokens, tools: aiSdkTools, toolChoice, - // Add thinking/reasoning configuration if present - // Cast to any to bypass strict JSONObject typing - the AI SDK accepts the correct runtime values - ...(thinkingConfig && { - providerOptions: { google: { thinkingConfig } } as any, - }), + ...(providerOptions ? ({ providerOptions } as Record) : {}), } try { @@ -246,6 +268,14 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl usage: { inputTokens?: number outputTokens?: number + cachedInputTokens?: number + reasoningTokens?: number + inputTokenDetails?: { + cacheReadTokens?: number + } + outputTokenDetails?: { + reasoningTokens?: number + } details?: { cachedInputTokens?: number reasoningTokens?: number @@ -254,23 +284,22 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl info: ModelInfo, providerMetadata?: Record, ): ApiStreamUsageChunk { - const inputTokens = usage.inputTokens || 0 - const outputTokens = usage.outputTokens || 0 - const cacheReadTokens = usage.details?.cachedInputTokens - const reasoningTokens = usage.details?.reasoningTokens + const normalized = normalizeProviderUsage({ + provider: "gemini", + apiProtocol: "openai", + usage: usage as any, + providerMetadata, + modelInfo: info, + }) return { - type: "usage", - inputTokens, - outputTokens, - cacheReadTokens, - reasoningTokens, + ...normalized.chunk, totalCost: this.calculateCost({ info, - inputTokens, - outputTokens, - cacheReadTokens, - reasoningTokens, + inputTokens: normalized.canonical.inputTokensTotal, + outputTokens: normalized.canonical.outputTokens, + cacheReadTokens: normalized.canonical.cacheReadTokens, + reasoningTokens: normalized.canonical.reasoningTokens ?? 0, }), } } diff --git a/src/api/providers/lite-llm.ts b/src/api/providers/lite-llm.ts index 8bca632f4c..e53e359e21 100644 --- a/src/api/providers/lite-llm.ts +++ b/src/api/providers/lite-llm.ts @@ -13,7 +13,7 @@ import { litellmDefaultModelId, litellmDefaultModelInfo, type ModelInfo, type Mo import { type ApiHandlerOptions, getModelMaxOutputTokens } from "../../shared/api" -import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" +import { ApiStream } from "../transform/stream" import { OpenAICompatibleHandler } from "./openai-compatible" import { getModels, getModelsFromCache } from "./fetchers/modelCache" @@ -34,6 +34,7 @@ export class LiteLLMHandler extends OpenAICompatibleHandler implements SingleCom modelInfo: litellmDefaultModelInfo, temperature: options.modelTemperature ?? 0, modelMaxTokens: options.modelMaxTokens, + cacheOverrideKey: "litellm", }) } @@ -92,22 +93,4 @@ export class LiteLLMHandler extends OpenAICompatibleHandler implements SingleCom await this.fetchModel() return super.completePrompt(prompt) } - - protected override processUsageMetrics(usage: { - inputTokens?: number - outputTokens?: number - details?: { - cachedInputTokens?: number - reasoningTokens?: number - } - raw?: Record - }): ApiStreamUsageChunk { - return { - type: "usage", - inputTokens: usage.inputTokens || 0, - outputTokens: usage.outputTokens || 0, - cacheReadTokens: usage.details?.cachedInputTokens, - reasoningTokens: usage.details?.reasoningTokens, - } - } } diff --git a/src/api/providers/lm-studio.ts b/src/api/providers/lm-studio.ts index 751c4ccde4..0ca89e7f6c 100644 --- a/src/api/providers/lm-studio.ts +++ b/src/api/providers/lm-studio.ts @@ -20,6 +20,7 @@ import { mapToolChoice, handleAiSdkError, } from "../transform/ai-sdk" +import { applyPromptCacheToMessages, mergeProviderOptions } from "../transform/prompt-cache" import { ApiStream } from "../transform/stream" import { OpenAICompatibleHandler, OpenAICompatibleConfig } from "./openai-compatible" @@ -43,6 +44,7 @@ export class LmStudioHandler extends OpenAICompatibleHandler implements SingleCo modelInfo, temperature: options.modelTemperature ?? LMSTUDIO_DEFAULT_TEMPERATURE, modelMaxTokens: options.modelMaxTokens ?? undefined, + cacheOverrideKey: "lmstudio", } super(options, config) @@ -66,12 +68,27 @@ export class LmStudioHandler extends OpenAICompatibleHandler implements SingleCo const aiSdkMessages = messages as ModelMessage[] + const promptCache = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: "lmstudio", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: model.info.supportsPromptCache, + promptCacheRetention: model.info.promptCacheRetention, + }, + settings: this.options, + }) + const openAiTools = this.convertToolsForOpenAI(metadata?.tools) - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + const aiSdkTools = convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined const requestOptions: Parameters[0] = { model: languageModel, - system: systemPrompt, + system: promptCache.systemProviderOptions + ? ({ role: "system", content: systemPrompt, providerOptions: promptCache.systemProviderOptions } as any) + : systemPrompt, messages: aiSdkMessages, temperature: model.temperature ?? this.config.temperature ?? LMSTUDIO_DEFAULT_TEMPERATURE, maxOutputTokens: this.getMaxOutputTokens(), @@ -79,11 +96,16 @@ export class LmStudioHandler extends OpenAICompatibleHandler implements SingleCo toolChoice: mapToolChoice(metadata?.tool_choice), } + let providerOptions: Record | undefined if (this.options.lmStudioSpeculativeDecodingEnabled && this.options.lmStudioDraftModelId) { - requestOptions.providerOptions = { + providerOptions = { lmstudio: { draft_model: this.options.lmStudioDraftModelId }, } } + providerOptions = mergeProviderOptions(providerOptions, promptCache.providerOptionsPatch) + if (providerOptions) { + requestOptions.providerOptions = providerOptions as any + } const result = streamText(requestOptions) diff --git a/src/api/providers/minimax.ts b/src/api/providers/minimax.ts index e5a83b7a8e..778dc75bbd 100644 --- a/src/api/providers/minimax.ts +++ b/src/api/providers/minimax.ts @@ -16,7 +16,7 @@ import { yieldResponseMessage, } from "../transform/ai-sdk" import { applyPromptCacheToMessages } from "../transform/prompt-cache" -import { calculateApiCostAnthropic } from "../../shared/cost" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import { DEFAULT_HEADERS } from "./constants" import { BaseProvider } from "./base-provider" @@ -73,8 +73,22 @@ export class MiniMaxHandler extends BaseProvider implements SingleCompletionHand const mergedMessages = mergeEnvironmentDetailsForMiniMax(messages as any) const aiSdkMessages = (mergedMessages as ModelMessage[]).map((message) => ({ ...message })) + + const promptCache = applyPromptCacheToMessages({ + adapter: "minimax", + overrideKey: "minimax", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: modelConfig.info.supportsPromptCache, + promptCacheRetention: (modelConfig.info as ModelInfo).promptCacheRetention, + }, + settings: this.options, + }) + const openAiTools = this.convertToolsForOpenAI(metadata?.tools) - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + const aiSdkTools = convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined const anthropicProviderOptions: Record = {} @@ -89,23 +103,11 @@ export class MiniMaxHandler extends BaseProvider implements SingleCompletionHand anthropicProviderOptions.disableParallelToolUse = true } - const promptCache = applyPromptCacheToMessages({ - adapter: "anthropic", - overrideKey: "minimax", - messages: aiSdkMessages, - modelInfo: { - supportsPromptCache: modelConfig.info.supportsPromptCache, - promptCacheRetention: (modelConfig.info as ModelInfo).promptCacheRetention, - }, - settings: this.options, - }) - const requestOptions = { model: this.client(modelConfig.id), - system: systemPrompt, - ...(promptCache.systemProviderOptions - ? ({ systemProviderOptions: promptCache.systemProviderOptions } as Record) - : {}), + system: promptCache.systemProviderOptions + ? ({ role: "system", content: systemPrompt, providerOptions: promptCache.systemProviderOptions } as any) + : systemPrompt, messages: aiSdkMessages, temperature: modelParams.temperature, maxOutputTokens: modelParams.maxTokens ?? modelConfig.info.maxTokens, @@ -150,35 +152,27 @@ export class MiniMaxHandler extends BaseProvider implements SingleCompletionHand } private processUsageMetrics( - usage: { inputTokens?: number; outputTokens?: number }, + usage: { + inputTokens?: number + outputTokens?: number + inputTokenDetails?: { + noCacheTokens?: number + cacheReadTokens?: number + cacheWriteTokens?: number + } + cachedInputTokens?: number + }, info: ModelInfo, providerMetadata?: Record>, ): ApiStreamUsageChunk { - const inputTokens = usage.inputTokens ?? 0 - const outputTokens = usage.outputTokens ?? 0 - - const anthropicMeta = providerMetadata?.anthropic as - | { cacheCreationInputTokens?: number; cacheReadInputTokens?: number } - | undefined - const cacheWriteTokens = anthropicMeta?.cacheCreationInputTokens ?? 0 - const cacheReadTokens = anthropicMeta?.cacheReadInputTokens ?? 0 - - const { totalCost } = calculateApiCostAnthropic( - info, - inputTokens, - outputTokens, - cacheWriteTokens, - cacheReadTokens, - ) - - return { - type: "usage", - inputTokens, - outputTokens, - cacheWriteTokens: cacheWriteTokens > 0 ? cacheWriteTokens : undefined, - cacheReadTokens: cacheReadTokens > 0 ? cacheReadTokens : undefined, - totalCost, - } + const { chunk } = normalizeProviderUsage({ + provider: "minimax", + apiProtocol: "anthropic", + usage: usage as any, + providerMetadata: providerMetadata as Record | undefined, + modelInfo: info, + }) + return chunk } getModel() { diff --git a/src/api/providers/mistral.ts b/src/api/providers/mistral.ts index e11ac0a79e..526f477266 100644 --- a/src/api/providers/mistral.ts +++ b/src/api/providers/mistral.ts @@ -13,8 +13,10 @@ import { import type { ApiHandlerOptions } from "../../shared/api" import { convertToAiSdkMessages, convertToolsForAiSdk, consumeAiSdkStream, handleAiSdkError } from "../transform/ai-sdk" +import { applyPromptCacheToMessages, mergeProviderOptions } from "../transform/prompt-cache" import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" import { getModelParams } from "../transform/model-params" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import { DEFAULT_HEADERS } from "./constants" import { BaseProvider } from "./base-provider" @@ -73,21 +75,26 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand /** * Process usage metrics from the AI SDK response. */ - protected processUsageMetrics(usage: { - inputTokens?: number - outputTokens?: number - details?: { - cachedInputTokens?: number - reasoningTokens?: number - } - }): ApiStreamUsageChunk { - return { - type: "usage", - inputTokens: usage.inputTokens || 0, - outputTokens: usage.outputTokens || 0, - cacheReadTokens: usage.details?.cachedInputTokens, - reasoningTokens: usage.details?.reasoningTokens, - } + protected processUsageMetrics( + usage: { + inputTokens?: number + outputTokens?: number + details?: { + cachedInputTokens?: number + reasoningTokens?: number + } + raw?: Record + }, + providerMetadata?: Record, + ): ApiStreamUsageChunk { + const { chunk } = normalizeProviderUsage({ + provider: "mistral", + apiProtocol: "openai", + usage: usage as any, + providerMetadata, + modelInfo: this.getModel().info, + }) + return chunk } /** @@ -142,24 +149,43 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand metadata?: ApiHandlerCreateMessageMetadata, ): ApiStream { const languageModel = this.getLanguageModel() + const { info } = this.getModel() // Convert messages to AI SDK format const aiSdkMessages = messages as ModelMessage[] + const promptCache = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: "mistral", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: info.supportsPromptCache, + promptCacheRetention: info.promptCacheRetention, + }, + settings: this.options, + }) + // Convert tools to OpenAI format first, then to AI SDK format const openAiTools = this.convertToolsForOpenAI(metadata?.tools) - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + const aiSdkTools = convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined + + const providerOptions = mergeProviderOptions(undefined, promptCache.providerOptionsPatch) // Build the request options // Use MISTRAL_DEFAULT_TEMPERATURE (1) as fallback to match original behavior const requestOptions: Parameters[0] = { model: languageModel, - system: systemPrompt, + system: promptCache.systemProviderOptions + ? ({ role: "system", content: systemPrompt, providerOptions: promptCache.systemProviderOptions } as any) + : systemPrompt, messages: aiSdkMessages, temperature: this.options.modelTemperature ?? MISTRAL_DEFAULT_TEMPERATURE, maxOutputTokens: this.getMaxOutputTokens(), tools: aiSdkTools, toolChoice: this.mapToolChoice(metadata?.tool_choice), + ...(providerOptions ? ({ providerOptions } as Record) : {}), } // Use streamText for streaming responses @@ -168,8 +194,8 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand try { const processUsage = this.processUsageMetrics.bind(this) yield* consumeAiSdkStream(result, async function* () { - const usage = await result.usage - yield processUsage(usage) + const [usage, providerMetadata] = await Promise.all([result.usage, result.providerMetadata]) + yield processUsage(usage as any, providerMetadata as Record | undefined) }) } catch (error) { throw handleAiSdkError(error, "Mistral") diff --git a/src/api/providers/moonshot.ts b/src/api/providers/moonshot.ts index 3e90e48f7a..7e5f7d7733 100644 --- a/src/api/providers/moonshot.ts +++ b/src/api/providers/moonshot.ts @@ -2,7 +2,6 @@ import { moonshotModels, moonshotDefaultModelId, type ModelInfo } from "@roo-cod import type { ApiHandlerOptions } from "../../shared/api" -import type { ApiStreamUsageChunk } from "../transform/stream" import { getModelParams } from "../transform/model-params" import { OpenAICompatibleHandler, OpenAICompatibleConfig } from "./openai-compatible" @@ -21,6 +20,7 @@ export class MoonshotHandler extends OpenAICompatibleHandler { modelInfo, modelMaxTokens: options.modelMaxTokens ?? undefined, temperature: options.modelTemperature ?? undefined, + cacheOverrideKey: "moonshot", } super(options, config) @@ -39,31 +39,6 @@ export class MoonshotHandler extends OpenAICompatibleHandler { return { id, info, ...params } } - /** - * Override to handle Moonshot's usage metrics, including caching. - * Moonshot returns cached_tokens in a different location than standard OpenAI. - */ - protected override processUsageMetrics(usage: { - inputTokens?: number - outputTokens?: number - details?: { - cachedInputTokens?: number - reasoningTokens?: number - } - raw?: Record - }): ApiStreamUsageChunk { - // Moonshot uses cached_tokens at the top level of raw usage data - const rawUsage = usage.raw as { cached_tokens?: number } | undefined - - return { - type: "usage", - inputTokens: usage.inputTokens || 0, - outputTokens: usage.outputTokens || 0, - cacheWriteTokens: 0, - cacheReadTokens: rawUsage?.cached_tokens ?? usage.details?.cachedInputTokens, - } - } - /** * Override to always include max_tokens for Moonshot (not max_completion_tokens). * Moonshot requires max_tokens parameter to be sent. diff --git a/src/api/providers/native-ollama.ts b/src/api/providers/native-ollama.ts index f1da6a01bd..bff3f377ed 100644 --- a/src/api/providers/native-ollama.ts +++ b/src/api/providers/native-ollama.ts @@ -14,7 +14,9 @@ import { handleAiSdkError, yieldResponseMessage, } from "../transform/ai-sdk" +import { applyPromptCacheToMessages, mergeProviderOptions } from "../transform/prompt-cache" import { ApiStream } from "../transform/stream" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import { BaseProvider } from "./base-provider" import { getOllamaModels } from "./fetchers/ollama" @@ -89,7 +91,7 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio metadata?: ApiHandlerCreateMessageMetadata, ): ApiStream { await this.fetchModel() - const { id: modelId } = this.getModel() + const { id: modelId, info } = this.getModel() const useR1Format = modelId.toLowerCase().includes("deepseek-r1") const temperature = this.options.modelTemperature ?? (useR1Format ? DEEP_SEEK_DEFAULT_TEMPERATURE : 0) @@ -97,19 +99,37 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio const aiSdkMessages = messages as ModelMessage[] - const openAiTools = this.convertToolsForOpenAI(metadata?.tools) - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + const promptCache = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: "ollama", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: info.supportsPromptCache, + promptCacheRetention: info.promptCacheRetention, + }, + settings: this.options, + }) - const providerOptions = this.buildProviderOptions(useR1Format) + const openAiTools = this.convertToolsForOpenAI(metadata?.tools) + const aiSdkTools = convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined + + const providerOptions = mergeProviderOptions( + this.buildProviderOptions(useR1Format), + promptCache.providerOptionsPatch, + ) const requestOptions: Parameters[0] = { model: languageModel, - system: systemPrompt, + system: promptCache.systemProviderOptions + ? ({ role: "system", content: systemPrompt, providerOptions: promptCache.systemProviderOptions } as any) + : systemPrompt, messages: aiSdkMessages, temperature, tools: aiSdkTools, toolChoice: mapToolChoice(metadata?.tool_choice), - ...(providerOptions && { providerOptions }), + ...(providerOptions ? ({ providerOptions } as Record) : {}), } const result = streamText(requestOptions) @@ -123,11 +143,13 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio const usage = await result.usage if (usage) { - yield { - type: "usage", - inputTokens: usage.inputTokens || 0, - outputTokens: usage.outputTokens || 0, - } + const { chunk } = normalizeProviderUsage({ + provider: "native-ollama", + apiProtocol: "openai", + usage: usage as any, + modelInfo: info, + }) + yield chunk } yield* yieldResponseMessage(result) diff --git a/src/api/providers/openai-codex.ts b/src/api/providers/openai-codex.ts index d96e744661..3c4351246f 100644 --- a/src/api/providers/openai-codex.ts +++ b/src/api/providers/openai-codex.ts @@ -23,8 +23,10 @@ import { handleAiSdkError, yieldResponseMessage, } from "../transform/ai-sdk" +import { applyPromptCacheToMessages, mergeProviderOptions } from "../transform/prompt-cache" import { ApiStream } from "../transform/stream" import { getModelParams } from "../transform/model-params" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import { BaseProvider } from "./base-provider" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" @@ -186,19 +188,42 @@ export class OpenAiCodexHandler extends BaseProvider implements SingleCompletion injectEncryptedReasoning(aiSdkMessages, encryptedReasoningItems, messages as RooMessage[]) } + const promptCache = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: "openai-codex", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: model.info.supportsPromptCache, + promptCacheRetention: model.info.promptCacheRetention, + }, + settings: this.options, + }) + // Convert tools to OpenAI format first, then to AI SDK format const openAiTools = this.convertToolsForOpenAI(metadata?.tools) - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + const aiSdkTools = convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined - const providerOptions = this.buildProviderOptions(model, metadata, systemPrompt) + const providerOptions = mergeProviderOptions( + this.buildProviderOptions(model, metadata, systemPrompt), + promptCache.providerOptionsPatch, + ) // Note: maxOutputTokens is intentionally omitted — Codex backend rejects it. const result = streamText({ model: languageModel, + system: promptCache.systemProviderOptions + ? ({ + role: "system", + content: systemPrompt, + providerOptions: promptCache.systemProviderOptions, + } as any) + : systemPrompt, messages: aiSdkMessages, tools: aiSdkTools, toolChoice: mapToolChoice(metadata?.tool_choice), - providerOptions, + ...(providerOptions ? ({ providerOptions } as Record) : {}), ...(model.info.supportsTemperature !== false && { temperature: this.options.modelTemperature ?? 0, }), @@ -250,26 +275,15 @@ export class OpenAiCodexHandler extends BaseProvider implements SingleCompletion // Yield usage — subscription pricing means totalCost is always 0 const usage = await result.usage if (usage) { - const inputTokens = usage.inputTokens || 0 - const outputTokens = usage.outputTokens || 0 - const details = (usage as any).details as - | { cachedInputTokens?: number; reasoningTokens?: number } - | undefined - const cacheReadTokens = details?.cachedInputTokens ?? 0 - // The OpenAI Responses API does not report cache write tokens separately; - // only cached (read) tokens are available via usage.details.cachedInputTokens. - const cacheWriteTokens = 0 - const reasoningTokens = details?.reasoningTokens - - yield { - type: "usage", - inputTokens, - outputTokens, - cacheWriteTokens: cacheWriteTokens || undefined, - cacheReadTokens: cacheReadTokens || undefined, - ...(typeof reasoningTokens === "number" ? { reasoningTokens } : {}), - totalCost: 0, // Subscription-based pricing - } + const { chunk } = normalizeProviderUsage({ + provider: "openai-codex", + apiProtocol: "openai", + usage: usage as any, + providerMetadata: providerMeta as Record | undefined, + modelInfo: model.info, + totalCostOverride: 0, // Subscription-based pricing + }) + yield chunk } } catch (usageError) { if (lastStreamError) { diff --git a/src/api/providers/openai-compatible.ts b/src/api/providers/openai-compatible.ts index 56ca670756..1e0c2c172e 100644 --- a/src/api/providers/openai-compatible.ts +++ b/src/api/providers/openai-compatible.ts @@ -8,7 +8,7 @@ import OpenAI from "openai" import { createOpenAICompatible } from "@ai-sdk/openai-compatible" import { streamText, generateText, LanguageModel, ToolSet, ModelMessage } from "ai" -import type { ModelInfo } from "@roo-code/types" +import type { ModelInfo, ProviderName } from "@roo-code/types" import type { ApiHandlerOptions } from "../../shared/api" @@ -19,7 +19,9 @@ import { mapToolChoice, handleAiSdkError, } from "../transform/ai-sdk" +import { applyPromptCacheToMessages, mergeProviderOptions } from "../transform/prompt-cache" import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import { DEFAULT_HEADERS } from "./constants" import { BaseProvider } from "./base-provider" @@ -48,6 +50,10 @@ export interface OpenAICompatibleConfig { modelMaxTokens?: number /** Temperature setting */ temperature?: number + /** Canonical provider key used for prompt caching overrides. */ + cacheOverrideKey: ProviderName + /** Optional usage profile key for shared usage normalization. Defaults to cacheOverrideKey. */ + usageProfileKey?: string } /** @@ -92,22 +98,34 @@ export abstract class OpenAICompatibleHandler extends BaseProvider implements Si * Process usage metrics from the AI SDK response. * Can be overridden by subclasses to handle provider-specific usage formats. */ - protected processUsageMetrics(usage: { - inputTokens?: number - outputTokens?: number - details?: { + protected processUsageMetrics( + usage: { + inputTokens?: number + outputTokens?: number + inputTokenDetails?: { + cacheReadTokens?: number + } + outputTokenDetails?: { + reasoningTokens?: number + } cachedInputTokens?: number reasoningTokens?: number - } - raw?: Record - }): ApiStreamUsageChunk { - return { - type: "usage", - inputTokens: usage.inputTokens || 0, - outputTokens: usage.outputTokens || 0, - cacheReadTokens: usage.details?.cachedInputTokens, - reasoningTokens: usage.details?.reasoningTokens, - } + details?: { + cachedInputTokens?: number + reasoningTokens?: number + } + raw?: Record + }, + providerMetadata?: Record, + ): ApiStreamUsageChunk { + const { chunk } = normalizeProviderUsage({ + provider: this.config.usageProfileKey ?? this.config.cacheOverrideKey ?? "openai-compatible", + apiProtocol: "openai", + usage: usage as any, + providerMetadata, + modelInfo: this.getModel().info, + }) + return chunk } /** @@ -134,19 +152,37 @@ export abstract class OpenAICompatibleHandler extends BaseProvider implements Si // Convert messages to AI SDK format const aiSdkMessages = messages as ModelMessage[] + const promptCache = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: this.config.cacheOverrideKey, + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: model.info.supportsPromptCache, + promptCacheRetention: model.info.promptCacheRetention, + }, + settings: this.options, + }) + // Convert tools to OpenAI format first, then to AI SDK format const openAiTools = this.convertToolsForOpenAI(metadata?.tools) - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + const aiSdkTools = convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined + + const providerOptions = mergeProviderOptions(undefined, promptCache.providerOptionsPatch) // Build the request options const requestOptions: Parameters[0] = { model: languageModel, - system: systemPrompt, + system: promptCache.systemProviderOptions + ? ({ role: "system", content: systemPrompt, providerOptions: promptCache.systemProviderOptions } as any) + : systemPrompt, messages: aiSdkMessages, temperature: model.temperature ?? this.config.temperature ?? 0, maxOutputTokens: this.getMaxOutputTokens(), tools: aiSdkTools, toolChoice: mapToolChoice(metadata?.tool_choice), + ...(providerOptions ? ({ providerOptions } as Record) : {}), } // Use streamText for streaming responses @@ -155,8 +191,8 @@ export abstract class OpenAICompatibleHandler extends BaseProvider implements Si try { const processUsage = this.processUsageMetrics.bind(this) yield* consumeAiSdkStream(result, async function* () { - const usage = await result.usage - yield processUsage(usage) + const [usage, providerMetadata] = await Promise.all([result.usage, result.providerMetadata]) + yield processUsage(usage, providerMetadata as Record | undefined) }) } catch (error) { // Handle AI SDK errors (AI_RetryError, AI_APICallError, etc.) diff --git a/src/api/providers/openai-native.ts b/src/api/providers/openai-native.ts index b99103a857..c780d0bf12 100644 --- a/src/api/providers/openai-native.ts +++ b/src/api/providers/openai-native.ts @@ -17,12 +17,12 @@ import { } from "@roo-code/types" import type { ApiHandlerOptions } from "../../shared/api" -import { calculateApiCostOpenAI } from "../../shared/cost" import { convertToolsForAiSdk, consumeAiSdkStream, mapToolChoice, handleAiSdkError } from "../transform/ai-sdk" import { applyPromptCacheToMessages } from "../transform/prompt-cache" import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" import { getModelParams } from "../transform/model-params" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import { BaseProvider } from "./base-provider" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" @@ -339,44 +339,41 @@ export class OpenAiNativeHandler extends BaseProvider implements SingleCompletio usage: { inputTokens?: number outputTokens?: number + inputTokenDetails?: { + cacheReadTokens?: number + cacheWriteTokens?: number + } + outputTokenDetails?: { + reasoningTokens?: number + } + cachedInputTokens?: number + reasoningTokens?: number details?: { cachedInputTokens?: number reasoningTokens?: number } + raw?: { + input_tokens_details?: { + cached_tokens?: number | null + } + } }, model: OpenAiNativeModel, providerMetadata?: Record, ): ApiStreamUsageChunk { - const inputTokens = usage.inputTokens || 0 - const outputTokens = usage.outputTokens || 0 - - const cacheReadTokens = usage.details?.cachedInputTokens ?? 0 - // The OpenAI Responses API does not report cache write tokens separately; - // only cached (read) tokens are available via usage.details.cachedInputTokens. - const cacheWriteTokens = 0 - const reasoningTokens = usage.details?.reasoningTokens - const effectiveTier = this.lastServiceTier || (this.options.openAiNativeServiceTier as ServiceTier | undefined) || undefined const effectiveInfo = this.applyServiceTierPricing(model.info, effectiveTier) - const { totalCost } = calculateApiCostOpenAI( - effectiveInfo, - inputTokens, - outputTokens, - cacheWriteTokens, - cacheReadTokens, - ) + const { chunk } = normalizeProviderUsage({ + provider: "openai-native", + apiProtocol: "openai", + usage: usage as any, + providerMetadata, + modelInfo: effectiveInfo, + }) - return { - type: "usage", - inputTokens, - outputTokens, - cacheWriteTokens: cacheWriteTokens || undefined, - cacheReadTokens: cacheReadTokens || undefined, - ...(typeof reasoningTokens === "number" ? { reasoningTokens } : {}), - totalCost, - } + return chunk } /** diff --git a/src/api/providers/openai.ts b/src/api/providers/openai.ts index a6483ad6aa..051eb678d5 100644 --- a/src/api/providers/openai.ts +++ b/src/api/providers/openai.ts @@ -24,8 +24,10 @@ import { handleAiSdkError, yieldResponseMessage, } from "../transform/ai-sdk" +import { applyPromptCacheToMessages, mergeProviderOptions } from "../transform/prompt-cache" import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" import { getModelParams } from "../transform/model-params" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import { DEFAULT_HEADERS } from "./constants" import { BaseProvider } from "./base-provider" @@ -109,7 +111,6 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl const aiSdkMessages = messages as ModelMessage[] const openAiTools = this.convertToolsForOpenAI(metadata?.tools) - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined let effectiveSystemPrompt: string | undefined = systemPrompt let effectiveTemperature: number | undefined = @@ -144,26 +145,43 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl aiSdkMessages.unshift({ role: "user", content: systemPrompt }) } + const promptCache = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: "openai", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: modelInfo.supportsPromptCache, + promptCacheRetention: modelInfo.promptCacheRetention, + }, + settings: this.options, + }) + + const mergedProviderOptions = mergeProviderOptions(providerOptions, promptCache.providerOptionsPatch) + if (this.options.openAiStreamingEnabled ?? true) { yield* this.handleStreaming( languageModel, - effectiveSystemPrompt, + this.buildSystemPromptWithProviderOptions(effectiveSystemPrompt, promptCache.systemProviderOptions), aiSdkMessages, effectiveTemperature, - aiSdkTools, + convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined, metadata, - providerOptions, + (mergedProviderOptions as Record) ?? {}, modelInfo, ) } else { yield* this.handleNonStreaming( languageModel, - effectiveSystemPrompt, + this.buildSystemPromptWithProviderOptions(effectiveSystemPrompt, promptCache.systemProviderOptions), aiSdkMessages, effectiveTemperature, - aiSdkTools, + convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined, metadata, - providerOptions, + (mergedProviderOptions as Record) ?? {}, modelInfo, ) } @@ -171,7 +189,7 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl private async *handleStreaming( languageModel: LanguageModel, - systemPrompt: string | undefined, + system: string | Record | undefined, messages: ModelMessage[], temperature: number | undefined, tools: ToolSet | undefined, @@ -181,7 +199,7 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl ): ApiStream { const result = streamText({ model: languageModel, - system: systemPrompt, + system: system as any, messages, temperature, maxOutputTokens: this.getMaxOutputTokens(), @@ -242,7 +260,7 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl private async *handleNonStreaming( languageModel: LanguageModel, - systemPrompt: string | undefined, + system: string | Record | undefined, messages: ModelMessage[], temperature: number | undefined, tools: ToolSet | undefined, @@ -253,7 +271,7 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl try { const { text, toolCalls, usage, providerMetadata } = await generateText({ model: languageModel, - system: systemPrompt, + system: system as any, messages, temperature, maxOutputTokens: this.getMaxOutputTokens(), @@ -286,10 +304,31 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl } } + private buildSystemPromptWithProviderOptions( + systemPrompt: string | undefined, + systemProviderOptions: Record | undefined, + ): string | Record | undefined { + if (!systemPrompt) { + return systemPrompt + } + + return systemProviderOptions + ? ({ role: "system", content: systemPrompt, providerOptions: systemProviderOptions } as const) + : systemPrompt + } + protected processUsageMetrics( usage: { inputTokens?: number outputTokens?: number + inputTokenDetails?: { + cacheReadTokens?: number + } + outputTokenDetails?: { + reasoningTokens?: number + } + cachedInputTokens?: number + reasoningTokens?: number details?: { cachedInputTokens?: number reasoningTokens?: number @@ -303,18 +342,15 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl } }, ): ApiStreamUsageChunk { - // Extract cache and reasoning metrics from OpenAI's providerMetadata when available, - // falling back to usage.details for standard AI SDK fields. - const cacheReadTokens = providerMetadata?.openai?.cachedPromptTokens ?? usage.details?.cachedInputTokens - const reasoningTokens = providerMetadata?.openai?.reasoningTokens ?? usage.details?.reasoningTokens - - return { - type: "usage", - inputTokens: usage.inputTokens || 0, - outputTokens: usage.outputTokens || 0, - cacheReadTokens, - reasoningTokens, - } + const providerKey = this.isAzureAiInference || this.isAzureOpenAi ? "azure" : "openai" + const { chunk } = normalizeProviderUsage({ + provider: providerKey, + apiProtocol: "openai", + usage: usage as any, + providerMetadata: providerMetadata as Record | undefined, + modelInfo: this.getModel().info, + }) + return chunk } override getModel() { diff --git a/src/api/providers/openrouter.ts b/src/api/providers/openrouter.ts index 8462a4ac03..3fcace84c8 100644 --- a/src/api/providers/openrouter.ts +++ b/src/api/providers/openrouter.ts @@ -14,7 +14,6 @@ import { import { TelemetryService } from "@roo-code/telemetry" import type { ApiHandlerOptions } from "../../shared/api" -import { calculateApiCostOpenAI } from "../../shared/cost" import { getModelParams } from "../transform/model-params" import { @@ -23,15 +22,17 @@ import { processAiSdkStreamPart, yieldResponseMessage, } from "../transform/ai-sdk" +import { applyPromptCacheToMessages, mergeProviderOptions } from "../transform/prompt-cache" import { BaseProvider } from "./base-provider" import { getModels, getModelsFromCache } from "./fetchers/modelCache" import { getModelEndpoints } from "./fetchers/modelEndpointCache" import { applyRouterToolPreferences } from "./utils/router-tool-preferences" import { generateImageWithProvider, ImageGenerationResult } from "./utils/image-generation" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import type { ApiHandlerCreateMessageMetadata, SingleCompletionHandler } from "../index" -import type { ApiStreamChunk, ApiStreamUsageChunk } from "../transform/stream" +import type { ApiStreamChunk } from "../transform/stream" import type { RooMessage } from "../../core/task-persistence/rooMessage" export class OpenRouterHandler extends BaseProvider implements SingleCompletionHandler { @@ -86,48 +87,6 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH }) } - private normalizeUsage( - usage: { inputTokens: number; outputTokens: number }, - providerMetadata: Record | undefined, - modelInfo: ModelInfo, - ): ApiStreamUsageChunk { - const inputTokens = usage.inputTokens ?? 0 - const outputTokens = usage.outputTokens ?? 0 - const openrouterMeta = providerMetadata?.openrouter ?? {} - const cacheReadTokens = - openrouterMeta.cachedInputTokens ?? - openrouterMeta.cache_read_input_tokens ?? - openrouterMeta.cacheReadTokens ?? - openrouterMeta.cached_tokens ?? - 0 - const cacheWriteTokens = - openrouterMeta.cacheCreationInputTokens ?? - openrouterMeta.cache_creation_input_tokens ?? - openrouterMeta.cacheWriteTokens ?? - 0 - const reasoningTokens = - openrouterMeta.reasoningOutputTokens ?? - openrouterMeta.reasoning_tokens ?? - openrouterMeta.output_tokens_details?.reasoning_tokens ?? - undefined - const { totalCost } = calculateApiCostOpenAI( - modelInfo, - inputTokens, - outputTokens, - cacheWriteTokens, - cacheReadTokens, - ) - return { - type: "usage", - inputTokens, - outputTokens, - ...(cacheWriteTokens > 0 ? { cacheWriteTokens } : {}), - ...(cacheReadTokens > 0 ? { cacheReadTokens } : {}), - ...(typeof reasoningTokens === "number" && reasoningTokens > 0 ? { reasoningTokens } : {}), - totalCost, - } - } - override async *createMessage( systemPrompt: string, messages: RooMessage[], @@ -150,11 +109,24 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH const aiSdkMessages = messages as ModelMessage[] + const promptCache = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: "openrouter", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: model.info.supportsPromptCache, + promptCacheRetention: model.info.promptCacheRetention, + }, + settings: this.options, + }) + const openrouter = this.createOpenRouterProvider({ reasoning, headers }) - const tools = convertToolsForAiSdk(metadata?.tools) + const tools = convertToolsForAiSdk(metadata?.tools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) - const providerOptions: + const openRouterProviderOptions: | { openrouter?: { provider?: { order: string[]; only: string[]; allow_fallbacks: boolean } @@ -174,17 +146,28 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH } : undefined + const providerOptions = mergeProviderOptions( + openRouterProviderOptions as Record | undefined, + promptCache.providerOptionsPatch, + ) + try { const result = streamText({ model: openrouter.chat(modelId), - system: systemPrompt, + system: promptCache.systemProviderOptions + ? ({ + role: "system", + content: systemPrompt, + providerOptions: promptCache.systemProviderOptions, + } as any) + : systemPrompt, messages: aiSdkMessages, maxOutputTokens: maxTokens && maxTokens > 0 ? maxTokens : undefined, temperature, topP, tools, toolChoice: metadata?.tool_choice as any, - providerOptions, + ...(providerOptions ? ({ providerOptions } as Record) : {}), }) for await (const part of result.fullStream) { @@ -196,15 +179,19 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH const usage = await result.usage const totalUsage = await result.totalUsage - const usageChunk = this.normalizeUsage( - { - inputTokens: totalUsage.inputTokens ?? usage.inputTokens ?? 0, - outputTokens: totalUsage.outputTokens ?? usage.outputTokens ?? 0, - }, - providerMetadata, - model.info, - ) - yield usageChunk + const usageRecord = { + ...(usage as any), + inputTokens: totalUsage.inputTokens ?? usage.inputTokens ?? 0, + outputTokens: totalUsage.outputTokens ?? usage.outputTokens ?? 0, + } + const { chunk } = normalizeProviderUsage({ + provider: "openrouter", + apiProtocol: "openai", + usage: usageRecord as any, + providerMetadata: providerMetadata as Record | undefined, + modelInfo: model.info, + }) + yield chunk yield* yieldResponseMessage(result) } catch (error: any) { diff --git a/src/api/providers/qwen-code.ts b/src/api/providers/qwen-code.ts index c90854d630..b59819dfa0 100644 --- a/src/api/providers/qwen-code.ts +++ b/src/api/providers/qwen-code.ts @@ -80,6 +80,7 @@ export class QwenCodeHandler extends OpenAICompatibleHandler implements SingleCo modelInfo, modelMaxTokens: options.modelMaxTokens ?? undefined, temperature: options.modelTemperature ?? undefined, + cacheOverrideKey: "qwen-code", } super(options, config) diff --git a/src/api/providers/requesty.ts b/src/api/providers/requesty.ts index 29a43acb9b..4432d76cbb 100644 --- a/src/api/providers/requesty.ts +++ b/src/api/providers/requesty.ts @@ -5,7 +5,6 @@ import { streamText, generateText, ToolSet, ModelMessage } from "ai" import { type ModelInfo, type ModelRecord, requestyDefaultModelId, requestyDefaultModelInfo } from "@roo-code/types" import type { ApiHandlerOptions } from "../../shared/api" -import { calculateApiCostOpenAI } from "../../shared/cost" import { convertToAiSdkMessages, @@ -14,6 +13,7 @@ import { mapToolChoice, handleAiSdkError, } from "../transform/ai-sdk" +import { applyPromptCacheToMessages, mergeProviderOptions } from "../transform/prompt-cache" import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" import { getModelParams } from "../transform/model-params" @@ -23,6 +23,7 @@ import { BaseProvider } from "./base-provider" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" import { toRequestyServiceUrl } from "../../shared/utils/requesty" import { applyRouterToolPreferences } from "./utils/router-tool-preferences" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import type { RooMessage } from "../../core/task-persistence/rooMessage" /** @@ -148,24 +149,14 @@ export class RequestyHandler extends BaseProvider implements SingleCompletionHan modelInfo?: ModelInfo, providerMetadata?: RequestyProviderMetadata, ): ApiStreamUsageChunk { - const inputTokens = usage.inputTokens || 0 - const outputTokens = usage.outputTokens || 0 - const cacheWriteTokens = providerMetadata?.requesty?.usage?.cachingTokens ?? 0 - const cacheReadTokens = providerMetadata?.requesty?.usage?.cachedTokens ?? usage.details?.cachedInputTokens ?? 0 - - const { totalCost } = modelInfo - ? calculateApiCostOpenAI(modelInfo, inputTokens, outputTokens, cacheWriteTokens, cacheReadTokens) - : { totalCost: 0 } - - return { - type: "usage", - inputTokens, - outputTokens, - cacheWriteTokens, - cacheReadTokens, - reasoningTokens: usage.details?.reasoningTokens, - totalCost, - } + const { chunk } = normalizeProviderUsage({ + provider: "requesty", + apiProtocol: "openai", + usage: usage as any, + providerMetadata: providerMetadata as Record | undefined, + modelInfo, + }) + return chunk } /** @@ -181,20 +172,39 @@ export class RequestyHandler extends BaseProvider implements SingleCompletionHan const aiSdkMessages = messages as ModelMessage[] + const promptCache = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: "requesty", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: info.supportsPromptCache, + promptCacheRetention: info.promptCacheRetention, + }, + settings: this.options, + }) + const openAiTools = this.convertToolsForOpenAI(metadata?.tools) - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + const aiSdkTools = convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined const requestyOptions = this.getRequestyProviderOptions(metadata) + const providerOptions = mergeProviderOptions( + requestyOptions ? ({ requesty: requestyOptions } as Record) : undefined, + promptCache.providerOptionsPatch, + ) const requestOptions: Parameters[0] = { model: languageModel, - system: systemPrompt, + system: promptCache.systemProviderOptions + ? ({ role: "system", content: systemPrompt, providerOptions: promptCache.systemProviderOptions } as any) + : systemPrompt, messages: aiSdkMessages, temperature: this.options.modelTemperature ?? temperature ?? 0, maxOutputTokens: this.getMaxOutputTokens(), tools: aiSdkTools, toolChoice: mapToolChoice(metadata?.tool_choice), - ...(requestyOptions ? { providerOptions: { requesty: requestyOptions } } : {}), + ...(providerOptions ? ({ providerOptions } as Record) : {}), } const result = streamText(requestOptions) diff --git a/src/api/providers/roo.ts b/src/api/providers/roo.ts index 9e52de109e..7341a627d8 100644 --- a/src/api/providers/roo.ts +++ b/src/api/providers/roo.ts @@ -1,13 +1,12 @@ import { Anthropic } from "@anthropic-ai/sdk" import { createOpenAICompatible } from "@ai-sdk/openai-compatible" -import { streamText, generateText, type ModelMessage } from "ai" +import { createGateway, streamText, generateText, type ModelMessage } from "ai" import { rooDefaultModelId, getApiProtocol, type ImageGenerationApiMethod } from "@roo-code/types" import { CloudService } from "@roo-code/cloud" import { Package } from "../../shared/package" import type { ApiHandlerOptions } from "../../shared/api" -import { calculateApiCostOpenAI } from "../../shared/cost" import { ApiStream } from "../transform/stream" import { getModelParams } from "../transform/model-params" import { @@ -17,6 +16,7 @@ import { mapToolChoice, yieldResponseMessage, } from "../transform/ai-sdk" +import { applyPromptCacheToMessages, mergeProviderOptions } from "../transform/prompt-cache" import type { RooReasoningParams } from "../transform/reasoning" import { getRooReasoning } from "../transform/reasoning" @@ -24,6 +24,7 @@ import type { ApiHandlerCreateMessageMetadata, SingleCompletionHandler } from ". import { BaseProvider } from "./base-provider" import { getModels, getModelsFromCache } from "./fetchers/modelCache" import { generateImageWithProvider, generateImageWithImagesApi, ImageGenerationResult } from "./utils/image-generation" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import { t } from "../../i18n" import type { RooMessage } from "../../core/task-persistence/rooMessage" @@ -57,6 +58,15 @@ export class RooHandler extends BaseProvider implements SingleCompletionHandler }) } + private shouldUseGatewaySdk(): boolean { + const envValue = process.env.ROO_CODE_ROUTER_USE_GATEWAY_SDK + if (!envValue) { + return false + } + + return ["1", "true", "yes", "on"].includes(envValue.toLowerCase()) + } + /** * Per-request provider factory. Creates a fresh provider instance * to ensure the latest session token is used for each request. @@ -84,6 +94,22 @@ export class RooHandler extends BaseProvider implements SingleCompletionHandler }) } + private createRooGatewayProvider(options?: { taskId?: string }) { + const token = this.options.rooApiKey ?? getSessionToken() + const headers: Record = { + "X-Roo-App-Version": Package.version, + } + if (options?.taskId) { + headers["X-Roo-Task-ID"] = options.taskId + } + + return createGateway({ + apiKey: token || "not-provided", + baseURL: `${this.fetcherBaseURL}/v3/ai`, + headers, + }) + } + override isAiSdkProvider() { return true as const } @@ -116,24 +142,47 @@ export class RooHandler extends BaseProvider implements SingleCompletionHandler const maxTokens = params.maxTokens ?? undefined const temperature = params.temperature ?? 0 - // Create per-request provider with fresh session token - const provider = this.createRooProvider({ reasoning, taskId: metadata?.taskId }) + // Create per-request provider with fresh session token. + // Optional gateway mode can be enabled via ROO_CODE_ROUTER_USE_GATEWAY_SDK. + const provider = this.shouldUseGatewaySdk() + ? this.createRooGatewayProvider({ taskId: metadata?.taskId }) + : this.createRooProvider({ reasoning, taskId: metadata?.taskId }) // RooMessage[] is already AI SDK-compatible, cast directly const aiSdkMessages = messages as ModelMessage[] - const tools = convertToolsForAiSdk(this.convertToolsForOpenAI(metadata?.tools)) + const promptCache = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: "roo", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: info.supportsPromptCache, + promptCacheRetention: "promptCacheRetention" in info ? info.promptCacheRetention : undefined, + }, + settings: this.options, + }) + const tools = convertToolsForAiSdk(this.convertToolsForOpenAI(metadata?.tools), { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) + const providerOptions = mergeProviderOptions(undefined, promptCache.providerOptionsPatch) let lastStreamError: string | undefined try { const result = streamText({ model: provider(modelId), - system: systemPrompt, + system: promptCache.systemProviderOptions + ? ({ + role: "system", + content: systemPrompt, + providerOptions: promptCache.systemProviderOptions, + } as any) + : systemPrompt, messages: aiSdkMessages, maxOutputTokens: maxTokens && maxTokens > 0 ? maxTokens : undefined, temperature, tools, toolChoice: mapToolChoice(metadata?.tool_choice), + ...(providerOptions ? ({ providerOptions } as Record) : {}), }) for await (const part of result.fullStream) { @@ -148,42 +197,38 @@ export class RooHandler extends BaseProvider implements SingleCompletionHandler // Check provider metadata for usage details const providerMetadata = (await result.providerMetadata) ?? (await (result as any).experimental_providerMetadata) - const rooMeta = providerMetadata?.roo as Record | undefined - - // Process usage with protocol-aware normalization + // Process usage with shared protocol-aware normalization const usage = await result.usage - const promptTokens = usage.inputTokens ?? 0 - const completionTokens = usage.outputTokens ?? 0 - - // Extract cache tokens from provider metadata - const cacheCreation = (rooMeta?.cache_creation_input_tokens as number) ?? 0 - const cacheRead = (rooMeta?.cache_read_input_tokens as number) ?? (rooMeta?.cached_tokens as number) ?? 0 - - // Protocol-aware token normalization: - // - OpenAI protocol expects TOTAL input tokens (cached + non-cached) - // - Anthropic protocol expects NON-CACHED input tokens (caches passed separately) + type UsageLike = typeof usage & { + details?: { cachedInputTokens?: number } + cacheCreationInputTokens?: number + cache_creation_input_tokens?: number + cachedInputTokens?: number + cached_tokens?: number + raw?: Record + } + const usageLike = usage as UsageLike const apiProtocol = getApiProtocol("roo", modelId) - const nonCached = Math.max(0, promptTokens - cacheCreation - cacheRead) - const inputTokens = apiProtocol === "anthropic" ? nonCached : promptTokens + const normalizedUsage = normalizeProviderUsage({ + provider: "roo", + usage: usageLike, + apiProtocol, + providerMetadata: providerMetadata as Record | undefined, + modelInfo: info, + emitZeroCacheTokens: true, + }) - // Cost: prefer server-side cost, fall back to client-side calculation const isFreeModel = info.isFree === true - const serverCost = rooMeta?.cost as number | undefined - const { totalCost: calculatedCost } = calculateApiCostOpenAI( - info, - promptTokens, - completionTokens, - cacheCreation, - cacheRead, - ) - const totalCost = isFreeModel ? 0 : (serverCost ?? calculatedCost) + const totalCost = isFreeModel ? 0 : normalizedUsage.chunk.totalCost yield { type: "usage" as const, - inputTokens, - outputTokens: completionTokens, - cacheWriteTokens: cacheCreation, - cacheReadTokens: cacheRead, + inputTokens: normalizedUsage.chunk.inputTokens, + nonCachedInputTokens: normalizedUsage.chunk.nonCachedInputTokens, + outputTokens: normalizedUsage.chunk.outputTokens, + cacheWriteTokens: normalizedUsage.chunk.cacheWriteTokens, + cacheReadTokens: normalizedUsage.chunk.cacheReadTokens, + reasoningTokens: normalizedUsage.chunk.reasoningTokens, totalCost, } @@ -208,7 +253,7 @@ export class RooHandler extends BaseProvider implements SingleCompletionHandler async completePrompt(prompt: string): Promise { const { id: modelId } = this.getModel() - const provider = this.createRooProvider() + const provider = this.shouldUseGatewaySdk() ? this.createRooGatewayProvider() : this.createRooProvider() try { const result = await generateText({ diff --git a/src/api/providers/sambanova.ts b/src/api/providers/sambanova.ts index 8f12d757ca..6647c037ee 100644 --- a/src/api/providers/sambanova.ts +++ b/src/api/providers/sambanova.ts @@ -14,8 +14,10 @@ import { handleAiSdkError, flattenAiSdkMessagesToStringContent, } from "../transform/ai-sdk" +import { applyPromptCacheToMessages, mergeProviderOptions } from "../transform/prompt-cache" import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" import { getModelParams } from "../transform/model-params" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import { DEFAULT_HEADERS } from "./constants" import { BaseProvider } from "./base-provider" @@ -84,18 +86,14 @@ export class SambaNovaHandler extends BaseProvider implements SingleCompletionHa } }, ): ApiStreamUsageChunk { - // Extract cache metrics from SambaNova's providerMetadata if available - const cacheReadTokens = providerMetadata?.sambanova?.promptCacheHitTokens ?? usage.details?.cachedInputTokens - const cacheWriteTokens = providerMetadata?.sambanova?.promptCacheMissTokens - - return { - type: "usage", - inputTokens: usage.inputTokens || 0, - outputTokens: usage.outputTokens || 0, - cacheReadTokens, - cacheWriteTokens, - reasoningTokens: usage.details?.reasoningTokens, - } + const { chunk } = normalizeProviderUsage({ + provider: "sambanova", + apiProtocol: "openai", + usage: usage as any, + providerMetadata: providerMetadata as Record | undefined, + modelInfo: this.getModel().info, + }) + return chunk } /** @@ -122,19 +120,37 @@ export class SambaNovaHandler extends BaseProvider implements SingleCompletionHa const castMessages = messages as ModelMessage[] const aiSdkMessages = info.supportsImages ? castMessages : flattenAiSdkMessagesToStringContent(castMessages) + const promptCache = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: "sambanova", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: info.supportsPromptCache, + promptCacheRetention: info.promptCacheRetention, + }, + settings: this.options, + }) + // Convert tools to OpenAI format first, then to AI SDK format const openAiTools = this.convertToolsForOpenAI(metadata?.tools) - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + const aiSdkTools = convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined + + const providerOptions = mergeProviderOptions(undefined, promptCache.providerOptionsPatch) // Build the request options const requestOptions: Parameters[0] = { model: languageModel, - system: systemPrompt, + system: promptCache.systemProviderOptions + ? ({ role: "system", content: systemPrompt, providerOptions: promptCache.systemProviderOptions } as any) + : systemPrompt, messages: aiSdkMessages, temperature: this.options.modelTemperature ?? temperature ?? SAMBANOVA_DEFAULT_TEMPERATURE, maxOutputTokens: this.getMaxOutputTokens(), tools: aiSdkTools, toolChoice: mapToolChoice(metadata?.tool_choice), + ...(providerOptions ? ({ providerOptions } as Record) : {}), } // Use streamText for streaming responses diff --git a/src/api/providers/utils/USAGE_NORMALIZATION.md b/src/api/providers/utils/USAGE_NORMALIZATION.md new file mode 100644 index 0000000000..38cd5ec616 --- /dev/null +++ b/src/api/providers/utils/USAGE_NORMALIZATION.md @@ -0,0 +1,54 @@ +# Usage Normalization (AI SDK Providers) + +This folder centralizes usage/cache normalization for AI-SDK-backed providers. + +## Canonical Contract + +The normalizer derives a canonical usage shape before emitting stream chunks: + +- `inputTokensTotal` +- `inputTokensNonCached` (optional) +- `outputTokens` +- `cacheWriteTokens` +- `cacheReadTokens` +- `reasoningTokens` (optional) +- `totalCostCandidate` (optional, provider-reported) + +`ApiStreamUsageChunk` then carries: + +- `inputTokens` (always total input tokens) +- `nonCachedInputTokens?` +- `outputTokens` +- `cacheWriteTokens?` +- `cacheReadTokens?` +- `reasoningTokens?` +- `totalCost?` + +## Precedence Policy + +For each metric, extraction precedence is: + +1. `providerMetadata` +2. AI SDK `usage` +3. `usage.raw` fallback + +This is implemented in `normalize-provider-usage.ts`. + +## Profiles + +`usage-profiles.ts` defines per-provider extraction profiles: + +- The profile specifies candidate paths per metric for `providerMetadata`, `usage`, and `raw`. +- Profiles should contain field mapping only (no provider-specific arithmetic). +- Exceptions (e.g., custom pricing behavior) should be applied at call sites via `totalCostOverride`. + +## New Provider Checklist + +1. Add/confirm a profile in `usage-profiles.ts`. +2. Wire provider usage handling to `normalizeProviderUsage(...)`. +3. Ensure `inputTokens` emitted to stream remains total input tokens. +4. Add/extend provider usage tests for: + - first-turn cache write heavy payload + - second-turn cache read heavy payload + - metadata vs usage conflicts +5. If provider has custom pricing logic, pass `totalCostOverride` explicitly and test it. diff --git a/src/api/providers/utils/normalize-provider-usage.ts b/src/api/providers/utils/normalize-provider-usage.ts new file mode 100644 index 0000000000..a14f3e1cde --- /dev/null +++ b/src/api/providers/utils/normalize-provider-usage.ts @@ -0,0 +1,195 @@ +import type { ModelInfo } from "@roo-code/types" + +import { calculateApiCostAnthropic, calculateApiCostOpenAI } from "../../../shared/cost" +import type { ApiStreamUsageChunk } from "../../transform/stream" +import { type CanonicalUsageMetrics, type UsageLike, toFiniteNumber, getValueAtPath } from "./usage-metrics" +import { type UsageMetricPaths, type UsageProtocol, getUsageProfile } from "./usage-profiles" + +interface UsageSources { + providerMetadata?: Record + usage?: Record + raw?: Record +} + +export interface NormalizeProviderUsageOptions { + provider: string + apiProtocol?: UsageProtocol + usage: UsageLike + providerMetadata?: Record + modelInfo?: ModelInfo + totalCostOverride?: number + emitZeroCacheTokens?: boolean + deriveNonCachedInputFromTotalMinusCache?: boolean +} + +export interface NormalizeProviderUsageResult { + canonical: CanonicalUsageMetrics + chunk: ApiStreamUsageChunk +} + +function toRecord(value: unknown): Record | undefined { + return typeof value === "object" && value !== null && !Array.isArray(value) + ? (value as Record) + : undefined +} + +function readMetricFromPaths(sources: UsageSources, paths?: UsageMetricPaths): number | undefined { + if (!paths) return undefined + + // Precedence policy: providerMetadata > usage > raw + const providerMetadataNumber = + paths.providerMetadata && paths.providerMetadata.length > 0 + ? firstNumberFromPathsWithArraySupport(sources.providerMetadata, paths.providerMetadata) + : undefined + if (providerMetadataNumber !== undefined) { + return providerMetadataNumber + } + + const usageNumber = + paths.usage && paths.usage.length > 0 + ? firstNumberFromPathsWithArraySupport(sources.usage, paths.usage) + : undefined + if (usageNumber !== undefined) { + return usageNumber + } + + const rawNumber = + paths.raw && paths.raw.length > 0 ? firstNumberFromPathsWithArraySupport(sources.raw, paths.raw) : undefined + if (rawNumber !== undefined) { + return rawNumber + } + + return undefined +} + +function firstNumberFromPathsWithArraySupport(source: unknown, paths: string[]): number | undefined { + if (!source) return undefined + + for (const path of paths) { + const value = getValueAtPath(source, path) + const num = toFiniteNumber(value as number | string | null | undefined) + if (num !== undefined) { + return num + } + + if (Array.isArray(value)) { + const sum = value.reduce((acc, item) => { + if (typeof item !== "object" || item === null) return acc + const maybeTokenCount = toFiniteNumber((item as Record).tokenCount as any) ?? 0 + return acc + maybeTokenCount + }, 0) + if (sum > 0) { + return sum + } + } + } + + return undefined +} + +function buildCanonicalMetrics(options: NormalizeProviderUsageOptions): CanonicalUsageMetrics { + const profile = getUsageProfile(options.provider, options.apiProtocol) + const usageRecord = toRecord(options.usage) ?? {} + + const sources: UsageSources = { + providerMetadata: options.providerMetadata, + usage: usageRecord, + raw: toRecord(options.usage.raw), + } + + const cacheWriteTokens = readMetricFromPaths(sources, profile.metrics.cacheWriteTokens) ?? 0 + const cacheReadTokens = readMetricFromPaths(sources, profile.metrics.cacheReadTokens) ?? 0 + const outputTokens = readMetricFromPaths(sources, profile.metrics.outputTokens) ?? 0 + const reasoningTokens = readMetricFromPaths(sources, profile.metrics.reasoningTokens) + const totalCostCandidate = readMetricFromPaths(sources, profile.metrics.totalCostCandidate) + + let inputTokensTotal = readMetricFromPaths(sources, profile.metrics.inputTokensTotal) + let inputTokensNonCached = readMetricFromPaths(sources, profile.metrics.inputTokensNonCached) + + const shouldDeriveNonCached = + options.deriveNonCachedInputFromTotalMinusCache ?? + profile.deriveNonCachedInputFromTotalMinusCache ?? + profile.apiProtocol === "openai" + + if (inputTokensNonCached === undefined && shouldDeriveNonCached) { + if (inputTokensTotal !== undefined) { + inputTokensNonCached = Math.max(0, inputTokensTotal - cacheWriteTokens - cacheReadTokens) + } + } + + if (inputTokensTotal === undefined) { + if (inputTokensNonCached !== undefined) { + inputTokensTotal = inputTokensNonCached + cacheWriteTokens + cacheReadTokens + } else { + inputTokensTotal = 0 + } + } + + return { + inputTokensTotal, + inputTokensNonCached, + outputTokens, + cacheWriteTokens, + cacheReadTokens, + reasoningTokens, + totalCostCandidate, + } +} + +function computeFallbackCost( + modelInfo: ModelInfo | undefined, + apiProtocol: UsageProtocol, + canonical: CanonicalUsageMetrics, +): number | undefined { + if (!modelInfo) return undefined + + if (apiProtocol === "anthropic") { + const nonCached = + canonical.inputTokensNonCached ?? + Math.max(0, canonical.inputTokensTotal - canonical.cacheWriteTokens - canonical.cacheReadTokens) + + return calculateApiCostAnthropic( + modelInfo, + nonCached, + canonical.outputTokens, + canonical.cacheWriteTokens, + canonical.cacheReadTokens, + ).totalCost + } + + return calculateApiCostOpenAI( + modelInfo, + canonical.inputTokensTotal, + canonical.outputTokens, + canonical.cacheWriteTokens, + canonical.cacheReadTokens, + ).totalCost +} + +export function normalizeProviderUsage(options: NormalizeProviderUsageOptions): NormalizeProviderUsageResult { + const profile = getUsageProfile(options.provider, options.apiProtocol) + const canonical = buildCanonicalMetrics(options) + + const totalCost = + options.totalCostOverride ?? + canonical.totalCostCandidate ?? + computeFallbackCost(options.modelInfo, profile.apiProtocol, canonical) + + const emitZeroCacheTokens = options.emitZeroCacheTokens ?? profile.emitZeroCacheTokens ?? false + + return { + canonical, + chunk: { + type: "usage", + inputTokens: canonical.inputTokensTotal, + nonCachedInputTokens: canonical.inputTokensNonCached, + outputTokens: canonical.outputTokens, + cacheWriteTokens: + emitZeroCacheTokens || canonical.cacheWriteTokens > 0 ? canonical.cacheWriteTokens : undefined, + cacheReadTokens: + emitZeroCacheTokens || canonical.cacheReadTokens > 0 ? canonical.cacheReadTokens : undefined, + reasoningTokens: canonical.reasoningTokens, + totalCost, + }, + } +} diff --git a/src/api/providers/utils/usage-metrics.ts b/src/api/providers/utils/usage-metrics.ts new file mode 100644 index 0000000000..a0472712f4 --- /dev/null +++ b/src/api/providers/utils/usage-metrics.ts @@ -0,0 +1,211 @@ +import type { ApiStreamUsageChunk } from "../../transform/stream" + +type NumberLike = number | string | null | undefined + +export type UsageLike = { + inputTokens?: number + outputTokens?: number + inputTokenDetails?: { + noCacheTokens?: number + cacheReadTokens?: number + cacheWriteTokens?: number + } + details?: { + cachedInputTokens?: number + } + cacheCreationInputTokens?: number + cache_creation_input_tokens?: number + cachedInputTokens?: number + cached_tokens?: number + raw?: Record +} + +export interface CanonicalUsageMetrics { + inputTokensTotal: number + inputTokensNonCached?: number + outputTokens: number + cacheWriteTokens: number + cacheReadTokens: number + reasoningTokens?: number + totalCostCandidate?: number +} + +export interface NormalizeUsageMetricsOptions { + usage: UsageLike + apiProtocol?: "anthropic" | "openai" + cacheWriteCandidates?: NumberLike[] + cacheReadCandidates?: NumberLike[] + inputNonCachedCandidates?: NumberLike[] + reasoningCandidates?: NumberLike[] + costCandidates?: NumberLike[] + emitZeroCacheTokens?: boolean + deriveAnthropicNoCacheFromTotalWhenMissing?: boolean +} + +export function normalizeUsageMetrics({ + usage, + apiProtocol, + cacheWriteCandidates = [], + cacheReadCandidates = [], + inputNonCachedCandidates = [], + reasoningCandidates = [], + costCandidates = [], + emitZeroCacheTokens = false, + deriveAnthropicNoCacheFromTotalWhenMissing = false, +}: NormalizeUsageMetricsOptions): Pick< + ApiStreamUsageChunk, + | "inputTokens" + | "nonCachedInputTokens" + | "outputTokens" + | "cacheWriteTokens" + | "cacheReadTokens" + | "reasoningTokens" + | "totalCost" +> { + const canonical = toCanonicalUsageMetrics({ + usage, + apiProtocol, + cacheWriteCandidates, + cacheReadCandidates, + inputNonCachedCandidates, + reasoningCandidates, + costCandidates, + deriveAnthropicNoCacheFromTotalWhenMissing, + }) + + return { + inputTokens: canonical.inputTokensTotal, + nonCachedInputTokens: canonical.inputTokensNonCached, + outputTokens: canonical.outputTokens, + cacheWriteTokens: + emitZeroCacheTokens || canonical.cacheWriteTokens > 0 ? canonical.cacheWriteTokens : undefined, + cacheReadTokens: emitZeroCacheTokens || canonical.cacheReadTokens > 0 ? canonical.cacheReadTokens : undefined, + reasoningTokens: canonical.reasoningTokens, + totalCost: canonical.totalCostCandidate, + } +} + +export function toCanonicalUsageMetrics({ + usage, + apiProtocol, + cacheWriteCandidates = [], + cacheReadCandidates = [], + inputNonCachedCandidates = [], + reasoningCandidates = [], + costCandidates = [], + deriveAnthropicNoCacheFromTotalWhenMissing = false, +}: NormalizeUsageMetricsOptions): CanonicalUsageMetrics { + const promptTokens = toFiniteNumber(usage.inputTokens) ?? 0 + const outputTokens = toFiniteNumber(usage.outputTokens) ?? 0 + const raw = usage.raw + const rawInputTokenDetails = asRecord(raw?.input_tokens_details) + const rawOutputTokenDetails = asRecord(raw?.output_tokens_details) + + const cacheWriteTokens = + firstNumber([ + ...cacheWriteCandidates, + usage.inputTokenDetails?.cacheWriteTokens, + usage.cacheCreationInputTokens, + usage.cache_creation_input_tokens, + raw?.cache_creation_input_tokens as NumberLike, + raw?.cacheCreationInputTokens as NumberLike, + ]) ?? 0 + + const cacheReadTokens = + firstNumber([ + ...cacheReadCandidates, + usage.inputTokenDetails?.cacheReadTokens, + usage.details?.cachedInputTokens, + usage.cachedInputTokens, + usage.cached_tokens, + raw?.cache_read_input_tokens as NumberLike, + raw?.cacheReadInputTokens as NumberLike, + raw?.cached_tokens as NumberLike, + rawInputTokenDetails?.cached_tokens as NumberLike, + ]) ?? 0 + + const reasoningTokens = firstNumber([ + ...reasoningCandidates, + (rawOutputTokenDetails?.reasoning_tokens as NumberLike) ?? undefined, + ]) + + const explicitNoCacheTokens = firstNumber([...inputNonCachedCandidates, usage.inputTokenDetails?.noCacheTokens]) + + let inputTokensNonCached = explicitNoCacheTokens + if (inputTokensNonCached === undefined) { + if (apiProtocol === "anthropic" && deriveAnthropicNoCacheFromTotalWhenMissing) { + inputTokensNonCached = Math.max(0, promptTokens - cacheWriteTokens - cacheReadTokens) + } else if (apiProtocol === "openai") { + inputTokensNonCached = Math.max(0, promptTokens - cacheWriteTokens - cacheReadTokens) + } + } + + return { + inputTokensTotal: promptTokens, + inputTokensNonCached, + outputTokens, + cacheWriteTokens, + cacheReadTokens, + reasoningTokens, + totalCostCandidate: firstNumber(costCandidates), + } +} + +export function toFiniteNumber(value: NumberLike): number | undefined { + if (typeof value === "number" && Number.isFinite(value)) { + return value + } + if (typeof value === "string") { + const parsed = Number(value) + if (Number.isFinite(parsed)) { + return parsed + } + } + return undefined +} + +export function getValueAtPath(source: unknown, path: string): unknown { + if (!path) return source + const normalized = path.replace(/\[(\d+)\]/g, ".$1") + const segments = normalized.split(".").filter(Boolean) + + let cursor: unknown = source + for (const segment of segments) { + if (cursor === null || cursor === undefined) { + return undefined + } + if (typeof cursor !== "object") { + return undefined + } + cursor = (cursor as Record)[segment] + } + + return cursor +} + +export function firstNumberFromPaths(source: unknown, paths: string[]): number | undefined { + for (const path of paths) { + const value = getValueAtPath(source, path) + const asNumber = toFiniteNumber(value as NumberLike) + if (asNumber !== undefined) { + return asNumber + } + } + return undefined +} + +function firstNumber(values: NumberLike[]): number | undefined { + for (const value of values) { + const asNumber = toFiniteNumber(value) + if (asNumber !== undefined) { + return asNumber + } + } + return undefined +} + +function asRecord(value: unknown): Record | undefined { + return typeof value === "object" && value !== null && !Array.isArray(value) + ? (value as Record) + : undefined +} diff --git a/src/api/providers/utils/usage-profiles.ts b/src/api/providers/utils/usage-profiles.ts new file mode 100644 index 0000000000..082a04ef95 --- /dev/null +++ b/src/api/providers/utils/usage-profiles.ts @@ -0,0 +1,547 @@ +export type UsageProtocol = "anthropic" | "openai" + +export interface UsageMetricPaths { + providerMetadata?: string[] + usage?: string[] + raw?: string[] +} + +export interface UsageProfile { + key: string + apiProtocol: UsageProtocol + emitZeroCacheTokens?: boolean + deriveNonCachedInputFromTotalMinusCache?: boolean + metrics: { + inputTokensTotal?: UsageMetricPaths + inputTokensNonCached?: UsageMetricPaths + outputTokens?: UsageMetricPaths + cacheWriteTokens?: UsageMetricPaths + cacheReadTokens?: UsageMetricPaths + reasoningTokens?: UsageMetricPaths + totalCostCandidate?: UsageMetricPaths + } +} + +const OPENAI_BASE_PROFILE: UsageProfile = { + key: "default-openai", + apiProtocol: "openai", + emitZeroCacheTokens: false, + deriveNonCachedInputFromTotalMinusCache: true, + metrics: { + inputTokensTotal: { + usage: ["inputTokens"], + }, + outputTokens: { + usage: ["outputTokens"], + }, + cacheWriteTokens: { + usage: ["inputTokenDetails.cacheWriteTokens", "cacheCreationInputTokens", "cache_creation_input_tokens"], + raw: ["cache_creation_input_tokens", "cacheCreationInputTokens"], + }, + cacheReadTokens: { + usage: [ + "inputTokenDetails.cacheReadTokens", + "cachedInputTokens", + "details.cachedInputTokens", + "cached_tokens", + ], + raw: [ + "input_tokens_details.cached_tokens", + "cache_read_input_tokens", + "cacheReadInputTokens", + "cached_tokens", + ], + }, + reasoningTokens: { + usage: ["outputTokenDetails.reasoningTokens", "reasoningTokens", "details.reasoningTokens"], + raw: ["output_tokens_details.reasoning_tokens"], + }, + totalCostCandidate: { + providerMetadata: ["gateway.cost"], + }, + }, +} + +const ANTHROPIC_BASE_PROFILE: UsageProfile = { + key: "default-anthropic", + apiProtocol: "anthropic", + emitZeroCacheTokens: false, + deriveNonCachedInputFromTotalMinusCache: true, + metrics: { + inputTokensTotal: { + usage: ["inputTokens"], + }, + inputTokensNonCached: { + usage: ["inputTokenDetails.noCacheTokens"], + providerMetadata: ["anthropic.usage.input_tokens"], + raw: ["input_tokens"], + }, + outputTokens: { + usage: ["outputTokens"], + providerMetadata: ["anthropic.usage.output_tokens"], + raw: ["output_tokens"], + }, + cacheWriteTokens: { + providerMetadata: ["anthropic.usage.cache_creation_input_tokens", "anthropic.cacheCreationInputTokens"], + usage: ["inputTokenDetails.cacheWriteTokens", "cacheCreationInputTokens", "cache_creation_input_tokens"], + raw: ["cache_creation_input_tokens", "cacheCreationInputTokens"], + }, + cacheReadTokens: { + providerMetadata: [ + "anthropic.usage.cache_read_input_tokens", + "anthropic.cacheReadInputTokens", + "anthropic.usage.cached_tokens", + ], + usage: [ + "inputTokenDetails.cacheReadTokens", + "cachedInputTokens", + "details.cachedInputTokens", + "cached_tokens", + ], + raw: [ + "cache_read_input_tokens", + "cacheReadInputTokens", + "input_tokens_details.cached_tokens", + "cached_tokens", + ], + }, + reasoningTokens: { + usage: ["outputTokenDetails.reasoningTokens", "reasoningTokens", "details.reasoningTokens"], + raw: ["output_tokens_details.reasoning_tokens"], + }, + totalCostCandidate: { + providerMetadata: ["gateway.cost"], + }, + }, +} + +const PROFILES: Record> = { + anthropic: { + key: "anthropic", + apiProtocol: "anthropic", + }, + "anthropic-vertex": { + key: "anthropic-vertex", + apiProtocol: "anthropic", + }, + bedrock: { + key: "bedrock", + apiProtocol: "anthropic", + metrics: { + cacheWriteTokens: { + providerMetadata: ["bedrock.usage.cacheWriteInputTokens"], + usage: [ + "inputTokenDetails.cacheWriteTokens", + "cacheCreationInputTokens", + "cache_creation_input_tokens", + ], + }, + cacheReadTokens: { + providerMetadata: ["bedrock.usage.cacheReadInputTokens"], + usage: [ + "inputTokenDetails.cacheReadTokens", + "cachedInputTokens", + "details.cachedInputTokens", + "cached_tokens", + ], + raw: ["input_tokens_details.cached_tokens", "cache_read_input_tokens", "cached_tokens"], + }, + }, + }, + baseten: { + key: "baseten", + apiProtocol: "openai", + }, + mistral: { + key: "mistral", + apiProtocol: "openai", + }, + "native-ollama": { + key: "native-ollama", + apiProtocol: "openai", + }, + lmstudio: { + key: "lmstudio", + apiProtocol: "openai", + }, + "lm-studio": { + key: "lm-studio", + apiProtocol: "openai", + }, + litellm: { + key: "litellm", + apiProtocol: "openai", + }, + "lite-llm": { + key: "lite-llm", + apiProtocol: "openai", + }, + moonshot: { + key: "moonshot", + apiProtocol: "openai", + }, + "qwen-code": { + key: "qwen-code", + apiProtocol: "openai", + }, + minimax: { + key: "minimax", + apiProtocol: "anthropic", + }, + openai: { + key: "openai", + apiProtocol: "openai", + metrics: { + cacheReadTokens: { + providerMetadata: ["openai.cachedPromptTokens"], + usage: [ + "inputTokenDetails.cacheReadTokens", + "cachedInputTokens", + "details.cachedInputTokens", + "cached_tokens", + ], + raw: ["input_tokens_details.cached_tokens", "cached_tokens"], + }, + reasoningTokens: { + providerMetadata: ["openai.reasoningTokens"], + usage: ["outputTokenDetails.reasoningTokens", "reasoningTokens", "details.reasoningTokens"], + raw: ["output_tokens_details.reasoning_tokens"], + }, + }, + }, + "openai-native": { + key: "openai-native", + apiProtocol: "openai", + metrics: { + cacheReadTokens: { + providerMetadata: ["openai.cachedPromptTokens"], + usage: ["inputTokenDetails.cacheReadTokens", "cachedInputTokens", "details.cachedInputTokens"], + raw: ["input_tokens_details.cached_tokens", "cached_tokens"], + }, + reasoningTokens: { + providerMetadata: ["openai.reasoningTokens"], + usage: ["outputTokenDetails.reasoningTokens", "reasoningTokens", "details.reasoningTokens"], + raw: ["output_tokens_details.reasoning_tokens"], + }, + }, + }, + "openai-codex": { + key: "openai-codex", + apiProtocol: "openai", + metrics: { + cacheReadTokens: { + providerMetadata: ["openai.cachedPromptTokens"], + usage: ["inputTokenDetails.cacheReadTokens", "cachedInputTokens", "details.cachedInputTokens"], + raw: ["input_tokens_details.cached_tokens", "cached_tokens"], + }, + reasoningTokens: { + providerMetadata: ["openai.reasoningTokens"], + usage: ["outputTokenDetails.reasoningTokens", "reasoningTokens", "details.reasoningTokens"], + raw: ["output_tokens_details.reasoning_tokens"], + }, + }, + }, + openrouter: { + key: "openrouter", + apiProtocol: "openai", + metrics: { + cacheWriteTokens: { + providerMetadata: [ + "openrouter.cacheCreationInputTokens", + "openrouter.cache_creation_input_tokens", + "openrouter.cacheWriteTokens", + ], + usage: [ + "inputTokenDetails.cacheWriteTokens", + "cacheCreationInputTokens", + "cache_creation_input_tokens", + ], + }, + cacheReadTokens: { + providerMetadata: [ + "openrouter.cachedInputTokens", + "openrouter.cache_read_input_tokens", + "openrouter.cacheReadTokens", + "openrouter.cached_tokens", + ], + usage: [ + "inputTokenDetails.cacheReadTokens", + "cachedInputTokens", + "details.cachedInputTokens", + "cached_tokens", + ], + raw: ["input_tokens_details.cached_tokens", "cached_tokens"], + }, + reasoningTokens: { + providerMetadata: [ + "openrouter.reasoningOutputTokens", + "openrouter.reasoning_tokens", + "openrouter.output_tokens_details.reasoning_tokens", + ], + usage: ["outputTokenDetails.reasoningTokens", "reasoningTokens", "details.reasoningTokens"], + raw: ["output_tokens_details.reasoning_tokens"], + }, + }, + }, + requesty: { + key: "requesty", + apiProtocol: "openai", + metrics: { + cacheWriteTokens: { + providerMetadata: ["requesty.usage.cachingTokens"], + usage: [ + "inputTokenDetails.cacheWriteTokens", + "cacheCreationInputTokens", + "cache_creation_input_tokens", + ], + }, + cacheReadTokens: { + providerMetadata: ["requesty.usage.cachedTokens"], + usage: ["inputTokenDetails.cacheReadTokens", "cachedInputTokens", "details.cachedInputTokens"], + raw: ["input_tokens_details.cached_tokens", "cached_tokens"], + }, + }, + }, + fireworks: { + key: "fireworks", + apiProtocol: "openai", + metrics: { + cacheWriteTokens: { + providerMetadata: ["fireworks.promptCacheMissTokens"], + usage: [ + "inputTokenDetails.cacheWriteTokens", + "cacheCreationInputTokens", + "cache_creation_input_tokens", + ], + }, + cacheReadTokens: { + providerMetadata: ["fireworks.promptCacheHitTokens"], + usage: ["inputTokenDetails.cacheReadTokens", "cachedInputTokens", "details.cachedInputTokens"], + raw: ["input_tokens_details.cached_tokens", "cached_tokens"], + }, + }, + }, + deepseek: { + key: "deepseek", + apiProtocol: "openai", + metrics: { + cacheWriteTokens: { + providerMetadata: ["deepseek.promptCacheMissTokens"], + usage: [ + "inputTokenDetails.cacheWriteTokens", + "cacheCreationInputTokens", + "cache_creation_input_tokens", + ], + }, + cacheReadTokens: { + providerMetadata: ["deepseek.promptCacheHitTokens"], + usage: ["inputTokenDetails.cacheReadTokens", "cachedInputTokens", "details.cachedInputTokens"], + raw: ["input_tokens_details.cached_tokens", "cached_tokens"], + }, + }, + }, + xai: { + key: "xai", + apiProtocol: "openai", + metrics: { + cacheReadTokens: { + providerMetadata: ["xai.cachedPromptTokens"], + usage: ["inputTokenDetails.cacheReadTokens", "cachedInputTokens", "details.cachedInputTokens"], + raw: ["input_tokens_details.cached_tokens", "cached_tokens"], + }, + }, + }, + azure: { + key: "azure", + apiProtocol: "openai", + metrics: { + cacheReadTokens: { + providerMetadata: ["azure.promptCacheHitTokens"], + usage: ["inputTokenDetails.cacheReadTokens", "cachedInputTokens", "details.cachedInputTokens"], + raw: ["input_tokens_details.cached_tokens", "cached_tokens"], + }, + cacheWriteTokens: { + usage: ["inputTokenDetails.cacheWriteTokens"], + }, + }, + }, + sambanova: { + key: "sambanova", + apiProtocol: "openai", + metrics: { + cacheWriteTokens: { + providerMetadata: ["sambanova.promptCacheMissTokens"], + usage: [ + "inputTokenDetails.cacheWriteTokens", + "cacheCreationInputTokens", + "cache_creation_input_tokens", + ], + }, + cacheReadTokens: { + providerMetadata: ["sambanova.promptCacheHitTokens"], + usage: ["inputTokenDetails.cacheReadTokens", "cachedInputTokens", "details.cachedInputTokens"], + raw: ["input_tokens_details.cached_tokens", "cached_tokens"], + }, + }, + }, + gemini: { + key: "gemini", + apiProtocol: "openai", + metrics: { + cacheReadTokens: { + providerMetadata: [ + "google.usageMetadata.cachedContentTokenCount", + "google.usageMetadata.cacheTokensDetails", + ], + usage: ["inputTokenDetails.cacheReadTokens", "cachedInputTokens", "details.cachedInputTokens"], + }, + }, + }, + vertex: { + key: "vertex", + apiProtocol: "openai", + metrics: { + cacheReadTokens: { + providerMetadata: [ + "vertex.usageMetadata.cachedContentTokenCount", + "vertex.usageMetadata.cacheTokensDetails", + "google.usageMetadata.cachedContentTokenCount", + "google.usageMetadata.cacheTokensDetails", + ], + usage: ["inputTokenDetails.cacheReadTokens", "cachedInputTokens", "details.cachedInputTokens"], + }, + }, + }, + roo: { + key: "roo", + emitZeroCacheTokens: true, + metrics: { + cacheWriteTokens: { + providerMetadata: [ + "roo.cache_creation_input_tokens", + "roo.cacheCreationInputTokens", + "openai.cacheCreationInputTokens", + "anthropic.usage.cache_creation_input_tokens", + "anthropic.cacheCreationInputTokens", + "gateway.cache_creation_input_tokens", + "gateway.cacheCreationInputTokens", + ], + usage: [ + "inputTokenDetails.cacheWriteTokens", + "cacheCreationInputTokens", + "cache_creation_input_tokens", + ], + }, + cacheReadTokens: { + providerMetadata: [ + "roo.cache_read_input_tokens", + "roo.cacheReadInputTokens", + "roo.cached_tokens", + "openai.cachedPromptTokens", + "anthropic.usage.cache_read_input_tokens", + "anthropic.usage.cached_tokens", + "anthropic.cacheReadInputTokens", + "gateway.cache_read_input_tokens", + "gateway.cacheReadInputTokens", + "gateway.cached_tokens", + ], + usage: [ + "inputTokenDetails.cacheReadTokens", + "cachedInputTokens", + "details.cachedInputTokens", + "cached_tokens", + ], + raw: ["input_tokens_details.cached_tokens", "cache_read_input_tokens", "cached_tokens"], + }, + totalCostCandidate: { + providerMetadata: ["roo.cost", "gateway.cost"], + }, + }, + }, + "vercel-ai-gateway": { + key: "vercel-ai-gateway", + metrics: { + cacheWriteTokens: { + providerMetadata: [ + "anthropic.usage.cache_creation_input_tokens", + "anthropic.cacheCreationInputTokens", + "gateway.cache_creation_input_tokens", + "gateway.cacheCreationInputTokens", + ], + usage: [ + "inputTokenDetails.cacheWriteTokens", + "cacheCreationInputTokens", + "cache_creation_input_tokens", + ], + }, + cacheReadTokens: { + providerMetadata: [ + "openai.cachedPromptTokens", + "anthropic.usage.cache_read_input_tokens", + "anthropic.usage.cached_tokens", + "anthropic.cacheReadInputTokens", + "gateway.cache_read_input_tokens", + "gateway.cacheReadInputTokens", + "gateway.cached_tokens", + ], + usage: [ + "inputTokenDetails.cacheReadTokens", + "cachedInputTokens", + "details.cachedInputTokens", + "cached_tokens", + ], + raw: ["input_tokens_details.cached_tokens", "cache_read_input_tokens", "cached_tokens"], + }, + totalCostCandidate: { + providerMetadata: ["gateway.cost"], + }, + }, + }, +} + +function mergeMetricPaths(base?: UsageMetricPaths, override?: UsageMetricPaths): UsageMetricPaths | undefined { + if (!base && !override) return undefined + return { + providerMetadata: override?.providerMetadata ?? base?.providerMetadata, + usage: override?.usage ?? base?.usage, + raw: override?.raw ?? base?.raw, + } +} + +function mergeProfiles(base: UsageProfile, override: Partial): UsageProfile { + return { + ...base, + ...override, + metrics: { + inputTokensTotal: mergeMetricPaths(base.metrics.inputTokensTotal, override.metrics?.inputTokensTotal), + inputTokensNonCached: mergeMetricPaths( + base.metrics.inputTokensNonCached, + override.metrics?.inputTokensNonCached, + ), + outputTokens: mergeMetricPaths(base.metrics.outputTokens, override.metrics?.outputTokens), + cacheWriteTokens: mergeMetricPaths(base.metrics.cacheWriteTokens, override.metrics?.cacheWriteTokens), + cacheReadTokens: mergeMetricPaths(base.metrics.cacheReadTokens, override.metrics?.cacheReadTokens), + reasoningTokens: mergeMetricPaths(base.metrics.reasoningTokens, override.metrics?.reasoningTokens), + totalCostCandidate: mergeMetricPaths(base.metrics.totalCostCandidate, override.metrics?.totalCostCandidate), + }, + } +} + +export function getUsageProfile(providerKey: string, apiProtocol?: UsageProtocol): UsageProfile { + const override = PROFILES[providerKey] + const protocol = override?.apiProtocol ?? apiProtocol ?? "openai" + const base = protocol === "anthropic" ? ANTHROPIC_BASE_PROFILE : OPENAI_BASE_PROFILE + + if (!override) { + return { + ...base, + key: providerKey || base.key, + } + } + + return mergeProfiles(base, { + ...override, + key: override.key ?? providerKey, + apiProtocol: protocol, + }) +} diff --git a/src/api/providers/vercel-ai-gateway.ts b/src/api/providers/vercel-ai-gateway.ts index ace066f962..ec79f7e39f 100644 --- a/src/api/providers/vercel-ai-gateway.ts +++ b/src/api/providers/vercel-ai-gateway.ts @@ -5,6 +5,7 @@ import { vercelAiGatewayDefaultModelId, vercelAiGatewayDefaultModelInfo, VERCEL_AI_GATEWAY_DEFAULT_TEMPERATURE, + getApiProtocol, type ModelInfo, type ModelRecord, } from "@roo-code/types" @@ -19,11 +20,13 @@ import { handleAiSdkError, yieldResponseMessage, } from "../transform/ai-sdk" +import { applyPromptCacheToMessages, mergeProviderOptions } from "../transform/prompt-cache" import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" import { DEFAULT_HEADERS } from "./constants" import { BaseProvider } from "./base-provider" import { getModels, getModelsFromCache } from "./fetchers/modelCache" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" import type { RooMessage } from "../../core/task-persistence/rooMessage" @@ -85,27 +88,32 @@ export class VercelAiGatewayHandler extends BaseProvider implements SingleComple usage: { inputTokens?: number outputTokens?: number + inputTokenDetails?: { + noCacheTokens?: number + cacheReadTokens?: number + cacheWriteTokens?: number + } details?: { cachedInputTokens?: number reasoningTokens?: number } + cacheCreationInputTokens?: number + cache_creation_input_tokens?: number + cachedInputTokens?: number + cached_tokens?: number + raw?: Record }, providerMetadata?: Record>, + apiProtocol?: "anthropic" | "openai", ): ApiStreamUsageChunk { - const gatewayMeta = providerMetadata?.gateway as Record | undefined + const { chunk } = normalizeProviderUsage({ + provider: "vercel-ai-gateway", + apiProtocol, + usage: usage as any, + providerMetadata, + }) - const cacheWriteTokens = (gatewayMeta?.cache_creation_input_tokens as number) ?? undefined - const cacheReadTokens = usage.details?.cachedInputTokens ?? (gatewayMeta?.cached_tokens as number) ?? undefined - const totalCost = (gatewayMeta?.cost as number) ?? 0 - - return { - type: "usage", - inputTokens: usage.inputTokens || 0, - outputTokens: usage.outputTokens || 0, - cacheWriteTokens, - cacheReadTokens, - totalCost, - } + return chunk } override async *createMessage( @@ -118,21 +126,39 @@ export class VercelAiGatewayHandler extends BaseProvider implements SingleComple const aiSdkMessages = messages as ModelMessage[] + const promptCache = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: "vercel-ai-gateway", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: info.supportsPromptCache, + promptCacheRetention: info.promptCacheRetention, + }, + settings: this.options, + }) + const openAiTools = this.convertToolsForOpenAI(metadata?.tools) - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + const aiSdkTools = convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined const temperature = this.supportsTemperature(modelId) ? (this.options.modelTemperature ?? VERCEL_AI_GATEWAY_DEFAULT_TEMPERATURE) : undefined + const providerOptions = mergeProviderOptions(undefined, promptCache.providerOptionsPatch) + const result = streamText({ model: languageModel, - system: systemPrompt, + system: promptCache.systemProviderOptions + ? ({ role: "system", content: systemPrompt, providerOptions: promptCache.systemProviderOptions } as any) + : systemPrompt, messages: aiSdkMessages, temperature, maxOutputTokens: info.maxTokens ?? undefined, tools: aiSdkTools, toolChoice: mapToolChoice(metadata?.tool_choice), + ...(providerOptions ? ({ providerOptions } as Record) : {}), }) try { @@ -151,7 +177,8 @@ export class VercelAiGatewayHandler extends BaseProvider implements SingleComple const usage = await result.usage const providerMetadata = await result.providerMetadata if (usage) { - yield this.processUsageMetrics(usage, providerMetadata as any) + const apiProtocol = getApiProtocol("vercel-ai-gateway", modelId) + yield this.processUsageMetrics(usage, providerMetadata as any, apiProtocol) } } catch (usageError) { if (lastStreamError) { diff --git a/src/api/providers/vertex.ts b/src/api/providers/vertex.ts index 4fb011bd68..52e0e7efcb 100644 --- a/src/api/providers/vertex.ts +++ b/src/api/providers/vertex.ts @@ -21,9 +21,11 @@ import { handleAiSdkError, yieldResponseMessage, } from "../transform/ai-sdk" +import { applyPromptCacheToMessages, mergeProviderOptions } from "../transform/prompt-cache" import { t } from "i18next" import type { ApiStream, ApiStreamUsageChunk, GroundingSource } from "../transform/stream" import { getModelParams } from "../transform/model-params" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" import { BaseProvider } from "./base-provider" @@ -107,6 +109,17 @@ export class VertexHandler extends BaseProvider implements SingleCompletionHandl // Convert messages to AI SDK format const aiSdkMessages = filteredMessages as ModelMessage[] + const promptCache = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: "vertex", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: info.supportsPromptCache, + promptCacheRetention: info.promptCacheRetention, + }, + settings: this.options, + }) + // Convert tools to OpenAI format first, then to AI SDK format let openAiTools = this.convertToolsForOpenAI(metadata?.tools) @@ -116,7 +129,9 @@ export class VertexHandler extends BaseProvider implements SingleCompletionHandl openAiTools = openAiTools.filter((tool) => tool.type === "function" && allowedSet.has(tool.function.name)) } - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + const aiSdkTools = convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined // Build tool choice - use 'required' when allowedFunctionNames restricts available tools const toolChoice = @@ -125,19 +140,26 @@ export class VertexHandler extends BaseProvider implements SingleCompletionHandl : mapToolChoice(metadata?.tool_choice) // Build the request options + const providerOptions = mergeProviderOptions( + thinkingConfig ? ({ vertex: { thinkingConfig } } as Record) : undefined, + promptCache.providerOptionsPatch, + ) + const requestOptions: Parameters[0] = { model: this.provider(modelId), - system: systemInstruction, + system: promptCache.systemProviderOptions + ? ({ + role: "system", + content: systemInstruction, + providerOptions: promptCache.systemProviderOptions, + } as any) + : systemInstruction, messages: aiSdkMessages, temperature: temperatureConfig, maxOutputTokens, tools: aiSdkTools, toolChoice, - // Add thinking/reasoning configuration if present - // Cast to any to bypass strict JSONObject typing - the AI SDK accepts the correct runtime values - ...(thinkingConfig && { - providerOptions: { vertex: { thinkingConfig } } as any, - }), + ...(providerOptions ? ({ providerOptions } as Record) : {}), } try { @@ -227,6 +249,14 @@ export class VertexHandler extends BaseProvider implements SingleCompletionHandl usage: { inputTokens?: number outputTokens?: number + cachedInputTokens?: number + reasoningTokens?: number + inputTokenDetails?: { + cacheReadTokens?: number + } + outputTokenDetails?: { + reasoningTokens?: number + } details?: { cachedInputTokens?: number reasoningTokens?: number @@ -235,23 +265,22 @@ export class VertexHandler extends BaseProvider implements SingleCompletionHandl info: ModelInfo, providerMetadata?: Record, ): ApiStreamUsageChunk { - const inputTokens = usage.inputTokens || 0 - const outputTokens = usage.outputTokens || 0 - const cacheReadTokens = usage.details?.cachedInputTokens - const reasoningTokens = usage.details?.reasoningTokens + const normalized = normalizeProviderUsage({ + provider: "vertex", + apiProtocol: "openai", + usage: usage as any, + providerMetadata, + modelInfo: info, + }) return { - type: "usage", - inputTokens, - outputTokens, - cacheReadTokens, - reasoningTokens, + ...normalized.chunk, totalCost: this.calculateCost({ info, - inputTokens, - outputTokens, - cacheReadTokens, - reasoningTokens, + inputTokens: normalized.canonical.inputTokensTotal, + outputTokens: normalized.canonical.outputTokens, + cacheReadTokens: normalized.canonical.cacheReadTokens, + reasoningTokens: normalized.canonical.reasoningTokens ?? 0, }), } } diff --git a/src/api/providers/xai.ts b/src/api/providers/xai.ts index 88b17a9044..6a33c3ddc0 100644 --- a/src/api/providers/xai.ts +++ b/src/api/providers/xai.ts @@ -1,6 +1,6 @@ import { Anthropic } from "@anthropic-ai/sdk" import { createXai } from "@ai-sdk/xai" -import { streamText, generateText, ToolSet } from "ai" +import { streamText, generateText, ToolSet, ModelMessage } from "ai" import { type XAIModelId, xaiDefaultModelId, xaiModels, type ModelInfo } from "@roo-code/types" @@ -13,8 +13,10 @@ import { mapToolChoice, handleAiSdkError, } from "../transform/ai-sdk" +import { applyPromptCacheToMessages, mergeProviderOptions } from "../transform/prompt-cache" import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" import { getModelParams } from "../transform/model-params" +import { normalizeProviderUsage } from "./utils/normalize-provider-usage" import { DEFAULT_HEADERS } from "./constants" import { BaseProvider } from "./base-provider" @@ -92,18 +94,14 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler } }, ): ApiStreamUsageChunk { - // Extract cache metrics from xAI's providerMetadata if available - // xAI supports prompt caching through prompt_tokens_details.cached_tokens - const cacheReadTokens = providerMetadata?.xai?.cachedPromptTokens ?? usage.details?.cachedInputTokens - - return { - type: "usage", - inputTokens: usage.inputTokens || 0, - outputTokens: usage.outputTokens || 0, - cacheReadTokens, - cacheWriteTokens: undefined, // xAI doesn't report cache write tokens separately - reasoningTokens: usage.details?.reasoningTokens, - } + const { chunk } = normalizeProviderUsage({ + provider: "xai", + apiProtocol: "openai", + usage: usage as any, + providerMetadata: providerMetadata as Record | undefined, + modelInfo: this.getModel().info, + }) + return chunk } /** @@ -122,26 +120,46 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler messages: RooMessage[], metadata?: ApiHandlerCreateMessageMetadata, ): ApiStream { - const { temperature, reasoning } = this.getModel() + const { temperature, reasoning, info } = this.getModel() const languageModel = this.getLanguageModel() // Convert messages to AI SDK format - const aiSdkMessages = messages + const aiSdkMessages = messages as ModelMessage[] + + const promptCache = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: "xai", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: info.supportsPromptCache, + promptCacheRetention: info.promptCacheRetention, + }, + settings: this.options, + }) // Convert tools to OpenAI format first, then to AI SDK format const openAiTools = this.convertToolsForOpenAI(metadata?.tools) - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + const aiSdkTools = convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined + + const providerOptions = mergeProviderOptions( + reasoning ? ({ xai: reasoning } as Record) : undefined, + promptCache.providerOptionsPatch, + ) // Build the request options const requestOptions: Parameters[0] = { model: languageModel, - system: systemPrompt, + system: promptCache.systemProviderOptions + ? ({ role: "system", content: systemPrompt, providerOptions: promptCache.systemProviderOptions } as any) + : systemPrompt, messages: aiSdkMessages, temperature: this.options.modelTemperature ?? temperature ?? XAI_DEFAULT_TEMPERATURE, maxOutputTokens: this.getMaxOutputTokens(), tools: aiSdkTools, toolChoice: mapToolChoice(metadata?.tool_choice), - ...(reasoning && { providerOptions: { xai: reasoning } }), + ...(providerOptions ? ({ providerOptions } as Record) : {}), } // Use streamText for streaming responses diff --git a/src/api/providers/zai.ts b/src/api/providers/zai.ts index 6c7abb6c16..a7d89312a1 100644 --- a/src/api/providers/zai.ts +++ b/src/api/providers/zai.ts @@ -21,6 +21,7 @@ import { mapToolChoice, handleAiSdkError, } from "../transform/ai-sdk" +import { applyPromptCacheToMessages, mergeProviderOptions } from "../transform/prompt-cache" import { ApiStream } from "../transform/stream" import { getModelParams } from "../transform/model-params" @@ -100,12 +101,27 @@ export class ZAiHandler extends BaseProvider implements SingleCompletionHandler const aiSdkMessages = messages as ModelMessage[] + const promptCache = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: "zai", + messages: aiSdkMessages, + modelInfo: { + supportsPromptCache: info.supportsPromptCache, + promptCacheRetention: info.promptCacheRetention, + }, + settings: this.options, + }) + const openAiTools = this.convertToolsForOpenAI(metadata?.tools) - const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined + const aiSdkTools = convertToolsForAiSdk(openAiTools, { + functionToolProviderOptions: promptCache.toolProviderOptions, + }) as ToolSet | undefined const requestOptions: Parameters[0] = { model: languageModel, - system: systemPrompt, + system: promptCache.systemProviderOptions + ? ({ role: "system", content: systemPrompt, providerOptions: promptCache.systemProviderOptions } as any) + : systemPrompt, messages: aiSdkMessages, temperature: this.options.modelTemperature ?? temperature ?? ZAI_DEFAULT_TEMPERATURE, maxOutputTokens: this.getMaxOutputTokens(), @@ -125,6 +141,11 @@ export class ZAiHandler extends BaseProvider implements SingleCompletionHandler } } + requestOptions.providerOptions = mergeProviderOptions( + requestOptions.providerOptions as Record | undefined, + promptCache.providerOptionsPatch, + ) as any + const result = streamText(requestOptions) try { diff --git a/src/api/transform/__tests__/ai-sdk.spec.ts b/src/api/transform/__tests__/ai-sdk.spec.ts index 6d30099b08..dac65990e7 100644 --- a/src/api/transform/__tests__/ai-sdk.spec.ts +++ b/src/api/transform/__tests__/ai-sdk.spec.ts @@ -748,10 +748,7 @@ describe("AI SDK conversion utilities", () => { expect(chunks[0]).toEqual({ type: "tool_call_end", id: "call_1" }) }) - it("ignores tool-call chunks to prevent duplicate tools in UI", () => { - // tool-call is intentionally ignored because tool-input-start/delta/end already - // provide complete tool call information. Emitting tool-call would cause duplicate - // tools in the UI for AI SDK providers (e.g., DeepSeek, Moonshot). + it("converts complete tool-call chunks into tool_call events", () => { const part = { type: "tool-call" as const, toolCallId: "call_1", @@ -760,7 +757,13 @@ describe("AI SDK conversion utilities", () => { } const chunks = [...processAiSdkStreamPart(part)] - expect(chunks).toHaveLength(0) + expect(chunks).toHaveLength(1) + expect(chunks[0]).toEqual({ + type: "tool_call", + id: "call_1", + name: "read_file", + arguments: '{"path":"test.ts"}', + }) }) it("processes source chunks with URL", () => { diff --git a/src/api/transform/__tests__/prompt-cache.spec.ts b/src/api/transform/__tests__/prompt-cache.spec.ts index d8724f454a..9185dbcf47 100644 --- a/src/api/transform/__tests__/prompt-cache.spec.ts +++ b/src/api/transform/__tests__/prompt-cache.spec.ts @@ -3,8 +3,13 @@ import type { ModelMessage } from "ai" import { applyPromptCacheToMessages, resolvePromptCachePolicy } from "../prompt-cache" describe("prompt-cache", () => { + const unifiedMarkers = { + anthropic: { cacheControl: { type: "ephemeral" } }, + bedrock: { cachePoint: { type: "default" } }, + } + describe("resolvePromptCachePolicy", () => { - it("defaults to enabled with aggressive strategy", () => { + it("defaults to enabled", () => { const policy = resolvePromptCachePolicy({ overrideKey: "bedrock", supportsPromptCache: true, @@ -12,7 +17,6 @@ describe("prompt-cache", () => { expect(policy).toEqual({ enabled: true, - strategy: "aggressive", }) }) @@ -93,7 +97,7 @@ describe("prompt-cache", () => { ] } - it("applies anthropic strategy and system marker", () => { + it("applies anthropic cache markers with system and last-two message checkpoints", () => { const messages = buildMessages() const result = applyPromptCacheToMessages({ adapter: "anthropic", @@ -102,24 +106,67 @@ describe("prompt-cache", () => { modelInfo: { supportsPromptCache: true, }, - settings: { - promptCachingStrategy: "aggressive", - }, }) expect(result.systemProviderOptions).toEqual({ + ...unifiedMarkers, + }) + expect(result.toolProviderOptions).toEqual({ anthropic: { cacheControl: { type: "ephemeral" } }, }) expect((messages[0] as any).providerOptions).toBeUndefined() expect((messages[2] as any).providerOptions).toEqual({ - anthropic: { cacheControl: { type: "ephemeral" } }, + ...unifiedMarkers, }) expect((messages[4] as any).providerOptions).toEqual({ - anthropic: { cacheControl: { type: "ephemeral" } }, + ...unifiedMarkers, }) }) - it("applies bedrock aggressive checkpoints to last three user messages", () => { + it("applies the same anthropic-family markers for anthropic-vertex and minimax adapters", () => { + for (const adapter of ["anthropic-vertex", "minimax"] as const) { + const messages = buildMessages() + const result = applyPromptCacheToMessages({ + adapter, + overrideKey: adapter, + messages, + modelInfo: { + supportsPromptCache: true, + }, + }) + + expect(result.systemProviderOptions).toEqual({ + ...unifiedMarkers, + }) + expect(result.toolProviderOptions).toEqual({ + anthropic: { cacheControl: { type: "ephemeral" } }, + }) + expect((messages[2] as any).providerOptions).toEqual({ + ...unifiedMarkers, + }) + expect((messages[4] as any).providerOptions).toEqual({ + ...unifiedMarkers, + }) + } + }) + + it("applies unified markers for generic ai-sdk adapter", () => { + const messages = buildMessages() + const result = applyPromptCacheToMessages({ + adapter: "ai-sdk", + overrideKey: "openrouter", + messages, + modelInfo: { + supportsPromptCache: true, + }, + }) + + expect(result.systemProviderOptions).toEqual(unifiedMarkers) + expect((messages[2] as any).providerOptions).toEqual(unifiedMarkers) + expect((messages[4] as any).providerOptions).toEqual(unifiedMarkers) + }) + + it("applies bedrock checkpoints to the last two non-assistant messages", () => { const messages = buildMessages() const result = applyPromptCacheToMessages({ adapter: "bedrock", @@ -128,17 +175,12 @@ describe("prompt-cache", () => { modelInfo: { supportsPromptCache: true, }, - settings: { - promptCachingStrategy: "aggressive", - }, }) expect(result.systemProviderOptions).toEqual({ bedrock: { cachePoint: { type: "default" } }, }) - expect((messages[0] as any).providerOptions).toEqual({ - bedrock: { cachePoint: { type: "default" } }, - }) + expect((messages[0] as any).providerOptions).toBeUndefined() expect((messages[2] as any).providerOptions).toEqual({ bedrock: { cachePoint: { type: "default" } }, }) @@ -147,8 +189,15 @@ describe("prompt-cache", () => { }) }) - it("applies balanced strategy with fewer checkpoints", () => { - const messages = buildMessages() + it("targets the last two non-assistant messages when tool messages are present", () => { + const messages: ModelMessage[] = [ + { role: "user", content: [{ type: "text", text: "u1" }] }, + { role: "assistant", content: [{ type: "text", text: "a1" }] }, + { role: "tool", content: [] as any }, + { role: "assistant", content: [{ type: "text", text: "a2" }] }, + { role: "user", content: [{ type: "text", text: "u2" }] }, + ] + applyPromptCacheToMessages({ adapter: "bedrock", overrideKey: "bedrock", @@ -156,9 +205,6 @@ describe("prompt-cache", () => { modelInfo: { supportsPromptCache: true, }, - settings: { - promptCachingStrategy: "balanced", - }, }) expect((messages[0] as any).providerOptions).toBeUndefined() diff --git a/src/api/transform/ai-sdk.ts b/src/api/transform/ai-sdk.ts index 3f879566a0..787bececbe 100644 --- a/src/api/transform/ai-sdk.ts +++ b/src/api/transform/ai-sdk.ts @@ -377,6 +377,9 @@ export function flattenAiSdkMessagesToStringContent( */ export function convertToolsForAiSdk( tools: OpenAI.Chat.ChatCompletionTool[] | undefined, + options?: { + functionToolProviderOptions?: Record + }, ): Record> | undefined { if (!tools || tools.length === 0) { return undefined @@ -389,6 +392,9 @@ export function convertToolsForAiSdk( toolSet[t.function.name] = createTool({ description: t.function.description, inputSchema: jsonSchema(t.function.parameters as any), + ...(options?.functionToolProviderOptions + ? { providerOptions: options.functionToolProviderOptions as any } + : {}), }) } } @@ -473,10 +479,28 @@ export function* processAiSdkStreamPart(part: ExtendedStreamPart): Generator -export type PromptCacheAdapter = "anthropic" | "bedrock" | "openai-native" +export type PromptCacheAdapter = "anthropic" | "anthropic-vertex" | "minimax" | "bedrock" | "openai-native" | "ai-sdk" export interface PromptCachePolicy { enabled: boolean - strategy: PromptCachingStrategy } export interface ApplyPromptCacheArgs { @@ -14,43 +12,34 @@ export interface ApplyPromptCacheArgs { overrideKey: string messages: ModelMessage[] modelInfo: Pick - settings?: Pick< - ProviderSettings, - "promptCachingEnabled" | "promptCachingStrategy" | "promptCachingProviderOverrides" - > + settings?: Pick } export interface AppliedPromptCache { enabled: boolean - strategy: PromptCachingStrategy systemProviderOptions?: Record providerOptionsPatch?: Record> + toolProviderOptions?: Record } -const DEFAULT_PROMPT_CACHING_STRATEGY: PromptCachingStrategy = "aggressive" - export function resolvePromptCachePolicy({ overrideKey, settings, supportsPromptCache, }: { overrideKey: string - settings?: Pick< - ProviderSettings, - "promptCachingEnabled" | "promptCachingStrategy" | "promptCachingProviderOverrides" - > + settings?: Pick supportsPromptCache: boolean }): PromptCachePolicy { - const strategy = settings?.promptCachingStrategy ?? DEFAULT_PROMPT_CACHING_STRATEGY if (!supportsPromptCache) { - return { enabled: false, strategy } + return { enabled: false } } const globalEnabled = settings?.promptCachingEnabled ?? true const providerOverride = settings?.promptCachingProviderOverrides?.[overrideKey] const enabled = providerOverride ?? globalEnabled - return { enabled, strategy } + return { enabled } } export function applyPromptCacheToMessages({ @@ -69,51 +58,44 @@ export function applyPromptCacheToMessages({ if (!policy.enabled) { return { enabled: false, - strategy: policy.strategy, } } if (adapter === "openai-native") { - if (modelInfo.promptCacheRetention === "24h") { - return { - enabled: true, - strategy: policy.strategy, - providerOptionsPatch: { - openai: { - promptCacheRetention: "24h", - }, - }, - } - } + const providerOptionsPatch = + modelInfo.promptCacheRetention === "24h" + ? ({ + openai: { + promptCacheRetention: "24h", + }, + } as const) + : undefined return { enabled: true, - strategy: policy.strategy, + providerOptionsPatch, } } const adapterConfig = getMessageAdapterConfig(adapter) - const checkpointCount = resolveCheckpointCount(policy.strategy, adapterConfig.maxUserCheckpoints) - const userIndices = getUserMessageIndices(messages) - const targetIndices = userIndices.slice(-checkpointCount) + const targetIndices = getNonAssistantMessageIndices(messages).slice(-2) applyProviderOptionAtIndices(messages, targetIndices, adapterConfig.messageProviderOption) return { enabled: true, - strategy: policy.strategy, systemProviderOptions: adapterConfig.systemProviderOptions, + toolProviderOptions: adapterConfig.toolProviderOptions, } } function getMessageAdapterConfig(adapter: Exclude): { - maxUserCheckpoints: number systemProviderOptions: Record messageProviderOption: Record> + toolProviderOptions?: Record } { if (adapter === "bedrock") { return { - maxUserCheckpoints: 3, systemProviderOptions: { bedrock: { cachePoint: { type: "default" } }, }, @@ -123,37 +105,57 @@ function getMessageAdapterConfig(adapter: Exclude | undefined, + patch: Record> | undefined, +): Record | undefined { + if (!patch) { + return base } - if (strategy === "conservative") { - return 1 + const next: Record = { ...(base ?? {}) } + + for (const [providerName, providerPatch] of Object.entries(patch)) { + const existingProviderOptions = next[providerName] + if ( + typeof existingProviderOptions === "object" && + existingProviderOptions !== null && + !Array.isArray(existingProviderOptions) + ) { + next[providerName] = { + ...(existingProviderOptions as Record), + ...providerPatch, + } + } else { + next[providerName] = providerPatch + } } - if (strategy === "balanced") { - return Math.max(1, Math.ceil(maxUserCheckpoints / 2)) - } - - return maxUserCheckpoints + return Object.keys(next).length > 0 ? next : undefined } -function getUserMessageIndices(messages: ModelMessage[]): number[] { +function getNonAssistantMessageIndices(messages: ModelMessage[]): number[] { const indices: number[] = [] for (let i = 0; i < messages.length; i++) { - if (messages[i].role === "user") { + if (messages[i].role !== "assistant") { indices.push(i) } } diff --git a/src/api/transform/stream.ts b/src/api/transform/stream.ts index bae4f5164e..794cd16e35 100644 --- a/src/api/transform/stream.ts +++ b/src/api/transform/stream.ts @@ -60,7 +60,10 @@ export interface ApiStreamThinkingCompleteChunk { export interface ApiStreamUsageChunk { type: "usage" + /** Total input tokens (cached + non-cached). */ inputTokens: number + /** Non-cached input tokens, when available/derivable from provider usage. */ + nonCachedInputTokens?: number outputTokens: number cacheWriteTokens?: number cacheReadTokens?: number diff --git a/src/core/config/ContextProxy.ts b/src/core/config/ContextProxy.ts index 588bcf250a..2825d1c945 100644 --- a/src/core/config/ContextProxy.ts +++ b/src/core/config/ContextProxy.ts @@ -22,7 +22,6 @@ import { TelemetryService } from "@roo-code/telemetry" import { logger } from "../../utils/logging" import { supportPrompt } from "../../shared/support-prompt" -import { migrateLegacyPromptCacheSettings } from "./migrateLegacyPromptCacheSettings" type GlobalStateKey = keyof GlobalState type SecretStateKey = keyof SecretState @@ -95,9 +94,6 @@ export class ContextProxy { // Migration: Sanitize invalid/removed API providers await this.migrateInvalidApiProvider() - // Migration: one-time read of removed provider-level prompt cache keys. - await this.migrateLegacyPromptCacheKeys() - // Migration: Move legacy customCondensingPrompt to customSupportPrompts await this.migrateLegacyCondensingPrompt() @@ -107,45 +103,6 @@ export class ContextProxy { this._isInitialized = true } - /** - * Migrates removed provider-level prompt cache toggles to the unified - * `promptCachingProviderOverrides` map and clears the legacy keys. - */ - private async migrateLegacyPromptCacheKeys() { - try { - const rawAwsToggle = this.originalContext.globalState.get("awsUsePromptCache" as any) - const rawLiteLlmToggle = this.originalContext.globalState.get("litellmUsePromptCache" as any) - - if (rawAwsToggle === undefined && rawLiteLlmToggle === undefined) { - return - } - - const migrationInput: Record = { - ...this.stateCache, - awsUsePromptCache: rawAwsToggle, - litellmUsePromptCache: rawLiteLlmToggle, - } - - const migration = migrateLegacyPromptCacheSettings(migrationInput) - if (!migration.changed) { - return - } - - const overrides = migration.config.promptCachingProviderOverrides as Record | undefined - await this.originalContext.globalState.update("promptCachingProviderOverrides", overrides) - this.stateCache.promptCachingProviderOverrides = overrides - - await this.originalContext.globalState.update("awsUsePromptCache" as any, undefined) - await this.originalContext.globalState.update("litellmUsePromptCache" as any, undefined) - delete (this.stateCache as Record).awsUsePromptCache - delete (this.stateCache as Record).litellmUsePromptCache - } catch (error) { - logger.error( - `Error during legacy prompt cache migration: ${error instanceof Error ? error.message : String(error)}`, - ) - } - } - /** * Migrates the legacy customCondensingPrompt to the new customSupportPrompts structure * and removes the legacy field. @@ -502,13 +459,6 @@ export class ContextProxy { } } - const promptCacheMigration = migrateLegacyPromptCacheSettings({ - ...(sanitizedValues as unknown as Record), - awsUsePromptCache: (this.stateCache as Record).awsUsePromptCache, - litellmUsePromptCache: (this.stateCache as Record).litellmUsePromptCache, - }) - sanitizedValues = promptCacheMigration.config as RooCodeSettings - const isKnownProvider = typeof values.apiProvider === "string" && (isProviderName(values.apiProvider) || isRetiredProvider(values.apiProvider)) diff --git a/src/core/config/ProviderSettingsManager.ts b/src/core/config/ProviderSettingsManager.ts index 7fdf55cebd..f5c753c46f 100644 --- a/src/core/config/ProviderSettingsManager.ts +++ b/src/core/config/ProviderSettingsManager.ts @@ -18,7 +18,6 @@ import { TelemetryService } from "@roo-code/telemetry" import { Mode, modes } from "../../shared/modes" import { buildApiHandler } from "../../api" -import { migrateLegacyPromptCacheSettings } from "./migrateLegacyPromptCacheSettings" // Type-safe model migrations mapping type ModelMigrations = { @@ -129,10 +128,6 @@ export class ProviderSettingsManager { isDirty = true } - if (this.applyLegacyPromptCacheMigration(providerProfiles)) { - isDirty = true - } - // Ensure all configs have IDs. for (const [_name, apiConfig] of Object.entries(providerProfiles.apiConfigs)) { if (!apiConfig.id) { @@ -318,22 +313,6 @@ export class ProviderSettingsManager { return migrated } - private applyLegacyPromptCacheMigration(providerProfiles: ProviderProfiles): boolean { - let migrated = false - - for (const [name, apiConfig] of Object.entries(providerProfiles.apiConfigs)) { - const migrationResult = migrateLegacyPromptCacheSettings(apiConfig as unknown as Record) - if (!migrationResult.changed) { - continue - } - - providerProfiles.apiConfigs[name] = migrationResult.config as ProviderSettingsWithId - migrated = true - } - - return migrated - } - /** * Clean model ID by removing prefix before "/" */ @@ -666,9 +645,7 @@ export class ProviderSettingsManager { return apiConfig } - const migrationResult = migrateLegacyPromptCacheSettings(apiConfig as Record) - const config = migrationResult.config as Record - + const config = apiConfig as Record const apiProvider = config.apiProvider // Check if apiProvider is set and if it's still recognized (active or retired) diff --git a/src/core/config/__tests__/ContextProxy.spec.ts b/src/core/config/__tests__/ContextProxy.spec.ts index f11d80c111..7c1d2a6e3c 100644 --- a/src/core/config/__tests__/ContextProxy.spec.ts +++ b/src/core/config/__tests__/ContextProxy.spec.ts @@ -70,20 +70,16 @@ describe("ContextProxy", () => { describe("constructor", () => { it("should initialize state cache with all global state keys", () => { - // +5 for the migration checks: + // +3 for the migration checks: // 1. openRouterImageGenerationSettings - // 2. awsUsePromptCache - // 3. litellmUsePromptCache - // 4. customCondensingPrompt - // 5. customSupportPrompts (for migrateOldDefaultCondensingPrompt) - expect(mockGlobalState.get).toHaveBeenCalledTimes(GLOBAL_STATE_KEYS.length + 5) + // 2. customCondensingPrompt + // 3. customSupportPrompts (for migrateOldDefaultCondensingPrompt) + expect(mockGlobalState.get).toHaveBeenCalledTimes(GLOBAL_STATE_KEYS.length + 3) for (const key of GLOBAL_STATE_KEYS) { expect(mockGlobalState.get).toHaveBeenCalledWith(key) } // Also check for migration calls expect(mockGlobalState.get).toHaveBeenCalledWith("openRouterImageGenerationSettings") - expect(mockGlobalState.get).toHaveBeenCalledWith("awsUsePromptCache") - expect(mockGlobalState.get).toHaveBeenCalledWith("litellmUsePromptCache") expect(mockGlobalState.get).toHaveBeenCalledWith("customCondensingPrompt") expect(mockGlobalState.get).toHaveBeenCalledWith("customSupportPrompts") }) @@ -108,8 +104,8 @@ describe("ContextProxy", () => { const result = proxy.getGlobalState("apiProvider") expect(result).toBe("deepseek") - // Original context should be called once during updateGlobalState (+5 for migration checks) - expect(mockGlobalState.get).toHaveBeenCalledTimes(GLOBAL_STATE_KEYS.length + 5) // From initialization + migration checks + // Original context should be called once during updateGlobalState (+3 for migration checks) + expect(mockGlobalState.get).toHaveBeenCalledTimes(GLOBAL_STATE_KEYS.length + 3) // From initialization + migration checks }) it("should handle default values correctly", async () => { @@ -557,23 +553,6 @@ describe("ContextProxy", () => { // Should not throw and should return undefined expect(settings.apiProvider).toBeUndefined() }) - - it("should migrate legacy prompt cache toggles to provider overrides", async () => { - await proxy.setValues({ - apiProvider: "bedrock", - awsUsePromptCache: false as any, - litellmUsePromptCache: false as any, - } as any) - - const settings = proxy.getProviderSettings() - - expect((settings as any).awsUsePromptCache).toBeUndefined() - expect((settings as any).litellmUsePromptCache).toBeUndefined() - expect(settings.promptCachingProviderOverrides).toEqual({ - bedrock: false, - litellm: false, - }) - }) }) describe("old default condensing prompt migration", () => { diff --git a/src/core/config/__tests__/ProviderSettingsManager.spec.ts b/src/core/config/__tests__/ProviderSettingsManager.spec.ts index 1fe891a69b..3f6b4f7847 100644 --- a/src/core/config/__tests__/ProviderSettingsManager.spec.ts +++ b/src/core/config/__tests__/ProviderSettingsManager.spec.ts @@ -336,45 +336,6 @@ describe("ProviderSettingsManager", () => { expect(storedConfig.apiConfigs.default.apiModelId).toEqual("roo/code-supernova-1-million") }) - it("should migrate legacy prompt cache toggles to provider overrides", async () => { - mockSecrets.get.mockResolvedValue( - JSON.stringify({ - currentApiConfigName: "default", - apiConfigs: { - default: { - id: "default", - apiProvider: "bedrock", - awsUsePromptCache: false, - }, - lite: { - id: "lite", - apiProvider: "litellm", - litellmUsePromptCache: false, - }, - }, - migrations: { - rateLimitSecondsMigrated: true, - openAiHeadersMigrated: true, - consecutiveMistakeLimitMigrated: true, - todoListEnabledMigrated: true, - claudeCodeLegacySettingsMigrated: true, - }, - }), - ) - - await providerSettingsManager.initialize() - - expect(mockSecrets.store).toHaveBeenCalled() - const calls = mockSecrets.store.mock.calls - const storedConfig = JSON.parse(calls[calls.length - 1][1]) - - expect(storedConfig.apiConfigs.default.awsUsePromptCache).toBeUndefined() - expect(storedConfig.apiConfigs.default.promptCachingProviderOverrides).toEqual({ bedrock: false }) - - expect(storedConfig.apiConfigs.lite.litellmUsePromptCache).toBeUndefined() - expect(storedConfig.apiConfigs.lite.promptCachingProviderOverrides).toEqual({ litellm: false }) - }) - it("should throw error if secrets storage fails", async () => { mockSecrets.get.mockRejectedValue(new Error("Storage failed")) diff --git a/src/core/config/__tests__/migrateLegacyPromptCacheSettings.spec.ts b/src/core/config/__tests__/migrateLegacyPromptCacheSettings.spec.ts deleted file mode 100644 index 58f0ba8277..0000000000 --- a/src/core/config/__tests__/migrateLegacyPromptCacheSettings.spec.ts +++ /dev/null @@ -1,67 +0,0 @@ -import { migrateLegacyPromptCacheSettings } from "../migrateLegacyPromptCacheSettings" - -describe("migrateLegacyPromptCacheSettings", () => { - it("maps legacy false toggles to provider overrides and removes legacy keys", () => { - const input = { - apiProvider: "bedrock", - awsUsePromptCache: false, - litellmUsePromptCache: false, - } - - const result = migrateLegacyPromptCacheSettings(input) - - expect(result.changed).toBe(true) - expect(result.config).toEqual({ - apiProvider: "bedrock", - promptCachingProviderOverrides: { - bedrock: false, - litellm: false, - }, - }) - }) - - it("does not create overrides for legacy true toggles and still removes legacy keys", () => { - const input = { - apiProvider: "bedrock", - awsUsePromptCache: true, - litellmUsePromptCache: true, - } - - const result = migrateLegacyPromptCacheSettings(input) - - expect(result.changed).toBe(true) - expect(result.config).toEqual({ - apiProvider: "bedrock", - }) - }) - - it("does not overwrite explicit new-format overrides", () => { - const input = { - awsUsePromptCache: false, - promptCachingProviderOverrides: { - bedrock: true, - }, - } - - const result = migrateLegacyPromptCacheSettings(input) - - expect(result.changed).toBe(true) - expect(result.config).toEqual({ - promptCachingProviderOverrides: { - bedrock: true, - }, - }) - }) - - it("returns unchanged when no legacy keys exist", () => { - const input = { - apiProvider: "anthropic", - promptCachingEnabled: true, - } - - const result = migrateLegacyPromptCacheSettings(input) - - expect(result.changed).toBe(false) - expect(result.config).toEqual(input) - }) -}) diff --git a/src/core/config/migrateLegacyPromptCacheSettings.ts b/src/core/config/migrateLegacyPromptCacheSettings.ts deleted file mode 100644 index b722ffce43..0000000000 --- a/src/core/config/migrateLegacyPromptCacheSettings.ts +++ /dev/null @@ -1,50 +0,0 @@ -type PromptCacheMigrationInput = Record - -export interface PromptCacheMigrationResult { - config: T - changed: boolean -} - -/** - * One-time migration helper for legacy provider-specific prompt cache toggles. - * - Maps legacy `false` values to provider overrides - * - Drops legacy keys from the config object - */ -export function migrateLegacyPromptCacheSettings( - config: T, -): PromptCacheMigrationResult { - let changed = false - const next = { ...config } as Record - - const currentOverrides = next.promptCachingProviderOverrides - const overrides = - typeof currentOverrides === "object" && currentOverrides !== null && !Array.isArray(currentOverrides) - ? { ...(currentOverrides as Record) } - : {} - - if (next.awsUsePromptCache === false && overrides.bedrock === undefined) { - overrides.bedrock = false - changed = true - } - - if (next.litellmUsePromptCache === false && overrides.litellm === undefined) { - overrides.litellm = false - changed = true - } - - if (Object.keys(overrides).length > 0) { - next.promptCachingProviderOverrides = overrides - } - - if ("awsUsePromptCache" in next) { - delete next.awsUsePromptCache - changed = true - } - - if ("litellmUsePromptCache" in next) { - delete next.litellmUsePromptCache - changed = true - } - - return { config: next as T, changed } -} diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index ee9f15e4e3..2b994a188f 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -2886,6 +2886,7 @@ export class Task extends EventEmitter implements TaskLike { let cacheWriteTokens = 0 let cacheReadTokens = 0 let inputTokens = 0 + let nonCachedInputTokens: number | undefined let outputTokens = 0 let totalCost: number | undefined @@ -2915,7 +2916,8 @@ export class Task extends EventEmitter implements TaskLike { apiProtocol === "anthropic" ? calculateApiCostAnthropic( streamModelInfo, - inputTokens, + nonCachedInputTokens ?? + Math.max(0, inputTokens - cacheWriteTokens - cacheReadTokens), outputTokens, cacheWriteTokens, cacheReadTokens, @@ -3060,6 +3062,9 @@ export class Task extends EventEmitter implements TaskLike { } case "usage": inputTokens += chunk.inputTokens + if (typeof chunk.nonCachedInputTokens === "number") { + nonCachedInputTokens = (nonCachedInputTokens ?? 0) + chunk.nonCachedInputTokens + } outputTokens += chunk.outputTokens cacheWriteTokens += chunk.cacheWriteTokens ?? 0 cacheReadTokens += chunk.cacheReadTokens ?? 0 @@ -3098,6 +3103,20 @@ export class Task extends EventEmitter implements TaskLike { case "tool_call": { // Legacy: Handle complete tool calls (for backward compatibility) + // Skip duplicate complete tool calls when start/delta/end flow already + // constructed a tool_use with the same ID. + const alreadyProcessed = this.assistantMessageContent.some( + (block) => + (block.type === "tool_use" || block.type === "mcp_tool_use") && + block.id === chunk.id, + ) + if (alreadyProcessed) { + console.debug( + `[Task#${this.taskId}] Ignoring duplicate tool_call for ID: ${chunk.id}`, + ) + break + } + // Convert native tool call to ToolUse format const toolUse = NativeToolCallParser.parseToolCall({ id: chunk.id, @@ -3184,6 +3203,7 @@ export class Task extends EventEmitter implements TaskLike { // Create a copy of current token values to avoid race conditions const currentTokens = { input: inputTokens, + nonCachedInput: nonCachedInputTokens, output: outputTokens, cacheWrite: cacheWriteTokens, cacheRead: cacheReadTokens, @@ -3197,6 +3217,7 @@ export class Task extends EventEmitter implements TaskLike { // Local variables to accumulate usage data without affecting the main flow let bgInputTokens = currentTokens.input + let bgNonCachedInputTokens = currentTokens.nonCachedInput let bgOutputTokens = currentTokens.output let bgCacheWriteTokens = currentTokens.cacheWrite let bgCacheReadTokens = currentTokens.cacheRead @@ -3206,6 +3227,7 @@ export class Task extends EventEmitter implements TaskLike { const captureUsageData = async ( tokens: { input: number + nonCachedInput?: number output: number cacheWrite: number cacheRead: number @@ -3215,12 +3237,14 @@ export class Task extends EventEmitter implements TaskLike { ) => { if ( tokens.input > 0 || + (tokens.nonCachedInput ?? 0) > 0 || tokens.output > 0 || tokens.cacheWrite > 0 || tokens.cacheRead > 0 ) { // Update the shared variables atomically inputTokens = tokens.input + nonCachedInputTokens = tokens.nonCachedInput outputTokens = tokens.output cacheWriteTokens = tokens.cacheWrite cacheReadTokens = tokens.cacheRead @@ -3249,7 +3273,8 @@ export class Task extends EventEmitter implements TaskLike { apiProtocol === "anthropic" ? calculateApiCostAnthropic( streamModelInfo, - tokens.input, + tokens.nonCachedInput ?? + Math.max(0, tokens.input - tokens.cacheWrite - tokens.cacheRead), tokens.output, tokens.cacheWrite, tokens.cacheRead, @@ -3298,6 +3323,10 @@ export class Task extends EventEmitter implements TaskLike { if (chunk && chunk.type === "usage") { usageFound = true bgInputTokens += chunk.inputTokens + if (typeof chunk.nonCachedInputTokens === "number") { + bgNonCachedInputTokens = + (bgNonCachedInputTokens ?? 0) + chunk.nonCachedInputTokens + } bgOutputTokens += chunk.outputTokens bgCacheWriteTokens += chunk.cacheWriteTokens ?? 0 bgCacheReadTokens += chunk.cacheReadTokens ?? 0 @@ -3316,6 +3345,7 @@ export class Task extends EventEmitter implements TaskLike { await captureUsageData( { input: bgInputTokens, + nonCachedInput: bgNonCachedInputTokens, output: bgOutputTokens, cacheWrite: bgCacheWriteTokens, cacheRead: bgCacheReadTokens, @@ -3340,6 +3370,7 @@ export class Task extends EventEmitter implements TaskLike { await captureUsageData( { input: bgInputTokens, + nonCachedInput: bgNonCachedInputTokens, output: bgOutputTokens, cacheWrite: bgCacheWriteTokens, cacheRead: bgCacheReadTokens, diff --git a/webview-ui/src/components/settings/ApiOptions.tsx b/webview-ui/src/components/settings/ApiOptions.tsx index 1725778021..5a059ca0d8 100644 --- a/webview-ui/src/components/settings/ApiOptions.tsx +++ b/webview-ui/src/components/settings/ApiOptions.tsx @@ -574,6 +574,7 @@ const ApiOptions = ({ )} @@ -806,28 +807,6 @@ const ApiOptions = ({
{t("settings:providers.enablePromptCachingTitle")}
- {(apiConfiguration.promptCachingEnabled ?? true) && ( -
- - -
- )} {selectedProvider === "openrouter" && openRouterModelProviders && Object.keys(openRouterModelProviders).length > 0 && ( diff --git a/webview-ui/src/components/settings/__tests__/ApiOptions.spec.tsx b/webview-ui/src/components/settings/__tests__/ApiOptions.spec.tsx index 3464fd424a..7895c44328 100644 --- a/webview-ui/src/components/settings/__tests__/ApiOptions.spec.tsx +++ b/webview-ui/src/components/settings/__tests__/ApiOptions.spec.tsx @@ -310,10 +310,10 @@ describe("ApiOptions", () => { }) expect(screen.getByText("settings:providers.enablePromptCaching")).toBeInTheDocument() - expect(screen.getByText("Prompt caching strategy")).toBeInTheDocument() + expect(screen.queryByText("Prompt caching strategy")).not.toBeInTheDocument() }) - it("updates prompt caching fields from advanced settings controls", () => { + it("updates prompt caching toggle from advanced settings controls", () => { const mockSetApiConfigurationField = vi.fn() renderApiOptions({ apiConfiguration: {}, @@ -324,21 +324,6 @@ describe("ApiOptions", () => { const enablePromptCachingInput = enablePromptCachingLabel?.querySelector("input") as HTMLInputElement fireEvent.click(enablePromptCachingInput) expect(mockSetApiConfigurationField).toHaveBeenCalledWith("promptCachingEnabled", false) - - const conservativeOption = screen.getByText("Conservative") - const strategySelect = conservativeOption.closest("select") as HTMLSelectElement - fireEvent.change(strategySelect, { target: { value: "conservative" } }) - expect(mockSetApiConfigurationField).toHaveBeenCalledWith("promptCachingStrategy", expect.any(String)) - }) - - it("hides prompt caching strategy selector when prompt caching is disabled", () => { - renderApiOptions({ - apiConfiguration: { - promptCachingEnabled: false, - }, - }) - - expect(screen.queryByText("Prompt caching strategy")).not.toBeInTheDocument() }) it("hides all controls when fromWelcomeView is true", () => { diff --git a/webview-ui/src/components/settings/providers/Bedrock.tsx b/webview-ui/src/components/settings/providers/Bedrock.tsx index ea0d50783f..b5b86dcd39 100644 --- a/webview-ui/src/components/settings/providers/Bedrock.tsx +++ b/webview-ui/src/components/settings/providers/Bedrock.tsx @@ -4,6 +4,7 @@ import { VSCodeTextField } from "@vscode/webview-ui-toolkit/react" import { type ProviderSettings, + type ModelInfo, type BedrockServiceTier, BEDROCK_REGIONS, BEDROCK_1M_CONTEXT_MODEL_IDS, @@ -12,17 +13,18 @@ import { } from "@roo-code/types" import { useAppTranslation } from "@src/i18n/TranslationContext" -import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@src/components/ui" +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue, StandardTooltip } from "@src/components/ui" import { inputEventTransform } from "../transforms" type BedrockProps = { apiConfiguration: ProviderSettings setApiConfigurationField: (field: keyof ProviderSettings, value: ProviderSettings[keyof ProviderSettings]) => void + selectedModelInfo?: ModelInfo simplifySettings?: boolean } -export const Bedrock = ({ apiConfiguration, setApiConfigurationField }: BedrockProps) => { +export const Bedrock = ({ apiConfiguration, setApiConfigurationField, selectedModelInfo }: BedrockProps) => { const { t } = useAppTranslation() const [awsEndpointSelected, setAwsEndpointSelected] = useState(!!apiConfiguration?.awsBedrockEndpointEnabled) @@ -55,6 +57,35 @@ export const Bedrock = ({ apiConfiguration, setApiConfigurationField }: BedrockP [setApiConfigurationField], ) + const handlePromptCacheOverrideChange = useCallback( + (enabled: boolean) => { + const globalEnabled = apiConfiguration.promptCachingEnabled ?? true + const nextOverrides = { ...(apiConfiguration.promptCachingProviderOverrides ?? {}) } + + if (enabled === globalEnabled) { + delete nextOverrides.bedrock + } else { + nextOverrides.bedrock = enabled + } + + setApiConfigurationField( + "promptCachingProviderOverrides", + (Object.keys(nextOverrides).length > 0 + ? nextOverrides + : undefined) as ProviderSettings["promptCachingProviderOverrides"], + ) + }, + [ + apiConfiguration.promptCachingEnabled, + apiConfiguration.promptCachingProviderOverrides, + setApiConfigurationField, + ], + ) + + const globalPromptCachingEnabled = apiConfiguration.promptCachingEnabled ?? true + const bedrockPromptCachingEnabled = + apiConfiguration.promptCachingProviderOverrides?.bedrock ?? globalPromptCachingEnabled + return ( <>
@@ -193,6 +224,24 @@ export const Bedrock = ({ apiConfiguration, setApiConfigurationField }: BedrockP }}> {t("settings:providers.awsCrossRegion")} + {selectedModelInfo?.supportsPromptCache && ( + <> + +
+ {t("settings:providers.enablePromptCaching")} + + + +
+ +
+ {t("settings:providers.cacheUsageNote")} +
+ + )} {supports1MContextBeta && (
{ + const globalEnabled = apiConfiguration.promptCachingEnabled ?? true + const nextOverrides = { ...(apiConfiguration.promptCachingProviderOverrides ?? {}) } + + if (enabled === globalEnabled) { + delete nextOverrides.litellm + } else { + nextOverrides.litellm = enabled + } + + setApiConfigurationField( + "promptCachingProviderOverrides", + (Object.keys(nextOverrides).length > 0 + ? nextOverrides + : undefined) as ProviderSettings["promptCachingProviderOverrides"], + ) + }, + [ + apiConfiguration.promptCachingEnabled, + apiConfiguration.promptCachingProviderOverrides, + setApiConfigurationField, + ], + ) + + const selectedModelId = apiConfiguration.litellmModelId || litellmDefaultModelId + const selectedModel = routerModels?.litellm?.[selectedModelId] + const globalPromptCachingEnabled = apiConfiguration.promptCachingEnabled ?? true + const litellmPromptCachingEnabled = + apiConfiguration.promptCachingProviderOverrides?.litellm ?? globalPromptCachingEnabled + return ( <> + {selectedModel?.supportsPromptCache && ( +
+ +
+ {t("settings:providers.enablePromptCaching")} +
+ +
+ {t("settings:providers.enablePromptCachingTitle")} +
+
+ )} ) } diff --git a/webview-ui/src/components/settings/providers/__tests__/Bedrock.spec.tsx b/webview-ui/src/components/settings/providers/__tests__/Bedrock.spec.tsx index 461a7ca61a..7f3c48bed2 100644 --- a/webview-ui/src/components/settings/providers/__tests__/Bedrock.spec.tsx +++ b/webview-ui/src/components/settings/providers/__tests__/Bedrock.spec.tsx @@ -260,7 +260,7 @@ describe("Bedrock Component", () => { // Test Scenario 3: UI Elements Tests describe("UI Elements", () => { - it("does not render legacy provider-level prompt caching controls", () => { + it("renders provider-level prompt caching controls for models that support caching", () => { const apiConfiguration: Partial = { apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0", awsUseProfile: true, @@ -270,12 +270,84 @@ describe("Bedrock Component", () => { , ) + expect(screen.getByText("settings:providers.enablePromptCaching")).toBeInTheDocument() + expect(screen.getByText("settings:providers.cacheUsageNote")).toBeInTheDocument() + }) + + it("does not render provider-level prompt caching controls for unsupported models", () => { + const apiConfiguration: Partial = { + apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0", + awsUseProfile: true, + } + + render( + , + ) + + expect(screen.queryByText("settings:providers.enablePromptCaching")).not.toBeInTheDocument() expect(screen.queryByText("settings:providers.cacheUsageNote")).not.toBeInTheDocument() }) + it("writes bedrock override when provider toggle diverges from global", () => { + const apiConfiguration: Partial = { + apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0", + awsUseProfile: true, + promptCachingEnabled: true, + } + + render( + , + ) + + const cacheLabel = screen.getByText("settings:providers.enablePromptCaching").closest("label") + const cacheInput = cacheLabel?.querySelector("input") as HTMLInputElement + fireEvent.click(cacheInput) + + expect(mockSetApiConfigurationField).toHaveBeenCalledWith("promptCachingProviderOverrides", { + bedrock: false, + }) + }) + + it("removes bedrock override when provider toggle matches global", () => { + const apiConfiguration: Partial = { + apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0", + awsUseProfile: true, + promptCachingEnabled: false, + promptCachingProviderOverrides: { + bedrock: true, + litellm: true, + }, + } + + render( + , + ) + + const cacheLabel = screen.getByText("settings:providers.enablePromptCaching").closest("label") + const cacheInput = cacheLabel?.querySelector("input") as HTMLInputElement + fireEvent.click(cacheInput) + + expect(mockSetApiConfigurationField).toHaveBeenCalledWith("promptCachingProviderOverrides", { + litellm: true, + }) + }) + it("should display example URLs when VPC endpoint checkbox is checked", () => { const apiConfiguration: Partial = { awsBedrockEndpoint: "https://example.com", diff --git a/webview-ui/src/components/settings/providers/__tests__/LiteLLM.spec.tsx b/webview-ui/src/components/settings/providers/__tests__/LiteLLM.spec.tsx new file mode 100644 index 0000000000..7aa6c89dd1 --- /dev/null +++ b/webview-ui/src/components/settings/providers/__tests__/LiteLLM.spec.tsx @@ -0,0 +1,135 @@ +import { render, screen, fireEvent } from "@/utils/test-utils" + +import { LiteLLM } from "../LiteLLM" +import type { ProviderSettings } from "@roo-code/types" + +const mockUseExtensionState = vi.fn() + +vi.mock("@src/context/ExtensionStateContext", () => ({ + useExtensionState: () => mockUseExtensionState(), +})) + +vi.mock("@src/i18n/TranslationContext", () => ({ + useAppTranslation: () => ({ + t: (key: string) => key, + }), +})) + +vi.mock("@src/utils/vscode", () => ({ + vscode: { + postMessage: vi.fn(), + }, +})) + +vi.mock("@vscode/webview-ui-toolkit/react", () => ({ + VSCodeTextField: ({ children, value, onInput, placeholder, className, type }: any) => ( +
+ {children} + onInput?.(e)} placeholder={placeholder} /> +
+ ), +})) + +vi.mock("vscrui", () => ({ + Checkbox: ({ children, checked, onChange }: any) => ( + + ), +})) + +vi.mock("@src/components/ui", () => ({ + Button: ({ children, onClick, disabled }: any) => ( + + ), +})) + +vi.mock("../../ModelPicker", () => ({ + ModelPicker: () =>
model picker
, +})) + +describe("LiteLLM prompt caching toggle", () => { + const mockSetApiConfigurationField = vi.fn() + + beforeEach(() => { + vi.clearAllMocks() + mockUseExtensionState.mockReturnValue({ + routerModels: { + litellm: { + "model-with-cache": { supportsPromptCache: true }, + }, + }, + }) + }) + + function renderLiteLLM(apiConfiguration: Partial) { + render( + , + ) + } + + it("renders provider-level prompt caching toggle when selected model supports prompt cache", () => { + renderLiteLLM({ + litellmModelId: "model-with-cache", + }) + + expect(screen.getByText("settings:providers.enablePromptCaching")).toBeInTheDocument() + }) + + it("does not render provider-level prompt caching toggle when selected model does not support prompt cache", () => { + mockUseExtensionState.mockReturnValue({ + routerModels: { + litellm: { + "model-without-cache": { supportsPromptCache: false }, + }, + }, + }) + + renderLiteLLM({ + litellmModelId: "model-without-cache", + }) + + expect(screen.queryByText("settings:providers.enablePromptCaching")).not.toBeInTheDocument() + }) + + it("writes litellm override when provider toggle diverges from global", () => { + renderLiteLLM({ + litellmModelId: "model-with-cache", + promptCachingEnabled: true, + }) + + const cacheLabel = screen.getByText("settings:providers.enablePromptCaching").closest("label") + const cacheInput = cacheLabel?.querySelector("input") as HTMLInputElement + fireEvent.click(cacheInput) + + expect(mockSetApiConfigurationField).toHaveBeenCalledWith("promptCachingProviderOverrides", { + litellm: false, + }) + }) + + it("removes litellm override when provider toggle matches global", () => { + renderLiteLLM({ + litellmModelId: "model-with-cache", + promptCachingEnabled: false, + promptCachingProviderOverrides: { + bedrock: true, + litellm: true, + }, + }) + + const cacheLabel = screen.getByText("settings:providers.enablePromptCaching").closest("label") + const cacheInput = cacheLabel?.querySelector("input") as HTMLInputElement + fireEvent.click(cacheInput) + + expect(mockSetApiConfigurationField).toHaveBeenCalledWith("promptCachingProviderOverrides", { + bedrock: true, + }) + }) +})