From 6dd3940567e04da5bfaf503650adf2ca9fa1a93b Mon Sep 17 00:00:00 2001
From: Pesto <120122957+AlapinEnjoyer@users.noreply.github.com>
Date: Thu, 25 Jun 2026 19:19:58 +0200
Subject: [PATCH] feat: enhance token usage resolution for OpenAI and Anthropic
 compatibility

---
 internal/llm/usage_resolver.go      | 20 +++++---
 internal/llm/usage_resolver_test.go | 80 +++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+), 6 deletions(-)
 create mode 100644 internal/llm/usage_resolver_test.go

diff --git a/internal/llm/usage_resolver.go b/internal/llm/usage_resolver.go
index 13454207..661d6591 100644
--- a/internal/llm/usage_resolver.go
+++ b/internal/llm/usage_resolver.go
@@ -27,15 +27,23 @@ var completionTokensPaths = []string{
 }
 
 var cacheReadTokensPaths = []string{
-	"usage.cache_read_input_tokens",                // Anthropic
-	"cache_read_input_tokens",                      // flat at root
-	"usage.prompt_tokens_details.cache_tokens_hit", // some providers
-	"usage.prompt_tokens_details.cache_tokens",     // some providers
+	"usage.cache_read_input_tokens",                  // Anthropic
+	"cache_read_input_tokens",                        // flat at root
+	"data.usage.cache_read_input_tokens",             // wrapped Anthropic-compatible proxy
+	"usage.prompt_tokens_details.cached_tokens",      // OpenAI-compatible providers
+	"data.usage.prompt_tokens_details.cached_tokens", // wrapped OpenAI-compatible providers
+	"usage.prompt_tokens_details.cache_tokens_hit",   // some providers
+	"data.usage.prompt_tokens_details.cache_tokens_hit",
+	"usage.prompt_tokens_details.cache_tokens", // some providers
+	"data.usage.prompt_tokens_details.cache_tokens",
 }
 
 var cacheWriteTokensPaths = []string{
-	"usage.cache_creation_input_tokens", // Anthropic / proxy
-	"cache_creation_input_tokens",       // flat at root
+	"usage.cache_creation_input_tokens",      // Anthropic / proxy
+	"cache_creation_input_tokens",            // flat at root
+	"data.usage.cache_creation_input_tokens", // wrapped Anthropic-compatible proxy
+	"usage.prompt_tokens_details.cache_creation_tokens",
+	"data.usage.prompt_tokens_details.cache_creation_tokens",
 }
 
 // totalTokensPaths is an ordered list of JSON paths to try when extracting
diff --git a/internal/llm/usage_resolver_test.go b/internal/llm/usage_resolver_test.go
new file mode 100644
index 00000000..5b670995
--- /dev/null
+++ b/internal/llm/usage_resolver_test.go
@@ -0,0 +1,80 @@
+package llm
+
+import "testing"
+
+func TestResolveUsageOpenAICompatibleCachedTokens(t *testing.T) {
+	usage := resolveUsage([]byte(`{
+		"usage": {
+			"prompt_tokens": 100,
+			"completion_tokens": 20,
+			"total_tokens": 120,
+			"prompt_tokens_details": {
+				"cached_tokens": 75
+			}
+		}
+	}`))
+
+	if usage == nil {
+		t.Fatal("resolveUsage returned nil")
+	}
+	if usage.CacheReadTokens != 75 {
+		t.Errorf("CacheReadTokens = %d, want 75", usage.CacheReadTokens)
+	}
+	if usage.PromptTokens != 100 {
+		t.Errorf("PromptTokens = %d, want 100", usage.PromptTokens)
+	}
+	if usage.CompletionTokens != 20 {
+		t.Errorf("CompletionTokens = %d, want 20", usage.CompletionTokens)
+	}
+}
+
+func TestResolveUsageWrappedCachedTokens(t *testing.T) {
+	usage := resolveUsage([]byte(`{
+		"data": {
+			"usage": {
+				"prompt_tokens": 100,
+				"completion_tokens": 20,
+				"prompt_tokens_details": {
+					"cached_tokens": 75,
+					"cache_creation_tokens": 10
+				}
+			}
+		}
+	}`))
+
+	if usage == nil {
+		t.Fatal("resolveUsage returned nil")
+	}
+	if usage.CacheReadTokens != 75 {
+		t.Errorf("CacheReadTokens = %d, want 75", usage.CacheReadTokens)
+	}
+	if usage.CacheWriteTokens != 10 {
+		t.Errorf("CacheWriteTokens = %d, want 10", usage.CacheWriteTokens)
+	}
+	if usage.TotalTokens != 205 {
+		t.Errorf("TotalTokens = %d, want 205", usage.TotalTokens)
+	}
+}
+
+func TestResolveUsageWrappedAnthropicCompatibleCacheTokens(t *testing.T) {
+	usage := resolveUsage([]byte(`{
+		"data": {
+			"usage": {
+				"prompt_tokens": 100,
+				"completion_tokens": 20,
+				"cache_read_input_tokens": 40,
+				"cache_creation_input_tokens": 15
+			}
+		}
+	}`))
+
+	if usage == nil {
+		t.Fatal("resolveUsage returned nil")
+	}
+	if usage.CacheReadTokens != 40 {
+		t.Errorf("CacheReadTokens = %d, want 40", usage.CacheReadTokens)
+	}
+	if usage.CacheWriteTokens != 15 {
+		t.Errorf("CacheWriteTokens = %d, want 15", usage.CacheWriteTokens)
+	}
+}