diff --git a/src/api/providers/__tests__/claude-code-caching.spec.ts b/src/api/providers/__tests__/claude-code-caching.spec.ts index 96b19964fc..a0996ab244 100644 --- a/src/api/providers/__tests__/claude-code-caching.spec.ts +++ b/src/api/providers/__tests__/claude-code-caching.spec.ts @@ -1,79 +1,57 @@ -import type { Anthropic } from "@anthropic-ai/sdk" - import { ClaudeCodeHandler } from "../claude-code" -import { runClaudeCode } from "../../../integrations/claude-code/run" import type { ApiHandlerOptions } from "../../../shared/api" -import type { ClaudeCodeMessage } from "../../../integrations/claude-code/types" +import type { StreamChunk } from "../../../integrations/claude-code/streaming-client" import type { ApiStreamUsageChunk } from "../../transform/stream" -// Mock the runClaudeCode function -vi.mock("../../../integrations/claude-code/run", () => ({ - runClaudeCode: vi.fn(), +// Mock the OAuth manager +vi.mock("../../../integrations/claude-code/oauth", () => ({ + claudeCodeOAuthManager: { + getAccessToken: vi.fn(), + getEmail: vi.fn(), + loadCredentials: vi.fn(), + saveCredentials: vi.fn(), + clearCredentials: vi.fn(), + isAuthenticated: vi.fn(), + }, + generateUserId: vi.fn(() => "user_abc123_account_def456_session_ghi789"), })) +// Mock the streaming client +vi.mock("../../../integrations/claude-code/streaming-client", () => ({ + createStreamingMessage: vi.fn(), +})) + +const { claudeCodeOAuthManager } = await import("../../../integrations/claude-code/oauth") +const { createStreamingMessage } = await import("../../../integrations/claude-code/streaming-client") + +const mockGetAccessToken = vi.mocked(claudeCodeOAuthManager.getAccessToken) +const mockCreateStreamingMessage = vi.mocked(createStreamingMessage) + describe("ClaudeCodeHandler - Caching Support", () => { let handler: ClaudeCodeHandler const mockOptions: ApiHandlerOptions = { - apiKey: "test-key", - apiModelId: "claude-3-5-sonnet-20241022", - claudeCodePath: "/test/path", + apiModelId: "claude-sonnet-4-5", } beforeEach(() => { handler = new ClaudeCodeHandler(mockOptions) vi.clearAllMocks() + mockGetAccessToken.mockResolvedValue("test-access-token") }) it("should collect cache read tokens from API response", async () => { - const mockStream = async function* (): AsyncGenerator { - // Initial system message + const mockStream = async function* (): AsyncGenerator { + yield { type: "text", text: "Hello!" } yield { - type: "system", - subtype: "init", - session_id: "test-session", - tools: [], - mcp_servers: [], - apiKeySource: "user", - } as ClaudeCodeMessage - - // Assistant message with cache tokens - const message: Anthropic.Messages.Message = { - id: "msg_123", - type: "message", - role: "assistant", - model: "claude-3-5-sonnet-20241022", - content: [{ type: "text", text: "Hello!", citations: [] }], - usage: { - input_tokens: 100, - output_tokens: 50, - cache_read_input_tokens: 80, // 80 tokens read from cache - cache_creation_input_tokens: 20, // 20 new tokens cached - }, - stop_reason: "end_turn", - stop_sequence: null, + type: "usage", + inputTokens: 100, + outputTokens: 50, + cacheReadTokens: 80, + cacheWriteTokens: 20, } - - yield { - type: "assistant", - message, - session_id: "test-session", - } as ClaudeCodeMessage - - // Result with cost - yield { - type: "result", - subtype: "success", - result: "success", - total_cost_usd: 0.001, - is_error: false, - duration_ms: 1000, - duration_api_ms: 900, - num_turns: 1, - session_id: "test-session", - } as ClaudeCodeMessage } - vi.mocked(runClaudeCode).mockReturnValue(mockStream()) + mockCreateStreamingMessage.mockReturnValue(mockStream()) const stream = handler.createMessage("System prompt", [{ role: "user", content: "Hello" }]) @@ -92,76 +70,29 @@ describe("ClaudeCodeHandler - Caching Support", () => { }) it("should accumulate cache tokens across multiple messages", async () => { - const mockStream = async function* (): AsyncGenerator { + // Note: The streaming client handles accumulation internally. + // Each usage chunk represents the accumulated totals for that point in the stream. + // This test verifies that we correctly pass through the accumulated values. + const mockStream = async function* (): AsyncGenerator { + yield { type: "text", text: "Part 1" } yield { - type: "system", - subtype: "init", - session_id: "test-session", - tools: [], - mcp_servers: [], - apiKeySource: "user", - } as ClaudeCodeMessage - - // First message chunk - const message1: Anthropic.Messages.Message = { - id: "msg_1", - type: "message", - role: "assistant", - model: "claude-3-5-sonnet-20241022", - content: [{ type: "text", text: "Part 1", citations: [] }], - usage: { - input_tokens: 50, - output_tokens: 25, - cache_read_input_tokens: 40, - cache_creation_input_tokens: 10, - }, - stop_reason: null, - stop_sequence: null, + type: "usage", + inputTokens: 50, + outputTokens: 25, + cacheReadTokens: 40, + cacheWriteTokens: 10, } - + yield { type: "text", text: "Part 2" } yield { - type: "assistant", - message: message1, - session_id: "test-session", - } as ClaudeCodeMessage - - // Second message chunk - const message2: Anthropic.Messages.Message = { - id: "msg_2", - type: "message", - role: "assistant", - model: "claude-3-5-sonnet-20241022", - content: [{ type: "text", text: "Part 2", citations: [] }], - usage: { - input_tokens: 50, - output_tokens: 25, - cache_read_input_tokens: 30, - cache_creation_input_tokens: 20, - }, - stop_reason: "end_turn", - stop_sequence: null, + type: "usage", + inputTokens: 100, // Accumulated: 50 + 50 + outputTokens: 50, // Accumulated: 25 + 25 + cacheReadTokens: 70, // Accumulated: 40 + 30 + cacheWriteTokens: 30, // Accumulated: 10 + 20 } - - yield { - type: "assistant", - message: message2, - session_id: "test-session", - } as ClaudeCodeMessage - - yield { - type: "result", - subtype: "success", - result: "success", - total_cost_usd: 0.002, - is_error: false, - duration_ms: 2000, - duration_api_ms: 1800, - num_turns: 1, - session_id: "test-session", - } as ClaudeCodeMessage } - vi.mocked(runClaudeCode).mockReturnValue(mockStream()) + mockCreateStreamingMessage.mockReturnValue(mockStream()) const stream = handler.createMessage("System prompt", [{ role: "user", content: "Hello" }]) @@ -170,62 +101,29 @@ describe("ClaudeCodeHandler - Caching Support", () => { chunks.push(chunk) } - const usageChunk = chunks.find((c) => c.type === "usage" && "totalCost" in c) as ApiStreamUsageChunk | undefined - expect(usageChunk).toBeDefined() - expect(usageChunk!.inputTokens).toBe(100) // 50 + 50 - expect(usageChunk!.outputTokens).toBe(50) // 25 + 25 - expect(usageChunk!.cacheReadTokens).toBe(70) // 40 + 30 - expect(usageChunk!.cacheWriteTokens).toBe(30) // 10 + 20 + // Get the last usage chunk which should have accumulated totals + const usageChunks = chunks.filter((c) => c.type === "usage" && "totalCost" in c) as ApiStreamUsageChunk[] + expect(usageChunks.length).toBe(2) + + const lastUsageChunk = usageChunks[usageChunks.length - 1] + expect(lastUsageChunk.inputTokens).toBe(100) // 50 + 50 + expect(lastUsageChunk.outputTokens).toBe(50) // 25 + 25 + expect(lastUsageChunk.cacheReadTokens).toBe(70) // 40 + 30 + expect(lastUsageChunk.cacheWriteTokens).toBe(30) // 10 + 20 }) it("should handle missing cache token fields gracefully", async () => { - const mockStream = async function* (): AsyncGenerator { + const mockStream = async function* (): AsyncGenerator { + yield { type: "text", text: "Hello!" } yield { - type: "system", - subtype: "init", - session_id: "test-session", - tools: [], - mcp_servers: [], - apiKeySource: "user", - } as ClaudeCodeMessage - - // Message without cache tokens - const message: Anthropic.Messages.Message = { - id: "msg_123", - type: "message", - role: "assistant", - model: "claude-3-5-sonnet-20241022", - content: [{ type: "text", text: "Hello!", citations: [] }], - usage: { - input_tokens: 100, - output_tokens: 50, - cache_read_input_tokens: null, - cache_creation_input_tokens: null, - }, - stop_reason: "end_turn", - stop_sequence: null, + type: "usage", + inputTokens: 100, + outputTokens: 50, + // No cache tokens provided } - - yield { - type: "assistant", - message, - session_id: "test-session", - } as ClaudeCodeMessage - - yield { - type: "result", - subtype: "success", - result: "success", - total_cost_usd: 0.001, - is_error: false, - duration_ms: 1000, - duration_api_ms: 900, - num_turns: 1, - session_id: "test-session", - } as ClaudeCodeMessage } - vi.mocked(runClaudeCode).mockReturnValue(mockStream()) + mockCreateStreamingMessage.mockReturnValue(mockStream()) const stream = handler.createMessage("System prompt", [{ role: "user", content: "Hello" }]) @@ -238,58 +136,24 @@ describe("ClaudeCodeHandler - Caching Support", () => { expect(usageChunk).toBeDefined() expect(usageChunk!.inputTokens).toBe(100) expect(usageChunk!.outputTokens).toBe(50) - expect(usageChunk!.cacheReadTokens).toBe(0) - expect(usageChunk!.cacheWriteTokens).toBe(0) + expect(usageChunk!.cacheReadTokens).toBeUndefined() + expect(usageChunk!.cacheWriteTokens).toBeUndefined() }) it("should report zero cost for subscription usage", async () => { - const mockStream = async function* (): AsyncGenerator { - // Subscription usage has apiKeySource: "none" + // Claude Code is always subscription-based, cost should always be 0 + const mockStream = async function* (): AsyncGenerator { + yield { type: "text", text: "Hello!" } yield { - type: "system", - subtype: "init", - session_id: "test-session", - tools: [], - mcp_servers: [], - apiKeySource: "none", - } as ClaudeCodeMessage - - const message: Anthropic.Messages.Message = { - id: "msg_123", - type: "message", - role: "assistant", - model: "claude-3-5-sonnet-20241022", - content: [{ type: "text", text: "Hello!", citations: [] }], - usage: { - input_tokens: 100, - output_tokens: 50, - cache_read_input_tokens: 80, - cache_creation_input_tokens: 20, - }, - stop_reason: "end_turn", - stop_sequence: null, + type: "usage", + inputTokens: 100, + outputTokens: 50, + cacheReadTokens: 80, + cacheWriteTokens: 20, } - - yield { - type: "assistant", - message, - session_id: "test-session", - } as ClaudeCodeMessage - - yield { - type: "result", - subtype: "success", - result: "success", - total_cost_usd: 0.001, // This should be ignored for subscription usage - is_error: false, - duration_ms: 1000, - duration_api_ms: 900, - num_turns: 1, - session_id: "test-session", - } as ClaudeCodeMessage } - vi.mocked(runClaudeCode).mockReturnValue(mockStream()) + mockCreateStreamingMessage.mockReturnValue(mockStream()) const stream = handler.createMessage("System prompt", [{ role: "user", content: "Hello" }]) @@ -300,6 +164,6 @@ describe("ClaudeCodeHandler - Caching Support", () => { const usageChunk = chunks.find((c) => c.type === "usage" && "totalCost" in c) as ApiStreamUsageChunk | undefined expect(usageChunk).toBeDefined() - expect(usageChunk!.totalCost).toBe(0) // Should be 0 for subscription usage + expect(usageChunk!.totalCost).toBe(0) // Should always be 0 for Claude Code (subscription-based) }) })