History

feat(translator): support cache write tokens in claude responses

- Extract cache write tokens from OpenAI and Codex usage details. - Map cache write tokens to Claude `cache_creation_input_tokens` for streaming and non-streaming responses.

Luis Pater committed 2 months ago

Commit 893abba1 parent 17a65ee4 files · +394 −14Browse files

Unified · Split

4 files
  1. internal/translator/codex/claude/codex_claude_response.go +15 −5
  2. internal/translator/codex/claude/codex_claude_response_test.go +194 −0
  3. internal/translator/openai/claude/openai_claude_response.go +22 −9
  4. internal/translator/openai/claude/openai_claude_response_test.go +163 −0
internal/translator/codex/claude/codex_claude_response.go+15−5
@@ -157,12 +157,15 @@ func ConvertCodexResponseToClaude(_ context.Context, _ string, originalRequestRa · 156 unmodified lines
157 output = append(output, stopCodexTextBlock(params)...)157 output = append(output, stopCodexTextBlock(params)...)
158 template, _ = sjson.SetBytes(template, "delta.stop_reason", mapCodexStopReasonToClaude(codexStopReason(responseData), params.HasEmittedToolUse))158 template, _ = sjson.SetBytes(template, "delta.stop_reason", mapCodexStopReasonToClaude(codexStopReason(responseData), params.HasEmittedToolUse))
159 template = setClaudeStopSequence(template, "delta.stop_sequence", responseData)159 template = setClaudeStopSequence(template, "delta.stop_sequence", responseData)
160 inputTokens, outputTokens, cachedTokens := extractResponsesUsage(responseData.Get("usage"))160 inputTokens, outputTokens, cachedTokens, cacheWriteTokens := extractResponsesUsage(responseData.Get("usage"))
161 template, _ = sjson.SetBytes(template, "usage.input_tokens", inputTokens)161 template, _ = sjson.SetBytes(template, "usage.input_tokens", inputTokens)
162 template, _ = sjson.SetBytes(template, "usage.output_tokens", outputTokens)162 template, _ = sjson.SetBytes(template, "usage.output_tokens", outputTokens)
163 if cachedTokens > 0 {163 if cachedTokens > 0 {
164 template, _ = sjson.SetBytes(template, "usage.cache_read_input_tokens", cachedTokens)164 template, _ = sjson.SetBytes(template, "usage.cache_read_input_tokens", cachedTokens)
165 }165 }
166 if cacheWriteTokens > 0 {
167 template, _ = sjson.SetBytes(template, "usage.cache_creation_input_tokens", cacheWriteTokens)
168 }
166 169
167 output = translatorcommon.AppendSSEEventBytes(output, "message_delta", template, 2)170 output = translatorcommon.AppendSSEEventBytes(output, "message_delta", template, 2)
168 output = translatorcommon.AppendSSEEventBytes(output, "message_stop", []byte(`{"type":"message_stop"}`), 2)171 output = translatorcommon.AppendSSEEventBytes(output, "message_stop", []byte(`{"type":"message_stop"}`), 2)
@@ -361,12 +364,15 @@ func ConvertCodexResponseToClaudeNonStream(_ context.Context, _ string, original · 192 unmodified lines
361 out := []byte(`{"id":"","type":"message","role":"assistant","model":"","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":0,"output_tokens":0}}`)364 out := []byte(`{"id":"","type":"message","role":"assistant","model":"","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":0,"output_tokens":0}}`)
362 out, _ = sjson.SetBytes(out, "id", responseData.Get("id").String())365 out, _ = sjson.SetBytes(out, "id", responseData.Get("id").String())
363 out, _ = sjson.SetBytes(out, "model", responseData.Get("model").String())366 out, _ = sjson.SetBytes(out, "model", responseData.Get("model").String())
364 inputTokens, outputTokens, cachedTokens := extractResponsesUsage(responseData.Get("usage"))367 inputTokens, outputTokens, cachedTokens, cacheWriteTokens := extractResponsesUsage(responseData.Get("usage"))
365 out, _ = sjson.SetBytes(out, "usage.input_tokens", inputTokens)368 out, _ = sjson.SetBytes(out, "usage.input_tokens", inputTokens)
366 out, _ = sjson.SetBytes(out, "usage.output_tokens", outputTokens)369 out, _ = sjson.SetBytes(out, "usage.output_tokens", outputTokens)
367 if cachedTokens > 0 {370 if cachedTokens > 0 {
368 out, _ = sjson.SetBytes(out, "usage.cache_read_input_tokens", cachedTokens)371 out, _ = sjson.SetBytes(out, "usage.cache_read_input_tokens", cachedTokens)
369 }372 }
373 if cacheWriteTokens > 0 {
374 out, _ = sjson.SetBytes(out, "usage.cache_creation_input_tokens", cacheWriteTokens)
375 }
370 376
371 hasToolCall := false377 hasToolCall := false
372 webSearchSeen := make(map[string]struct{})378 webSearchSeen := make(map[string]struct{})
@@ -794,14 +800,18 @@ func resolveCodexClaudeToolUseName(originalRequestRawJSON []byte, name string) s · 421 unmodified lines
794 return name800 return name
795}801}
796 802
797func extractResponsesUsage(usage gjson.Result) (int64, int64, int64) {803func extractResponsesUsage(usage gjson.Result) (int64, int64, int64, int64) {
798 if !usage.Exists() || usage.Type == gjson.Null {804 if !usage.Exists() || usage.Type == gjson.Null {
799 return 0, 0, 0805 return 0, 0, 0, 0
800 }806 }
801 807
802 inputTokens := usage.Get("input_tokens").Int()808 inputTokens := usage.Get("input_tokens").Int()
803 outputTokens := usage.Get("output_tokens").Int()809 outputTokens := usage.Get("output_tokens").Int()
804 cachedTokens := usage.Get("input_tokens_details.cached_tokens").Int()810 cachedTokens := usage.Get("input_tokens_details.cached_tokens").Int()
811 cacheWriteTokens := usage.Get("input_tokens_details.cache_write_tokens").Int()
812 if cacheWriteTokens == 0 {
813 cacheWriteTokens = usage.Get("input_tokens_details.cache_creation_tokens").Int()
814 }
805 815
806 if cachedTokens > 0 {816 if cachedTokens > 0 {
807 if inputTokens >= cachedTokens {817 if inputTokens >= cachedTokens {
@@ -811,7 +821,7 @@ func extractResponsesUsage(usage gjson.Result) (int64, int64, int64) { · 3 unmodified lines
811 }821 }
812 }822 }
813 823
814 return inputTokens, outputTokens, cachedTokens824 return inputTokens, outputTokens, cachedTokens, cacheWriteTokens
815}825}
816 826
817// buildReverseMapFromClaudeOriginalShortToOriginal builds a map[short]original from original Claude request tools.827// buildReverseMapFromClaudeOriginalShortToOriginal builds a map[short]original from original Claude request tools.
internal/translator/codex/claude/codex_claude_response_test.go+194−0
@@ -3,6 +3,7 @@ package claude · 2 unmodified lines
3import (3import (
4 "bytes"4 "bytes"
5 "context"5 "context"
6 "fmt"
6 "strings"7 "strings"
7 "testing"8 "testing"
8 9
@@ -1328,3 +1329,196 @@ func firstClaudeStreamPayloadForEvent(output, event string) (gjson.Result, bool) · 1319 unmodified lines
1328 }1329 }
1329 return gjson.Result{}, false1330 return gjson.Result{}, false
1330}1331}
1332
1333func TestConvertCodexResponseToClaude_StreamPreservesCacheWriteUsage(t *testing.T) {
1334 tests := []struct {
1335 name string
1336 terminalUsageJSON string
1337 wantInputTokens int64
1338 wantOutputTokens int64
1339 wantCacheReadTokens int64
1340 wantCacheWriteTokens int64
1341 }{
1342 {
1343 name: "cache_write_tokens field",
1344 terminalUsageJSON: `{"input_tokens":1000,"output_tokens":200,"input_tokens_details":{"cached_tokens":800,"cache_write_tokens":150}}`,
1345 wantInputTokens: 200,
1346 wantOutputTokens: 200,
1347 wantCacheReadTokens: 800,
1348 wantCacheWriteTokens: 150,
1349 },
1350 {
1351 name: "cache_creation_tokens field alias",
1352 terminalUsageJSON: `{"input_tokens":1000,"output_tokens":200,"input_tokens_details":{"cached_tokens":800,"cache_creation_tokens":150}}`,
1353 wantInputTokens: 200,
1354 wantOutputTokens: 200,
1355 wantCacheReadTokens: 800,
1356 wantCacheWriteTokens: 150,
1357 },
1358 {
1359 name: "cached_tokens greater than input_tokens clamps input_tokens to zero",
1360 terminalUsageJSON: `{"input_tokens":500,"output_tokens":100,"input_tokens_details":{"cached_tokens":800,"cache_write_tokens":50}}`,
1361 wantInputTokens: 0,
1362 wantOutputTokens: 100,
1363 wantCacheReadTokens: 800,
1364 wantCacheWriteTokens: 50,
1365 },
1366 {
1367 name: "zero cache_write_tokens does not emit cache_creation_input_tokens",
1368 terminalUsageJSON: `{"input_tokens":1000,"output_tokens":200,"input_tokens_details":{"cached_tokens":800,"cache_write_tokens":0}}`,
1369 wantInputTokens: 200,
1370 wantOutputTokens: 200,
1371 wantCacheReadTokens: 800,
1372 wantCacheWriteTokens: 0,
1373 },
1374 }
1375
1376 for _, tt := range tests {
1377 t.Run(tt.name, func(t *testing.T) {
1378 ctx := context.Background()
1379 originalRequest := []byte(`{"messages":[]}`)
1380 var param any
1381
1382 chunks := [][]byte{
1383 []byte(`data: {"type":"response.created","response":{"id":"resp_1","model":"gpt-5"}}`),
1384 []byte(`data: {"type":"response.output_item.done","item":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"ok"}]}}`),
1385 []byte(fmt.Sprintf(`data: {"type":"response.completed","response":{"stop_reason":"stop","usage":%s}}`, tt.terminalUsageJSON)),
1386 }
1387
1388 var outputs [][]byte
1389 for _, chunk := range chunks {
1390 outputs = append(outputs, ConvertCodexResponseToClaude(ctx, "", originalRequest, nil, chunk, &param)...)
1391 }
1392
1393 delta, ok := findClaudeStreamMessageDelta(outputs)
1394 if !ok {
1395 t.Fatalf("missing message_delta event; outputs=%q", outputs)
1396 }
1397
1398 usage := delta.Get("usage")
1399 if got := usage.Get("input_tokens").Int(); got != tt.wantInputTokens {
1400 t.Fatalf("input_tokens = %d, want %d", got, tt.wantInputTokens)
1401 }
1402 if got := usage.Get("output_tokens").Int(); got != tt.wantOutputTokens {
1403 t.Fatalf("output_tokens = %d, want %d", got, tt.wantOutputTokens)
1404 }
1405 if got := usage.Get("cache_read_input_tokens").Int(); got != tt.wantCacheReadTokens {
1406 t.Fatalf("cache_read_input_tokens = %d, want %d", got, tt.wantCacheReadTokens)
1407 }
1408 if tt.wantCacheWriteTokens == 0 {
1409 if usage.Get("cache_creation_input_tokens").Exists() {
1410 t.Fatalf("cache_creation_input_tokens should not be emitted when zero; got %v", usage.Get("cache_creation_input_tokens").Raw)
1411 }
1412 } else if got := usage.Get("cache_creation_input_tokens").Int(); got != tt.wantCacheWriteTokens {
1413 t.Fatalf("cache_creation_input_tokens = %d, want %d", got, tt.wantCacheWriteTokens)
1414 }
1415 })
1416 }
1417}
1418
1419func TestConvertCodexResponseToClaudeNonStream_PreservesCacheWriteUsage(t *testing.T) {
1420 tests := []struct {
1421 name string
1422 responseJSON string
1423 wantInputTokens int64
1424 wantOutputTokens int64
1425 wantCacheReadTokens int64
1426 wantCacheWriteTokens int64
1427 }{
1428 {
1429 name: "cache_write_tokens field",
1430 responseJSON: `{
1431 "type":"response.completed",
1432 "response":{
1433 "id":"resp_1",
1434 "model":"gpt-5",
1435 "stop_reason":"stop",
1436 "usage":{"input_tokens":1000,"output_tokens":200,"input_tokens_details":{"cached_tokens":800,"cache_write_tokens":150}},
1437 "output":[{"type":"message","content":[{"type":"output_text","text":"ok"}]}]
1438 }
1439 }`,
1440 wantInputTokens: 200,
1441 wantOutputTokens: 200,
1442 wantCacheReadTokens: 800,
1443 wantCacheWriteTokens: 150,
1444 },
1445 {
1446 name: "cache_creation_tokens alias",
1447 responseJSON: `{
1448 "type":"response.completed",
1449 "response":{
1450 "id":"resp_1",
1451 "model":"gpt-5",
1452 "stop_reason":"stop",
1453 "usage":{"input_tokens":1000,"output_tokens":200,"input_tokens_details":{"cached_tokens":800,"cache_creation_tokens":150}},
1454 "output":[{"type":"message","content":[{"type":"output_text","text":"ok"}]}]
1455 }
1456 }`,
1457 wantInputTokens: 200,
1458 wantOutputTokens: 200,
1459 wantCacheReadTokens: 800,
1460 wantCacheWriteTokens: 150,
1461 },
1462 {
1463 name: "cached_tokens greater than input_tokens clamps input_tokens to zero",
1464 responseJSON: `{
1465 "type":"response.completed",
1466 "response":{
1467 "id":"resp_1",
1468 "model":"gpt-5",
1469 "stop_reason":"stop",
1470 "usage":{"input_tokens":500,"output_tokens":100,"input_tokens_details":{"cached_tokens":800,"cache_write_tokens":50}},
1471 "output":[{"type":"message","content":[{"type":"output_text","text":"ok"}]}]
1472 }
1473 }`,
1474 wantInputTokens: 0,
1475 wantOutputTokens: 100,
1476 wantCacheReadTokens: 800,
1477 wantCacheWriteTokens: 50,
1478 },
1479 {
1480 name: "zero cache_write_tokens does not emit cache_creation_input_tokens",
1481 responseJSON: `{
1482 "type":"response.completed",
1483 "response":{
1484 "id":"resp_1",
1485 "model":"gpt-5",
1486 "stop_reason":"stop",
1487 "usage":{"input_tokens":1000,"output_tokens":200,"input_tokens_details":{"cached_tokens":800,"cache_write_tokens":0}},
1488 "output":[{"type":"message","content":[{"type":"output_text","text":"ok"}]}]
1489 }
1490 }`,
1491 wantInputTokens: 200,
1492 wantOutputTokens: 200,
1493 wantCacheReadTokens: 800,
1494 wantCacheWriteTokens: 0,
1495 },
1496 }
1497
1498 for _, tt := range tests {
1499 t.Run(tt.name, func(t *testing.T) {
1500 ctx := context.Background()
1501 originalRequest := []byte(`{"messages":[]}`)
1502 out := ConvertCodexResponseToClaudeNonStream(ctx, "", originalRequest, nil, []byte(tt.responseJSON), nil)
1503 parsed := gjson.ParseBytes(out)
1504
1505 usage := parsed.Get("usage")
1506 if got := usage.Get("input_tokens").Int(); got != tt.wantInputTokens {
1507 t.Fatalf("input_tokens = %d, want %d", got, tt.wantInputTokens)
1508 }
1509 if got := usage.Get("output_tokens").Int(); got != tt.wantOutputTokens {
1510 t.Fatalf("output_tokens = %d, want %d", got, tt.wantOutputTokens)
1511 }
1512 if got := usage.Get("cache_read_input_tokens").Int(); got != tt.wantCacheReadTokens {
1513 t.Fatalf("cache_read_input_tokens = %d, want %d", got, tt.wantCacheReadTokens)
1514 }
1515 if tt.wantCacheWriteTokens == 0 {
1516 if usage.Get("cache_creation_input_tokens").Exists() {
1517 t.Fatalf("cache_creation_input_tokens should not be emitted when zero; got %v", usage.Get("cache_creation_input_tokens").Raw)
1518 }
1519 } else if got := usage.Get("cache_creation_input_tokens").Int(); got != tt.wantCacheWriteTokens {
1520 t.Fatalf("cache_creation_input_tokens = %d, want %d", got, tt.wantCacheWriteTokens)
1521 }
1522 })
1523 }
1524}
internal/translator/openai/claude/openai_claude_response.go+22−9
@@ -310,8 +310,8 @@ func convertOpenAIStreamingChunkToAnthropic(rawJSON []byte, param *ConvertOpenAI · 309 unmodified lines
310 usage := root.Get("usage")310 usage := root.Get("usage")
311 if usage.Exists() && usage.Type != gjson.Null {311 if usage.Exists() && usage.Type != gjson.Null {
312 finalizeOpenAIAnthropicContentBlocks(param, &results)312 finalizeOpenAIAnthropicContentBlocks(param, &results)
313 inputTokens, outputTokens, cachedTokens := extractOpenAIUsage(usage)313 inputTokens, outputTokens, cachedTokens, cacheWriteTokens := extractOpenAIUsage(usage)
314 emitAnthropicMessageDelta(param, &results, inputTokens, outputTokens, cachedTokens)314 emitAnthropicMessageDelta(param, &results, inputTokens, outputTokens, cachedTokens, cacheWriteTokens)
315 emitMessageStopIfNeeded(param, &results)315 emitMessageStopIfNeeded(param, &results)
316 }316 }
317 }317 }
@@ -326,7 +326,7 @@ func convertOpenAIDoneToAnthropic(param *ConvertOpenAIResponseToAnthropicParams) · 8 unmodified lines
326 finalizeOpenAIAnthropicContentBlocks(param, &results)326 finalizeOpenAIAnthropicContentBlocks(param, &results)
327 327
328 if !param.MessageDeltaSent {328 if !param.MessageDeltaSent {
329 emitAnthropicMessageDelta(param, &results, 0, 0, 0)329 emitAnthropicMessageDelta(param, &results, 0, 0, 0, 0)
330 }330 }
331 331
332 emitMessageStopIfNeeded(param, &results)332 emitMessageStopIfNeeded(param, &results)
@@ -400,12 +400,15 @@ func convertOpenAINonStreamingToAnthropic(rawJSON []byte) [][]byte { · 67 unmodified lines
400 400
401 // Set usage information401 // Set usage information
402 if usage := root.Get("usage"); usage.Exists() {402 if usage := root.Get("usage"); usage.Exists() {
403 inputTokens, outputTokens, cachedTokens := extractOpenAIUsage(usage)403 inputTokens, outputTokens, cachedTokens, cacheWriteTokens := extractOpenAIUsage(usage)
404 out, _ = sjson.SetBytes(out, "usage.input_tokens", inputTokens)404 out, _ = sjson.SetBytes(out, "usage.input_tokens", inputTokens)
405 out, _ = sjson.SetBytes(out, "usage.output_tokens", outputTokens)405 out, _ = sjson.SetBytes(out, "usage.output_tokens", outputTokens)
406 if cachedTokens > 0 {406 if cachedTokens > 0 {
407 out, _ = sjson.SetBytes(out, "usage.cache_read_input_tokens", cachedTokens)407 out, _ = sjson.SetBytes(out, "usage.cache_read_input_tokens", cachedTokens)
408 }408 }
409 if cacheWriteTokens > 0 {
410 out, _ = sjson.SetBytes(out, "usage.cache_creation_input_tokens", cacheWriteTokens)
411 }
409 }412 }
410 413
411 return [][]byte{out}414 return [][]byte{out}
@@ -570,7 +573,7 @@ func finalizeOpenAIAnthropicContentBlocks(param *ConvertOpenAIResponseToAnthropi · 158 unmodified lines
570 }573 }
571}574}
572 575
573func emitAnthropicMessageDelta(param *ConvertOpenAIResponseToAnthropicParams, results *[][]byte, inputTokens, outputTokens, cachedTokens int64) {576func emitAnthropicMessageDelta(param *ConvertOpenAIResponseToAnthropicParams, results *[][]byte, inputTokens, outputTokens, cachedTokens, cacheWriteTokens int64) {
574 if param == nil || param.MessageDeltaSent {577 if param == nil || param.MessageDeltaSent {
575 return578 return
576 }579 }
@@ -581,6 +584,9 @@ func emitAnthropicMessageDelta(param *ConvertOpenAIResponseToAnthropicParams, re · 4 unmodified lines
581 if cachedTokens > 0 {584 if cachedTokens > 0 {
582 messageDeltaJSON, _ = sjson.SetBytes(messageDeltaJSON, "usage.cache_read_input_tokens", cachedTokens)585 messageDeltaJSON, _ = sjson.SetBytes(messageDeltaJSON, "usage.cache_read_input_tokens", cachedTokens)
583 }586 }
587 if cacheWriteTokens > 0 {
588 messageDeltaJSON, _ = sjson.SetBytes(messageDeltaJSON, "usage.cache_creation_input_tokens", cacheWriteTokens)
589 }
584 *results = append(*results, translatorcommon.AppendSSEEventBytes(nil, "message_delta", messageDeltaJSON, 2))590 *results = append(*results, translatorcommon.AppendSSEEventBytes(nil, "message_delta", messageDeltaJSON, 2))
585 param.MessageDeltaSent = true591 param.MessageDeltaSent = true
586}592}
@@ -748,12 +754,15 @@ func ConvertOpenAIResponseToClaudeNonStream(_ context.Context, _ string, origina · 161 unmodified lines
748 }754 }
749 755
750 if respUsage := root.Get("usage"); respUsage.Exists() {756 if respUsage := root.Get("usage"); respUsage.Exists() {
751 inputTokens, outputTokens, cachedTokens := extractOpenAIUsage(respUsage)757 inputTokens, outputTokens, cachedTokens, cacheWriteTokens := extractOpenAIUsage(respUsage)
752 out, _ = sjson.SetBytes(out, "usage.input_tokens", inputTokens)758 out, _ = sjson.SetBytes(out, "usage.input_tokens", inputTokens)
753 out, _ = sjson.SetBytes(out, "usage.output_tokens", outputTokens)759 out, _ = sjson.SetBytes(out, "usage.output_tokens", outputTokens)
754 if cachedTokens > 0 {760 if cachedTokens > 0 {
755 out, _ = sjson.SetBytes(out, "usage.cache_read_input_tokens", cachedTokens)761 out, _ = sjson.SetBytes(out, "usage.cache_read_input_tokens", cachedTokens)
756 }762 }
763 if cacheWriteTokens > 0 {
764 out, _ = sjson.SetBytes(out, "usage.cache_creation_input_tokens", cacheWriteTokens)
765 }
757 }766 }
758 767
759 if !stopReasonSet {768 if !stopReasonSet {
@@ -771,14 +780,18 @@ func ClaudeTokenCount(ctx context.Context, count int64) []byte { · 11 unmodified lines
771 return translatorcommon.ClaudeInputTokensJSON(count)780 return translatorcommon.ClaudeInputTokensJSON(count)
772}781}
773 782
774func extractOpenAIUsage(usage gjson.Result) (int64, int64, int64) {783func extractOpenAIUsage(usage gjson.Result) (int64, int64, int64, int64) {
775 if !usage.Exists() || usage.Type == gjson.Null {784 if !usage.Exists() || usage.Type == gjson.Null {
776 return 0, 0, 0785 return 0, 0, 0, 0
777 }786 }
778 787
779 inputTokens := usage.Get("prompt_tokens").Int()788 inputTokens := usage.Get("prompt_tokens").Int()
780 outputTokens := usage.Get("completion_tokens").Int()789 outputTokens := usage.Get("completion_tokens").Int()
781 cachedTokens := usage.Get("prompt_tokens_details.cached_tokens").Int()790 cachedTokens := usage.Get("prompt_tokens_details.cached_tokens").Int()
791 cacheWriteTokens := usage.Get("prompt_tokens_details.cache_write_tokens").Int()
792 if cacheWriteTokens == 0 {
793 cacheWriteTokens = usage.Get("prompt_tokens_details.cache_creation_tokens").Int()
794 }
782 795
783 if cachedTokens > 0 {796 if cachedTokens > 0 {
784 if inputTokens >= cachedTokens {797 if inputTokens >= cachedTokens {
@@ -788,5 +801,5 @@ func extractOpenAIUsage(usage gjson.Result) (int64, int64, int64) { · 3 unmodified lines
788 }801 }
789 }802 }
790 803
791 return inputTokens, outputTokens, cachedTokens804 return inputTokens, outputTokens, cachedTokens, cacheWriteTokens
792}805}
internal/translator/openai/claude/openai_claude_response_test.go+163−0
@@ -3,6 +3,7 @@ package claude · 2 unmodified lines
3import (3import (
4 "bytes"4 "bytes"
5 "context"5 "context"
6 "fmt"
6 "strings"7 "strings"
7 "testing"8 "testing"
8 9
@@ -577,3 +578,165 @@ func TestStreamingTool_OmittedToolCallIndexPreservesParallelCalls(t *testing.T) · 568 unmodified lines
577 t.Fatalf("stop_reason = %q, want %q", got, "tool_use")578 t.Fatalf("stop_reason = %q, want %q", got, "tool_use")
578 }579 }
579}580}
581
582func TestStreamingUsage_PreservesCacheWriteTokens(t *testing.T) {
583 tests := []struct {
584 name string
585 usageJSON string
586 wantInputTokens int64
587 wantOutputTokens int64
588 wantCacheReadTokens int64
589 wantCacheWriteTokens int64
590 }{
591 {
592 name: "cache_write_tokens field",
593 usageJSON: `{"prompt_tokens":1000,"completion_tokens":200,"prompt_tokens_details":{"cached_tokens":800,"cache_write_tokens":150}}`,
594 wantInputTokens: 200,
595 wantOutputTokens: 200,
596 wantCacheReadTokens: 800,
597 wantCacheWriteTokens: 150,
598 },
599 {
600 name: "cache_creation_tokens alias",
601 usageJSON: `{"prompt_tokens":1000,"completion_tokens":200,"prompt_tokens_details":{"cached_tokens":800,"cache_creation_tokens":150}}`,
602 wantInputTokens: 200,
603 wantOutputTokens: 200,
604 wantCacheReadTokens: 800,
605 wantCacheWriteTokens: 150,
606 },
607 {
608 name: "cached_tokens greater than prompt_tokens clamps input_tokens to zero",
609 usageJSON: `{"prompt_tokens":500,"completion_tokens":100,"prompt_tokens_details":{"cached_tokens":800,"cache_write_tokens":50}}`,
610 wantInputTokens: 0,
611 wantOutputTokens: 100,
612 wantCacheReadTokens: 800,
613 wantCacheWriteTokens: 50,
614 },
615 {
616 name: "zero cache_write_tokens does not emit cache_creation_input_tokens",
617 usageJSON: `{"prompt_tokens":1000,"completion_tokens":200,"prompt_tokens_details":{"cached_tokens":800,"cache_write_tokens":0}}`,
618 wantInputTokens: 200,
619 wantOutputTokens: 200,
620 wantCacheReadTokens: 800,
621 wantCacheWriteTokens: 0,
622 },
623 }
624
625 for _, tt := range tests {
626 t.Run(tt.name, func(t *testing.T) {
627 events := runStream(t, streamReq,
628 `{"id":"c1","model":"m","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"}}]}`,
629 `{"id":"c1","model":"m","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}`,
630 fmt.Sprintf(`{"id":"c1","model":"m","choices":[],"usage":%s}`, tt.usageJSON),
631 )
632
633 var deltaEvent *sseEvent
634 for _, e := range events {
635 if e.Type == "message_delta" {
636 deltaEvent = &e
637 break
638 }
639 }
640 if deltaEvent == nil {
641 t.Fatalf("missing message_delta event")
642 }
643 if input := gjson.Get(deltaEvent.Payload, "usage.input_tokens").Int(); input != tt.wantInputTokens {
644 t.Fatalf("input_tokens = %d, want %d", input, tt.wantInputTokens)
645 }
646 if output := gjson.Get(deltaEvent.Payload, "usage.output_tokens").Int(); output != tt.wantOutputTokens {
647 t.Fatalf("output_tokens = %d, want %d", output, tt.wantOutputTokens)
648 }
649 if cacheRead := gjson.Get(deltaEvent.Payload, "usage.cache_read_input_tokens").Int(); cacheRead != tt.wantCacheReadTokens {
650 t.Fatalf("cache_read_input_tokens = %d, want %d", cacheRead, tt.wantCacheReadTokens)
651 }
652 if tt.wantCacheWriteTokens == 0 {
653 if gjson.Get(deltaEvent.Payload, "usage.cache_creation_input_tokens").Exists() {
654 t.Fatalf("cache_creation_input_tokens should not be emitted when zero; got %v", gjson.Get(deltaEvent.Payload, "usage.cache_creation_input_tokens").Raw)
655 }
656 } else if cacheWrite := gjson.Get(deltaEvent.Payload, "usage.cache_creation_input_tokens").Int(); cacheWrite != tt.wantCacheWriteTokens {
657 t.Fatalf("cache_creation_input_tokens = %d, want %d", cacheWrite, tt.wantCacheWriteTokens)
658 }
659 })
660 }
661}
662
663func TestNonStreamingUsage_PreservesCacheWriteTokens(t *testing.T) {
664 tests := []struct {
665 name string
666 usageJSON string
667 wantInputTokens int64
668 wantOutputTokens int64
669 wantCacheReadTokens int64
670 wantCacheWriteTokens int64
671 }{
672 {
673 name: "cache_write_tokens field",
674 usageJSON: `{"prompt_tokens":1000,"completion_tokens":200,"prompt_tokens_details":{"cached_tokens":800,"cache_write_tokens":150}}`,
675 wantInputTokens: 200,
676 wantOutputTokens: 200,
677 wantCacheReadTokens: 800,
678 wantCacheWriteTokens: 150,
679 },
680 {
681 name: "cache_creation_tokens alias",
682 usageJSON: `{"prompt_tokens":1000,"completion_tokens":200,"prompt_tokens_details":{"cached_tokens":800,"cache_creation_tokens":150}}`,
683 wantInputTokens: 200,
684 wantOutputTokens: 200,
685 wantCacheReadTokens: 800,
686 wantCacheWriteTokens: 150,
687 },
688 {
689 name: "cached_tokens greater than prompt_tokens clamps input_tokens to zero",
690 usageJSON: `{"prompt_tokens":500,"completion_tokens":100,"prompt_tokens_details":{"cached_tokens":800,"cache_write_tokens":50}}`,
691 wantInputTokens: 0,
692 wantOutputTokens: 100,
693 wantCacheReadTokens: 800,
694 wantCacheWriteTokens: 50,
695 },
696 {
697 name: "zero cache_write_tokens does not emit cache_creation_input_tokens",
698 usageJSON: `{"prompt_tokens":1000,"completion_tokens":200,"prompt_tokens_details":{"cached_tokens":800,"cache_write_tokens":0}}`,
699 wantInputTokens: 200,
700 wantOutputTokens: 200,
701 wantCacheReadTokens: 800,
702 wantCacheWriteTokens: 0,
703 },
704 }
705
706 for _, tt := range tests {
707 t.Run(tt.name, func(t *testing.T) {
708 rawJSON := []byte(fmt.Sprintf(`{
709 "id":"chatcmpl-123",
710 "object":"chat.completion",
711 "created":1677652288,
712 "model":"gpt-5.4",
713 "choices":[{"index":0,"message":{"role":"assistant","content":"Hello world"},"finish_reason":"stop"}],
714 "usage":%s
715 }`, tt.usageJSON))
716
717 ctx := context.Background()
718 reqJSON := []byte(`{"model":"claude-3-5-sonnet-20241022","messages":[{"role":"user","content":"Hello"}]}`)
719
720 out := ConvertOpenAIResponseToClaudeNonStream(ctx, "", reqJSON, reqJSON, rawJSON, nil)
721 parsed := gjson.ParseBytes(out)
722
723 usage := parsed.Get("usage")
724 if got := usage.Get("input_tokens").Int(); got != tt.wantInputTokens {
725 t.Fatalf("input_tokens = %d, want %d", got, tt.wantInputTokens)
726 }
727 if got := usage.Get("output_tokens").Int(); got != tt.wantOutputTokens {
728 t.Fatalf("output_tokens = %d, want %d", got, tt.wantOutputTokens)
729 }
730 if got := usage.Get("cache_read_input_tokens").Int(); got != tt.wantCacheReadTokens {
731 t.Fatalf("cache_read_input_tokens = %d, want %d", got, tt.wantCacheReadTokens)
732 }
733 if tt.wantCacheWriteTokens == 0 {
734 if usage.Get("cache_creation_input_tokens").Exists() {
735 t.Fatalf("cache_creation_input_tokens should not be emitted when zero; got %v", usage.Get("cache_creation_input_tokens").Raw)
736 }
737 } else if got := usage.Get("cache_creation_input_tokens").Int(); got != tt.wantCacheWriteTokens {
738 t.Fatalf("cache_creation_input_tokens = %d, want %d", got, tt.wantCacheWriteTokens)
739 }
740 })
741 }
742}