From fbe116071a0e46525edddd6398043eca8fd90a99 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Sat, 18 Jul 2026 00:40:56 +0800 Subject: [PATCH] feat(websockets): implement active connection tracking for websocket sessions - Introduced per-connection active channel management for Codex and XAI websocket sessions. - Enhanced session behavior to support clearing and reassigning active channels tied to specific connections. - Added support for detecting and detaching stale connections upon session target changes. - Refactored websocket message handling logic to ensure connection-specific channel consistency. - Extended test coverage for websocket session behaviors, including reconnections, target switching, and pinned auth transitions. --- .../executor/codex_websockets_executor.go | 146 +++++++++---- .../executor/websocket_session_target_test.go | 198 ++++++++++++++++++ .../executor/xai_websockets_executor.go | 66 +++--- .../executor/xai_websockets_executor_test.go | 119 +++++++++++ .../openai/openai_responses_websocket.go | 101 ++++++++- .../openai/openai_responses_websocket_test.go | 173 +++++++++++++++ 6 files changed, 720 insertions(+), 83 deletions(-) create mode 100644 internal/runtime/executor/websocket_session_target_test.go diff --git a/internal/runtime/executor/codex_websockets_executor.go b/internal/runtime/executor/codex_websockets_executor.go index 275148f1..af0e415c 100644 --- a/internal/runtime/executor/codex_websockets_executor.go +++ b/internal/runtime/executor/codex_websockets_executor.go @@ -72,6 +72,7 @@ type codexWebsocketSession struct { writeMu sync.Mutex activeMu sync.Mutex + activeConn *websocket.Conn activeCh chan codexWebsocketRead activeDone <-chan struct{} activeCancel context.CancelFunc @@ -96,7 +97,7 @@ type codexWebsocketRead struct { err error } -func (s *codexWebsocketSession) setActive(ch chan codexWebsocketRead) { +func (s *codexWebsocketSession) setActive(conn *websocket.Conn, ch chan codexWebsocketRead) { if s == nil { return } @@ -106,8 +107,9 @@ func (s *codexWebsocketSession) setActive(ch chan codexWebsocketRead) { s.activeCancel = nil s.activeDone = nil } + s.activeConn = conn s.activeCh = ch - if ch != nil { + if conn != nil && ch != nil { activeCtx, activeCancel := context.WithCancel(context.Background()) s.activeDone = activeCtx.Done() s.activeCancel = activeCancel @@ -115,20 +117,44 @@ func (s *codexWebsocketSession) setActive(ch chan codexWebsocketRead) { s.activeMu.Unlock() } -func (s *codexWebsocketSession) clearActive(ch chan codexWebsocketRead) { - if s == nil { - return +func (s *codexWebsocketSession) activate(conn *websocket.Conn) chan codexWebsocketRead { + if s == nil || conn == nil { + return nil + } + ch := make(chan codexWebsocketRead, 4096) + s.setActive(conn, ch) + return ch +} + +func (s *codexWebsocketSession) activeForConn(conn *websocket.Conn) (chan codexWebsocketRead, <-chan struct{}) { + if s == nil || conn == nil { + return nil, nil } s.activeMu.Lock() - if s.activeCh == ch { - s.activeCh = nil - if s.activeCancel != nil { - s.activeCancel() - } - s.activeCancel = nil - s.activeDone = nil + defer s.activeMu.Unlock() + if s.activeConn != conn { + return nil, nil } - s.activeMu.Unlock() + return s.activeCh, s.activeDone +} + +func (s *codexWebsocketSession) clearActive(conn *websocket.Conn, ch chan codexWebsocketRead) bool { + if s == nil { + return false + } + s.activeMu.Lock() + defer s.activeMu.Unlock() + if s.activeConn != conn || s.activeCh != ch { + return false + } + s.activeConn = nil + s.activeCh = nil + if s.activeCancel != nil { + s.activeCancel() + } + s.activeCancel = nil + s.activeDone = nil + return true } func (s *codexWebsocketSession) writeMessage(conn *websocket.Conn, msgType int, payload []byte) error { @@ -155,6 +181,40 @@ func (s *codexWebsocketSession) configureConn(conn *websocket.Conn) { }) } +func websocketSessionTargetChanged(sess *codexWebsocketSession, authID string, wsURL string) bool { + if sess == nil { + return false + } + + sess.connMu.Lock() + defer sess.connMu.Unlock() + if strings.TrimSpace(sess.authID) == "" && strings.TrimSpace(sess.wsURL) == "" { + return false + } + return strings.TrimSpace(sess.authID) != strings.TrimSpace(authID) || strings.TrimSpace(sess.wsURL) != strings.TrimSpace(wsURL) +} + +func detachMismatchedWebsocketSessionConn(sess *codexWebsocketSession, authID string, wsURL string) (*websocket.Conn, string, string) { + if sess == nil { + return nil, "", "" + } + + sess.connMu.Lock() + defer sess.connMu.Unlock() + conn := sess.conn + if conn == nil || (strings.TrimSpace(sess.authID) == strings.TrimSpace(authID) && strings.TrimSpace(sess.wsURL) == strings.TrimSpace(wsURL)) { + return nil, "", "" + } + + previousAuthID := sess.authID + previousWSURL := sess.wsURL + sess.conn = nil + if sess.readerConn == conn { + sess.readerConn = nil + } + return conn, previousAuthID, previousWSURL +} + func (s *codexWebsocketSession) notifyUpstreamDisconnect(err error) { if s == nil { return @@ -296,9 +356,10 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut var readCh chan codexWebsocketRead if sess != nil { - readCh = make(chan codexWebsocketRead, 4096) - sess.setActive(readCh) - defer sess.clearActive(readCh) + readCh = sess.activate(conn) + defer func() { + sess.clearActive(conn, readCh) + }() } if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil { @@ -310,6 +371,8 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut // execution session. connRetry, respHSRetry, errDialRetry := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) if errDialRetry == nil && connRetry != nil { + conn = connRetry + readCh = sess.activate(conn) wsReqBodyRetry := buildCodexWebsocketRequestBody(upstreamBody) helps.RecordAPIWebsocketRequest(ctx, e.cfg, helps.UpstreamRequestLog{ URL: wsURL, @@ -324,8 +387,7 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut }) recordAPIWebsocketHandshake(ctx, e.cfg, respHSRetry) reporter.StartResponseTTFT() - if errSendRetry := writeCodexWebsocketMessage(sess, connRetry, wsReqBodyRetry); errSendRetry == nil { - conn = connRetry + if errSendRetry := writeCodexWebsocketMessage(sess, conn, wsReqBodyRetry); errSendRetry == nil { wsReqBody = wsReqBodyRetry } else { e.invalidateUpstreamConn(sess, connRetry, "send_error", errSendRetry) @@ -515,8 +577,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr var readCh chan codexWebsocketRead if sess != nil { - readCh = make(chan codexWebsocketRead, 4096) - sess.setActive(readCh) + readCh = sess.activate(conn) } if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil { @@ -529,10 +590,12 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr if errDialRetry != nil || connRetry == nil { closeHTTPResponseBody(respHSRetry, "codex websockets executor: close handshake response body error") helps.RecordAPIWebsocketError(ctx, e.cfg, "dial_retry", errDialRetry) - sess.clearActive(readCh) + sess.clearActive(conn, readCh) sess.reqMu.Unlock() return nil, errDialRetry } + conn = connRetry + readCh = sess.activate(conn) wsReqBodyRetry := buildCodexWebsocketRequestBody(upstreamBody) helps.RecordAPIWebsocketRequest(ctx, e.cfg, helps.UpstreamRequestLog{ URL: wsURL, @@ -547,14 +610,13 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr }) recordAPIWebsocketHandshake(ctx, e.cfg, respHSRetry) reporter.StartResponseTTFT() - if errSendRetry := writeCodexWebsocketMessage(sess, connRetry, wsReqBodyRetry); errSendRetry != nil { + if errSendRetry := writeCodexWebsocketMessage(sess, conn, wsReqBodyRetry); errSendRetry != nil { helps.RecordAPIWebsocketError(ctx, e.cfg, "send_retry", errSendRetry) - e.invalidateUpstreamConn(sess, connRetry, "send_error", errSendRetry) - sess.clearActive(readCh) + e.invalidateUpstreamConn(sess, conn, "send_error", errSendRetry) + sess.clearActive(conn, readCh) sess.reqMu.Unlock() return nil, errSendRetry } - conn = connRetry wsReqBody = wsReqBodyRetry } else { logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, "send_error", errSend) @@ -573,7 +635,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr defer close(out) defer func() { if sess != nil { - sess.clearActive(readCh) + sess.clearActive(conn, readCh) sess.reqMu.Unlock() return } @@ -1399,6 +1461,13 @@ func (e *CodexWebsocketsExecutor) ensureUpstreamConn(ctx context.Context, auth * return e.dialCodexWebsocket(ctx, auth, wsURL, headers) } + if staleConn, staleAuthID, staleWSURL := detachMismatchedWebsocketSessionConn(sess, authID, wsURL); staleConn != nil { + logCodexWebsocketDisconnected(sess.sessionID, staleAuthID, staleWSURL, "target_changed", nil) + if errClose := staleConn.Close(); errClose != nil { + log.Errorf("codex websockets executor: close stale websocket error: %v", errClose) + } + } + sess.connMu.Lock() conn := sess.conn readerConn := sess.readerConn @@ -1448,18 +1517,16 @@ func (e *CodexWebsocketsExecutor) readUpstreamLoop(sess *codexWebsocketSession, _ = conn.SetReadDeadline(time.Now().Add(codexResponsesWebsocketIdleTimeout)) msgType, payload, errRead := conn.ReadMessage() if errRead != nil { - sess.activeMu.Lock() - ch := sess.activeCh - done := sess.activeDone - sess.activeMu.Unlock() + ch, done := sess.activeForConn(conn) if ch != nil { select { case ch <- codexWebsocketRead{conn: conn, err: errRead}: case <-done: default: } - sess.clearActive(ch) - close(ch) + if sess.clearActive(conn, ch) { + close(ch) + } } e.invalidateUpstreamConn(sess, conn, "upstream_disconnected", errRead) return @@ -1468,18 +1535,16 @@ func (e *CodexWebsocketsExecutor) readUpstreamLoop(sess *codexWebsocketSession, if msgType != websocket.TextMessage { if msgType == websocket.BinaryMessage { errBinary := fmt.Errorf("codex websockets executor: unexpected binary message") - sess.activeMu.Lock() - ch := sess.activeCh - done := sess.activeDone - sess.activeMu.Unlock() + ch, done := sess.activeForConn(conn) if ch != nil { select { case ch <- codexWebsocketRead{conn: conn, err: errBinary}: case <-done: default: } - sess.clearActive(ch) - close(ch) + if sess.clearActive(conn, ch) { + close(ch) + } } e.invalidateUpstreamConn(sess, conn, "unexpected_binary", errBinary) return @@ -1487,10 +1552,7 @@ func (e *CodexWebsocketsExecutor) readUpstreamLoop(sess *codexWebsocketSession, continue } - sess.activeMu.Lock() - ch := sess.activeCh - done := sess.activeDone - sess.activeMu.Unlock() + ch, done := sess.activeForConn(conn) if ch == nil { continue } diff --git a/internal/runtime/executor/websocket_session_target_test.go b/internal/runtime/executor/websocket_session_target_test.go new file mode 100644 index 00000000..324fe8e5 --- /dev/null +++ b/internal/runtime/executor/websocket_session_target_test.go @@ -0,0 +1,198 @@ +package executor + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gorilla/websocket" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" +) + +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") + } +} + +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, *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, *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 +} diff --git a/internal/runtime/executor/xai_websockets_executor.go b/internal/runtime/executor/xai_websockets_executor.go index 72a43d42..7fd8cdb5 100644 --- a/internal/runtime/executor/xai_websockets_executor.go +++ b/internal/runtime/executor/xai_websockets_executor.go @@ -414,22 +414,15 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox if err != nil { return nil, err } - if idMapper != nil { - prepared.body = idMapper.upstreamRequestPayload(prepared.body) - } reporter := helps.NewExecutorUsageReporter(ctx, e, prepared.baseModel, auth) defer reporter.TrackFailure(ctx, &err) - reporter.SetTranslatedReasoningEffort(prepared.body, e.Identifier()) httpURL := strings.TrimSuffix(baseURL, "/") + "/responses" wsURL, err := buildXAIResponsesWebsocketURL(httpURL) if err != nil { return nil, err } - wsHeaders := applyXAIWebsocketHeaders(http.Header{}, auth, token, prepared.sessionID) - wsReqBody := buildXAIWebsocketRequestBody(prepared.body) - warmupRequest := xaiWebsocketGenerateFalse(wsReqBody) var authID, authLabel, authType, authValue string if auth != nil { @@ -445,6 +438,17 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox sess.reqMu.Lock() } } + if idMapper != nil { + if websocketSessionTargetChanged(sess, authID, wsURL) { + idMapper.upstreamPreviousID = "" + } + prepared.body = idMapper.upstreamRequestPayload(prepared.body) + } + reporter.SetTranslatedReasoningEffort(prepared.body, e.Identifier()) + + wsHeaders := applyXAIWebsocketHeaders(http.Header{}, auth, token, prepared.sessionID) + wsReqBody := buildXAIWebsocketRequestBody(prepared.body) + warmupRequest := xaiWebsocketGenerateFalse(wsReqBody) wsReqLog := helps.UpstreamRequestLog{ URL: wsURL, @@ -491,8 +495,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox var readCh chan codexWebsocketRead if sess != nil { - readCh = make(chan codexWebsocketRead, 4096) - sess.setActive(readCh) + readCh = sess.activate(conn) } if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil { @@ -504,13 +507,15 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox bodyErrRetry := websocketHandshakeBody(respHSRetry) closeHTTPResponseBody(respHSRetry, "xai websockets executor: close handshake response body error") helps.RecordAPIWebsocketError(ctx, e.cfg, "dial_retry", errDialRetry) - sess.clearActive(readCh) + sess.clearActive(conn, readCh) sess.reqMu.Unlock() if respHSRetry != nil && respHSRetry.StatusCode > 0 { return nil, xaiStatusErr(respHSRetry.StatusCode, bodyErrRetry) } return nil, errDialRetry } + conn = connRetry + readCh = sess.activate(conn) wsReqBodyRetry := buildXAIWebsocketRequestBody(prepared.body) helps.RecordAPIWebsocketRequest(ctx, e.cfg, helps.UpstreamRequestLog{ URL: wsURL, @@ -526,14 +531,13 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox logXAIWebsocketRequest(executionSessionID, authID, wsURL, wsReqBodyRetry) recordAPIWebsocketHandshake(ctx, e.cfg, respHSRetry) reporter.StartResponseTTFT() - if errSendRetry := writeCodexWebsocketMessage(sess, connRetry, wsReqBodyRetry); errSendRetry != nil { + if errSendRetry := writeCodexWebsocketMessage(sess, conn, wsReqBodyRetry); errSendRetry != nil { helps.RecordAPIWebsocketError(ctx, e.cfg, "send_retry", errSendRetry) - e.invalidateUpstreamConn(sess, connRetry, "send_error", errSendRetry) - sess.clearActive(readCh) + e.invalidateUpstreamConn(sess, conn, "send_error", errSendRetry) + sess.clearActive(conn, readCh) sess.reqMu.Unlock() return nil, errSendRetry } - conn = connRetry wsReqBody = wsReqBodyRetry } else { logXAIWebsocketDisconnected(executionSessionID, authID, wsURL, "send_error", errSend) @@ -552,7 +556,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox defer close(out) defer func() { if sess != nil { - sess.clearActive(readCh) + sess.clearActive(conn, readCh) sess.reqMu.Unlock() return } @@ -938,6 +942,13 @@ func (e *XAIWebsocketsExecutor) ensureUpstreamConn(ctx context.Context, auth *cl return e.dialXAIWebsocket(ctx, auth, wsURL, headers) } + if staleConn, staleAuthID, staleWSURL := detachMismatchedWebsocketSessionConn(sess, authID, wsURL); staleConn != nil { + logXAIWebsocketDisconnected(sess.sessionID, staleAuthID, staleWSURL, "target_changed", nil) + if errClose := staleConn.Close(); errClose != nil { + log.Errorf("xai websockets executor: close stale websocket error: %v", errClose) + } + } + sess.connMu.Lock() conn := sess.conn readerConn := sess.readerConn @@ -1033,18 +1044,16 @@ func (e *XAIWebsocketsExecutor) readUpstreamLoop(sess *codexWebsocketSession, co for { msgType, payload, errRead := conn.ReadMessage() if errRead != nil { - sess.activeMu.Lock() - ch := sess.activeCh - done := sess.activeDone - sess.activeMu.Unlock() + ch, done := sess.activeForConn(conn) if ch != nil { select { case ch <- codexWebsocketRead{conn: conn, err: errRead}: case <-done: default: } - sess.clearActive(ch) - close(ch) + if sess.clearActive(conn, ch) { + close(ch) + } } e.invalidateUpstreamConn(sess, conn, "upstream_disconnected", errRead) return @@ -1053,18 +1062,16 @@ func (e *XAIWebsocketsExecutor) readUpstreamLoop(sess *codexWebsocketSession, co if msgType != websocket.TextMessage { if msgType == websocket.BinaryMessage { errBinary := fmt.Errorf("xai websockets executor: unexpected binary message") - sess.activeMu.Lock() - ch := sess.activeCh - done := sess.activeDone - sess.activeMu.Unlock() + ch, done := sess.activeForConn(conn) if ch != nil { select { case ch <- codexWebsocketRead{conn: conn, err: errBinary}: case <-done: default: } - sess.clearActive(ch) - close(ch) + if sess.clearActive(conn, ch) { + close(ch) + } } e.invalidateUpstreamConn(sess, conn, "unexpected_binary", errBinary) return @@ -1072,10 +1079,7 @@ func (e *XAIWebsocketsExecutor) readUpstreamLoop(sess *codexWebsocketSession, co continue } - sess.activeMu.Lock() - ch := sess.activeCh - done := sess.activeDone - sess.activeMu.Unlock() + ch, done := sess.activeForConn(conn) if ch == nil { continue } diff --git a/internal/runtime/executor/xai_websockets_executor_test.go b/internal/runtime/executor/xai_websockets_executor_test.go index 114042b5..edc694d8 100644 --- a/internal/runtime/executor/xai_websockets_executor_test.go +++ b/internal/runtime/executor/xai_websockets_executor_test.go @@ -820,6 +820,125 @@ func TestXAIWebsocketsExecuteStreamRewritesRepeatedResponseIDWithoutPreviousResp } } +func TestXAIWebsocketsExecuteStreamReplaysTranscriptWhenAuthChanges(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + type capturedRequest struct { + authorization string + payload []byte + } + captured := make(chan capturedRequest, 2) + 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 + } + defer func() { + if errClose := conn.Close(); errClose != nil { + t.Errorf("close upstream websocket: %v", errClose) + } + }() + + for { + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + return + } + authorization := r.Header.Get("Authorization") + captured <- capturedRequest{authorization: authorization, payload: bytes.Clone(payload)} + responseID := "resp-auth-a" + if strings.Contains(authorization, "token-c") { + responseID = "resp-auth-c" + } + completed := []byte(fmt.Sprintf(`{"type":"response.completed","response":{"id":%q,"output":[{"type":"message","id":%q,"role":"assistant","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}`, responseID, "msg-"+responseID)) + if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { + t.Errorf("write completed websocket message: %v", errWrite) + return + } + } + })) + defer server.Close() + rejectedServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + http.Error(w, "rejected", http.StatusUnauthorized) + })) + defer rejectedServer.Close() + + exec := NewXAIWebsocketsExecutor(&config.Config{}) + exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + exec.idStore = &xaiWebsocketIDStateStore{sessions: make(map[string]*xaiWebsocketIDState)} + defer exec.CloseExecutionSession("xai-auth-switch-session") + + newAuth := func(id string, token string, baseURL string) *cliproxyauth.Auth { + return &cliproxyauth.Auth{ + ID: id, + Provider: "xai", + Attributes: map[string]string{ + "base_url": baseURL, + "websockets": "true", + }, + Metadata: map[string]any{"access_token": token}, + } + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "xai-auth-switch-session", + }, + } + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + + runRequest := func(auth *cliproxyauth.Auth, body []byte) []byte { + result, errExecute := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{Model: "grok-4.3", Payload: body}, opts) + if errExecute != nil { + t.Fatalf("ExecuteStream() error = %v", errExecute) + } + var completed []byte + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatalf("stream chunk error = %v", chunk.Err) + } + if gjson.GetBytes(chunk.Payload, "type").String() == "response.completed" { + completed = bytes.Clone(chunk.Payload) + } + } + if len(completed) == 0 { + t.Fatal("stream did not return response.completed") + } + return completed + } + + firstCompleted := runRequest(newAuth("auth-a", "token-a", server.URL), []byte(`{"model":"grok-4.3","input":[{"type":"message","id":"user-1","role":"user","content":"first"}]}`)) + firstResponseID := gjson.GetBytes(firstCompleted, "response.id").String() + if firstResponseID != "resp-auth-a" { + t.Fatalf("first response ID = %q, want resp-auth-a", firstResponseID) + } + firstUpstream := <-captured + if firstUpstream.authorization != "Bearer token-a" { + t.Fatalf("first Authorization = %q, want Bearer token-a", firstUpstream.authorization) + } + + secondBody := []byte(fmt.Sprintf(`{"model":"grok-4.3","previous_response_id":%q,"input":[{"type":"message","id":"user-2","role":"user","content":"second"}]}`, firstResponseID)) + if _, errExecute := exec.ExecuteStream(ctx, newAuth("auth-b", "token-b", rejectedServer.URL), cliproxyexecutor.Request{Model: "grok-4.3", Payload: secondBody}, opts); errExecute == nil { + t.Fatal("expected auth B websocket handshake to fail") + } + runRequest(newAuth("auth-c", "token-c", server.URL), secondBody) + secondUpstream := <-captured + if secondUpstream.authorization != "Bearer token-c" { + t.Fatalf("second successful Authorization = %q, want Bearer token-c", secondUpstream.authorization) + } + if gjson.GetBytes(secondUpstream.payload, "previous_response_id").Exists() { + t.Fatalf("previous_response_id was sent after auth switch: %s", secondUpstream.payload) + } + input := gjson.GetBytes(secondUpstream.payload, "input").Array() + if len(input) != 3 { + t.Fatalf("replayed input len = %d, want 3: %s", len(input), secondUpstream.payload) + } + if input[0].Get("id").String() != "user-1" || input[1].Get("id").String() != "msg-resp-auth-a" || input[2].Get("id").String() != "user-2" { + t.Fatalf("unexpected replayed input: %s", secondUpstream.payload) + } +} + func TestXAIWebsocketsExecuteStreamCompactionTriggerUsesHTTPCompactWithRecordedContext(t *testing.T) { nativeEncryptedContent := testValidGrokEncryptedContent() upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} diff --git a/sdk/api/handlers/openai/openai_responses_websocket.go b/sdk/api/handlers/openai/openai_responses_websocket.go index d9cd1190..3c63a543 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket.go +++ b/sdk/api/handlers/openai/openai_responses_websocket.go @@ -53,6 +53,11 @@ type websocketTimelineAppender interface { Append(eventType string, payload []byte, timestamp time.Time) } +type responsesWebsocketPinnedAuthState struct { + authID string + modelKey string +} + type websocketTimelineLog struct { enabled bool source *requestlogging.FileBodySource @@ -278,16 +283,43 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { lastResponseID := "" var lastResponsePendingToolCallIDs []string pinnedAuthID := "" - lastAttemptedAuthID := "" + // Preserve independent upstream auth affinity when a downstream session switches providers. + pinnedAuthByProvider := make(map[string]responsesWebsocketPinnedAuthState) passthroughModelName := "" - sessionAuthByID := func(authID string) (*coreauth.Auth, bool) { + sessionAuthByIDWithSource := func(authID string) (*coreauth.Auth, bool, bool) { if h == nil || h.AuthManager == nil { - return nil, false + return nil, false, false } if auth, ok := h.AuthManager.GetExecutionSessionAuthByID(passthroughSessionID, authID); ok { - return auth, true + return auth, true, true } - return h.AuthManager.GetByID(authID) + auth, ok := h.AuthManager.GetByID(authID) + return auth, false, ok + } + sessionAuthByID := func(authID string) (*coreauth.Auth, bool) { + auth, _, ok := sessionAuthByIDWithSource(authID) + return auth, ok + } + rememberPinnedAuth := func(authID string, modelName string) { + authID = strings.TrimSpace(authID) + auth, ok := sessionAuthByID(authID) + if authID == "" || !ok || auth == nil { + return + } + pinnedAuthID = authID + providerKey := strings.ToLower(strings.TrimSpace(auth.Provider)) + _, modelKey := responsesWebsocketProviderSetForModel(responsesWebsocketResolvedModelName(modelName)) + if providerKey != "" { + pinnedAuthByProvider[providerKey] = responsesWebsocketPinnedAuthState{authID: authID, modelKey: modelKey} + } + } + forgetPinnedAuth := func() { + for providerKey, state := range pinnedAuthByProvider { + if state.authID == pinnedAuthID { + delete(pinnedAuthByProvider, providerKey) + } + } + pinnedAuthID = "" } forceTranscriptReplayNextRequest := false @@ -315,14 +347,43 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { wsTimelineLog.BeginRequest() wsTimelineLog.Append("request", payload, time.Now()) - requestModelName := strings.TrimSpace(gjson.GetBytes(payload, "model").String()) + explicitRequestModelName := strings.TrimSpace(gjson.GetBytes(payload, "model").String()) + requestModelName := explicitRequestModelName if requestModelName == "" { requestModelName = passthroughModelName } if requestModelName == "" { requestModelName = strings.TrimSpace(gjson.GetBytes(lastRequest, "model").String()) } + if pinnedAuthID != "" { + pinnedAuth, homeRuntime, ok := sessionAuthByIDWithSource(pinnedAuthID) + providerKey := "" + if pinnedAuth != nil { + providerKey = strings.ToLower(strings.TrimSpace(pinnedAuth.Provider)) + } + state, hasState := pinnedAuthByProvider[providerKey] + if !ok || !hasState || state.authID != pinnedAuthID || !responsesWebsocketPinnedAuthMatchesModel(pinnedAuth, requestModelName, state.modelKey, homeRuntime) { + pinnedAuthID = "" + } + } + if pinnedAuthID == "" { + providerSet, _ := responsesWebsocketProviderSetForModel(responsesWebsocketResolvedModelName(requestModelName)) + if len(providerSet) == 1 { + for providerKey := range providerSet { + state, ok := pinnedAuthByProvider[providerKey] + candidateAuth, homeRuntime, okAuth := sessionAuthByIDWithSource(state.authID) + if ok && okAuth && responsesWebsocketPinnedAuthMatchesModel(candidateAuth, requestModelName, state.modelKey, homeRuntime) { + pinnedAuthID = state.authID + } else { + delete(pinnedAuthByProvider, providerKey) + } + } + } + } useUpstreamWebsocketPassthrough := h.responsesWebsocketUsesUpstreamWebsocketPassthrough(requestModelName) + if explicitRequestModelName != "" && !useUpstreamWebsocketPassthrough { + passthroughModelName = "" + } allowIncrementalInputWithPreviousResponseID := false allowCompactionReplayBypass := false if !useUpstreamWebsocketPassthrough { @@ -417,6 +478,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { } modelName := gjson.GetBytes(requestJSON, "model").String() + lastAttemptedAuthID := pinnedAuthID cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background()) cliCtx = cliproxyexecutor.WithDownstreamWebsocket(cliCtx) cliCtx = handlers.WithExecutionSessionID(cliCtx, passthroughSessionID) @@ -434,7 +496,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { return } if websocketUpstreamSupportsIncrementalInput(selectedAuth.Attributes, selectedAuth.Metadata) { - pinnedAuthID = authID + rememberPinnedAuth(authID, modelName) } }) } @@ -449,16 +511,16 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { if forwardErrMsg == nil && !useUpstreamWebsocketPassthrough && lastAttemptedAuthID != "" { if selectedAuth, ok := sessionAuthByID(lastAttemptedAuthID); ok && selectedAuth != nil { if websocketUpstreamSupportsIncrementalInput(selectedAuth.Attributes, selectedAuth.Metadata) { - pinnedAuthID = lastAttemptedAuthID + rememberPinnedAuth(lastAttemptedAuthID, modelName) } else if pinnedAuthID != "" { if pinnedAuth, ok := sessionAuthByID(pinnedAuthID); ok && pinnedAuth != nil && websocketUpstreamSupportsIncrementalInput(pinnedAuth.Attributes, pinnedAuth.Metadata) { - pinnedAuthID = lastAttemptedAuthID + rememberPinnedAuth(lastAttemptedAuthID, modelName) } } } } if shouldReleaseResponsesWebsocketPinnedAuth(forwardErrMsg) { - pinnedAuthID = "" + forgetPinnedAuth() forceTranscriptReplayNextRequest = true if useUpstreamWebsocketPassthrough { passthroughModelName = "" @@ -1077,6 +1139,25 @@ func responsesWebsocketAuthSupportsIncrementalInput(auth *coreauth.Auth) bool { return websocketUpstreamSupportsIncrementalInput(auth.Attributes, auth.Metadata) } +func responsesWebsocketPinnedAuthMatchesModel(auth *coreauth.Auth, modelName string, pinnedModelKey string, homeRuntime bool) bool { + if auth == nil { + return false + } + providerSet, modelKey := responsesWebsocketProviderSetForModel(responsesWebsocketResolvedModelName(modelName)) + providerKey := strings.ToLower(strings.TrimSpace(auth.Provider)) + if _, ok := providerSet[providerKey]; !ok { + return false + } + if !responsesWebsocketAuthAvailableForModel(auth, modelKey, time.Now()) { + return false + } + + if homeRuntime { + return strings.EqualFold(strings.TrimSpace(pinnedModelKey), strings.TrimSpace(modelKey)) + } + return registry.GetGlobalRegistry().ClientSupportsModel(auth.ID, modelKey) +} + func normalizeResponsesWebsocketPassthroughRequest(rawJSON []byte, modelName string) ([]byte, *interfaces.ErrorMessage) { if !json.Valid(rawJSON) { return nil, &interfaces.ErrorMessage{ diff --git a/sdk/api/handlers/openai/openai_responses_websocket_test.go b/sdk/api/handlers/openai/openai_responses_websocket_test.go index 4cd522e4..c5f54b06 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket_test.go +++ b/sdk/api/handlers/openai/openai_responses_websocket_test.go @@ -1951,6 +1951,179 @@ func TestResponsesWebsocketXAIWebsocketPassthroughCarriesPreviousResponseID(t *t } } +func TestResponsesWebsocketSwitchesPinnedAuthAcrossProviders(t *testing.T) { + for _, testCase := range []struct { + name string + xaiWebsockets bool + returnToDifferentXAIModel bool + }{ + {name: "xai SSE", xaiWebsockets: false}, + {name: "xai websocket", xaiWebsockets: true}, + {name: "xai websocket different model", xaiWebsockets: true, returnToDifferentXAIModel: true}, + } { + t.Run(testCase.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + + xaiModel := "xai-provider-switch-" + strings.ReplaceAll(testCase.name, " ", "-") + returnXAIModel := xaiModel + if testCase.returnToDifferentXAIModel { + returnXAIModel += "-return" + } + codexModel := "codex-provider-switch-" + strings.ReplaceAll(testCase.name, " ", "-") + xaiExecutor := &websocketDirectCaptureExecutor{provider: "xai"} + codexExecutor := &websocketDirectCaptureExecutor{provider: "codex"} + + xaiAuth := &coreauth.Auth{ + ID: "auth-" + xaiModel, + Provider: "xai", + Status: coreauth.StatusActive, + } + if testCase.xaiWebsockets { + xaiAuth.Attributes = map[string]string{"websockets": "true"} + } + codexAuth := &coreauth.Auth{ + ID: "auth-" + codexModel, + Provider: "codex", + Status: coreauth.StatusActive, + Attributes: map[string]string{"websockets": "true"}, + } + selector := &orderedWebsocketSelector{order: []string{xaiAuth.ID, codexAuth.ID}} + manager := coreauth.NewManager(nil, selector, nil) + manager.RegisterExecutor(xaiExecutor) + manager.RegisterExecutor(codexExecutor) + if _, errRegister := manager.Register(context.Background(), xaiAuth); errRegister != nil { + t.Fatalf("Register xAI auth: %v", errRegister) + } + if _, errRegister := manager.Register(context.Background(), codexAuth); errRegister != nil { + t.Fatalf("Register Codex auth: %v", errRegister) + } + + registry.GetGlobalRegistry().RegisterClient(xaiAuth.ID, xaiAuth.Provider, []*registry.ModelInfo{{ID: xaiModel}}) + registry.GetGlobalRegistry().RegisterClient(codexAuth.ID, codexAuth.Provider, []*registry.ModelInfo{{ID: codexModel}}) + registeredAuthIDs := []string{xaiAuth.ID, codexAuth.ID} + xaiAlternateAuthID := "" + if testCase.xaiWebsockets { + xaiAlternateAuth := &coreauth.Auth{ + ID: "auth-alternate-" + xaiModel, + Provider: "xai", + Status: coreauth.StatusActive, + Attributes: map[string]string{"websockets": "true"}, + } + xaiAlternateAuthID = xaiAlternateAuth.ID + selector.order = append(selector.order, xaiAlternateAuth.ID) + if _, errRegister := manager.Register(context.Background(), xaiAlternateAuth); errRegister != nil { + t.Fatalf("Register alternate xAI auth: %v", errRegister) + } + alternateModels := []*registry.ModelInfo{{ID: xaiModel}} + if testCase.returnToDifferentXAIModel { + alternateModels = []*registry.ModelInfo{{ID: returnXAIModel}} + } + registry.GetGlobalRegistry().RegisterClient(xaiAlternateAuth.ID, xaiAlternateAuth.Provider, alternateModels) + registeredAuthIDs = append(registeredAuthIDs, xaiAlternateAuth.ID) + } + t.Cleanup(func() { + for _, authID := range registeredAuthIDs { + registry.GetGlobalRegistry().UnregisterClient(authID) + } + }) + + base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager) + h := NewOpenAIResponsesAPIHandler(base) + router := gin.New() + router.GET("/v1/responses/ws", h.ResponsesWebsocket) + + server := httptest.NewServer(router) + defer server.Close() + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses/ws" + conn, _, errDial := websocket.DefaultDialer.Dial(wsURL, nil) + if errDial != nil { + t.Fatalf("dial websocket: %v", errDial) + } + defer func() { + if errClose := conn.Close(); errClose != nil { + t.Errorf("close websocket: %v", errClose) + } + }() + + requests := []string{ + fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-xai-1"}]}`, xaiModel), + fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-codex-1"}]}`, codexModel), + fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-xai-2"}]}`, returnXAIModel), + `{"type":"response.create","input":[{"type":"message","id":"msg-xai-3"}]}`, + } + for index, request := range requests { + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(request)); errWrite != nil { + t.Fatalf("write websocket message %d: %v", index+1, errWrite) + } + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + t.Fatalf("read websocket response %d: %v", index+1, errRead) + } + if got := gjson.GetBytes(payload, "type").String(); got != wsEventTypeCompleted { + t.Fatalf("response %d type = %s, want %s: %s", index+1, got, wsEventTypeCompleted, payload) + } + } + + wantReturnAuthID := xaiAuth.ID + if testCase.returnToDifferentXAIModel { + wantReturnAuthID = xaiAlternateAuthID + } + if got := xaiExecutor.AuthIDs(); len(got) != 3 || got[0] != xaiAuth.ID || got[1] != wantReturnAuthID || got[2] != wantReturnAuthID { + t.Fatalf("xAI auth IDs = %v, want [%s %s %s]", got, xaiAuth.ID, wantReturnAuthID, wantReturnAuthID) + } + if got := codexExecutor.AuthIDs(); len(got) != 1 || got[0] != codexAuth.ID { + t.Fatalf("Codex auth IDs = %v, want [%s]", got, codexAuth.ID) + } + }) + } +} + +func TestResponsesWebsocketPinnedAuthMatchesModel(t *testing.T) { + modelA := "xai-pinned-auth-model-a" + modelB := "xai-pinned-auth-model-b" + auth := &coreauth.Auth{ID: "xai-pinned-auth", Provider: "xai", Status: coreauth.StatusActive} + otherAuthID := "xai-pinned-auth-other" + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: modelA}}) + registry.GetGlobalRegistry().RegisterClient(otherAuthID, auth.Provider, []*registry.ModelInfo{{ID: modelB}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(auth.ID) + registry.GetGlobalRegistry().UnregisterClient(otherAuthID) + }) + + if !responsesWebsocketPinnedAuthMatchesModel(auth, modelA, modelA, false) { + t.Fatal("expected registered auth to match its supported model") + } + if responsesWebsocketPinnedAuthMatchesModel(auth, modelB, modelA, false) { + t.Fatal("registered auth matched an unsupported model from the same provider") + } + + disabledAuth := auth.Clone() + disabledAuth.Disabled = true + if responsesWebsocketPinnedAuthMatchesModel(disabledAuth, modelA, modelA, false) { + t.Fatal("disabled auth matched a model") + } + + cooldownAuth := auth.Clone() + cooldownAuth.ModelStates = map[string]*coreauth.ModelState{ + modelA: {Unavailable: true, NextRetryAfter: time.Now().Add(time.Minute)}, + } + if responsesWebsocketPinnedAuthMatchesModel(cooldownAuth, modelA, modelA, false) { + t.Fatal("auth in model cooldown matched a model") + } + + unregisteredAuth := &coreauth.Auth{ID: "unregistered-auth", Provider: "xai", Status: coreauth.StatusActive} + if responsesWebsocketPinnedAuthMatchesModel(unregisteredAuth, modelA, modelA, false) { + t.Fatal("unregistered ordinary auth matched a model") + } + if !responsesWebsocketPinnedAuthMatchesModel(unregisteredAuth, modelA, modelA, true) { + t.Fatal("expected Home runtime auth to match its pinned model") + } + if responsesWebsocketPinnedAuthMatchesModel(unregisteredAuth, modelB, modelA, true) { + t.Fatal("Home runtime auth matched a different model") + } +} + func TestWebsocketUpstreamSupportsIncrementalInputForModel(t *testing.T) { manager := coreauth.NewManager(nil, nil, nil) auth := &coreauth.Auth{