Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
42 kB · 1116 lines
Go
at commit e8e39526
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117package executor
import ( "context" "encoding/json" "fmt" "net" "net/http" "net/http/httptest" "net/url" "reflect" "strconv" "strings" "sync" "sync/atomic" "testing" "time"
"github.com/gorilla/websocket" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" internalhome "github.com/router-for-me/CLIProxyAPI/v7/internal/home" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator")
type rejectSecondBindLifecycle struct { binds atomic.Int32}
func (l *rejectSecondBindLifecycle) Bind(func() error) error { if l.binds.Add(1) > 1 { return fmt.Errorf("retry lifecycle bind rejected") } return nil}
func (*rejectSecondBindLifecycle) End(string) {}
func TestCodexWebsocketSessionActiveChannelBelongsToConnection(t *testing.T) { sess := &codexWebsocketSession{} oldConn := &websocket.Conn{} newConn := &websocket.Conn{} oldCh := make(chan codexWebsocketRead, 1) newCh := make(chan codexWebsocketRead, 1)
sess.setActive(oldConn, oldCh) if ch, _ := sess.activeForConn(oldConn); ch != oldCh { t.Fatal("old connection did not own its active channel") }
sess.setActive(newConn, newCh) if sess.clearActive(oldConn, oldCh) { t.Fatal("old connection cleared the new active channel") } if ch, _ := sess.activeForConn(oldConn); ch != nil { t.Fatal("old connection retained access to an active channel") } if ch, _ := sess.activeForConn(newConn); ch != newCh { t.Fatal("new connection lost its active channel") } if !sess.clearActive(newConn, newCh) { t.Fatal("new connection could not clear its active channel") }
closedOldCh := sess.activate(oldConn) if !sess.clearActive(oldConn, closedOldCh) { t.Fatal("old connection could not clear its active channel before retry") } close(closedOldCh) retryCh := sess.activate(newConn) if retryCh == closedOldCh { t.Fatal("retry reused the old connection's read channel") } select { case retryCh <- codexWebsocketRead{conn: newConn}: default: t.Fatal("retry read channel was not writable") }}
type trackedWebsocketLifecycle struct { mu sync.Mutex close func() error once sync.Once ends atomic.Int32}
type drainDuringBindWebsocketLifecycle struct{}
func (drainDuringBindWebsocketLifecycle) Bind(closeFn func() error) error { if errClose := closeFn(); errClose != nil { return errClose } return fmt.Errorf("execution lifecycle drained during Bind")}
func (drainDuringBindWebsocketLifecycle) End(string) {}
func (l *trackedWebsocketLifecycle) Bind(closeFn func() error) error { l.mu.Lock() l.close = closeFn l.mu.Unlock() return nil}
func (l *trackedWebsocketLifecycle) End(string) { l.once.Do(func() { l.ends.Add(1) l.mu.Lock() closeFn := l.close l.mu.Unlock() if closeFn != nil { _ = closeFn() } })}
func TestClearRetryActiveStateClearsOriginalConnection(t *testing.T) { sess := &codexWebsocketSession{} originalConn := &websocket.Conn{} originalCh := sess.activate(originalConn) if !clearRetryActiveState(sess, originalConn, originalCh) { t.Fatal("clearRetryActiveState() = false, want true") } if ch, done := sess.activeForConn(originalConn); ch != nil || done != nil { t.Fatalf("original active state = %v/%v, want nil", ch, done) }}
func TestWebsocketRetryBindFailureClearsActiveSessionState(t *testing.T) { tests := []struct { name string run func(t *testing.T, baseURL string) (func(cliproxyexecutor.Options) error, *codexWebsocketSession) }{ { name: "Codex nonstream", run: func(t *testing.T, baseURL string) (func(cliproxyexecutor.Options) error, *codexWebsocketSession) { executor := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) executor.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} auth := &cliproxyauth.Auth{ID: "retry-bind-codex", Provider: "codex", Attributes: map[string]string{"api_key": "test-key", "base_url": baseURL}} req := cliproxyexecutor.Request{Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`)} primed := false return func(runOpts cliproxyexecutor.Options) error { if !primed { wsURL := "ws" + strings.TrimPrefix(baseURL, "http") + "/responses" conn, _, _, errEnsure := executor.ensureUpstreamConn(context.Background(), auth, executor.getOrCreateSession("retry-bind"), auth.ID, wsURL, http.Header{}) if errEnsure != nil { return errEnsure } if errDeadline := conn.SetWriteDeadline(time.Now().Add(-time.Second)); errDeadline != nil { return errDeadline } primed = true } _, errExecute := executor.Execute(context.Background(), auth, req, runOpts) return errExecute }, executor.getOrCreateSession("retry-bind") }, }, { name: "Codex stream", run: func(t *testing.T, baseURL string) (func(cliproxyexecutor.Options) error, *codexWebsocketSession) { executor := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) executor.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} auth := &cliproxyauth.Auth{ID: "retry-bind-codex", Provider: "codex", Attributes: map[string]string{"api_key": "test-key", "base_url": baseURL}} req := cliproxyexecutor.Request{Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`)} primed := false return func(runOpts cliproxyexecutor.Options) error { if !primed { wsURL := "ws" + strings.TrimPrefix(baseURL, "http") + "/responses" conn, _, _, errEnsure := executor.ensureUpstreamConn(context.Background(), auth, executor.getOrCreateSession("retry-bind"), auth.ID, wsURL, http.Header{}) if errEnsure != nil { return errEnsure } if errDeadline := conn.SetWriteDeadline(time.Now().Add(-time.Second)); errDeadline != nil { return errDeadline } primed = true } result, errExecute := executor.ExecuteStream(context.Background(), auth, req, runOpts) if errExecute != nil { return errExecute } for chunk := range result.Chunks { if chunk.Err != nil { return chunk.Err } } return nil }, executor.getOrCreateSession("retry-bind") }, }, { name: "xAI stream", run: func(t *testing.T, baseURL string) (func(cliproxyexecutor.Options) error, *codexWebsocketSession) { executor := NewXAIWebsocketsExecutor(&config.Config{}) executor.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} auth := &cliproxyauth.Auth{ID: "retry-bind-xai", Provider: "xai", Attributes: map[string]string{"base_url": baseURL, "websockets": "true"}, Metadata: map[string]any{"access_token": "test-token"}} req := cliproxyexecutor.Request{Model: "grok-4", Payload: []byte(`{"model":"grok-4","input":[{"type":"message","role":"user","content":"hello"}]}`)} primed := false return func(runOpts cliproxyexecutor.Options) error { if !primed { wsURL := "ws" + strings.TrimPrefix(baseURL, "http") + "/responses" conn, _, _, errEnsure := executor.ensureUpstreamConn(context.Background(), auth, executor.getOrCreateSession("retry-bind"), auth.ID, wsURL, http.Header{}) if errEnsure != nil { return errEnsure } if errDeadline := conn.SetWriteDeadline(time.Now().Add(-time.Second)); errDeadline != nil { return errDeadline } primed = true } result, errExecute := executor.ExecuteStream(context.Background(), auth, req, runOpts) if errExecute != nil { return errExecute } for chunk := range result.Chunks { if chunk.Err != nil { return chunk.Err } } return nil }, executor.getOrCreateSession("retry-bind") }, }, }
for _, test := range tests { t.Run(test.name, func(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} var connections atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { conn, errUpgrade := upgrader.Upgrade(w, r, nil) if errUpgrade != nil { t.Errorf("upgrade websocket: %v", errUpgrade) return } connection := connections.Add(1) defer func() { _ = conn.Close() }() if connection == 1 { _, _, _ = conn.ReadMessage() return } if connection == 2 { return } if _, _, errRead := conn.ReadMessage(); errRead != nil { return } completed := []byte(`{"type":"response.completed","response":{"id":"response-1","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { t.Errorf("write websocket completion: %v", errWrite) } })) defer server.Close()
lifecycle := &rejectSecondBindLifecycle{} opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse, ResponseFormat: sdktranslator.FormatOpenAIResponse, ExecutionLifecycle: lifecycle, Metadata: map[string]any{cliproxyexecutor.ExecutionSessionMetadataKey: "retry-bind"}} run, sess := test.run(t, server.URL) if errRun := run(opts); errRun == nil { t.Fatal("first request error = nil, want retry lifecycle bind rejection") } if got := lifecycle.binds.Load(); got != 2 { t.Fatalf("lifecycle binds = %d, want 2", got) } sess.activeMu.Lock() active := sess.activeConn != nil || sess.activeCh != nil || sess.activeDone != nil || sess.activeCancel != nil sess.activeMu.Unlock() if active { t.Fatal("retry bind failure left the old active websocket state") }
opts.ExecutionLifecycle = nil if errRun := run(opts); errRun != nil { t.Fatalf("second request error = %v", errRun) } if got := connections.Load(); got != 3 { t.Fatalf("websocket connections = %d, want 3 after retry bind failure", got) } }) }}
func TestWebsocketSessionCloseEndsRetainedLifecycleOnce(t *testing.T) { exec := NewCodexWebsocketsExecutor(&config.Config{}) exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} server, closed := newWebsocketTargetServer(t) defer server.Close()
sess := exec.getOrCreateSession("retained-lifecycle") auth := &cliproxyauth.Auth{ID: "auth-a"} wsURL := "ws" + strings.TrimPrefix(server.URL, "http") conn := ensureWebsocketTargetConn(t, exec.ensureUpstreamConn, auth, sess, auth.ID, wsURL) lifecycle := &trackedWebsocketLifecycle{} if errBind := sess.bindExecutionLifecycle(cliproxyexecutor.Options{ExecutionLifecycle: lifecycle}, conn, sess.connCloser, "model-a"); errBind != nil { t.Fatalf("bind execution lifecycle: %v", errBind) }
exec.CloseExecutionSession("retained-lifecycle") lifecycle.End("duplicate_close") if got := lifecycle.ends.Load(); got != 1 { t.Fatalf("lifecycle End calls = %d, want 1", got) } if got := <-closed; got != auth.ID { t.Fatalf("closed server auth = %q, want %q", got, auth.ID) }}
type closeCountingNetConn struct { net.Conn closes atomic.Int32}
func (c *closeCountingNetConn) Close() error { c.closes.Add(1) return c.Conn.Close()}
func newCloseCountingWebsocketConn(t *testing.T, rawURL string) (*websocket.Conn, *closeCountingNetConn) { t.Helper() parsed, errParse := url.Parse(rawURL) if errParse != nil { t.Fatalf("parse websocket URL: %v", errParse) } conn, errDial := net.Dial("tcp", parsed.Host) if errDial != nil { t.Fatalf("dial websocket: %v", errDial) } counting := &closeCountingNetConn{Conn: conn} wsConn, _, errClient := websocket.NewClient(counting, parsed, nil, 1024, 1024) if errClient != nil { _ = counting.Close() t.Fatalf("create websocket client: %v", errClient) } return wsConn, counting}
func TestSessionlessWebsocketSelectionEndAndDirectCloseRaceClosesOnce(t *testing.T) { server, _ := newWebsocketTargetServer(t) defer server.Close()
wsURL := "ws" + strings.TrimPrefix(server.URL, "http") conn, physical := newCloseCountingWebsocketConn(t, wsURL) closer := newWebsocketConnectionCloser(conn) lifecycle := &trackedWebsocketLifecycle{} if errBind := (*codexWebsocketSession)(nil).bindExecutionLifecycle(cliproxyexecutor.Options{ExecutionLifecycle: lifecycle}, conn, closer, "model-a"); errBind != nil { t.Fatalf("bind sessionless lifecycle: %v", errBind) }
var wait sync.WaitGroup wait.Add(2) go func() { defer wait.Done() lifecycle.End("selection_ended") }() go func() { defer wait.Done() if errClose := closer.Close(); errClose != nil { t.Errorf("direct close: %v", errClose) } }() wait.Wait()
if got := physical.closes.Load(); got != 1 { t.Fatalf("physical websocket closes = %d, want 1", got) }}
func TestWebsocketDrainDuringBindClosesOwnedConnectionOnce(t *testing.T) { server, _ := newWebsocketTargetServer(t) defer server.Close()
wsURL := "ws" + strings.TrimPrefix(server.URL, "http") conn, physical := newCloseCountingWebsocketConn(t, wsURL) closer := newWebsocketConnectionCloser(conn) sess := &codexWebsocketSession{conn: conn, connCloser: closer, wsURL: wsURL, authID: "auth-a", readerConn: conn}
errBind := sess.bindExecutionLifecycle(cliproxyexecutor.Options{ExecutionLifecycle: drainDuringBindWebsocketLifecycle{}}, conn, closer, "model-a") if errBind == nil { t.Fatal("bind execution lifecycle error = nil, want drain error") } closeWebsocketAfterBindFailure(sess, conn, closer)
if got := physical.closes.Load(); got != 1 { t.Fatalf("physical websocket closes = %d, want 1", got) } sess.connMu.Lock() defer sess.connMu.Unlock() if sess.conn != nil || sess.connCloser != nil || sess.lifecycle != nil { t.Fatalf("drained session state = conn:%v closer:%v lifecycle:%v, want detached", sess.conn, sess.connCloser, sess.lifecycle) }}
func TestWebsocketTargetReplacementPhysicallyClosesOwnedConnectionOnce(t *testing.T) { tests := []struct { name string }{ {name: "Codex"}, {name: "xAI"}, }
for _, test := range tests { t.Run(test.name, func(t *testing.T) { serverA, _ := newWebsocketTargetServer(t) defer serverA.Close() serverB, _ := newWebsocketTargetServer(t) defer serverB.Close()
var ensure func(context.Context, *cliproxyauth.Auth, *codexWebsocketSession, string, string, http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error) var closeSession func(string) var sess *codexWebsocketSession switch test.name { case "Codex": exec := NewCodexWebsocketsExecutor(&config.Config{}) exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} ensure = exec.ensureUpstreamConn closeSession = exec.CloseExecutionSession sess = exec.getOrCreateSession("counted-target-change") case "xAI": exec := NewXAIWebsocketsExecutor(&config.Config{}) exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} ensure = exec.ensureUpstreamConn closeSession = exec.CloseExecutionSession sess = exec.getOrCreateSession("counted-target-change") } defer closeSession("counted-target-change")
wsURLA := "ws" + strings.TrimPrefix(serverA.URL, "http") wsURLB := "ws" + strings.TrimPrefix(serverB.URL, "http") connA, physical := newCloseCountingWebsocketConn(t, wsURLA) sess.connMu.Lock() sess.conn = connA sess.connCloser = newWebsocketConnectionCloser(connA) sess.wsURL = wsURLA sess.authID = "auth-a" sess.readerConn = connA sess.connMu.Unlock() lifecycle := &trackedWebsocketLifecycle{} if errBind := sess.bindExecutionLifecycle(cliproxyexecutor.Options{ExecutionLifecycle: lifecycle}, connA, sess.connCloser, "model-a"); errBind != nil { t.Fatalf("bind execution lifecycle: %v", errBind) }
if _, _, _, errEnsure := ensure(context.Background(), &cliproxyauth.Auth{ID: "auth-b"}, sess, "auth-b", wsURLB, nil); errEnsure != nil { t.Fatalf("replace websocket target: %v", errEnsure) } if got := physical.closes.Load(); got != 1 { t.Fatalf("physical websocket closes = %d, want 1", got) } }) }}
func TestWebsocketLifecycleEndThenInvalidateAndCloseAllPhysicallyClosesOnce(t *testing.T) { tests := []struct { name string run func(*codexWebsocketSession, *websocket.Conn, *trackedWebsocketLifecycle) }{ { name: "Codex", run: func(sess *codexWebsocketSession, conn *websocket.Conn, lifecycle *trackedWebsocketLifecycle) { exec := NewCodexWebsocketsExecutor(&config.Config{}) exec.store = &codexWebsocketSessionStore{sessions: map[string]*codexWebsocketSession{sess.sessionID: sess}} lifecycle.End("lifecycle_ended") exec.invalidateUpstreamConn(sess, conn, "invalidated", nil) exec.CloseExecutionSession(cliproxyauth.CloseAllExecutionSessionsID) }, }, { name: "xAI", run: func(sess *codexWebsocketSession, conn *websocket.Conn, lifecycle *trackedWebsocketLifecycle) { exec := NewXAIWebsocketsExecutor(&config.Config{}) exec.store = &codexWebsocketSessionStore{sessions: map[string]*codexWebsocketSession{sess.sessionID: sess}} lifecycle.End("lifecycle_ended") exec.invalidateUpstreamConn(sess, conn, "invalidated", nil) exec.CloseExecutionSession(cliproxyauth.CloseAllExecutionSessionsID) }, }, }
for _, test := range tests { t.Run(test.name, func(t *testing.T) { server, _ := newWebsocketTargetServer(t) defer server.Close() wsURL := "ws" + strings.TrimPrefix(server.URL, "http") conn, physical := newCloseCountingWebsocketConn(t, wsURL) sess := &codexWebsocketSession{sessionID: "counted-lifecycle", conn: conn, connCloser: newWebsocketConnectionCloser(conn), wsURL: wsURL, authID: "auth-a", readerConn: conn} lifecycle := &trackedWebsocketLifecycle{} if errBind := sess.bindExecutionLifecycle(cliproxyexecutor.Options{ExecutionLifecycle: lifecycle}, conn, sess.connCloser, "model-a"); errBind != nil { t.Fatalf("bind execution lifecycle: %v", errBind) }
test.run(sess, conn, lifecycle) if got := physical.closes.Load(); got != 1 { t.Fatalf("physical websocket closes = %d, want 1", got) } }) }}
func TestWebsocketExecutorsReconnectWhenSessionTargetChanges(t *testing.T) { t.Run("Codex", func(t *testing.T) { exec := NewCodexWebsocketsExecutor(&config.Config{}) exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} testWebsocketExecutorReconnectsWhenSessionTargetChanges( t, exec.UpstreamDisconnectChan, exec.getOrCreateSession, exec.ensureUpstreamConn, exec.CloseExecutionSession, ) })
t.Run("xAI", func(t *testing.T) { exec := NewXAIWebsocketsExecutor(&config.Config{}) exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} testWebsocketExecutorReconnectsWhenSessionTargetChanges( t, exec.UpstreamDisconnectChan, exec.getOrCreateSession, exec.ensureUpstreamConn, exec.CloseExecutionSession, ) })}
func testWebsocketExecutorReconnectsWhenSessionTargetChanges( t *testing.T, disconnectChan func(string) <-chan error, getSession func(string) *codexWebsocketSession, ensureConn func(context.Context, *cliproxyauth.Auth, *codexWebsocketSession, string, string, http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error), closeSession func(string),) { t.Helper()
serverA, closedA := newWebsocketTargetServer(t) defer serverA.Close() serverB, closedB := newWebsocketTargetServer(t) defer serverB.Close()
sessionID := "target-switch-session" disconnectCh := disconnectChan(sessionID) sess := getSession(sessionID) if sess == nil { t.Fatal("expected websocket session") } defer closeSession(sessionID)
authA := &cliproxyauth.Auth{ID: "auth-a"} authB := &cliproxyauth.Auth{ID: "auth-b"} wsURLA := "ws" + strings.TrimPrefix(serverA.URL, "http") wsURLB := "ws" + strings.TrimPrefix(serverB.URL, "http")
connA := ensureWebsocketTargetConn(t, ensureConn, authA, sess, authA.ID, wsURLA) connAReused := ensureWebsocketTargetConn(t, ensureConn, authA, sess, authA.ID, wsURLA) if connAReused != connA { t.Fatal("matching websocket target did not reuse the existing connection") }
connURLB := ensureWebsocketTargetConn(t, ensureConn, authA, sess, authA.ID, wsURLB) if connURLB == connA { t.Fatal("websocket URL change reused the existing connection") } if got := <-closedA; got != authA.ID { t.Fatalf("closed server A auth = %q, want %q", got, authA.ID) }
connAuthB := ensureWebsocketTargetConn(t, ensureConn, authB, sess, authB.ID, wsURLB) if connAuthB == connURLB { t.Fatal("websocket auth change reused the existing connection") } if got := <-closedB; got != authA.ID { t.Fatalf("first closed server B auth = %q, want %q", got, authA.ID) }
sess.connMu.Lock() gotAuthID := sess.authID gotURL := sess.wsURL sess.connMu.Unlock() if gotAuthID != authB.ID || gotURL != wsURLB { t.Fatalf("session target = {%q %q}, want {%q %q}", gotAuthID, gotURL, authB.ID, wsURLB) }
select { case errDisconnect := <-disconnectCh: t.Fatalf("controlled websocket target switch notified downstream: %v", errDisconnect) default: }}
func ensureWebsocketTargetConn( t *testing.T, ensureConn func(context.Context, *cliproxyauth.Auth, *codexWebsocketSession, string, string, http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error), auth *cliproxyauth.Auth, sess *codexWebsocketSession, authID string, wsURL string,) *websocket.Conn { t.Helper() headers := http.Header{"X-Test-Auth": []string{authID}} conn, _, resp, errEnsure := ensureConn(context.Background(), auth, sess, authID, wsURL, headers) if resp != nil && resp.Body != nil { defer func() { if errClose := resp.Body.Close(); errClose != nil { t.Errorf("close handshake response body: %v", errClose) } }() } if errEnsure != nil { t.Fatalf("ensure websocket connection: %v", errEnsure) } if conn == nil { t.Fatal("ensure websocket connection returned nil") } return conn}
func newWebsocketTargetServer(t *testing.T) (*httptest.Server, <-chan string) { t.Helper() closed := make(chan string, 4) upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { authID := r.Header.Get("X-Test-Auth") conn, errUpgrade := upgrader.Upgrade(w, r, nil) if errUpgrade != nil { t.Errorf("upgrade websocket: %v", errUpgrade) return } defer func() { if errClose := conn.Close(); errClose != nil { t.Errorf("close upstream websocket: %v", errClose) } closed <- authID }() for { if _, _, errRead := conn.ReadMessage(); errRead != nil { return } } })) return server, closed}
type registryDrainWebsocketLifecycle struct { scope *executionregistry.Scope ends atomic.Int32}
func (l *registryDrainWebsocketLifecycle) Bind(closeFn func() error) error { return l.scope.Bind(closeFn)}
func (l *registryDrainWebsocketLifecycle) End(string) { l.ends.Add(1) l.scope.End("websocket_closed")}
func (l *registryDrainWebsocketLifecycle) Retain() {}
type websocketHomeDispatcher struct { provider string}
func (d websocketHomeDispatcher) HeartbeatOK() bool { return true }
func (d websocketHomeDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { return json.Marshal(map[string]any{"auth": map[string]any{ "id": "home-websocket-auth", "provider": d.provider, "status": "active", "attributes": map[string]string{ "api_key": "home-key", }, }})}
func (websocketHomeDispatcher) AbortAmbiguousDispatch() {}
type accountedWebsocketHomeDispatcher struct { provider string baseURL string calls atomic.Int32 releases atomic.Int32 before atomic.Bool}
func (*accountedWebsocketHomeDispatcher) HeartbeatOK() bool { return true }
func (d *accountedWebsocketHomeDispatcher) RPopAuth(_ context.Context, model string, _ string, _ http.Header, _ int) ([]byte, error) { call := d.calls.Add(1) if call > 1 && d.releases.Load() != call-1 { d.before.Store(false) } else if call > 1 { d.before.Store(true) } upstreamModel := "model-a" if strings.Contains(strings.ToLower(model), "(custom)") { upstreamModel = "model-a(custom)" } return json.Marshal(map[string]any{ "model": upstreamModel, "auth_index": "accounted-websocket-auth", "auth": map[string]any{ "id": "accounted-websocket-auth", "provider": d.provider, "status": "active", "attributes": map[string]string{ "api_key": "test-key", "base_url": d.baseURL, "websockets": "true", }, }, "concurrency": map[string]any{ "accounted": true, "credential_id": "accounted-websocket-auth", "model": upstreamModel, }, })}
func (*accountedWebsocketHomeDispatcher) AbortAmbiguousDispatch() {}
func TestAuditAccountedCodexXAIReconnectReuseAndTargetChange(t *testing.T) { tests := []struct { name string provider string newExecutor func() cliproxyauth.ProviderExecutor }{ { name: "Codex", provider: "codex", newExecutor: func() cliproxyauth.ProviderExecutor { executor := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) executor.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} return executor }, }, { name: "xAI", provider: "xai", newExecutor: func() cliproxyauth.ProviderExecutor { executor := NewXAIWebsocketsExecutor(&config.Config{}) executor.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} executor.idStore = &xaiWebsocketIDStateStore{sessions: make(map[string]*xaiWebsocketIDState)} return executor }, }, }
for _, test := range tests { t.Run(test.name, func(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} var connections atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { conn, errUpgrade := upgrader.Upgrade(w, r, nil) if errUpgrade != nil { t.Errorf("upgrade websocket: %v", errUpgrade) return } connections.Add(1) defer func() { _ = conn.Close() }() for { if _, _, errRead := conn.ReadMessage(); errRead != nil { return } completed := []byte(`{"type":"response.completed","response":{"id":"response-1","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { return } } })) defer server.Close()
registry := executionregistry.New() dispatcher := &accountedWebsocketHomeDispatcher{provider: test.provider, baseURL: server.URL} var releaseGroups []executionregistry.ReleaseGroup registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, _ int64) { dispatcher.releases.Add(1) releaseGroups = append(releaseGroups, group) }) manager := cliproxyauth.NewManager(nil, nil, nil) manager.SetConfig(&config.Config{Home: config.HomeConfig{Enabled: true}}) manager.PublishHomeDispatch(dispatcher, registry, 1) manager.RegisterExecutor(test.newExecutor()) t.Cleanup(func() { manager.CloseExecutionSession("accounted-websocket-session") })
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) opts := cliproxyexecutor.Options{ Stream: true, SourceFormat: sdktranslator.FormatOpenAIResponse, ResponseFormat: sdktranslator.FormatOpenAIResponse, Metadata: map[string]any{ cliproxyexecutor.ExecutionSessionMetadataKey: "accounted-websocket-session", cliproxyexecutor.PinnedAuthMetadataKey: "accounted-websocket-auth", }, } execute := func(model string) { t.Helper() result, errExecute := manager.ExecuteStream(ctx, []string{test.provider}, cliproxyexecutor.Request{Model: model, Payload: []byte(`{"model":"model-a","input":[]}`)}, opts) if errExecute != nil { t.Fatalf("ExecuteStream(%q) error = %v", model, errExecute) } for chunk := range result.Chunks { if chunk.Err != nil { t.Fatalf("ExecuteStream(%q) chunk error = %v", model, chunk.Err) } } }
execute(" MODEL-A(HIGH) ") execute("model-a") if got := dispatcher.calls.Load(); got != 1 { t.Fatalf("Home RPOP calls = %d, want 1 for canonical retained reuse", got) } manager.CloseExecutionSession("accounted-websocket-session") execute("model-a") execute("model-a(custom)") if got := dispatcher.calls.Load(); got != 3 { t.Fatalf("Home RPOP calls = %d, want 3 after reconnect and target change", got) } if !dispatcher.before.Load() { t.Fatal("previous accounted selection was not released before redispatch") } manager.CloseExecutionSession("accounted-websocket-session") wantGroups := []executionregistry.ReleaseGroup{ {CredentialID: "accounted-websocket-auth", Model: "model-a"}, {CredentialID: "accounted-websocket-auth", Model: "model-a"}, {CredentialID: "accounted-websocket-auth", Model: "model-a(custom)"}, } if !reflect.DeepEqual(releaseGroups, wantGroups) { t.Fatalf("release groups = %#v, want %#v", releaseGroups, wantGroups) } if got := connections.Load(); got != 3 { t.Fatalf("upstream websocket connections = %d, want 3", got) } }) }}
func TestHomeSelectionRegistryDrainClosesRealWebsocketSessions(t *testing.T) { tests := []struct { name string provider string newExecutor func() (cliproxyauth.ProviderExecutor, func(string) *codexWebsocketSession, func(context.Context, *cliproxyauth.Auth, *codexWebsocketSession, string, string, http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error)) }{ { name: "Codex", provider: "codex", newExecutor: func() (cliproxyauth.ProviderExecutor, func(string) *codexWebsocketSession, func(context.Context, *cliproxyauth.Auth, *codexWebsocketSession, string, string, http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error)) { executor := NewCodexWebsocketsExecutor(&config.Config{}) executor.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} return executor, executor.getOrCreateSession, executor.ensureUpstreamConn }, }, { name: "xAI", provider: "xai", newExecutor: func() (cliproxyauth.ProviderExecutor, func(string) *codexWebsocketSession, func(context.Context, *cliproxyauth.Auth, *codexWebsocketSession, string, string, http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error)) { executor := NewXAIWebsocketsExecutor(&config.Config{}) executor.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} return executor, executor.getOrCreateSession, executor.ensureUpstreamConn }, }, }
for _, test := range tests { t.Run(test.name, func(t *testing.T) { server, closed := newWebsocketTargetServer(t) defer server.Close()
executor, getSession, ensureConn := test.newExecutor() registry := executionregistry.New() manager := cliproxyauth.NewManager(nil, nil, nil) manager.SetConfig(&config.Config{Home: config.HomeConfig{Enabled: true}}) manager.PublishHomeDispatch(websocketHomeDispatcher{provider: test.provider}, registry, 1) manager.RegisterExecutor(executor) selection, errSelect := manager.SelectHomeAuthByKind(context.Background(), test.provider, "model-a", cliproxyauth.AuthKindAPIKey, cliproxyexecutor.Options{}) if errSelect != nil { t.Fatalf("SelectHomeAuthByKind() error = %v", errSelect) } auth := selection.CloneAuth() sess := getSession("real-home-drain") wsURL := "ws" + strings.TrimPrefix(server.URL, "http") conn := ensureWebsocketTargetConn(t, ensureConn, auth, sess, auth.ID, wsURL) if errBind := sess.bindExecutionLifecycle(cliproxyexecutor.Options{ExecutionLifecycle: selection}, conn, sess.connCloser, "model-a"); errBind != nil { t.Fatalf("bind execution lifecycle: %v", errBind) }
drainCtx, cancelDrain := context.WithTimeout(context.Background(), time.Second) defer cancelDrain() if errDrain := registry.Drain(drainCtx); errDrain != nil { t.Fatalf("Drain() error = %v", errDrain) } if selection.Active() { t.Fatal("registry drain did not end the Home dispatch selection") } if got := <-closed; got != auth.ID { t.Fatalf("closed server auth = %q, want %q", got, auth.ID) } }) }}
type codex426RetryDispatcher struct { calls atomic.Int32 baseURLs []string websockets []bool releases atomic.Int32 releasedBeforeSecondRPop atomic.Bool}
func (d *codex426RetryDispatcher) HeartbeatOK() bool { return true }
func (d *codex426RetryDispatcher) RPopAuth(_ context.Context, model string, _ string, _ http.Header, _ int) ([]byte, error) { call := int(d.calls.Add(1)) if call > len(d.baseURLs) { return nil, fmt.Errorf("unexpected Home dispatch %d", call) } if call == 2 { d.releasedBeforeSecondRPop.Store(d.releases.Load() == 1) } credentialID := "codex-home-" + strconv.Itoa(call) attributes := map[string]string{ "api_key": "home-key", "base_url": d.baseURLs[call-1], } if call <= len(d.websockets) && d.websockets[call-1] { attributes["websockets"] = "true" } return json.Marshal(map[string]any{ "model": model, "auth_index": credentialID, "auth": map[string]any{ "id": credentialID, "provider": "codex", "status": "active", "attributes": attributes, }, "concurrency": map[string]any{ "accounted": true, "credential_id": credentialID, "model": model, }, })}
func (*codex426RetryDispatcher) AbortAmbiguousDispatch() {}
func TestAuditHomeCodex426WebsocketToHTTPFreshSelection(t *testing.T) { upgradeRequired := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { http.Error(w, "websocket upgrade required", http.StatusUpgradeRequired) })) defer upgradeRequired.Close()
var httpFallbackCalls atomic.Int32 httpFallback := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodPost || r.URL.Path != "/responses" { http.Error(w, "unexpected fallback request", http.StatusBadRequest) return } httpFallbackCalls.Add(1) w.Header().Set("Content-Type", "text/event-stream") _, _ = w.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"response-1\",\"output\":[],\"usage\":{\"input_tokens\":0,\"output_tokens\":0,\"total_tokens\":0}}}\n\n")) })) defer httpFallback.Close()
executor := NewCodexAutoExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) executor.wsExec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} dispatcher := &codex426RetryDispatcher{ baseURLs: []string{upgradeRequired.URL, httpFallback.URL}, websockets: []bool{true, false}, } registry := executionregistry.New() var releaseGroups []executionregistry.ReleaseGroup var releaseGroupsMu sync.Mutex releaseFlusher := internalhome.NewReleaseFlusher(func() config.CredentialConcurrencyConfig { return config.CredentialConcurrencyConfig{ ReleaseFlushInterval: time.Millisecond, ReleaseMaxBackoff: 10 * time.Millisecond, } }, func(_ context.Context, frame internalhome.ConcurrencyReleaseFrame) error { dispatcher.releases.Add(1) releaseGroupsMu.Lock() releaseGroups = append(releaseGroups, executionregistry.ReleaseGroup{CredentialID: frame.CredentialID, Model: frame.Model}) releaseGroupsMu.Unlock() return nil }) registry.SetReleaseSink(releaseFlusher.MarkDirty) releaseCtx, cancelRelease := context.WithCancel(context.Background()) releaseDone := make(chan struct{}) go func() { defer close(releaseDone) releaseFlusher.Run(releaseCtx) }() defer func() { cancelRelease() <-releaseDone }() manager := cliproxyauth.NewManager(nil, nil, nil) manager.SetConfig(&config.Config{Home: config.HomeConfig{Enabled: true}}) manager.PublishHomeDispatch(dispatcher, registry, 1) manager.RegisterExecutor(executor)
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) result, errExecute := manager.ExecuteStream(ctx, []string{"codex"}, cliproxyexecutor.Request{Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`)}, cliproxyexecutor.Options{Stream: true, SourceFormat: sdktranslator.FormatOpenAIResponse, ResponseFormat: sdktranslator.FormatOpenAIResponse, Metadata: map[string]any{cliproxyexecutor.ExecutionSessionMetadataKey: "home-426"}}) if errExecute != nil { t.Fatalf("ExecuteStream() error = %v", errExecute) } if !dispatcher.releasedBeforeSecondRPop.Load() { t.Fatal("first accounted selection was not released before the 426 retry RPOP") } if got := dispatcher.releases.Load(); got != 1 { t.Fatalf("accounted releases before response completion = %d, want 1", got) } for chunk := range result.Chunks { if chunk.Err != nil { t.Fatalf("stream chunk error = %v", chunk.Err) } } if got := dispatcher.calls.Load(); got != 2 { t.Fatalf("Home RPOP calls = %d, want 2 after 426", got) } if got := httpFallbackCalls.Load(); got != 1 { t.Fatalf("HTTP fallback calls = %d, want 1 on the fresh Home selection", got) } deadline := time.NewTimer(time.Second) defer deadline.Stop() for dispatcher.releases.Load() != 2 { select { case <-deadline.C: t.Fatalf("accounted releases after response completion = %d, want 2", dispatcher.releases.Load()) case <-time.After(time.Millisecond): } } releaseGroupsMu.Lock() gotReleaseGroups := append([]executionregistry.ReleaseGroup(nil), releaseGroups...) releaseGroupsMu.Unlock() wantReleaseGroups := []executionregistry.ReleaseGroup{ {CredentialID: "codex-home-1", Model: "gpt-5-codex"}, {CredentialID: "codex-home-2", Model: "gpt-5-codex"}, } if !reflect.DeepEqual(gotReleaseGroups, wantReleaseGroups) { t.Fatalf("accounted release groups = %#v, want %#v", gotReleaseGroups, wantReleaseGroups) }}
func TestWebsocketRegistryDrainClosesAndEndsRetainedSession(t *testing.T) { tests := []struct { name string getSession func(string) *codexWebsocketSession ensureConn func(context.Context, *cliproxyauth.Auth, *codexWebsocketSession, string, string, http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error) }{ { name: "Codex", getSession: func(sessionID string) *codexWebsocketSession { executor := NewCodexWebsocketsExecutor(&config.Config{}) executor.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} return executor.getOrCreateSession(sessionID) }, ensureConn: func(ctx context.Context, auth *cliproxyauth.Auth, sess *codexWebsocketSession, authID, wsURL string, headers http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error) { executor := NewCodexWebsocketsExecutor(&config.Config{}) return executor.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, headers) }, }, { name: "xAI shared session", getSession: func(sessionID string) *codexWebsocketSession { executor := NewXAIWebsocketsExecutor(&config.Config{}) executor.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} return executor.getOrCreateSession(sessionID) }, ensureConn: func(ctx context.Context, auth *cliproxyauth.Auth, sess *codexWebsocketSession, authID, wsURL string, headers http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error) { executor := NewXAIWebsocketsExecutor(&config.Config{}) return executor.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, headers) }, }, }
for _, test := range tests { t.Run(test.name, func(t *testing.T) { server, closed := newWebsocketTargetServer(t) defer server.Close()
registry := executionregistry.New() pending, errBegin := registry.BeginDispatch() if errBegin != nil { t.Fatalf("BeginDispatch() error = %v", errBegin) } scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{Kind: "websocket"}) if errInstall != nil { t.Fatalf("Install() error = %v", errInstall) } lifecycle := ®istryDrainWebsocketLifecycle{scope: scope} auth := &cliproxyauth.Auth{ID: "auth-a"} wsURL := "ws" + strings.TrimPrefix(server.URL, "http") sess := test.getSession("drain-retained-session") conn := ensureWebsocketTargetConn(t, test.ensureConn, auth, sess, auth.ID, wsURL) if errBind := sess.bindExecutionLifecycle(cliproxyexecutor.Options{ExecutionLifecycle: lifecycle}, conn, sess.connCloser, "model-a"); errBind != nil { t.Fatalf("bind execution lifecycle: %v", errBind) }
drainCtx, cancelDrain := context.WithTimeout(context.Background(), time.Second) defer cancelDrain() if errDrain := registry.Drain(drainCtx); errDrain != nil { t.Fatalf("Drain() error = %v", errDrain) } if got := lifecycle.ends.Load(); got != 1 { t.Fatalf("lifecycle End calls = %d, want 1", got) } if got := <-closed; got != auth.ID { t.Fatalf("closed server auth = %q, want %q", got, auth.ID) } }) }}