From dea4787970a6e87f5e7f043766c529b083275f66 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Wed, 8 Jul 2026 04:01:22 +0800 Subject: [PATCH] refactor(executor): centralize OpenAI stream usage handling with `StreamUsageBuffer` - Introduced `StreamUsageBuffer` to encapsulate stream usage observation and publishing logic. - Updated OpenAI-compatible executors and helpers to use `StreamUsageBuffer` for consistent stream usage tracking. - Added unit tests to validate `StreamUsageBuffer` behavior, ensuring proper handling of usage details. Closes: #4053 --- .../runtime/executor/codex_openai_images.go | 6 +-- .../runtime/executor/helps/usage_helpers.go | 32 ++++++++++++++++ .../executor/helps/usage_helpers_test.go | 37 +++++++++++++++++++ internal/runtime/executor/kimi_executor.go | 6 +-- .../executor/openai_compat_executor.go | 9 +++-- 5 files changed, 80 insertions(+), 10 deletions(-) diff --git a/internal/runtime/executor/codex_openai_images.go b/internal/runtime/executor/codex_openai_images.go index 53981799..4ca89827 100644 --- a/internal/runtime/executor/codex_openai_images.go +++ b/internal/runtime/executor/codex_openai_images.go @@ -428,10 +428,12 @@ func (e *CodexExecutor) executeDirectOpenAIImageStream(ctx context.Context, auth out := make(chan cliproxyexecutor.StreamChunk) go func() { defer close(out) + var streamUsage helps.StreamUsageBuffer defer func() { if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("codex executor: close response body error: %v", errClose) } + streamUsage.Publish(ctx, reporter) reporter.EnsurePublished(ctx) }() @@ -443,9 +445,7 @@ func (e *CodexExecutor) executeDirectOpenAIImageStream(ctx context.Context, auth chunk = applyCodexIdentityConfuseResponsePayload(chunk, identityState) helps.AppendAPIResponseChunk(ctx, e.cfg, chunk) for _, line := range bytes.Split(chunk, []byte("\n")) { - if detail, ok := helps.ParseOpenAIStreamUsage(bytes.TrimSpace(line)); ok { - reporter.Publish(ctx, detail) - } + streamUsage.Observe(helps.ParseOpenAIStreamUsage(bytes.TrimSpace(line))) } select { case out <- cliproxyexecutor.StreamChunk{Payload: chunk}: diff --git a/internal/runtime/executor/helps/usage_helpers.go b/internal/runtime/executor/helps/usage_helpers.go index 67d4205b..687bced1 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -447,6 +447,38 @@ func resolveUsageAuthType(auth *cliproxyauth.Auth) string { return auth.AuthKind() } +// StreamUsageBuffer keeps the latest usage detail observed in a stream. +type StreamUsageBuffer struct { + detail usage.Detail + ok bool +} + +// Observe records detail when ok is true, allowing the final stream usage to win. +func (b *StreamUsageBuffer) Observe(detail usage.Detail, ok bool) { + if b == nil || !ok { + return + } + b.detail = detail + b.ok = true +} + +// Publish emits the latest observed usage detail, if any. +func (b *StreamUsageBuffer) Publish(ctx context.Context, reporter *UsageReporter) bool { + if b == nil || !b.ok || reporter == nil { + return false + } + reporter.Publish(ctx, b.detail) + return true +} + +// Detail returns the latest observed usage detail. +func (b *StreamUsageBuffer) Detail() (usage.Detail, bool) { + if b == nil || !b.ok { + return usage.Detail{}, false + } + return b.detail, true +} + func ParseCodexUsage(data []byte) (usage.Detail, bool) { usageNode := gjson.ParseBytes(data).Get("response.usage") if !hasOpenAIStyleUsageTokenFields(usageNode) { diff --git a/internal/runtime/executor/helps/usage_helpers_test.go b/internal/runtime/executor/helps/usage_helpers_test.go index 03520744..a1ff29e0 100644 --- a/internal/runtime/executor/helps/usage_helpers_test.go +++ b/internal/runtime/executor/helps/usage_helpers_test.go @@ -89,6 +89,43 @@ func TestParseOpenAIStreamUsageResponsesFields(t *testing.T) { } } +func TestStreamUsageBufferKeepsLastUsage(t *testing.T) { + var buffer StreamUsageBuffer + buffer.Observe(usage.Detail{}, true) + buffer.Observe(usage.Detail{InputTokens: 1, OutputTokens: 1, TotalTokens: 2}, false) + buffer.Observe(usage.Detail{InputTokens: 39320, OutputTokens: 26, TotalTokens: 39346, CachedTokens: 33280}, true) + + detail, ok := buffer.Detail() + if !ok { + t.Fatal("buffer detail ok = false, want true") + } + if detail.InputTokens != 39320 { + t.Fatalf("input tokens = %d, want %d", detail.InputTokens, 39320) + } + if detail.OutputTokens != 26 { + t.Fatalf("output tokens = %d, want %d", detail.OutputTokens, 26) + } + if detail.TotalTokens != 39346 { + t.Fatalf("total tokens = %d, want %d", detail.TotalTokens, 39346) + } + if detail.CachedTokens != 33280 { + t.Fatalf("cached tokens = %d, want %d", detail.CachedTokens, 33280) + } +} + +func TestStreamUsageBufferPreservesOnlyZeroUsage(t *testing.T) { + var buffer StreamUsageBuffer + buffer.Observe(usage.Detail{}, true) + + detail, ok := buffer.Detail() + if !ok { + t.Fatal("buffer detail ok = false, want true") + } + if detail != (usage.Detail{}) { + t.Fatalf("detail = %+v, want zero detail", detail) + } +} + func TestParseClaudeUsageIncludesCacheTokensInTotal(t *testing.T) { data := []byte(`{"usage":{"input_tokens":3085,"output_tokens":253,"cache_read_input_tokens":7,"cache_creation_input_tokens":19514}}`) detail := ParseClaudeUsage(data) diff --git a/internal/runtime/executor/kimi_executor.go b/internal/runtime/executor/kimi_executor.go index f296687f..439a92c7 100644 --- a/internal/runtime/executor/kimi_executor.go +++ b/internal/runtime/executor/kimi_executor.go @@ -288,12 +288,12 @@ func (e *KimiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut scanner := bufio.NewScanner(httpResp.Body) scanner.Buffer(nil, 1_048_576) // 1MB var param any + var streamUsage helps.StreamUsageBuffer + defer streamUsage.Publish(ctx, reporter) for scanner.Scan() { line := scanner.Bytes() helps.AppendAPIResponseChunk(ctx, e.cfg, line) - if detail, ok := helps.ParseOpenAIStreamUsage(line); ok { - reporter.Publish(ctx, detail) - } + streamUsage.Observe(helps.ParseOpenAIStreamUsage(line)) chunks := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, body, bytes.Clone(line), ¶m) for i := range chunks { select { diff --git a/internal/runtime/executor/openai_compat_executor.go b/internal/runtime/executor/openai_compat_executor.go index 5bfba83d..f67ddaf4 100644 --- a/internal/runtime/executor/openai_compat_executor.go +++ b/internal/runtime/executor/openai_compat_executor.go @@ -393,12 +393,12 @@ func (e *OpenAICompatExecutor) ExecuteStream(ctx context.Context, auth *cliproxy scanner := bufio.NewScanner(httpResp.Body) scanner.Buffer(nil, 52_428_800) // 50MB var param any + var streamUsage helps.StreamUsageBuffer + defer streamUsage.Publish(ctx, reporter) for scanner.Scan() { line := scanner.Bytes() helps.AppendAPIResponseChunk(ctx, e.cfg, line) - if detail, ok := helps.ParseOpenAIStreamUsage(line); ok { - reporter.Publish(ctx, detail) - } + streamUsage.Observe(helps.ParseOpenAIStreamUsage(line)) trimmedLine := bytes.TrimSpace(line) if len(trimmedLine) == 0 { continue @@ -452,7 +452,8 @@ func (e *OpenAICompatExecutor) ExecuteStream(ctx context.Context, auth *cliproxy } } } - // Ensure we record the request if no usage chunk was ever seen + // Ensure we record the request if no usage chunk was ever seen. + streamUsage.Publish(ctx, reporter) reporter.EnsurePublished(ctx) }() return &cliproxyexecutor.StreamResult{Headers: httpResp.Header.Clone(), Chunks: out}, nil