From 768b4c49fd44e1c6e21981d96dfcb8dd31490bf0 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Wed, 15 Jul 2026 21:30:09 +0800 Subject: [PATCH] feat(usage): track generate flag with backward-compatible defaults Propagate client generate metadata through usage records while keeping legacy callers that omit the field enabled by default. Preserve an existing context generate=false when metadata does not set the value, and normalize omission to true only at publishing boundaries. --- internal/pluginhost/adapters.go | 1 + internal/pluginhost/adapters_test.go | 49 +++++++++++++++++ internal/redisqueue/plugin.go | 2 + internal/redisqueue/plugin_test.go | 42 ++++++++++++++ .../runtime/executor/helps/usage_helpers.go | 7 ++- .../executor/helps/usage_helpers_test.go | 19 +++++++ sdk/api/handlers/handlers.go | 17 ++++++ sdk/api/handlers/handlers_metadata_test.go | 30 ++++++++++ sdk/cliproxy/auth/conductor.go | 19 +++++++ sdk/cliproxy/auth/conductor_usage_test.go | 29 ++++++++++ sdk/cliproxy/executor/types.go | 4 ++ sdk/cliproxy/usage/manager.go | 55 +++++++++++++++++-- sdk/cliproxy/usage/manager_test.go | 52 ++++++++++++++++++ sdk/pluginapi/types.go | 3 + 14 files changed, 321 insertions(+), 8 deletions(-) create mode 100644 sdk/cliproxy/usage/manager_test.go diff --git a/internal/pluginhost/adapters.go b/internal/pluginhost/adapters.go index 403a8c1f..385ffa16 100644 --- a/internal/pluginhost/adapters.go +++ b/internal/pluginhost/adapters.go @@ -1829,6 +1829,7 @@ func (a *usageAdapter) HandleUsage(ctx context.Context, record coreusage.Record) Source: record.Source, ReasoningEffort: record.ReasoningEffort, ServiceTier: record.ServiceTier, + Generate: coreusage.GenerateEnabled(record.Generate), RequestedAt: record.RequestedAt, Latency: record.Latency, TTFT: record.TTFT, diff --git a/internal/pluginhost/adapters_test.go b/internal/pluginhost/adapters_test.go index 6817d0a9..4512be3c 100644 --- a/internal/pluginhost/adapters_test.go +++ b/internal/pluginhost/adapters_test.go @@ -2138,6 +2138,55 @@ func TestUsageAdapterPanicFusesPlugin(t *testing.T) { } } +func TestUsageAdapterNormalizesOmittedGenerateToTrue(t *testing.T) { + var gotGenerate bool + plugin := usagePluginFunc(func(ctx context.Context, record pluginapi.UsageRecord) { + gotGenerate = record.Generate + }) + host := newHostWithRecords(capabilityRecord{ + id: "usage-generate", + plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + UsagePlugin: plugin, + }}, + }) + adapter := &usageAdapter{ + host: host, + pluginID: "usage-generate", + } + + // Legacy callers construct usage.Record without Generate; adapter must publish true. + adapter.HandleUsage(context.Background(), coreusage.Record{Provider: "provider", Model: "gpt-5.4"}) + if !gotGenerate { + t.Fatalf("plugin Generate = %v, want true for omitted field", gotGenerate) + } +} + +func TestUsageAdapterPreservesExplicitGenerateFalse(t *testing.T) { + var gotGenerate bool + plugin := usagePluginFunc(func(ctx context.Context, record pluginapi.UsageRecord) { + gotGenerate = record.Generate + }) + host := newHostWithRecords(capabilityRecord{ + id: "usage-generate-false", + plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + UsagePlugin: plugin, + }}, + }) + adapter := &usageAdapter{ + host: host, + pluginID: "usage-generate-false", + } + + adapter.HandleUsage(context.Background(), coreusage.Record{ + Provider: "provider", + Model: "gpt-5.4", + Generate: coreusage.GenerateFlag(false), + }) + if gotGenerate { + t.Fatalf("plugin Generate = %v, want false", gotGenerate) + } +} + func TestUsageManagerRegisterNamedReplacesWithoutDuplicateDispatch(t *testing.T) { manager := coreusage.NewManager(0) defer manager.Stop() diff --git a/internal/redisqueue/plugin.go b/internal/redisqueue/plugin.go index 78a17322..0968606f 100644 --- a/internal/redisqueue/plugin.go +++ b/internal/redisqueue/plugin.go @@ -96,6 +96,7 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec AuthIndex: record.AuthIndex, Tokens: tokens, Failed: failed, + Generate: coreusage.GenerateEnabled(record.Generate), Fail: fail, ResponseHeaders: record.ResponseHeaders, } @@ -143,6 +144,7 @@ type requestDetail struct { AuthIndex string `json:"auth_index"` Tokens tokenStats `json:"tokens"` Failed bool `json:"failed"` + Generate bool `json:"generate"` 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 32873ab8..3db9189f 100644 --- a/internal/redisqueue/plugin_test.go +++ b/internal/redisqueue/plugin_test.go @@ -36,6 +36,7 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) { ReasoningEffort: "medium", ServiceTier: "auto", ResponseServiceTier: "default", + Generate: coreusage.GenerateFlag(true), RequestedAt: time.Date(2026, 4, 25, 0, 0, 0, 0, time.UTC), Latency: 1500 * time.Millisecond, Detail: coreusage.Detail{ @@ -64,10 +65,51 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) { requireHeaderField(t, payload, "response_headers", "X-Upstream-Request-Id", []string{"upstream-req-1"}) requireHeaderField(t, payload, "response_headers", "Retry-After", []string{"30"}) requireBoolField(t, payload, "failed", false) + requireBoolField(t, payload, "generate", true) requireFailField(t, payload, http.StatusOK, "") }) } +func TestUsageQueuePluginPayloadIncludesGenerateFalse(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", + Generate: coreusage.GenerateFlag(false), + Detail: coreusage.Detail{ + InputTokens: 1, + TotalTokens: 1, + }, + }) + + payload := popSinglePayload(t) + requireBoolField(t, payload, "generate", false) + }) +} + +func TestUsageQueuePluginPayloadDefaultsGenerateTrueWhenOmitted(t *testing.T) { + withEnabledQueue(t, func() { + ctx := internallogging.WithResponseStatusHolder(context.Background()) + internallogging.SetResponseStatus(ctx, http.StatusOK) + + // Legacy callers construct usage.Record without Generate; omission must publish as true. + (&usageQueuePlugin{}).HandleUsage(ctx, coreusage.Record{ + Provider: "openai", + Model: "gpt-5.4", + Detail: coreusage.Detail{ + InputTokens: 1, + TotalTokens: 1, + }, + }) + + payload := popSinglePayload(t) + requireBoolField(t, payload, "generate", true) + }) +} + func TestUsageQueuePluginMarksCanonicalZeroCacheRead(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 d3d9caee..c4b8d827 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -33,6 +33,7 @@ type UsageReporter struct { source string reasoning string serviceTier string + generate bool requestedAt time.Time ttftMu sync.RWMutex ttft time.Duration @@ -71,6 +72,7 @@ func NewUsageReporter(ctx context.Context, provider, model string, auth *cliprox authType: resolveUsageAuthType(auth), reasoning: usage.ReasoningEffortFromContext(ctx), serviceTier: usage.ServiceTierFromContext(ctx), + generate: usage.GenerateFromContext(ctx), } if auth != nil { reporter.authID = auth.ID @@ -249,14 +251,14 @@ func (r *UsageReporter) buildRecord(detail usage.Detail, failed bool, failures . fail = failures[0] } if r == nil { - return usage.Record{Detail: detail, Failed: failed, Fail: fail} + return usage.Record{Detail: detail, Failed: failed, Fail: fail, Generate: usage.GenerateFlag(true)} } return r.buildRecordForModel(r.model, detail, failed, fail) } func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, failed bool, fail usage.Failure) usage.Record { if r == nil { - return usage.Record{Model: model, Detail: detail, Failed: failed, Fail: fail} + return usage.Record{Model: model, Detail: detail, Failed: failed, Fail: fail, Generate: usage.GenerateFlag(true)} } return usage.Record{ Provider: r.provider, @@ -271,6 +273,7 @@ func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, f ReasoningEffort: r.reasoning, ServiceTier: r.serviceTier, ResponseServiceTier: strings.TrimSpace(detail.ResponseServiceTier), + Generate: usage.GenerateFlag(r.generate), 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 a08cfbf2..c54b38e4 100644 --- a/internal/runtime/executor/helps/usage_helpers_test.go +++ b/internal/runtime/executor/helps/usage_helpers_test.go @@ -469,6 +469,25 @@ func TestUsageReporterBuildRecordIncludesServiceTier(t *testing.T) { } } +func TestUsageReporterBuildRecordDefaultsGenerateTrue(t *testing.T) { + reporter := NewUsageReporter(context.Background(), "openai", "gpt-5.4", nil) + + record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false) + if !usage.GenerateEnabled(record.Generate) { + t.Fatalf("generate = %v, want true", usage.GenerateEnabled(record.Generate)) + } +} + +func TestUsageReporterBuildRecordIncludesGenerateFalse(t *testing.T) { + ctx := usage.WithGenerate(context.Background(), false) + reporter := NewUsageReporter(ctx, "openai", "gpt-5.4", nil) + + record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false) + if usage.GenerateEnabled(record.Generate) { + t.Fatalf("generate = %v, want false", usage.GenerateEnabled(record.Generate)) + } +} + 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.go b/sdk/api/handlers/handlers.go index 06965042..f3ab7178 100644 --- a/sdk/api/handlers/handlers.go +++ b/sdk/api/handlers/handlers.go @@ -329,6 +329,19 @@ func setServiceTierMetadata(meta map[string]any, rawJSON []byte) { meta[coreexecutor.ServiceTierMetadataKey] = serviceTier } +func setGenerateMetadata(meta map[string]any, rawJSON []byte) { + if meta == nil { + return + } + // Missing or true means generation is enabled; only an explicit false disables generation. + generate := true + node := gjson.GetBytes(rawJSON, "generate") + if node.Exists() && node.IsBool() && !node.Bool() { + generate = false + } + meta[coreexecutor.GenerateMetadataKey] = generate +} + // headersFromContext extracts the original HTTP request headers from the gin context // embedded in the provided context. This allows session affinity selectors to read // client-provided session headers. @@ -739,6 +752,7 @@ func (h *BaseAPIHandler) executeWithAuthManagerFormats(ctx context.Context, entr addModelExecutionSourceMetadata(reqMeta, execOptions.InternalSource) setReasoningEffortMetadata(reqMeta, entryProtocol, normalizedModel, rawJSON) setServiceTierMetadata(reqMeta, rawJSON) + setGenerateMetadata(reqMeta, rawJSON) payload := rawJSON if len(payload) == 0 { payload = nil @@ -806,6 +820,7 @@ func (h *BaseAPIHandler) executeCountWithAuthManager(ctx context.Context, handle addAuthSelectionModelMetadata(reqMeta, execOptions.AuthSelectionModel) setReasoningEffortMetadata(reqMeta, handlerType, normalizedModel, rawJSON) setServiceTierMetadata(reqMeta, rawJSON) + setGenerateMetadata(reqMeta, rawJSON) payload := rawJSON if len(payload) == 0 { payload = nil @@ -893,6 +908,7 @@ func (h *BaseAPIHandler) pluginExecutorRequest(ctx context.Context, entryProtoco addModelExecutionSourceMetadata(reqMeta, execOptions.InternalSource) setReasoningEffortMetadata(reqMeta, entryProtocol, modelName, rawJSON) setServiceTierMetadata(reqMeta, rawJSON) + setGenerateMetadata(reqMeta, rawJSON) payload := rawJSON if len(payload) == 0 { payload = nil @@ -1140,6 +1156,7 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context addModelExecutionSourceMetadata(reqMeta, execOptions.InternalSource) setReasoningEffortMetadata(reqMeta, entryProtocol, normalizedModel, rawJSON) setServiceTierMetadata(reqMeta, rawJSON) + setGenerateMetadata(reqMeta, rawJSON) payload := rawJSON if len(payload) == 0 { payload = nil diff --git a/sdk/api/handlers/handlers_metadata_test.go b/sdk/api/handlers/handlers_metadata_test.go index 107c09f5..60022b56 100644 --- a/sdk/api/handlers/handlers_metadata_test.go +++ b/sdk/api/handlers/handlers_metadata_test.go @@ -70,3 +70,33 @@ func TestSetServiceTierMetadataPreservesExplicitDefault(t *testing.T) { t.Fatalf("ServiceTierMetadataKey = %v, want %q", gotServiceTier, "default") } } + +func TestSetGenerateMetadataDefaultsWhenMissing(t *testing.T) { + meta := make(map[string]any) + + setGenerateMetadata(meta, []byte(`{"model":"gpt-5.4"}`)) + + if got := meta[coreexecutor.GenerateMetadataKey]; got != true { + t.Fatalf("GenerateMetadataKey = %v, want true", got) + } +} + +func TestSetGenerateMetadataPreservesTrue(t *testing.T) { + meta := make(map[string]any) + + setGenerateMetadata(meta, []byte(`{"generate":true}`)) + + if got := meta[coreexecutor.GenerateMetadataKey]; got != true { + t.Fatalf("GenerateMetadataKey = %v, want true", got) + } +} + +func TestSetGenerateMetadataHonorsExplicitFalse(t *testing.T) { + meta := make(map[string]any) + + setGenerateMetadata(meta, []byte(`{"generate":false}`)) + + if got := meta[coreexecutor.GenerateMetadataKey]; got != false { + t.Fatalf("GenerateMetadataKey = %v, want false", got) + } +} diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 04e87f26..8164749b 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -3050,6 +3050,9 @@ func contextWithRequestedModelAlias(ctx context.Context, opts cliproxyexecutor.O if serviceTier != "" { ctx = coreusage.WithServiceTier(ctx, serviceTier) } + if generate, ok := generateFromOptions(opts); ok { + ctx = coreusage.WithGenerate(ctx, generate) + } return ctx } @@ -3100,6 +3103,22 @@ func serviceTierFromOptions(opts cliproxyexecutor.Options) string { return stringMetadataValue(opts.Metadata, cliproxyexecutor.ServiceTierMetadataKey) } +func generateFromOptions(opts cliproxyexecutor.Options) (bool, bool) { + if len(opts.Metadata) == 0 { + return false, false + } + raw, ok := opts.Metadata[cliproxyexecutor.GenerateMetadataKey] + if !ok || raw == nil { + return false, false + } + switch value := raw.(type) { + case bool: + return value, true + default: + return false, false + } +} + func stringMetadataValue(metadata map[string]any, key string) string { if len(metadata) == 0 { return "" diff --git a/sdk/cliproxy/auth/conductor_usage_test.go b/sdk/cliproxy/auth/conductor_usage_test.go index 90a55ceb..91aa237c 100644 --- a/sdk/cliproxy/auth/conductor_usage_test.go +++ b/sdk/cliproxy/auth/conductor_usage_test.go @@ -14,6 +14,7 @@ func TestContextWithRequestedModelAliasIncludesReasoningEffort(t *testing.T) { cliproxyexecutor.RequestedModelMetadataKey: "client-model", cliproxyexecutor.ReasoningEffortMetadataKey: "medium", cliproxyexecutor.ServiceTierMetadataKey: "auto", + cliproxyexecutor.GenerateMetadataKey: false, }, }, "fallback-model") @@ -27,4 +28,32 @@ func TestContextWithRequestedModelAliasIncludesReasoningEffort(t *testing.T) { if gotServiceTier != "auto" { t.Fatalf("service tier = %q, want %q", gotServiceTier, "auto") } + if got := coreusage.GenerateFromContext(ctx); got { + t.Fatalf("generate = %v, want false", got) + } +} + +func TestContextWithRequestedModelAliasDefaultsGenerateTrue(t *testing.T) { + ctx := contextWithRequestedModelAlias(context.Background(), cliproxyexecutor.Options{ + Metadata: map[string]any{ + cliproxyexecutor.RequestedModelMetadataKey: "client-model", + }, + }, "fallback-model") + + if got := coreusage.GenerateFromContext(ctx); !got { + t.Fatalf("generate = %v, want true", got) + } +} + +func TestContextWithRequestedModelAliasPreservesExistingGenerateFalse(t *testing.T) { + ctx := coreusage.WithGenerate(context.Background(), false) + ctx = contextWithRequestedModelAlias(ctx, cliproxyexecutor.Options{ + Metadata: map[string]any{ + cliproxyexecutor.RequestedModelMetadataKey: "client-model", + }, + }, "fallback-model") + + if got := coreusage.GenerateFromContext(ctx); got { + t.Fatalf("generate = %v, want false", got) + } } diff --git a/sdk/cliproxy/executor/types.go b/sdk/cliproxy/executor/types.go index ae3f1881..c9e9a346 100644 --- a/sdk/cliproxy/executor/types.go +++ b/sdk/cliproxy/executor/types.go @@ -27,6 +27,10 @@ const ReasoningEffortMetadataKey = "reasoning_effort" // ServiceTierMetadataKey stores the client-requested service tier for usage logs. const ServiceTierMetadataKey = "service_tier" +// GenerateMetadataKey stores whether the client requested actual generation for usage logs. +// Missing or true means generation is enabled; only an explicit false disables generation. +const GenerateMetadataKey = "generate" + const ( // PinnedAuthMetadataKey locks execution to a specific auth ID. PinnedAuthMetadataKey = "pinned_auth_id" diff --git a/sdk/cliproxy/usage/manager.go b/sdk/cliproxy/usage/manager.go index bbc83ff5..0c9529f1 100644 --- a/sdk/cliproxy/usage/manager.go +++ b/sdk/cliproxy/usage/manager.go @@ -39,12 +39,16 @@ type Record struct { RequestServiceTier string // ResponseServiceTier stores the final tier reported by the upstream response. ResponseServiceTier string - RequestedAt time.Time - Latency time.Duration - TTFT time.Duration - Failed bool - Fail Failure - Detail Detail + // 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 + RequestedAt time.Time + Latency time.Duration + TTFT time.Duration + Failed bool + Fail Failure + Detail Detail // ResponseHeaders stores a snapshot of upstream response headers for usage sinks. ResponseHeaders http.Header } @@ -70,6 +74,7 @@ type Detail struct { type requestedModelAliasContextKey struct{} type reasoningEffortContextKey struct{} type serviceTierContextKey struct{} +type generateContextKey struct{} // WithRequestedModelAlias stores the client-requested model name for usage sinks. func WithRequestedModelAlias(ctx context.Context, alias string) context.Context { @@ -163,6 +168,44 @@ func ServiceTierFromContext(ctx context.Context) string { } } +// WithGenerate stores whether the client requested actual generation for usage sinks. +// Missing context values default to true; only an explicit false disables generation. +func WithGenerate(ctx context.Context, generate bool) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, generateContextKey{}, generate) +} + +// GenerateFromContext returns whether the client requested actual generation. +// Missing values default to true. +func GenerateFromContext(ctx context.Context) bool { + if ctx == nil { + return true + } + raw := ctx.Value(generateContextKey{}) + switch value := raw.(type) { + case bool: + return value + default: + return true + } +} + +// GenerateFlag returns a pointer suitable for Record.Generate. +func GenerateFlag(generate bool) *bool { + return &generate +} + +// GenerateEnabled reports whether generation is enabled for the record field. +// A nil value defaults to true so legacy callers that omit Generate keep the historical behavior. +func GenerateEnabled(generate *bool) bool { + if generate == nil { + return true + } + return *generate +} + // Plugin consumes usage records emitted by the proxy runtime. type Plugin interface { HandleUsage(ctx context.Context, record Record) diff --git a/sdk/cliproxy/usage/manager_test.go b/sdk/cliproxy/usage/manager_test.go new file mode 100644 index 00000000..6f7b1fbb --- /dev/null +++ b/sdk/cliproxy/usage/manager_test.go @@ -0,0 +1,52 @@ +package usage + +import ( + "context" + "testing" +) + +func TestGenerateEnabledDefaultsNilToTrue(t *testing.T) { + if !GenerateEnabled(nil) { + t.Fatalf("GenerateEnabled(nil) = false, want true") + } +} + +func TestGenerateEnabledHonorsExplicitFalse(t *testing.T) { + if GenerateEnabled(GenerateFlag(false)) { + t.Fatalf("GenerateEnabled(false) = true, want false") + } +} + +func TestGenerateEnabledHonorsExplicitTrue(t *testing.T) { + if !GenerateEnabled(GenerateFlag(true)) { + t.Fatalf("GenerateEnabled(true) = false, want true") + } +} + +func TestGenerateFromContextDefaultsMissingToTrue(t *testing.T) { + if !GenerateFromContext(context.Background()) { + t.Fatalf("GenerateFromContext(background) = false, want true") + } +} + +func TestGenerateFromContextHonorsExplicitFalse(t *testing.T) { + ctx := WithGenerate(context.Background(), false) + if GenerateFromContext(ctx) { + t.Fatalf("GenerateFromContext(false) = true, want false") + } +} + +func TestRecordOmittedGenerateIsEnabled(t *testing.T) { + // Existing callers construct Record without setting Generate. + // Omission must remain distinguishable from explicit false and default to true. + record := Record{ + Provider: "openai", + Model: "gpt-5.4", + } + if record.Generate != nil { + t.Fatalf("Record.Generate = %v, want nil for omitted field", record.Generate) + } + if !GenerateEnabled(record.Generate) { + t.Fatalf("GenerateEnabled(omitted) = false, want true") + } +} diff --git a/sdk/pluginapi/types.go b/sdk/pluginapi/types.go index 5bd97508..ae59a210 100644 --- a/sdk/pluginapi/types.go +++ b/sdk/pluginapi/types.go @@ -1276,6 +1276,9 @@ type UsageRecord struct { ReasoningEffort string // ServiceTier records the requested or reported service tier. ServiceTier string + // Generate reports whether the client requested actual generation. + // The host normalizes omitted usage.Record values to true before delivery. + Generate bool // RequestedAt is the time the request was received. RequestedAt time.Time // Latency is the total request latency.