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:
Luis Pater
2026-07-14 14:59:15 +08:00
parent 03d58c4402
commit 160a55612e
2 changed files with 73 additions and 6 deletions
+39 -6
View File
@@ -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)
}
}