From ef7e8206d370ab49b5b9f20b9db03990d412e190 Mon Sep 17 00:00:00 2001
From: Luis Pater <webmaster@idotorg.org>
Date: Sun, 9 Nov 2025 17:24:47 +0800
Subject: [PATCH] fix(executor): ensure usage reporting for upstream responses
 lacking usage data

Add `ensurePublished` to guarantee request counting even when usage fields (e.g., tokens) are absent in OpenAI-compatible executor responses, particularly for streaming paths.
---
 .../executor/openai_compat_executor.go        |  4 ++++
 internal/runtime/executor/usage_helpers.go    | 22 +++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/internal/runtime/executor/openai_compat_executor.go b/internal/runtime/executor/openai_compat_executor.go
index e08196db..cf6af2dd 100644
--- a/internal/runtime/executor/openai_compat_executor.go
+++ b/internal/runtime/executor/openai_compat_executor.go
@@ -116,6 +116,8 @@ func (e *OpenAICompatExecutor) Execute(ctx context.Context, auth *cliproxyauth.A
 	}
 	appendAPIResponseChunk(ctx, e.cfg, body)
 	reporter.publish(ctx, parseOpenAIUsage(body))
+	// Ensure we at least record the request even if upstream doesn't return usage
+	reporter.ensurePublished(ctx)
 	// Translate response back to source format when needed
 	var param any
 	out := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), translated, body, &param)
@@ -225,6 +227,8 @@ func (e *OpenAICompatExecutor) ExecuteStream(ctx context.Context, auth *cliproxy
 			reporter.publishFailure(ctx)
 			out <- cliproxyexecutor.StreamChunk{Err: errScan}
 		}
+		// Ensure we record the request if no usage chunk was ever seen
+		reporter.ensurePublished(ctx)
 	}()
 	return stream, nil
 }
diff --git a/internal/runtime/executor/usage_helpers.go b/internal/runtime/executor/usage_helpers.go
index 4e7cab90..1f81c048 100644
--- a/internal/runtime/executor/usage_helpers.go
+++ b/internal/runtime/executor/usage_helpers.go
@@ -84,6 +84,28 @@ func (r *usageReporter) publishWithOutcome(ctx context.Context, detail usage.Det
 	})
 }
 
+// ensurePublished guarantees that a usage record is emitted exactly once.
+// It is safe to call multiple times; only the first call wins due to once.Do.
+// This is used to ensure request counting even when upstream responses do not
+// include any usage fields (tokens), especially for streaming paths.
+func (r *usageReporter) ensurePublished(ctx context.Context) {
+	if r == nil {
+		return
+	}
+	r.once.Do(func() {
+		usage.PublishRecord(ctx, usage.Record{
+			Provider:    r.provider,
+			Model:       r.model,
+			Source:      r.source,
+			APIKey:      r.apiKey,
+			AuthID:      r.authID,
+			RequestedAt: r.requestedAt,
+			Failed:      false,
+			Detail:      usage.Detail{},
+		})
+	})
+}
+
 func apiKeyFromContext(ctx context.Context) string {
 	if ctx == nil {
 		return ""