From 9721d9939ed5aeb20dd3db811f1fa9b9c9e66cba Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Mon, 31 Aug 2026 20:29:05 +0800 Subject: [PATCH] feat(usage): track streaming execution state in usage records - Add `Stream` field to usage records, context helpers, and Redis queue payloads. - Propagate streaming mode across execution handlers, conductors, and usage reporters. Closes: #5361 --- internal/redisqueue/plugin.go | 7 +++++ internal/redisqueue/plugin_test.go | 25 ++++++++++++++++ .../runtime/executor/helps/usage_helpers.go | 11 +++++++ .../executor/helps/usage_helpers_test.go | 29 +++++++++++++++++++ sdk/api/handlers/handlers_execution.go | 3 +- .../handlers_plugin_executor_usage_test.go | 6 ++++ sdk/api/handlers/handlers_stream.go | 3 +- sdk/cliproxy/auth/conductor_execution.go | 1 + sdk/cliproxy/auth/conductor_usage_test.go | 14 +++++++++ sdk/cliproxy/usage/manager.go | 28 +++++++++++++++++- sdk/cliproxy/usage/manager_test.go | 24 +++++++++++++++ 11 files changed, 148 insertions(+), 3 deletions(-) diff --git a/internal/redisqueue/plugin.go b/internal/redisqueue/plugin.go index d91c8a28..5a08e3ac 100644 --- a/internal/redisqueue/plugin.go +++ b/internal/redisqueue/plugin.go @@ -84,6 +84,11 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec } fail := resolveFail(ctx, record, failed) + stream := record.Stream + if !stream { + stream = coreusage.StreamFromContext(ctx) + } + detail := requestDetail{ Timestamp: timestamp, LatencyMs: record.Latency.Milliseconds(), @@ -97,6 +102,7 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec Tokens: tokens, Failed: failed, Generate: coreusage.GenerateEnabled(record.Generate), + Stream: stream, Fail: fail, ResponseHeaders: record.ResponseHeaders, } @@ -153,6 +159,7 @@ type requestDetail struct { Tokens tokenStats `json:"tokens"` Failed bool `json:"failed"` Generate bool `json:"generate"` + Stream bool `json:"stream"` Fail failDetail `json:"fail"` ResponseHeaders http.Header `json:"response_headers,omitempty"` } diff --git a/internal/redisqueue/plugin_test.go b/internal/redisqueue/plugin_test.go index c1a1f010..449c6583 100644 --- a/internal/redisqueue/plugin_test.go +++ b/internal/redisqueue/plugin_test.go @@ -78,6 +78,7 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) { requireHeaderField(t, payload, "response_headers", "Retry-After", []string{"30"}) requireBoolField(t, payload, "failed", false) requireBoolField(t, payload, "generate", true) + requireBoolField(t, payload, "stream", false) requireFailField(t, payload, http.StatusOK, "") }) } @@ -154,6 +155,30 @@ func TestUsageQueuePluginPayloadDefaultsGenerateTrueWhenOmitted(t *testing.T) { }) } +func TestUsageQueuePluginPublishesStreamFlag(t *testing.T) { + for _, stream := range []bool{true, false} { + t.Run(map[bool]string{true: "stream_true", false: "stream_false"}[stream], func(t *testing.T) { + withEnabledQueue(t, func() { + ctx := internallogging.WithResponseStatusHolder(context.Background()) + internallogging.SetResponseStatus(ctx, http.StatusOK) + + (&usageQueuePlugin{}).HandleUsage(ctx, coreusage.Record{ + Provider: "openai", + Model: "gpt-5.4", + Stream: stream, + Detail: coreusage.Detail{ + InputTokens: 1, + TotalTokens: 1, + }, + }) + + payload := popSinglePayload(t) + requireBoolField(t, payload, "stream", stream) + }) + }) + } +} + func TestUsageQueuePluginPreservesLegacyCachedOnlyUsage(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithResponseStatusHolder(context.Background()) diff --git a/internal/runtime/executor/helps/usage_helpers.go b/internal/runtime/executor/helps/usage_helpers.go index cfb2dd2d..bee80cc1 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -38,6 +38,7 @@ type UsageReporter struct { reasoning string serviceTier string generate bool + stream bool requestedAt time.Time ttftMu sync.RWMutex ttft time.Duration @@ -79,6 +80,7 @@ func NewUsageReporter(ctx context.Context, provider, model string, auth *cliprox reasoning: usage.ReasoningEffortFromContext(ctx), serviceTier: usage.ServiceTierFromContext(ctx), generate: usage.GenerateFromContext(ctx), + stream: usage.StreamFromContext(ctx), } if auth != nil { reporter.authID = auth.ID @@ -88,6 +90,14 @@ func NewUsageReporter(ctx context.Context, provider, model string, auth *cliprox return reporter } +// SetStream records whether the request was executed in streaming mode. +func (r *UsageReporter) SetStream(stream bool) { + if r == nil { + return + } + r.stream = stream +} + // UpdateAccessTokenFingerprint records the token version actually used upstream. func (r *UsageReporter) UpdateAccessTokenFingerprint(auth *cliproxyauth.Auth) { if r == nil { @@ -391,6 +401,7 @@ func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, f ServiceTier: r.serviceTier, ResponseServiceTier: strings.TrimSpace(detail.ResponseServiceTier), Generate: usage.GenerateFlag(r.generate), + Stream: r.stream, RequestedAt: r.requestedAt, Latency: r.latency(), TTFT: r.ttftDuration(), diff --git a/internal/runtime/executor/helps/usage_helpers_test.go b/internal/runtime/executor/helps/usage_helpers_test.go index 82eefcda..f97a7f96 100644 --- a/internal/runtime/executor/helps/usage_helpers_test.go +++ b/internal/runtime/executor/helps/usage_helpers_test.go @@ -789,6 +789,35 @@ func TestUsageReporterBuildRecordIncludesGenerateFalse(t *testing.T) { } } +func TestUsageReporterBuildRecordDefaultsStreamFalse(t *testing.T) { + reporter := NewUsageReporter(context.Background(), "openai", "gpt-5.4", nil) + + record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false) + if record.Stream { + t.Fatalf("stream = %v, want false", record.Stream) + } +} + +func TestUsageReporterBuildRecordIncludesStreamTrue(t *testing.T) { + ctx := usage.WithStream(context.Background(), true) + reporter := NewUsageReporter(ctx, "openai", "gpt-5.4", nil) + + record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false) + if !record.Stream { + t.Fatalf("stream = %v, want true", record.Stream) + } +} + +func TestUsageReporterSetStream(t *testing.T) { + reporter := NewUsageReporter(context.Background(), "openai", "gpt-5.4", nil) + reporter.SetStream(true) + + record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false) + if !record.Stream { + t.Fatalf("stream = %v, want true", record.Stream) + } +} + func TestUsageReporterSetTranslatedReasoningEffortPreservesClientServiceTier(t *testing.T) { ctx := usage.WithServiceTier(context.Background(), "auto") reporter := NewUsageReporter(ctx, "openai", "gpt-5.4", nil) diff --git a/sdk/api/handlers/handlers_execution.go b/sdk/api/handlers/handlers_execution.go index 79632acc..77340df7 100644 --- a/sdk/api/handlers/handlers_execution.go +++ b/sdk/api/handlers/handlers_execution.go @@ -10,6 +10,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" "golang.org/x/net/context" @@ -179,7 +180,7 @@ func (h *BaseAPIHandler) executeWithPluginExecutor(ctx context.Context, entryPro if host == nil { return nil, nil, &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("plugin executor host is unavailable")} } - execCtx, nestedTracker := withNestedExecutionTracker(ctx) + execCtx, nestedTracker := withNestedExecutionTracker(coreusage.WithStream(ctx, false)) req, opts := h.pluginExecutorRequest(execCtx, entryProtocol, responseProtocol, modelName, originalRequestedModel, rawJSON, alt, false, execOptions) lifecycle := h.newRequestLifecycleTracker(execCtx, entryProtocol, modelName, originalRequestedModel, false, opts.Metadata, execOptions.SkipInterceptorPluginID) var interceptErr *interfaces.ErrorMessage diff --git a/sdk/api/handlers/handlers_plugin_executor_usage_test.go b/sdk/api/handlers/handlers_plugin_executor_usage_test.go index 5d9c4f12..f5e3f022 100644 --- a/sdk/api/handlers/handlers_plugin_executor_usage_test.go +++ b/sdk/api/handlers/handlers_plugin_executor_usage_test.go @@ -103,6 +103,9 @@ func TestHandlerPluginExecutorPublishesUsageNonStreamOpenAI(t *testing.T) { if record.Provider != targetPluginID { t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID) } + if record.Stream { + t.Errorf("record.Stream = true, want false") + } if record.Detail.InputTokens != 12 || record.Detail.OutputTokens != 34 || record.Detail.TotalTokens != 46 { t.Errorf("record.Detail = %+v, want prompt=12 completion=34 total=46", record.Detail) } @@ -145,6 +148,9 @@ func TestHandlerPluginExecutorPublishesUsageStreamOpenAI(t *testing.T) { if record.Provider != targetPluginID { t.Errorf("record.Provider = %q, want %q", record.Provider, targetPluginID) } + if !record.Stream { + t.Errorf("record.Stream = false, want true") + } if record.Detail.InputTokens != 15 || record.Detail.OutputTokens != 25 || record.Detail.TotalTokens != 40 { t.Errorf("record.Detail = %+v, want prompt=15 completion=25 total=40", record.Detail) } diff --git a/sdk/api/handlers/handlers_stream.go b/sdk/api/handlers/handlers_stream.go index f476e104..f1cc74eb 100644 --- a/sdk/api/handlers/handlers_stream.go +++ b/sdk/api/handlers/handlers_stream.go @@ -9,6 +9,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" "golang.org/x/net/context" @@ -40,7 +41,7 @@ func (h *BaseAPIHandler) streamWithPluginExecutor(ctx context.Context, entryProt close(errChan) return nil, nil, errChan } - execCtx, nestedTracker := withNestedExecutionTracker(ctx) + execCtx, nestedTracker := withNestedExecutionTracker(coreusage.WithStream(ctx, true)) req, opts := h.pluginExecutorRequest(execCtx, entryProtocol, responseProtocol, modelName, originalRequestedModel, rawJSON, alt, true, execOptions) lifecycle := h.newRequestLifecycleTracker(execCtx, entryProtocol, modelName, originalRequestedModel, true, opts.Metadata, execOptions.SkipInterceptorPluginID) var interceptErr *interfaces.ErrorMessage diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index 3aa1f67d..188fb054 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -1421,6 +1421,7 @@ func contextWithRequestedModelAlias(ctx context.Context, opts cliproxyexecutor.O if generate, ok := generateFromOptions(opts); ok { ctx = coreusage.WithGenerate(ctx, generate) } + ctx = coreusage.WithStream(ctx, opts.Stream) return ctx } diff --git a/sdk/cliproxy/auth/conductor_usage_test.go b/sdk/cliproxy/auth/conductor_usage_test.go index 91aa237c..f3c32c2f 100644 --- a/sdk/cliproxy/auth/conductor_usage_test.go +++ b/sdk/cliproxy/auth/conductor_usage_test.go @@ -8,6 +8,20 @@ import ( coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" ) +func TestContextWithRequestedModelAliasIncludesStream(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(map[bool]string{false: "stream_false", true: "stream_true"}[stream], func(t *testing.T) { + ctx := contextWithRequestedModelAlias(context.Background(), cliproxyexecutor.Options{ + Stream: stream, + }, "fallback-model") + + if got := coreusage.StreamFromContext(ctx); got != stream { + t.Fatalf("stream = %v, want %v", got, stream) + } + }) + } +} + func TestContextWithRequestedModelAliasIncludesReasoningEffort(t *testing.T) { ctx := contextWithRequestedModelAlias(context.Background(), cliproxyexecutor.Options{ Metadata: map[string]any{ diff --git a/sdk/cliproxy/usage/manager.go b/sdk/cliproxy/usage/manager.go index ca36dc55..4338dba7 100644 --- a/sdk/cliproxy/usage/manager.go +++ b/sdk/cliproxy/usage/manager.go @@ -44,7 +44,9 @@ type Record struct { // Generate reports whether the client requested actual generation. // nil or true means generation is enabled; only an explicit false disables generation. // Use GenerateFlag to set the value and GenerateEnabled to read it with the default. - Generate *bool + Generate *bool + // Stream reports whether the request was executed in streaming mode. + Stream bool RequestedAt time.Time Latency time.Duration TTFT time.Duration @@ -78,6 +80,7 @@ type requestedModelAliasContextKey struct{} type reasoningEffortContextKey struct{} type serviceTierContextKey struct{} type generateContextKey struct{} +type streamContextKey struct{} // WithRequestedModelAlias stores the client-requested model name for usage sinks. func WithRequestedModelAlias(ctx context.Context, alias string) context.Context { @@ -195,6 +198,29 @@ func GenerateFromContext(ctx context.Context) bool { } } +// WithStream stores whether the request was executed in streaming mode for usage sinks. +func WithStream(ctx context.Context, stream bool) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, streamContextKey{}, stream) +} + +// StreamFromContext returns whether the request was executed in streaming mode. +// Missing values default to false. +func StreamFromContext(ctx context.Context) bool { + if ctx == nil { + return false + } + raw := ctx.Value(streamContextKey{}) + switch value := raw.(type) { + case bool: + return value + default: + return false + } +} + // GenerateFlag returns a pointer suitable for Record.Generate. func GenerateFlag(generate bool) *bool { return &generate diff --git a/sdk/cliproxy/usage/manager_test.go b/sdk/cliproxy/usage/manager_test.go index 6f7b1fbb..f1d19284 100644 --- a/sdk/cliproxy/usage/manager_test.go +++ b/sdk/cliproxy/usage/manager_test.go @@ -5,6 +5,30 @@ import ( "testing" ) +func TestStreamFromContextDefaultsMissingToFalse(t *testing.T) { + if StreamFromContext(context.Background()) { + t.Fatalf("StreamFromContext(background) = true, want false") + } +} + +func TestStreamFromContextHonorsExplicitTrue(t *testing.T) { + ctx := WithStream(context.Background(), true) + if !StreamFromContext(ctx) { + t.Fatalf("StreamFromContext(true) = false, want true") + } +} + +func TestRecordStreamField(t *testing.T) { + record := Record{ + Provider: "openai", + Model: "gpt-5.4", + Stream: true, + } + if !record.Stream { + t.Fatalf("Record.Stream = false, want true") + } +} + func TestGenerateEnabledDefaultsNilToTrue(t *testing.T) { if !GenerateEnabled(nil) { t.Fatalf("GenerateEnabled(nil) = false, want true") -- 2.51.2