mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-09-06 16:15:50 +08:00
fix(cliproxy): refine Codex model resolution and credential validation
- Adjusted `resolveConfigCodexStyleKey` to include `validateIndexCredentials` for stricter credential checks. - Enhanced logic in `buildCodexConfigModels` to handle empty model lists more effectively. - Standardized credential matching with helper functions for improved maintainability.
This commit is contained in:
188
sdk/cliproxy/service_codex_models_test.go
Normal file
188
sdk/cliproxy/service_codex_models_test.go
Normal file
@@ -0,0 +1,188 @@
|
||||
package cliproxy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
internalregistry "github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
)
|
||||
|
||||
func TestRegisterModelsForAuthCodexAPIKeyModels(t *testing.T) {
|
||||
defaultModels := internalregistry.GetCodexProModels()
|
||||
if len(defaultModels) == 0 {
|
||||
t.Fatal("expected Codex Pro default models")
|
||||
}
|
||||
|
||||
excludedModelID := defaultModels[0].ID
|
||||
tests := []struct {
|
||||
name string
|
||||
entry config.CodexKey
|
||||
wantIDs map[string]struct{}
|
||||
}{
|
||||
{
|
||||
name: "defaults without configuration",
|
||||
entry: config.CodexKey{APIKey: "default-key"},
|
||||
wantIDs: codexModelIDSet(defaultModels),
|
||||
},
|
||||
{
|
||||
name: "explicit configuration replaces defaults",
|
||||
entry: config.CodexKey{
|
||||
APIKey: "configured-key",
|
||||
Models: []internalconfig.CodexModel{{
|
||||
Name: "upstream-codex", Alias: "configured-codex",
|
||||
}},
|
||||
},
|
||||
wantIDs: map[string]struct{}{"configured-codex": {}},
|
||||
},
|
||||
{
|
||||
name: "exclusions apply to defaults",
|
||||
entry: config.CodexKey{
|
||||
APIKey: "excluded-key",
|
||||
ExcludedModels: []string{excludedModelID},
|
||||
},
|
||||
wantIDs: codexModelIDSet(defaultModels[1:]),
|
||||
},
|
||||
}
|
||||
|
||||
for index := range tests {
|
||||
testCase := tests[index]
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
authID := fmt.Sprintf("codex-api-key-models-%d", index)
|
||||
modelRegistry := internalregistry.GetGlobalRegistry()
|
||||
modelRegistry.UnregisterClient(authID)
|
||||
t.Cleanup(func() { modelRegistry.UnregisterClient(authID) })
|
||||
|
||||
service := &Service{cfg: &config.Config{CodexKey: []config.CodexKey{testCase.entry}}}
|
||||
auth := &coreauth.Auth{
|
||||
ID: authID,
|
||||
Provider: "codex",
|
||||
Status: coreauth.StatusActive,
|
||||
Attributes: map[string]string{
|
||||
coreauth.AttributeAPIKey: testCase.entry.APIKey,
|
||||
coreauth.AttributeConfigIndex: "0",
|
||||
coreauth.AttributeSource: "config:codex:test",
|
||||
},
|
||||
}
|
||||
|
||||
service.registerModelsForAuth(context.Background(), auth)
|
||||
gotIDs := codexModelIDSet(modelRegistry.GetModelsForClient(authID))
|
||||
if len(gotIDs) != len(testCase.wantIDs) {
|
||||
t.Fatalf("registered model IDs = %#v, want %#v", gotIDs, testCase.wantIDs)
|
||||
}
|
||||
for modelID := range testCase.wantIDs {
|
||||
if _, ok := gotIDs[modelID]; !ok {
|
||||
t.Errorf("missing registered model %q", modelID)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterModelsForAuthCodexAPIKeyDefaultRequiresConfigMatch(t *testing.T) {
|
||||
defaultIDs := codexModelIDSet(internalregistry.GetCodexProModels())
|
||||
tests := []struct {
|
||||
name string
|
||||
config config.Config
|
||||
attributes map[string]string
|
||||
wantIDs map[string]struct{}
|
||||
}{
|
||||
{
|
||||
name: "valid index with unmatched API key",
|
||||
config: config.Config{CodexKey: []config.CodexKey{{
|
||||
APIKey: "configured-key",
|
||||
}}},
|
||||
attributes: map[string]string{
|
||||
coreauth.AttributeAPIKey: "stale-key",
|
||||
coreauth.AttributeConfigIndex: "0",
|
||||
coreauth.AttributeSource: "config:codex:stale",
|
||||
},
|
||||
wantIDs: map[string]struct{}{},
|
||||
},
|
||||
{
|
||||
name: "valid index with unmatched base URL",
|
||||
config: config.Config{CodexKey: []config.CodexKey{{
|
||||
APIKey: "configured-key", BaseURL: "https://new.example.com",
|
||||
}}},
|
||||
attributes: map[string]string{
|
||||
coreauth.AttributeAPIKey: "configured-key",
|
||||
coreauth.AttributeConfigIndex: "0",
|
||||
coreauth.AttributeSource: "config:codex:stale",
|
||||
"base_url": "https://old.example.com",
|
||||
},
|
||||
wantIDs: map[string]struct{}{},
|
||||
},
|
||||
{
|
||||
name: "stale index falls back to matching credentials",
|
||||
config: config.Config{CodexKey: []config.CodexKey{
|
||||
{
|
||||
APIKey: "wrong-key",
|
||||
Models: []internalconfig.CodexModel{{Name: "wrong-model"}},
|
||||
},
|
||||
{APIKey: "configured-key"},
|
||||
}},
|
||||
attributes: map[string]string{
|
||||
coreauth.AttributeAPIKey: "configured-key",
|
||||
coreauth.AttributeConfigIndex: "0",
|
||||
coreauth.AttributeSource: "config:codex:stale",
|
||||
},
|
||||
wantIDs: defaultIDs,
|
||||
},
|
||||
{
|
||||
name: "API key ignores OAuth plan type",
|
||||
config: config.Config{CodexKey: []config.CodexKey{{
|
||||
APIKey: "configured-key",
|
||||
}}},
|
||||
attributes: map[string]string{
|
||||
coreauth.AttributeAPIKey: "configured-key",
|
||||
coreauth.AttributeConfigIndex: "0",
|
||||
coreauth.AttributeSource: "config:codex:test",
|
||||
"plan_type": "free",
|
||||
},
|
||||
wantIDs: defaultIDs,
|
||||
},
|
||||
}
|
||||
|
||||
for index := range tests {
|
||||
testCase := tests[index]
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
authID := fmt.Sprintf("codex-api-key-config-match-%d", index)
|
||||
modelRegistry := internalregistry.GetGlobalRegistry()
|
||||
modelRegistry.UnregisterClient(authID)
|
||||
modelRegistry.RegisterClient(authID, "codex", []*internalregistry.ModelInfo{{ID: "stale-model"}})
|
||||
t.Cleanup(func() { modelRegistry.UnregisterClient(authID) })
|
||||
|
||||
service := &Service{cfg: &testCase.config}
|
||||
auth := &coreauth.Auth{
|
||||
ID: authID,
|
||||
Provider: "codex",
|
||||
Status: coreauth.StatusActive,
|
||||
Attributes: testCase.attributes,
|
||||
}
|
||||
|
||||
service.registerModelsForAuth(context.Background(), auth)
|
||||
gotIDs := codexModelIDSet(modelRegistry.GetModelsForClient(authID))
|
||||
if len(gotIDs) != len(testCase.wantIDs) {
|
||||
t.Fatalf("registered model IDs = %#v, want %#v", gotIDs, testCase.wantIDs)
|
||||
}
|
||||
for modelID := range testCase.wantIDs {
|
||||
if _, ok := gotIDs[modelID]; !ok {
|
||||
t.Errorf("missing registered model %q", modelID)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func codexModelIDSet(models []*internalregistry.ModelInfo) map[string]struct{} {
|
||||
ids := make(map[string]struct{}, len(models))
|
||||
for _, model := range models {
|
||||
if model != nil && model.ID != "" {
|
||||
ids[model.ID] = struct{}{}
|
||||
}
|
||||
}
|
||||
return ids
|
||||
}
|
||||
@@ -117,10 +117,11 @@ func (s *Service) registerModelsForAuthWithCache(ctx context.Context, a *coreaut
|
||||
case "codex":
|
||||
if authKind == "apikey" {
|
||||
if entry := s.resolveConfigCodexKey(a); entry != nil {
|
||||
models = buildCodexConfigModels(entry)
|
||||
models = registry.GetCodexProModels()
|
||||
if len(entry.Models) > 0 {
|
||||
models = buildCodexConfigModels(entry)
|
||||
}
|
||||
excluded = entry.ExcludedModels
|
||||
} else {
|
||||
models = nil
|
||||
}
|
||||
models = applyExcludedModels(models, excluded)
|
||||
break
|
||||
@@ -486,39 +487,41 @@ func (s *Service) resolveConfigCodexKey(auth *coreauth.Auth) *config.CodexKey {
|
||||
if s == nil || s.cfg == nil {
|
||||
return nil
|
||||
}
|
||||
return resolveConfigCodexStyleKey(auth, s.cfg.CodexKey)
|
||||
return resolveConfigCodexStyleKey(auth, s.cfg.CodexKey, true)
|
||||
}
|
||||
|
||||
func (s *Service) resolveConfigXAIKey(auth *coreauth.Auth) *config.XAIKey {
|
||||
if s == nil || s.cfg == nil {
|
||||
return nil
|
||||
}
|
||||
return resolveConfigCodexStyleKey(auth, s.cfg.XAIKey)
|
||||
return resolveConfigCodexStyleKey(auth, s.cfg.XAIKey, false)
|
||||
}
|
||||
|
||||
func resolveConfigCodexStyleKey(auth *coreauth.Auth, entries []config.CodexKey) *config.CodexKey {
|
||||
func resolveConfigCodexStyleKey(auth *coreauth.Auth, entries []config.CodexKey, validateIndexCredentials bool) *config.CodexKey {
|
||||
if auth == nil {
|
||||
return nil
|
||||
}
|
||||
if entry := configEntryForAuthIndex(auth, entries); entry != nil {
|
||||
return entry
|
||||
}
|
||||
var attrKey, attrBase string
|
||||
if auth.Attributes != nil {
|
||||
attrKey = strings.TrimSpace(auth.Attributes["api_key"])
|
||||
attrBase = strings.TrimSpace(auth.Attributes["base_url"])
|
||||
}
|
||||
for i := range entries {
|
||||
entry := &entries[i]
|
||||
matchesCredentials := func(entry *config.CodexKey) bool {
|
||||
if entry == nil {
|
||||
return false
|
||||
}
|
||||
cfgKey := strings.TrimSpace(entry.APIKey)
|
||||
cfgBase := strings.TrimSpace(entry.BaseURL)
|
||||
if attrKey != "" && strings.EqualFold(cfgKey, attrKey) {
|
||||
if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) {
|
||||
return entry
|
||||
}
|
||||
continue
|
||||
if attrKey != "" {
|
||||
return strings.EqualFold(cfgKey, attrKey) && (cfgBase == "" || strings.EqualFold(cfgBase, attrBase))
|
||||
}
|
||||
if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) {
|
||||
return attrBase != "" && strings.EqualFold(cfgBase, attrBase)
|
||||
}
|
||||
if entry := configEntryForAuthIndex(auth, entries); entry != nil && (!validateIndexCredentials || matchesCredentials(entry)) {
|
||||
return entry
|
||||
}
|
||||
for i := range entries {
|
||||
if entry := &entries[i]; matchesCredentials(entry) {
|
||||
return entry
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user