mirror of
https://gh-proxy.org/https://github.com/klinsmaya/CLIProxyAPI.git
synced 2026-11-09 18:06:50 +08:00
feat(auth): add home force mapping logic with tests and alias result handling
- Implemented `homeForceMappingAliasResult` for handling forced alias mapping in the authentication process. - Updated `resolveExecutionAliasResultForRequested` to incorporate home force mapping. - Introduced new attributes (`home_force_mapping` and `home_original_alias`) for managing alias behavior. - Extended `homeAuthDispatchResponse` with fields for force mapping and original alias. - Added test cases to validate force mapping logic and edge cases.
This commit is contained in:
@@ -1248,12 +1248,34 @@ func (m *Manager) resolveExecutionAliasResult(auth *Auth, routeModel string) OAu
|
||||
}
|
||||
|
||||
func (m *Manager) resolveExecutionAliasResultForRequested(auth *Auth, requestedModel string) OAuthModelAliasResult {
|
||||
if result := homeForceMappingAliasResult(auth, requestedModel); result.ForceMapping {
|
||||
return result
|
||||
}
|
||||
if auth != nil && auth.AuthKind() == AuthKindAPIKey {
|
||||
return m.resolveAPIKeyModelAliasWithResult(auth, requestedModel)
|
||||
}
|
||||
return m.applyOAuthModelAliasWithResult(auth, requestedModel)
|
||||
}
|
||||
|
||||
func homeForceMappingAliasResult(auth *Auth, requestedModel string) OAuthModelAliasResult {
|
||||
if auth == nil || auth.Attributes == nil || !strings.EqualFold(strings.TrimSpace(auth.Attributes[homeForceMappingAttributeKey]), "true") {
|
||||
return OAuthModelAliasResult{}
|
||||
}
|
||||
originalAlias := strings.TrimSpace(auth.Attributes[homeOriginalAliasAttributeKey])
|
||||
if originalAlias == "" {
|
||||
return OAuthModelAliasResult{}
|
||||
}
|
||||
upstreamModel := strings.TrimSpace(auth.Attributes[homeUpstreamModelAttributeKey])
|
||||
if upstreamModel == "" {
|
||||
upstreamModel = strings.TrimSpace(requestedModel)
|
||||
}
|
||||
return OAuthModelAliasResult{
|
||||
UpstreamModel: upstreamModel,
|
||||
ForceMapping: true,
|
||||
OriginalAlias: originalAlias,
|
||||
}
|
||||
}
|
||||
|
||||
func executionAliasPoolModel(auth *Auth, requestedModel string, aliasResult OAuthModelAliasResult) string {
|
||||
if auth != nil && auth.AuthKind() == AuthKindAPIKey {
|
||||
if strings.TrimSpace(requestedModel) != "" {
|
||||
@@ -4899,6 +4921,8 @@ type homeErrorDetail struct {
|
||||
|
||||
const (
|
||||
homeUpstreamModelAttributeKey = "home_upstream_model"
|
||||
homeForceMappingAttributeKey = "home_force_mapping"
|
||||
homeOriginalAliasAttributeKey = "home_original_alias"
|
||||
homeRequestRetryExceededErrorCode = "request_retry_exceeded"
|
||||
)
|
||||
|
||||
@@ -4938,11 +4962,13 @@ func repeatedHomeAuthError() *Error {
|
||||
}
|
||||
|
||||
type homeAuthDispatchResponse struct {
|
||||
Model string `json:"model"`
|
||||
Provider string `json:"provider"`
|
||||
AuthIndex string `json:"auth_index"`
|
||||
UserAPIKey string `json:"user_api_key"`
|
||||
Auth Auth `json:"auth"`
|
||||
Model string `json:"model"`
|
||||
Provider string `json:"provider"`
|
||||
AuthIndex string `json:"auth_index"`
|
||||
UserAPIKey string `json:"user_api_key"`
|
||||
ForceMapping bool `json:"force_mapping"`
|
||||
OriginalAlias string `json:"original_alias"`
|
||||
Auth Auth `json:"auth"`
|
||||
}
|
||||
|
||||
type homeAuthDispatcher interface {
|
||||
@@ -5204,10 +5230,17 @@ func (m *Manager) pickNextViaHome(ctx context.Context, model string, opts clipro
|
||||
}
|
||||
if upstreamModel := strings.TrimSpace(dispatch.Model); upstreamModel != "" {
|
||||
if auth.Attributes == nil {
|
||||
auth.Attributes = make(map[string]string, 1)
|
||||
auth.Attributes = make(map[string]string, 3)
|
||||
}
|
||||
auth.Attributes[homeUpstreamModelAttributeKey] = upstreamModel
|
||||
}
|
||||
if originalAlias := strings.TrimSpace(dispatch.OriginalAlias); dispatch.ForceMapping && originalAlias != "" {
|
||||
if auth.Attributes == nil {
|
||||
auth.Attributes = make(map[string]string, 2)
|
||||
}
|
||||
auth.Attributes[homeForceMappingAttributeKey] = "true"
|
||||
auth.Attributes[homeOriginalAliasAttributeKey] = originalAlias
|
||||
}
|
||||
if strings.TrimSpace(auth.ID) == "" {
|
||||
return nil, nil, "", &Error{Code: "invalid_auth", Message: "home returned auth without id", HTTPStatus: http.StatusBadGateway}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
package auth
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestHomeForceMappingAliasResult(t *testing.T) {
|
||||
auth := &Auth{
|
||||
Provider: "xai",
|
||||
Attributes: map[string]string{
|
||||
homeUpstreamModelAttributeKey: "grok-4.5",
|
||||
homeForceMappingAttributeKey: "true",
|
||||
homeOriginalAliasAttributeKey: "grok-latest",
|
||||
},
|
||||
}
|
||||
|
||||
result := homeForceMappingAliasResult(auth, "grok-latest")
|
||||
if result.UpstreamModel != "grok-4.5" || !result.ForceMapping || result.OriginalAlias != "grok-latest" {
|
||||
t.Fatalf("homeForceMappingAliasResult() = %+v", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeForceMappingAliasResultRequiresExplicitFlag(t *testing.T) {
|
||||
auth := &Auth{
|
||||
Provider: "xai",
|
||||
Attributes: map[string]string{
|
||||
homeUpstreamModelAttributeKey: "grok-4.5",
|
||||
homeOriginalAliasAttributeKey: "grok-latest",
|
||||
},
|
||||
}
|
||||
|
||||
result := homeForceMappingAliasResult(auth, "grok-latest")
|
||||
if result.ForceMapping || result.OriginalAlias != "" {
|
||||
t.Fatalf("homeForceMappingAliasResult() = %+v, want no force mapping", result)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user