package 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)			}		})	}}

Lines 1–500 of 1116. Load more