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.
This commit is contained in:
Luis Pater
2026-05-05 01:47:53 +08:00
parent 28b4b19e7e
commit ba5d8ca733
8 changed files with 125 additions and 6 deletions
+35
View File
@@ -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 {
@@ -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)
}
}