From ba5d8ca7336e78ca2b7c6134638a4054b589c330 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Tue, 5 May 2026 01:47:53 +0800 Subject: [PATCH] feat(usage): add support for requested model alias handling - Introduced methods for setting and retrieving model aliases in execution and usage contexts. - Enhanced `UsageReporter` and related structures to include client-requested aliases. - Updated tests to validate alias propagation and ensure correct usage reporting. - Adjusted metadata handling in CLIProxyAPI executors to address alias integration. --- internal/redisqueue/plugin.go | 6 ++++ internal/redisqueue/plugin_test.go | 6 ++++ .../runtime/executor/helps/usage_helpers.go | 7 ++++ .../executor/helps/usage_helpers_test.go | 14 ++++++++ sdk/api/handlers/handlers.go | 6 ++-- sdk/cliproxy/auth/conductor.go | 35 +++++++++++++++++++ .../conductor_oauth_alias_suspension_test.go | 25 +++++++++++-- sdk/cliproxy/usage/manager.go | 32 +++++++++++++++++ 8 files changed, 125 insertions(+), 6 deletions(-) diff --git a/internal/redisqueue/plugin.go b/internal/redisqueue/plugin.go index 971684190..b33bc8fd9 100644 --- a/internal/redisqueue/plugin.go +++ b/internal/redisqueue/plugin.go @@ -33,6 +33,10 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec if modelName == "" { modelName = "unknown" } + aliasName := strings.TrimSpace(record.Alias) + if aliasName == "" { + aliasName = modelName + } provider := strings.TrimSpace(record.Provider) if provider == "" { provider = "unknown" @@ -76,6 +80,7 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec requestDetail: detail, Provider: provider, Model: modelName, + Alias: aliasName, Endpoint: resolveEndpoint(ctx), AuthType: authType, APIKey: apiKey, @@ -91,6 +96,7 @@ type queuedUsageDetail struct { requestDetail Provider string `json:"provider"` Model string `json:"model"` + Alias string `json:"alias"` Endpoint string `json:"endpoint"` AuthType string `json:"auth_type"` APIKey string `json:"api_key"` diff --git a/internal/redisqueue/plugin_test.go b/internal/redisqueue/plugin_test.go index 0cc8b9b9c..8dcade90e 100644 --- a/internal/redisqueue/plugin_test.go +++ b/internal/redisqueue/plugin_test.go @@ -24,6 +24,7 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) { plugin.HandleUsage(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4", + Alias: "client-gpt", APIKey: "test-key", AuthIndex: "0", AuthType: "apikey", @@ -40,6 +41,7 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) { payload := popSinglePayload(t) requireStringField(t, payload, "provider", "openai") requireStringField(t, payload, "model", "gpt-5.4") + requireStringField(t, payload, "alias", "client-gpt") requireStringField(t, payload, "endpoint", "POST /v1/chat/completions") requireStringField(t, payload, "auth_type", "apikey") requireStringField(t, payload, "request_id", "ctx-request-id") @@ -58,6 +60,7 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndFailureAndGinRequestID(t plugin.HandleUsage(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4-mini", + Alias: "client-mini", APIKey: "test-key", AuthIndex: "0", AuthType: "apikey", @@ -74,6 +77,7 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndFailureAndGinRequestID(t payload := popSinglePayload(t) requireStringField(t, payload, "provider", "openai") requireStringField(t, payload, "model", "gpt-5.4-mini") + requireStringField(t, payload, "alias", "client-mini") requireStringField(t, payload, "endpoint", "GET /v1/responses") requireStringField(t, payload, "auth_type", "apikey") requireStringField(t, payload, "request_id", "gin-request-id") @@ -102,6 +106,7 @@ func TestUsageQueuePluginAsyncIgnoresRecycledGinContext(t *testing.T) { mgr.Publish(ctx, coreusage.Record{ Provider: "openai", Model: "gpt-5.4", + Alias: "client-gpt", APIKey: "test-key", AuthIndex: "0", AuthType: "apikey", @@ -117,6 +122,7 @@ func TestUsageQueuePluginAsyncIgnoresRecycledGinContext(t *testing.T) { payload := waitForSinglePayload(t, 2*time.Second) requireStringField(t, payload, "endpoint", "POST /v1/chat/completions") + requireStringField(t, payload, "alias", "client-gpt") requireStringField(t, payload, "request_id", "ctx-request-id") requireBoolField(t, payload, "failed", true) }) diff --git a/internal/runtime/executor/helps/usage_helpers.go b/internal/runtime/executor/helps/usage_helpers.go index c5e258c86..312a1d35c 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -18,6 +18,7 @@ import ( type UsageReporter struct { provider string model string + alias string authID string authIndex string authType string @@ -29,9 +30,14 @@ type UsageReporter struct { func NewUsageReporter(ctx context.Context, provider, model string, auth *cliproxyauth.Auth) *UsageReporter { apiKey := APIKeyFromContext(ctx) + alias := usage.RequestedModelAliasFromContext(ctx) + if alias == "" { + alias = model + } reporter := &UsageReporter{ provider: provider, model: model, + alias: strings.TrimSpace(alias), requestedAt: time.Now(), apiKey: apiKey, source: resolveUsageSource(auth, apiKey), @@ -139,6 +145,7 @@ func (r *UsageReporter) buildRecordForModel(model string, detail usage.Detail, f return usage.Record{ Provider: r.provider, Model: model, + Alias: r.alias, Source: r.source, APIKey: r.apiKey, AuthID: r.authID, diff --git a/internal/runtime/executor/helps/usage_helpers_test.go b/internal/runtime/executor/helps/usage_helpers_test.go index c77335fd6..ef2c7de58 100644 --- a/internal/runtime/executor/helps/usage_helpers_test.go +++ b/internal/runtime/executor/helps/usage_helpers_test.go @@ -1,6 +1,7 @@ package helps import ( + "context" "testing" "time" @@ -107,6 +108,19 @@ func TestUsageReporterBuildRecordIncludesLatency(t *testing.T) { } } +func TestUsageReporterBuildRecordIncludesRequestedModelAlias(t *testing.T) { + ctx := usage.WithRequestedModelAlias(context.Background(), "client-gpt") + reporter := NewUsageReporter(ctx, "openai", "gpt-5.4", nil) + + record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false) + if record.Model != "gpt-5.4" { + t.Fatalf("model = %q, want %q", record.Model, "gpt-5.4") + } + if record.Alias != "client-gpt" { + t.Fatalf("alias = %q, want %q", record.Alias, "client-gpt") + } +} + func TestUsageReporterBuildAdditionalModelRecordSkipsZeroTokens(t *testing.T) { reporter := &UsageReporter{ provider: "codex", diff --git a/sdk/api/handlers/handlers.go b/sdk/api/handlers/handlers.go index 52b2a4fde..e89227aa7 100644 --- a/sdk/api/handlers/handlers.go +++ b/sdk/api/handlers/handlers.go @@ -539,7 +539,7 @@ func (h *BaseAPIHandler) ExecuteWithAuthManager(ctx context.Context, handlerType return nil, nil, errMsg } reqMeta := requestExecutionMetadata(ctx) - reqMeta[coreexecutor.RequestedModelMetadataKey] = normalizedModel + reqMeta[coreexecutor.RequestedModelMetadataKey] = modelName payload := rawJSON if len(payload) == 0 { payload = nil @@ -587,7 +587,7 @@ func (h *BaseAPIHandler) ExecuteCountWithAuthManager(ctx context.Context, handle return nil, nil, errMsg } reqMeta := requestExecutionMetadata(ctx) - reqMeta[coreexecutor.RequestedModelMetadataKey] = normalizedModel + reqMeta[coreexecutor.RequestedModelMetadataKey] = modelName payload := rawJSON if len(payload) == 0 { payload = nil @@ -639,7 +639,7 @@ func (h *BaseAPIHandler) ExecuteStreamWithAuthManager(ctx context.Context, handl return nil, nil, errChan } reqMeta := requestExecutionMetadata(ctx) - reqMeta[coreexecutor.RequestedModelMetadataKey] = normalizedModel + reqMeta[coreexecutor.RequestedModelMetadataKey] = modelName payload := rawJSON if len(payload) == 0 { payload = nil diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index d2a3db188..ab3eca495 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -22,6 +22,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v6/internal/thinking" "github.com/router-for-me/CLIProxyAPI/v6/internal/util" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor" + coreusage "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage" log "github.com/sirupsen/logrus" ) @@ -827,6 +828,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi if executor == nil { return nil, &Error{Code: "executor_not_found", Message: "executor not registered"} } + ctx = contextWithRequestedModelAlias(ctx, opts, routeModel) var lastErr error for idx, execModel := range execModels { resultModel := m.stateModelForExecution(auth, routeModel, execModel, pooled) @@ -1319,6 +1321,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) } + execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel) models, pooled := m.preparedExecutionModels(auth, routeModel) if len(models) == 0 { @@ -1397,6 +1400,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) } + execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel) models, pooled := m.preparedExecutionModels(auth, routeModel) if len(models) == 0 { @@ -1534,6 +1538,36 @@ func hasRequestedModelMetadata(meta map[string]any) bool { } } +func contextWithRequestedModelAlias(ctx context.Context, opts cliproxyexecutor.Options, fallback string) context.Context { + alias := requestedModelAliasFromOptions(opts, fallback) + return coreusage.WithRequestedModelAlias(ctx, alias) +} + +func requestedModelAliasFromOptions(opts cliproxyexecutor.Options, fallback string) string { + fallback = strings.TrimSpace(fallback) + if len(opts.Metadata) == 0 { + return fallback + } + raw, ok := opts.Metadata[cliproxyexecutor.RequestedModelMetadataKey] + if !ok || raw == nil { + return fallback + } + switch value := raw.(type) { + case string: + if strings.TrimSpace(value) == "" { + return fallback + } + return strings.TrimSpace(value) + case []byte: + if len(value) == 0 { + return fallback + } + return strings.TrimSpace(string(value)) + default: + return fallback + } +} + func pinnedAuthIDFromMetadata(meta map[string]any) string { if len(meta) == 0 { return "" @@ -3096,6 +3130,7 @@ func (m *Manager) tryAntigravityCreditsExecute(ctx context.Context, req cliproxy creditsCtx = context.WithValue(creditsCtx, "cliproxy.roundtripper", rt) } creditsOpts := ensureRequestedModelMetadata(opts, routeModel) + creditsCtx = contextWithRequestedModelAlias(creditsCtx, creditsOpts, routeModel) publishSelectedAuthMetadata(creditsOpts.Metadata, c.auth.ID) models := m.executionModelCandidates(c.auth, routeModel) if len(models) == 0 { diff --git a/sdk/cliproxy/auth/conductor_oauth_alias_suspension_test.go b/sdk/cliproxy/auth/conductor_oauth_alias_suspension_test.go index 8bc779e53..b4b72204c 100644 --- a/sdk/cliproxy/auth/conductor_oauth_alias_suspension_test.go +++ b/sdk/cliproxy/auth/conductor_oauth_alias_suspension_test.go @@ -10,20 +10,23 @@ import ( internalconfig "github.com/router-for-me/CLIProxyAPI/v6/internal/config" "github.com/router-for-me/CLIProxyAPI/v6/internal/registry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor" + coreusage "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage" ) type aliasRoutingExecutor struct { id string - mu sync.Mutex - executeModels []string + mu sync.Mutex + executeModels []string + executeAliases []string } func (e *aliasRoutingExecutor) Identifier() string { return e.id } -func (e *aliasRoutingExecutor) Execute(_ context.Context, _ *Auth, req cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { +func (e *aliasRoutingExecutor) Execute(ctx context.Context, _ *Auth, req cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { e.mu.Lock() e.executeModels = append(e.executeModels, req.Model) + e.executeAliases = append(e.executeAliases, coreusage.RequestedModelAliasFromContext(ctx)) e.mu.Unlock() return cliproxyexecutor.Response{Payload: []byte(req.Model)}, nil } @@ -52,6 +55,14 @@ func (e *aliasRoutingExecutor) ExecuteModels() []string { return out } +func (e *aliasRoutingExecutor) ExecuteAliases() []string { + e.mu.Lock() + defer e.mu.Unlock() + out := make([]string, len(e.executeAliases)) + copy(out, e.executeAliases) + return out +} + func TestManagerExecute_OAuthAliasBypassesBlockedRouteModel(t *testing.T) { const ( provider = "antigravity" @@ -108,4 +119,12 @@ func TestManagerExecute_OAuthAliasBypassesBlockedRouteModel(t *testing.T) { if gotModels[0] != targetModel { t.Fatalf("execute model = %q, want %q", gotModels[0], targetModel) } + + gotAliases := executor.ExecuteAliases() + if len(gotAliases) != 1 { + t.Fatalf("execute aliases len = %d, want 1", len(gotAliases)) + } + if gotAliases[0] != routeModel { + t.Fatalf("execute alias = %q, want %q", gotAliases[0], routeModel) + } } diff --git a/sdk/cliproxy/usage/manager.go b/sdk/cliproxy/usage/manager.go index c3d95f663..72405d758 100644 --- a/sdk/cliproxy/usage/manager.go +++ b/sdk/cliproxy/usage/manager.go @@ -2,6 +2,7 @@ package usage import ( "context" + "strings" "sync" "time" @@ -12,6 +13,7 @@ import ( type Record struct { Provider string Model string + Alias string APIKey string AuthID string AuthIndex string @@ -32,6 +34,36 @@ type Detail struct { TotalTokens int64 } +type requestedModelAliasContextKey struct{} + +// WithRequestedModelAlias stores the client-requested model name for usage sinks. +func WithRequestedModelAlias(ctx context.Context, alias string) context.Context { + if ctx == nil { + ctx = context.Background() + } + alias = strings.TrimSpace(alias) + if alias == "" { + return ctx + } + return context.WithValue(ctx, requestedModelAliasContextKey{}, alias) +} + +// RequestedModelAliasFromContext returns the client-requested model name stored in ctx. +func RequestedModelAliasFromContext(ctx context.Context) string { + if ctx == nil { + return "" + } + raw := ctx.Value(requestedModelAliasContextKey{}) + switch value := raw.(type) { + case string: + return strings.TrimSpace(value) + case []byte: + return strings.TrimSpace(string(value)) + default: + return "" + } +} + // Plugin consumes usage records emitted by the proxy runtime. type Plugin interface { HandleUsage(ctx context.Context, record Record)