package auth import ( "context" "errors" "net/http" "sync/atomic" "testing" "time" internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" ) func TestHomeDispatchSelectionOwnsScopeOutsideAuth(t *testing.T) { registry := executionregistry.New() pending, errBegin := registry.BeginDispatch() if errBegin != nil { t.Fatal(errBegin) } scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{RequestID: "req-1", CredentialID: "cred-1", Model: "gpt", Kind: "http", StartedAt: time.Now()}) if errInstall != nil { t.Fatal(errInstall) } selection, errSelection := newHomeDispatchSelection(&Auth{ID: "cred-1", Provider: "codex"}, nil, "codex", scope) if errSelection != nil { t.Fatal(errSelection) } clone := selection.CloneAuth() if clone == nil || clone.ID != "cred-1" || clone.Runtime != nil { t.Fatalf("clone = %#v", clone) } closed := atomic.Int32{} if errBind := selection.Bind(func() error { closed.Add(1); return nil }); errBind != nil { t.Fatal(errBind) } selection.End("completed") selection.End("duplicate") if closed.Load() != 1 { t.Fatalf("close calls = %d", closed.Load()) } } func TestHomeDispatchSelectionReplaceAuthPreservesRoutingAttributes(t *testing.T) { selection := &HomeDispatchSelection{Auth: &Auth{ ID: "cred-1", Provider: "codex", Attributes: map[string]string{ homeUpstreamModelAttributeKey: "gpt-5-upstream", homeForceMappingAttributeKey: "true", homeOriginalAliasAttributeKey: "team/gpt-5", }, Metadata: map[string]any{"access_token": "old"}, }} selection.ReplaceAuth(&Auth{ ID: "cred-1", Provider: "codex", Attributes: map[string]string{AttributeAuthKind: AuthKindOAuth}, Metadata: map[string]any{"access_token": "fresh"}, }) updated := selection.CloneAuth() if updated == nil || updated.Metadata["access_token"] != "fresh" { t.Fatalf("updated auth = %#v", updated) } if updated.Attributes[homeUpstreamModelAttributeKey] != "gpt-5-upstream" || updated.Attributes[homeForceMappingAttributeKey] != "true" || updated.Attributes[homeOriginalAliasAttributeKey] != "team/gpt-5" { t.Fatalf("routing attributes were not preserved: %#v", updated.Attributes) } } func TestHomeDispatchSelectionReplaceAuthConcurrentClone(t *testing.T) { selection := &HomeDispatchSelection{Auth: &Auth{ID: "cred-1", Metadata: map[string]any{"access_token": "old"}}} done := make(chan struct{}) go func() { defer close(done) for i := 0; i < 1000; i++ { selection.ReplaceAuth(&Auth{ID: "cred-1", Metadata: map[string]any{"access_token": "fresh"}}) } }() for i := 0; i < 1000; i++ { if auth := selection.CloneAuth(); auth == nil || auth.ID != "cred-1" { t.Fatalf("CloneAuth() = %#v", auth) } } <-done } func TestReplaceHomeSelectionAuthUpdatesRetainedRuntimeAuth(t *testing.T) { selection := &HomeDispatchSelection{Auth: &Auth{ID: "cred-1", Provider: "codex", Metadata: map[string]any{"access_token": "old"}}} manager := &Manager{ homeRuntimeAuths: map[string]map[string]*Auth{ "session-1": {"cred-1": selection.Auth.Clone()}, }, homeRuntimeAuthOwners: map[string]map[string]*HomeDispatchSelection{ "session-1": {"cred-1": selection}, }, } manager.replaceHomeSelectionAuth(selection, &Auth{ID: "cred-1", Provider: "codex", Metadata: map[string]any{"access_token": "fresh"}}) retained := manager.homeRuntimeAuths["session-1"]["cred-1"] if retained == nil || retained.Metadata["access_token"] != "fresh" { t.Fatalf("retained runtime auth = %#v, want fresh token", retained) } } func TestHomeDispatchSelectionDrainsResourcesAddedDuringEnd(t *testing.T) { registry := executionregistry.New() pending, errBegin := registry.BeginDispatch() if errBegin != nil { t.Fatal(errBegin) } scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{}) if errInstall != nil { t.Fatal(errInstall) } selection, errSelection := newHomeDispatchSelection(&Auth{ID: "cred-1"}, nil, "test", scope) if errSelection != nil { t.Fatal(errSelection) } started := make(chan struct{}) release := make(chan struct{}) if errBind := selection.Bind(func() error { close(started) <-release return nil }); errBind != nil { t.Fatal(errBind) } done := make(chan struct{}) go func() { selection.End("draining") close(done) }() <-started closedLate := atomic.Int32{} errLate := selection.Bind(func() error { closedLate.Add(1) return errors.New("late close") }) if !errors.Is(errLate, executionregistry.ErrRegistryNotAccepting) { t.Fatalf("late Bind() error = %v, want ErrRegistryNotAccepting", errLate) } if closedLate.Load() != 1 { t.Fatalf("late close calls = %d, want 1", closedLate.Load()) } close(release) select { case <-done: case <-time.After(time.Second): t.Fatal("End did not complete") } drainCtx, cancelDrain := context.WithTimeout(context.Background(), time.Second) defer cancelDrain() if errDrain := registry.Drain(drainCtx); errDrain != nil { t.Fatalf("Drain() error = %v", errDrain) } } type gatedHomeDispatcher struct { loaded chan struct{} release chan struct{} rpop atomic.Int32 } func (d *gatedHomeDispatcher) HeartbeatOK() bool { select { case <-d.loaded: default: close(d.loaded) } <-d.release return true } func (d *gatedHomeDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { d.rpop.Add(1) return nil, errors.New("old Home dispatcher was used") } func (*gatedHomeDispatcher) AbortAmbiguousDispatch() {} func TestManagerHomeDispatchBundleCompareAndClearDoesNotRemoveReplacement(t *testing.T) { manager := NewManager(nil, nil, nil) first := manager.PublishHomeDispatch(&gatedHomeDispatcher{loaded: make(chan struct{}), release: make(chan struct{})}, executionregistry.New(), 1) second := manager.PublishHomeDispatch(&gatedHomeDispatcher{loaded: make(chan struct{}), release: make(chan struct{})}, executionregistry.New(), 2) if manager.ClearHomeDispatchBundle(first) { t.Fatal("ClearHomeDispatchBundle() cleared a replacement bundle") } if got := manager.HomeDispatchBundle(); got != second { t.Fatalf("HomeDispatchBundle() = %p, want %p", got, second) } if !manager.ClearHomeDispatchBundle(second) { t.Fatal("ClearHomeDispatchBundle() = false, want true") } if got := manager.HomeDispatchBundle(); got != nil { t.Fatalf("HomeDispatchBundle() = %p, want nil", got) } } func TestPickHomeDispatchSelectionDoesNotMixDetachedBundleWithReplacement(t *testing.T) { manager := NewManager(nil, nil, nil) manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) oldDispatcher := &gatedHomeDispatcher{loaded: make(chan struct{}), release: make(chan struct{})} oldRegistry := executionregistry.New() oldBundle := manager.PublishHomeDispatch(oldDispatcher, oldRegistry, 1) result := make(chan error, 1) go func() { _, errSelect := manager.pickHomeDispatchSelection(context.Background(), "gpt-5.4", cliproxyexecutor.Options{}) result <- errSelect }() select { case <-oldDispatcher.loaded: case <-time.After(time.Second): t.Fatal("selection did not load the old dispatch bundle") } if !manager.ClearHomeDispatchBundle(oldBundle) { t.Fatal("ClearHomeDispatchBundle() = false, want true") } drainCtx, cancelDrain := context.WithTimeout(context.Background(), time.Second) defer cancelDrain() if errDrain := oldRegistry.Drain(drainCtx); errDrain != nil { t.Fatalf("old registry Drain() error = %v", errDrain) } manager.PublishHomeDispatch(&gatedHomeDispatcher{loaded: make(chan struct{}), release: make(chan struct{})}, executionregistry.New(), 2) close(oldDispatcher.release) select { case errSelect := <-result: var authErr *Error if !errors.As(errSelect, &authErr) || authErr.Code != "home_unavailable" { t.Fatalf("pickHomeDispatchSelection() error = %v, want home_unavailable", errSelect) } case <-time.After(time.Second): t.Fatal("selection did not resume after the old bundle was detached") } if got := oldDispatcher.rpop.Load(); got != 0 { t.Fatalf("old dispatcher RPopAuth() calls = %d, want 0", got) } }