From 04d83a6b7c29d7c7664906e056f9129c696e6c01 Mon Sep 17 00:00:00 2001
From: Luxferre <lux@ferre>
Date: Thu, 27 Aug 2026 14:42:44 +0300
Subject: [PATCH] feat: estimate and populate prompt cache tokens when upstream
 omits breakdown

---
 main.go        | 30 ++++++++++++++++++++++++++++--
 qorona_test.go | 18 ++++++++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/main.go b/main.go
index 6401aa2..2b4773a 100644
--- a/main.go
+++ b/main.go
@@ -135,6 +135,7 @@ type Usage struct {
   PromptTokens            int                      `json:"prompt_tokens"`
   CompletionTokens        int                      `json:"completion_tokens"`
   TotalTokens             int                      `json:"total_tokens"`
+  CachedTokens            int                      `json:"cached_tokens,omitempty"`
   PromptTokensDetails     *PromptTokensDetails     `json:"prompt_tokens_details,omitempty"`
   CompletionTokensDetails *CompletionTokensDetails `json:"completion_tokens_details,omitempty"`
 }
@@ -422,12 +423,35 @@ type FinalOutput struct {
   ReasoningTokens  int
 }
 
+func EstimateCachedTokens(promptTokens int, currentQuestion string) int {
+  if promptTokens <= 0 {
+    return 0
+  }
+  qLen := len(strings.TrimSpace(currentQuestion))
+  uncached := (qLen + 3) / 4
+  if uncached < 1 {
+    uncached = 1
+  }
+  if uncached >= promptTokens {
+    return 0
+  }
+  return promptTokens - uncached
+}
+
+func ResolveCachedTokens(upstreamCached int, promptTokens int, currentQuestion string) int {
+  if upstreamCached > 0 {
+    return upstreamCached
+  }
+  return EstimateCachedTokens(promptTokens, currentQuestion)
+}
+
 func BuildUsage(promptTokens, completionTokens, cachedTokens, reasoningTokens int) Usage {
   totalTokens := promptTokens + completionTokens
   return Usage{
     PromptTokens:     promptTokens,
     CompletionTokens: completionTokens,
     TotalTokens:      totalTokens,
+    CachedTokens:     cachedTokens,
     PromptTokensDetails: &PromptTokensDetails{
       CachedTokens: cachedTokens,
     },
@@ -1473,12 +1497,13 @@ func (gw *QoronaGateway) HandleChatCompletions(w http.ResponseWriter, r *http.Re
                 if cachedTokens == 0 {
                   cachedTokens = sResult.Result.FinalStatus.PromptTokensDetails.CachedTokens
                 }
+                cachedTokens = ResolveCachedTokens(cachedTokens, promptTokens, question)
 
                 if req.StreamOptions != nil && req.StreamOptions.IncludeUsage {
                   streamer.EmitUsage(BuildUsage(promptTokens, completionTokens, cachedTokens, 0))
                 }
                 streamer.Done()
-                log.Printf("[Stream] Completed run_id=%s finish_reason=%s tokens=%d/%d", runID, finishReason, promptTokens, completionTokens)
+                log.Printf("[Stream] Completed run_id=%s finish_reason=%s tokens=%d/%d cached=%d", runID, finishReason, promptTokens, completionTokens, cachedTokens)
                 return
               }
             }
@@ -1606,6 +1631,7 @@ func (gw *QoronaGateway) HandleChatCompletions(w http.ResponseWriter, r *http.Re
                 if cachedTokens == 0 {
                   cachedTokens = sResult.Result.FinalStatus.PromptTokensDetails.CachedTokens
                 }
+                cachedTokens = ResolveCachedTokens(cachedTokens, promptTokens, question)
 
                 fullContent := accumulatedContent.String()
                 rawThinking := accumulatedThinking.String()
@@ -1644,7 +1670,7 @@ func (gw *QoronaGateway) HandleChatCompletions(w http.ResponseWriter, r *http.Re
                 }
 
                 WriteCompletionResponse(w, completionID, createdTime, modelName, out)
-                log.Printf("[Non-Stream] Completed run_id=%s finish_reason=%s tokens=%d/%d", runID, finishReason, promptTokens, completionTokens)
+                log.Printf("[Non-Stream] Completed run_id=%s finish_reason=%s tokens=%d/%d cached=%d", runID, finishReason, promptTokens, completionTokens, cachedTokens)
                 return
               }
             }
diff --git a/qorona_test.go b/qorona_test.go
index 29e2c87..8c9721d 100644
--- a/qorona_test.go
+++ b/qorona_test.go
@@ -192,3 +192,21 @@ func TestStreamerUsage(t *testing.T) {
   }
 }
 
+func TestResolveCachedTokens(t *testing.T) {
+  // If upstream explicitly provides cached tokens (>0), preserve it
+  if got := ResolveCachedTokens(500, 1000, "hello"); got != 500 {
+    t.Fatalf("expected 500, got %d", got)
+  }
+
+  // If upstream returns 0, estimate based on promptTokens and current question
+  if got := ResolveCachedTokens(0, 1397, "Say hello in 1 word"); got <= 0 || got >= 1397 {
+    t.Fatalf("expected cached tokens between 1 and 1396, got %d", got)
+  }
+
+  // If promptTokens <= 0, return 0
+  if got := ResolveCachedTokens(0, 0, "hello"); got != 0 {
+    t.Fatalf("expected 0, got %d", got)
+  }
+}
+
+