From bc812e5f740d3285fb3bd7514841631e8eb5dd6c Mon Sep 17 00:00:00 2001 From: Supra4E8C Date: Sat, 11 Jul 2026 02:42:42 +0800 Subject: [PATCH] perf(usage): avoid parsing irrelevant stream chunks --- .../runtime/executor/codex_openai_images.go | 2 +- .../runtime/executor/helps/usage_helpers.go | 45 +++++++++++++++ .../executor/helps/usage_helpers_test.go | 55 ++++++++++++++++++- .../helps/usage_stream_benchmark_test.go | 31 +++++++++++ internal/runtime/executor/kimi_executor.go | 2 +- .../executor/openai_compat_executor.go | 2 +- 6 files changed, 132 insertions(+), 5 deletions(-) create mode 100644 internal/runtime/executor/helps/usage_stream_benchmark_test.go diff --git a/internal/runtime/executor/codex_openai_images.go b/internal/runtime/executor/codex_openai_images.go index 663acb64..10019f0c 100644 --- a/internal/runtime/executor/codex_openai_images.go +++ b/internal/runtime/executor/codex_openai_images.go @@ -449,7 +449,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")) { - streamUsage.Observe(helps.ParseOpenAIStreamUsage(bytes.TrimSpace(line))) + streamUsage.ObserveOpenAIStream(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 c04c2c83..8eaad5e8 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -455,6 +455,11 @@ type StreamUsageBuffer struct { ok bool } +var ( + openAIStreamUsageMarker = []byte(`"usage"`) + openAIStreamServiceTierMarker = []byte(`"service_tier"`) +) + // 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 { @@ -473,6 +478,42 @@ func (b *StreamUsageBuffer) Observe(detail usage.Detail, ok bool) { b.ok = true } +// ObserveOpenAIStream records response-tier state and the latest usage from an +// OpenAI-style stream while avoiding JSON parsing for irrelevant chunks. +func (b *StreamUsageBuffer) ObserveOpenAIStream(line []byte) { + if b == nil { + return + } + payload := jsonPayload(line) + if len(payload) == 0 { + return + } + + hasUsageCandidate := bytes.Contains(payload, openAIStreamUsageMarker) + needTier := b.detail.ResponseServiceTier == "" || hasUsageCandidate + hasTierCandidate := needTier && bytes.Contains(payload, openAIStreamServiceTierMarker) + if !hasUsageCandidate && !hasTierCandidate { + return + } + if !gjson.ValidBytes(payload) { + return + } + + detail := usage.Detail{} + usageOK := false + if hasUsageCandidate { + usageNode := gjson.GetBytes(payload, "usage") + if hasOpenAIStyleUsageTokenFields(usageNode) { + detail = parseOpenAIStyleUsageNode(usageNode) + usageOK = true + } + } + if hasTierCandidate { + detail.ResponseServiceTier = extractResponseServiceTierFromValidJSON(payload) + } + b.Observe(detail, usageOK || detail.ResponseServiceTier != "") +} + // 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 { @@ -682,6 +723,10 @@ func extractResponseServiceTier(payload []byte) string { if len(payload) == 0 || !gjson.ValidBytes(payload) { return "" } + return extractResponseServiceTierFromValidJSON(payload) +} + +func extractResponseServiceTierFromValidJSON(payload []byte) string { for _, path := range []string{"response.service_tier", "service_tier", "interaction.service_tier"} { if tier := strings.TrimSpace(gjson.GetBytes(payload, path).String()); tier != "" { return tier diff --git a/internal/runtime/executor/helps/usage_helpers_test.go b/internal/runtime/executor/helps/usage_helpers_test.go index 0d1ff671..0e863369 100644 --- a/internal/runtime/executor/helps/usage_helpers_test.go +++ b/internal/runtime/executor/helps/usage_helpers_test.go @@ -167,8 +167,8 @@ func TestStreamUsageBufferPreservesTierAcrossChunks(t *testing.T) { t.Parallel() var buffer StreamUsageBuffer - buffer.Observe(ParseOpenAIStreamUsage([]byte(`data: {"service_tier":"default"}`))) - buffer.Observe(ParseOpenAIStreamUsage([]byte(`data: {"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}`))) + buffer.ObserveOpenAIStream([]byte(`data: {"service_tier":"default"}`)) + buffer.ObserveOpenAIStream([]byte(`data: {"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}`)) detail, ok := buffer.Detail() if !ok { t.Fatal("Detail() ok = false, want true") @@ -178,6 +178,57 @@ func TestStreamUsageBufferPreservesTierAcrossChunks(t *testing.T) { } } +func TestStreamUsageBufferObserveOpenAIStreamStateTransitions(t *testing.T) { + t.Parallel() + + t.Run("same chunk", func(t *testing.T) { + var buffer StreamUsageBuffer + buffer.ObserveOpenAIStream([]byte(`data: {"service_tier":"flex","usage":{"input_tokens":2,"output_tokens":3,"total_tokens":5}}`)) + detail, ok := buffer.Detail() + if !ok || detail.InputTokens != 2 || detail.ResponseServiceTier != "flex" { + t.Fatalf("detail = %+v ok=%v", detail, ok) + } + }) + + t.Run("usage before tier", func(t *testing.T) { + var buffer StreamUsageBuffer + buffer.ObserveOpenAIStream([]byte(`data: {"usage":{"input_tokens":2,"output_tokens":3,"total_tokens":5}}`)) + buffer.ObserveOpenAIStream([]byte(`data: {"service_tier":"default"}`)) + detail, ok := buffer.Detail() + if !ok || detail.InputTokens != 2 || detail.ResponseServiceTier != "default" { + t.Fatalf("detail = %+v ok=%v", detail, ok) + } + }) + + t.Run("final usage tier overrides early tier", func(t *testing.T) { + var buffer StreamUsageBuffer + buffer.ObserveOpenAIStream([]byte(`data: {"service_tier":"default"}`)) + buffer.ObserveOpenAIStream([]byte(`data: {"service_tier":"priority","usage":{"input_tokens":2,"output_tokens":3,"total_tokens":5}}`)) + detail, ok := buffer.Detail() + if !ok || detail.ResponseServiceTier != "priority" { + t.Fatalf("detail = %+v ok=%v", detail, ok) + } + }) + + t.Run("irrelevant and invalid chunks do not change state", func(t *testing.T) { + var buffer StreamUsageBuffer + buffer.ObserveOpenAIStream([]byte(`data: {"content":"the word \"usage\" appears here"}`)) + buffer.ObserveOpenAIStream([]byte(`data: {"usage":`)) + buffer.ObserveOpenAIStream([]byte(`data: {"usage":null}`)) + if detail, ok := buffer.Detail(); ok { + t.Fatalf("detail = %+v ok=true, want empty buffer", detail) + } + }) + + t.Run("zero token usage is retained", func(t *testing.T) { + var buffer StreamUsageBuffer + buffer.ObserveOpenAIStream([]byte(`data: {"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}`)) + if _, ok := buffer.Detail(); !ok { + t.Fatal("Detail() ok = false, want true") + } + }) +} + func TestStreamUsageBufferPreservesOnlyZeroUsage(t *testing.T) { var buffer StreamUsageBuffer buffer.Observe(usage.Detail{}, true) diff --git a/internal/runtime/executor/helps/usage_stream_benchmark_test.go b/internal/runtime/executor/helps/usage_stream_benchmark_test.go new file mode 100644 index 00000000..1d5b8dd7 --- /dev/null +++ b/internal/runtime/executor/helps/usage_stream_benchmark_test.go @@ -0,0 +1,31 @@ +package helps + +import "testing" + +var ( + benchmarkOpenAIContentChunk = []byte(`data: {"choices":[{"delta":{"content":"hello"}}]}`) + benchmarkOpenAITierChunk = []byte(`data: {"service_tier":"default","choices":[]}`) + benchmarkOpenAIUsageChunk = []byte(`data: {"usage":{"input_tokens":10,"output_tokens":20,"total_tokens":30}}`) +) + +func BenchmarkStreamUsageBufferObserveOpenAIStreamContentChunk(b *testing.B) { + var buffer StreamUsageBuffer + buffer.ObserveOpenAIStream(benchmarkOpenAITierChunk) + b.ReportAllocs() + b.ResetTimer() + for index := 0; index < b.N; index++ { + buffer.ObserveOpenAIStream(benchmarkOpenAIContentChunk) + } +} + +func BenchmarkStreamUsageBufferObserveOpenAIStream100Chunks(b *testing.B) { + b.ReportAllocs() + for index := 0; index < b.N; index++ { + var buffer StreamUsageBuffer + buffer.ObserveOpenAIStream(benchmarkOpenAITierChunk) + for chunk := 0; chunk < 98; chunk++ { + buffer.ObserveOpenAIStream(benchmarkOpenAIContentChunk) + } + buffer.ObserveOpenAIStream(benchmarkOpenAIUsageChunk) + } +} diff --git a/internal/runtime/executor/kimi_executor.go b/internal/runtime/executor/kimi_executor.go index 439a92c7..f0fb2170 100644 --- a/internal/runtime/executor/kimi_executor.go +++ b/internal/runtime/executor/kimi_executor.go @@ -293,7 +293,7 @@ func (e *KimiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut for scanner.Scan() { line := scanner.Bytes() helps.AppendAPIResponseChunk(ctx, e.cfg, line) - streamUsage.Observe(helps.ParseOpenAIStreamUsage(line)) + streamUsage.ObserveOpenAIStream(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 f67ddaf4..75881614 100644 --- a/internal/runtime/executor/openai_compat_executor.go +++ b/internal/runtime/executor/openai_compat_executor.go @@ -398,7 +398,7 @@ func (e *OpenAICompatExecutor) ExecuteStream(ctx context.Context, auth *cliproxy for scanner.Scan() { line := scanner.Bytes() helps.AppendAPIResponseChunk(ctx, e.cfg, line) - streamUsage.Observe(helps.ParseOpenAIStreamUsage(line)) + streamUsage.ObserveOpenAIStream(line) trimmedLine := bytes.TrimSpace(line) if len(trimmedLine) == 0 { continue