diff --git a/internal/pluginhost/adapters_test.go b/internal/pluginhost/adapters_test.go index 302b1f38..610a3753 100644 --- a/internal/pluginhost/adapters_test.go +++ b/internal/pluginhost/adapters_test.go @@ -2266,6 +2266,54 @@ func TestUsageAdapterDetachesCanceledRequestContext(t *testing.T) { } } +func TestUsageAdapterIncludesRequestIdentityForSchemaFour(t *testing.T) { + var got pluginapi.UsageRecord + plugin := usagePluginFunc(func(ctx context.Context, record pluginapi.UsageRecord) { + got = record + }) + host := newHostWithRecords(capabilityRecord{ + id: "usage-identity", + plugin: pluginapi.Plugin{SchemaVersion: pluginabi.SchemaVersionUsageIdentity, Capabilities: pluginapi.Capabilities{ + UsagePlugin: plugin, + }}, + }) + adapter := &usageAdapter{ + host: host, + pluginID: "usage-identity", + schemaVersion: pluginabi.SchemaVersionUsageIdentity, + } + + ctx := coreusage.WithRequestIdentity(context.Background(), "request-1", "trace-1") + adapter.HandleUsage(ctx, coreusage.Record{Provider: "provider"}) + if got.RequestID != "request-1" || got.TraceID != "trace-1" { + t.Fatalf("usage identity = %q, %q", got.RequestID, got.TraceID) + } +} + +func TestUsageAdapterOmitsRequestIdentityForLegacySchema(t *testing.T) { + var got pluginapi.UsageRecord + plugin := usagePluginFunc(func(ctx context.Context, record pluginapi.UsageRecord) { + got = record + }) + host := newHostWithRecords(capabilityRecord{ + id: "usage-identity-legacy", + plugin: pluginapi.Plugin{SchemaVersion: pluginabi.SchemaVersionUsageIdentity - 1, Capabilities: pluginapi.Capabilities{ + UsagePlugin: plugin, + }}, + }) + adapter := &usageAdapter{ + host: host, + pluginID: "usage-identity-legacy", + schemaVersion: pluginabi.SchemaVersionUsageIdentity - 1, + } + + ctx := coreusage.WithRequestIdentity(context.Background(), "request-1", "trace-1") + adapter.HandleUsage(ctx, coreusage.Record{Provider: "provider"}) + if got.RequestID != "" || got.TraceID != "" { + t.Fatalf("legacy usage identity = %q, %q", got.RequestID, got.TraceID) + } +} + func TestUsageAdapterNormalizesOmittedGenerateToTrue(t *testing.T) { var gotGenerate bool plugin := usagePluginFunc(func(ctx context.Context, record pluginapi.UsageRecord) { diff --git a/internal/pluginhost/adapters_usage_translation.go b/internal/pluginhost/adapters_usage_translation.go index 93a4c8a5..ad046780 100644 --- a/internal/pluginhost/adapters_usage_translation.go +++ b/internal/pluginhost/adapters_usage_translation.go @@ -10,6 +10,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginabi" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" log "github.com/sirupsen/logrus" @@ -26,9 +27,10 @@ func (h *Host) RegisterUsagePlugins() { continue } coreusage.RegisterNamedPlugin("plugin:"+record.id, &usageAdapter{ - host: h, - pluginID: record.id, - plugin: plugin, + host: h, + pluginID: record.id, + plugin: plugin, + schemaVersion: record.plugin.SchemaVersion, }) } } @@ -114,9 +116,10 @@ func (h *Host) isPluginFused(id string) bool { } type usageAdapter struct { - host *Host - pluginID string - plugin pluginapi.UsagePlugin + host *Host + pluginID string + plugin pluginapi.UsagePlugin + schemaVersion uint32 } type thinkingAdapter struct { @@ -146,7 +149,13 @@ func (a *usageAdapter) HandleUsage(ctx context.Context, record coreusage.Record) a.host.fusePlugin(a.pluginID, "UsagePlugin.HandleUsage", recovered) } }() + requestID, traceID := "", "" + if a.schemaVersion >= pluginabi.SchemaVersionUsageIdentity { + requestID, traceID = coreusage.RequestIdentityFromContext(ctx) + } plugin.HandleUsage(ctx, pluginapi.UsageRecord{ + RequestID: requestID, + TraceID: traceID, Provider: record.Provider, ExecutorType: record.ExecutorType, Model: record.Model, diff --git a/sdk/api/handlers/handlers_execution.go b/sdk/api/handlers/handlers_execution.go index 994bd33a..3cc7743e 100644 --- a/sdk/api/handlers/handlers_execution.go +++ b/sdk/api/handlers/handlers_execution.go @@ -71,6 +71,7 @@ func (h *BaseAPIHandler) executeWithAuthManagerFormats(ctx context.Context, entr } afterAuthCapture := &requestAfterAuthCapture{} lifecycle := h.newRequestLifecycleTracker(ctx, entryProtocol, normalizedModel, originalRequestedModel, false, reqMeta, execOptions.SkipInterceptorPluginID) + ctx = lifecycle.executionContext() opts := coreexecutor.Options{ Stream: false, Alt: alt, @@ -136,6 +137,7 @@ func (h *BaseAPIHandler) executeCountWithAuthManager(ctx context.Context, handle } afterAuthCapture := &requestAfterAuthCapture{} lifecycle := h.newRequestLifecycleTracker(ctx, handlerType, normalizedModel, originalRequestedModel, false, reqMeta, execOptions.SkipInterceptorPluginID) + ctx = lifecycle.executionContext() opts := coreexecutor.Options{ Stream: false, Alt: alt, @@ -179,6 +181,7 @@ func (h *BaseAPIHandler) executeWithPluginExecutor(ctx context.Context, entryPro execCtx, nestedTracker := withNestedExecutionTracker(ctx) 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) + execCtx = lifecycle.executionContext() var interceptErr *interfaces.ErrorMessage req, opts, interceptErr = h.applyRequestInterceptorsBeforeAuth(execCtx, entryProtocol, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID) if interceptErr != nil { @@ -211,6 +214,7 @@ func (h *BaseAPIHandler) countWithPluginExecutor(ctx context.Context, handlerTyp } req, opts := h.pluginExecutorRequest(ctx, handlerType, handlerType, modelName, originalRequestedModel, rawJSON, alt, false, execOptions) lifecycle := h.newRequestLifecycleTracker(ctx, handlerType, modelName, originalRequestedModel, false, opts.Metadata, execOptions.SkipInterceptorPluginID) + ctx = lifecycle.executionContext() var interceptErr *interfaces.ErrorMessage req, opts, interceptErr = h.applyRequestInterceptorsBeforeAuth(ctx, handlerType, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID) if interceptErr != nil { diff --git a/sdk/api/handlers/handlers_interceptors.go b/sdk/api/handlers/handlers_interceptors.go index a8b35604..d27198c9 100644 --- a/sdk/api/handlers/handlers_interceptors.go +++ b/sdk/api/handlers/handlers_interceptors.go @@ -9,6 +9,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" 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" "golang.org/x/net/context" ) @@ -74,6 +75,7 @@ type requestLifecycleTracker struct { func (h *BaseAPIHandler) newRequestLifecycleTracker(ctx context.Context, sourceFormat, model, requestedModel string, stream bool, metadata map[string]any, skipPluginID string) *requestLifecycleTracker { requestID := uuid.NewString() traceID := logging.GetRequestID(ctx) + ctx = coreusage.WithRequestIdentity(ctx, requestID, traceID) return &requestLifecycleTracker{ ctx: ctx, host: h.interceptorHost(), @@ -91,6 +93,13 @@ func (h *BaseAPIHandler) newRequestLifecycleTracker(ctx context.Context, sourceF } } +func (t *requestLifecycleTracker) executionContext() context.Context { + if t == nil || t.ctx == nil { + return context.Background() + } + return t.ctx +} + func (t *requestLifecycleTracker) requestID() string { if t == nil { return "" diff --git a/sdk/api/handlers/handlers_interceptors_test.go b/sdk/api/handlers/handlers_interceptors_test.go index 0b328f3e..1ad85527 100644 --- a/sdk/api/handlers/handlers_interceptors_test.go +++ b/sdk/api/handlers/handlers_interceptors_test.go @@ -16,6 +16,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" ) @@ -230,6 +231,11 @@ func TestRequestLifecycleTrackerUsesUniqueExecutionIDs(t *testing.T) { if first.completion.TraceID != "trace-1" || second.completion.TraceID != "trace-1" { t.Fatalf("trace IDs = %q and %q", first.completion.TraceID, second.completion.TraceID) } + firstRequestID, firstTraceID := coreusage.RequestIdentityFromContext(first.executionContext()) + secondRequestID, secondTraceID := coreusage.RequestIdentityFromContext(second.executionContext()) + if firstRequestID != first.requestID() || secondRequestID != second.requestID() || firstTraceID != "trace-1" || secondTraceID != "trace-1" { + t.Fatalf("usage identities = (%q, %q) and (%q, %q)", firstRequestID, firstTraceID, secondRequestID, secondTraceID) + } } func TestHandlerRequestInterceptorTerminatesBeforeAuth(t *testing.T) { diff --git a/sdk/api/handlers/handlers_stream.go b/sdk/api/handlers/handlers_stream.go index 9cceec54..dcdc08df 100644 --- a/sdk/api/handlers/handlers_stream.go +++ b/sdk/api/handlers/handlers_stream.go @@ -43,6 +43,7 @@ func (h *BaseAPIHandler) streamWithPluginExecutor(ctx context.Context, entryProt execCtx, nestedTracker := withNestedExecutionTracker(ctx) 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) + execCtx = lifecycle.executionContext() var interceptErr *interfaces.ErrorMessage req, opts, interceptErr = h.applyRequestInterceptorsBeforeAuth(execCtx, entryProtocol, originalRequestedModel, lifecycle.requestID(), req, opts, execOptions.SkipInterceptorPluginID) if interceptErr != nil { @@ -305,6 +306,7 @@ func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context } afterAuthCapture := &requestAfterAuthCapture{} lifecycle := h.newRequestLifecycleTracker(ctx, entryProtocol, normalizedModel, originalRequestedModel, true, reqMeta, execOptions.SkipInterceptorPluginID) + ctx = lifecycle.executionContext() opts := coreexecutor.Options{ Stream: true, Alt: alt, diff --git a/sdk/cliproxy/usage/manager.go b/sdk/cliproxy/usage/manager.go index ca36dc55..13506fbb 100644 --- a/sdk/cliproxy/usage/manager.go +++ b/sdk/cliproxy/usage/manager.go @@ -78,6 +78,34 @@ type requestedModelAliasContextKey struct{} type reasoningEffortContextKey struct{} type serviceTierContextKey struct{} type generateContextKey struct{} +type requestIdentityContextKey struct{} + +type requestIdentity struct { + requestID string + traceID string +} + +// WithRequestIdentity stores the model execution and parent request IDs for usage sinks. +func WithRequestIdentity(ctx context.Context, requestID, traceID string) context.Context { + if ctx == nil { + ctx = context.Background() + } + requestID = strings.TrimSpace(requestID) + traceID = strings.TrimSpace(traceID) + if requestID == "" && traceID == "" { + return ctx + } + return context.WithValue(ctx, requestIdentityContextKey{}, requestIdentity{requestID: requestID, traceID: traceID}) +} + +// RequestIdentityFromContext returns the model execution and parent request IDs stored in ctx. +func RequestIdentityFromContext(ctx context.Context) (requestID, traceID string) { + if ctx == nil { + return "", "" + } + identity, _ := ctx.Value(requestIdentityContextKey{}).(requestIdentity) + return strings.TrimSpace(identity.requestID), strings.TrimSpace(identity.traceID) +} // WithRequestedModelAlias stores the client-requested model name for usage sinks. func WithRequestedModelAlias(ctx context.Context, alias string) context.Context { diff --git a/sdk/cliproxy/usage/manager_test.go b/sdk/cliproxy/usage/manager_test.go index 6f7b1fbb..fc2b3120 100644 --- a/sdk/cliproxy/usage/manager_test.go +++ b/sdk/cliproxy/usage/manager_test.go @@ -36,6 +36,21 @@ func TestGenerateFromContextHonorsExplicitFalse(t *testing.T) { } } +func TestRequestIdentityContextRoundTrip(t *testing.T) { + ctx := WithRequestIdentity(context.Background(), " request-1 ", " trace-1 ") + requestID, traceID := RequestIdentityFromContext(ctx) + if requestID != "request-1" || traceID != "trace-1" { + t.Fatalf("request identity = %q, %q", requestID, traceID) + } +} + +func TestRequestIdentityContextDefaultsEmpty(t *testing.T) { + requestID, traceID := RequestIdentityFromContext(nil) + if requestID != "" || traceID != "" { + t.Fatalf("empty request identity = %q, %q", requestID, traceID) + } +} + func TestRecordOmittedGenerateIsEnabled(t *testing.T) { // Existing callers construct Record without setting Generate. // Omission must remain distinguishable from explicit false and default to true. diff --git a/sdk/pluginabi/types.go b/sdk/pluginabi/types.go index 97c41a13..a8a3e470 100644 --- a/sdk/pluginabi/types.go +++ b/sdk/pluginabi/types.go @@ -10,10 +10,13 @@ const ( // Version 3 omits OriginalRequest/RequestBody on payload stream chunks // (ChunkIndex >= 0); those fields remain on StreamChunkHeaderInitIndex only. // Plugins that still need per-chunk request bodies should keep schema_version < 3. - SchemaVersion uint32 = 3 + SchemaVersion uint32 = 4 // SchemaVersionStreamChunkOmitRequestBody is the first schema version that omits // request bodies on payload stream-chunk interceptor calls. SchemaVersionStreamChunkOmitRequestBody uint32 = 3 + // SchemaVersionUsageIdentity is the first schema version that attaches + // request lifecycle identity to usage records. + SchemaVersionUsageIdentity uint32 = 4 ) const ( diff --git a/sdk/pluginabi/types_test.go b/sdk/pluginabi/types_test.go index 8fa63542..86253848 100644 --- a/sdk/pluginabi/types_test.go +++ b/sdk/pluginabi/types_test.go @@ -27,12 +27,15 @@ func TestEnvelopeRoundTrip(t *testing.T) { } func TestMethodNamesAreStable(t *testing.T) { - if SchemaVersion != 3 { - t.Fatalf("SchemaVersion = %d, want 3", SchemaVersion) + if SchemaVersion != 4 { + t.Fatalf("SchemaVersion = %d, want 4", SchemaVersion) } if SchemaVersionStreamChunkOmitRequestBody != 3 { t.Fatalf("SchemaVersionStreamChunkOmitRequestBody = %d, want 3", SchemaVersionStreamChunkOmitRequestBody) } + if SchemaVersionUsageIdentity != 4 { + t.Fatalf("SchemaVersionUsageIdentity = %d, want 4", SchemaVersionUsageIdentity) + } if MethodPluginRegister != "plugin.register" { t.Fatalf("MethodPluginRegister = %q", MethodPluginRegister) } diff --git a/sdk/pluginapi/types.go b/sdk/pluginapi/types.go index 6add5d69..894e3b77 100644 --- a/sdk/pluginapi/types.go +++ b/sdk/pluginapi/types.go @@ -1316,6 +1316,10 @@ type ManagementResponse struct { // UsageRecord describes request usage and billing metadata. type UsageRecord struct { + // RequestID identifies the request lifecycle associated with this usage record. + RequestID string + // TraceID identifies the parent inbound HTTP request when available. + TraceID string // Provider identifies the upstream provider. Provider string // ExecutorType identifies the executor implementation.