agent/cpa-plugin
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
954c0baeab5738f5bcaf54ebcf87ce626b2add52
cpa-plugin/patch/cli-proxy-api-usage-identity.patch
T

349 lines
16 KiB
Diff
Raw Blame History

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,
@@ -177,6 +179,7 @@ func (h *BaseAPIHandler) executeWithPluginExecutor(ctx context.Context, entryPro
}
req, opts := h.pluginExecutorRequest(ctx, entryProtocol, responseProtocol, modelName, originalRequestedModel, rawJSON, alt, false, execOptions)
lifecycle := h.newRequestLifecycleTracker(ctx, entryProtocol, modelName, originalRequestedModel, false, opts.Metadata, execOptions.SkipInterceptorPluginID)
+ ctx = lifecycle.executionContext()
var interceptErr *interfaces.ErrorMessage
req, opts, interceptErr = h.applyRequestInterceptorsBeforeAuth(ctx, 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
@@ -41,6 +41,7 @@ func (h *BaseAPIHandler) streamWithPluginExecutor(ctx context.Context, entryProt
}
req, opts := h.pluginExecutorRequest(ctx, entryProtocol, responseProtocol, modelName, originalRequestedModel, rawJSON, alt, true, execOptions)
lifecycle := h.newRequestLifecycleTracker(ctx, entryProtocol, modelName, originalRequestedModel, true, opts.Metadata, execOptions.SkipInterceptorPluginID)
+ ctx = lifecycle.executionContext()
var interceptErr *interfaces.ErrorMessage
req, opts, interceptErr = h.applyRequestInterceptorsBeforeAuth(ctx, 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.
Reference in New Issue View Git Blame Copy Permalink