Merge pull request #4522 from sususu98/fix/responses-ws-continuity

Luis Pater committed 3 months ago

Commit e5729872 parents 285322c a66117213 files · +2108 −359Browse files

Unified · Split

13 files

internal/runtime/executor/codex_executor.go +3 −0
internal/runtime/executor/codex_websockets_executor.go +80 −5
internal/runtime/executor/codex_websockets_executor_test.go +170 −0
internal/runtime/executor/xai_websockets_executor.go +149 −21
internal/runtime/executor/xai_websockets_executor_test.go +165 −0
sdk/api/handlers/handlers.go +25 −1
sdk/api/handlers/handlers_model_router_test.go +32 −0
sdk/api/handlers/openai/openai_responses_websocket.go +232 −75
sdk/api/handlers/openai/openai_responses_websocket_test.go +1061 −252
sdk/api/handlers/openai/openai_responses_websocket_toolcall_repair.go +118 −5
sdk/cliproxy/executor/context.go +19 −0
sdk/cliproxy/executor/websocket.go +29 −0
sdk/cliproxy/executor/websocket_test.go +25 −0

internal/runtime/executor/codex_executor.go+3−0

 	body = helps.SetStringIfDifferent(body, "model", baseModel)
 	body = helps.SetBoolIfDifferent(body, "stream", true)
 	body, _ = sjson.DeleteBytes(body, "previous_response_id")
+	body, _ = sjson.DeleteBytes(body, "generate")
 	body, _ = sjson.DeleteBytes(body, "prompt_cache_retention")
 	body, _ = sjson.DeleteBytes(body, "safety_identifier")
 	body, _ = sjson.DeleteBytes(body, "stream_options")
 	requestPath := helps.PayloadRequestPath(opts)
 	body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers)
 	body, _ = sjson.DeleteBytes(body, "previous_response_id")
+	body, _ = sjson.DeleteBytes(body, "generate")
 	body, _ = sjson.DeleteBytes(body, "prompt_cache_retention")
 	body, _ = sjson.DeleteBytes(body, "safety_identifier")
 	body, _ = sjson.DeleteBytes(body, "stream_options")
 	body = helps.SetStringIfDifferent(body, "model", baseModel)
 	body, _ = sjson.DeleteBytes(body, "previous_response_id")
+	body, _ = sjson.DeleteBytes(body, "generate")
 	body, _ = sjson.DeleteBytes(body, "prompt_cache_retention")
 	body, _ = sjson.DeleteBytes(body, "safety_identifier")
 	body, _ = sjson.DeleteBytes(body, "stream_options")

internal/runtime/executor/codex_websockets_executor.go+80−5

@@ -352,6 +352,23 @@ func websocketSessionTargetChanged(sess *codexWebsocketSession, authID string, w · 351 unmodified lines
352	return strings.TrimSpace(sess.authID) != strings.TrimSpace(authID) \|\| strings.TrimSpace(sess.wsURL) != strings.TrimSpace(wsURL)	352	return strings.TrimSpace(sess.authID) != strings.TrimSpace(authID) \|\| strings.TrimSpace(sess.wsURL) != strings.TrimSpace(wsURL)
353	}	353	}
354		354
		355	func existingWebsocketSessionConn(sess codexWebsocketSession, authID string, wsURL string) (websocket.Conn, *websocketConnectionCloser) {
		356	if sess == nil {
		357	return nil, nil
		358	}
		359	sess.connMu.Lock()
		360	conn := sess.conn
		361	closer := sess.connCloser
		362	matches := conn != nil && closer != nil &&
		363	strings.TrimSpace(sess.authID) == strings.TrimSpace(authID) &&
		364	strings.TrimSpace(sess.wsURL) == strings.TrimSpace(wsURL)
		365	sess.connMu.Unlock()
		366	if !matches \|\| sess.upstreamDisconnectError(conn) != nil {
		367	return nil, nil
		368	}
		369	return conn, closer
		370	}
		371
355	func detachMismatchedWebsocketSessionConn(sess codexWebsocketSession, authID string, wsURL string) (websocket.Conn, *websocketConnectionCloser, string, string, cliproxyexecutor.ExecutionLifecycle) {	372	func detachMismatchedWebsocketSessionConn(sess codexWebsocketSession, authID string, wsURL string) (websocket.Conn, *websocketConnectionCloser, string, string, cliproxyexecutor.ExecutionLifecycle) {
356	if sess == nil {	373	if sess == nil {
357	return nil, nil, "", "", nil	374	return nil, nil, "", "", nil
@@ -532,14 +549,25 @@ func (e CodexWebsocketsExecutor) Execute(ctx context.Context, auth cliproxyaut · 174 unmodified lines
532	}	549	}
533	helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog)	550	helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog)
534		551
535	conn, closer, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)	552	var conn *websocket.Conn
		553	var closer *websocketConnectionCloser
		554	var respHS *http.Response
		555	var errDial error
		556	if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
		557	conn, closer = existingWebsocketSessionConn(sess, authID, wsURL)
		558	if conn == nil {
		559	return resp, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
		560	}
		561	} else {
		562	conn, closer, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
		563	}
536	if errDial != nil {	564	if errDial != nil {
537	bodyErr := websocketHandshakeBody(respHS)	565	bodyErr := websocketHandshakeBody(respHS)
538	if respHS != nil {	566	if respHS != nil {
539	helps.RecordAPIWebsocketUpgradeRejection(ctx, e.cfg, websocketUpgradeRequestLog(wsReqLog), respHS.StatusCode, respHS.Header.Clone(), bodyErr)	567	helps.RecordAPIWebsocketUpgradeRejection(ctx, e.cfg, websocketUpgradeRequestLog(wsReqLog), respHS.StatusCode, respHS.Header.Clone(), bodyErr)
540	}	568	}
541	if respHS != nil && respHS.StatusCode == http.StatusUpgradeRequired {	569	if respHS != nil && respHS.StatusCode == http.StatusUpgradeRequired {
542	if opts.ExecutionLifecycle != nil {	570	if opts.ExecutionLifecycle != nil \|\| cliproxyexecutor.DownstreamWebsocket(ctx) {
543	return resp, statusErr{code: respHS.StatusCode, msg: string(bodyErr)}	571	return resp, statusErr{code: respHS.StatusCode, msg: string(bodyErr)}
544	}	572	}
545	return e.CodexExecutor.Execute(ctx, auth, req, opts)	573	return e.CodexExecutor.Execute(ctx, auth, req, opts)
@@ -582,6 +610,14 @@ func (e CodexWebsocketsExecutor) Execute(ctx context.Context, auth cliproxyaut · 36 unmodified lines
582	if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil {	610	if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil {
583	errSend = mapCodexWebsocketWriteError(sess, conn, errSend)	611	errSend = mapCodexWebsocketWriteError(sess, conn, errSend)
584	if sess != nil {	612	if sess != nil {
		613	if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
		614	e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "send_error", errSend)
		615	if !shouldRetryCodexWebsocketSend(errSend) {
		616	helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend)
		617	return resp, errSend
		618	}
		619	return resp, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
		620	}
585	e.invalidateUpstreamConn(sess, conn, "send_error", errSend)	621	e.invalidateUpstreamConn(sess, conn, "send_error", errSend)
586	if !shouldRetryCodexWebsocketSend(errSend) {	622	if !shouldRetryCodexWebsocketSend(errSend) {
587	helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend)	623	helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend)
@@ -810,7 +846,21 @@ func (e CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth clipr · 222 unmodified lines
810	}	846	}
811	helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog)	847	helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog)
812		848
813	conn, closer, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)	849	var conn *websocket.Conn
		850	var closer *websocketConnectionCloser
		851	var respHS *http.Response
		852	var errDial error
		853	if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
		854	conn, closer = existingWebsocketSessionConn(sess, authID, wsURL)
		855	if conn == nil {
		856	if sess != nil {
		857	sess.reqMu.Unlock()
		858	}
		859	return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
		860	}
		861	} else {
		862	conn, closer, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
		863	}
814	var upstreamHeaders http.Header	864	var upstreamHeaders http.Header
815	if respHS != nil {	865	if respHS != nil {
816	upstreamHeaders = respHS.Header.Clone()	866	upstreamHeaders = respHS.Header.Clone()
@@ -824,7 +874,7 @@ func (e CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth clipr · 7 unmodified lines
824	if sess != nil {	874	if sess != nil {
825	sess.reqMu.Unlock()	875	sess.reqMu.Unlock()
826	}	876	}
827	if opts.ExecutionLifecycle != nil {	877	if opts.ExecutionLifecycle != nil \|\| cliproxyexecutor.DownstreamWebsocket(ctx) {
828	return nil, statusErr{code: respHS.StatusCode, msg: string(bodyErr)}	878	return nil, statusErr{code: respHS.StatusCode, msg: string(bodyErr)}
829	}	879	}
830	return e.CodexExecutor.ExecuteStream(ctx, auth, req, opts)	880	return e.CodexExecutor.ExecuteStream(ctx, auth, req, opts)
@@ -864,6 +914,15 @@ func (e CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth clipr · 33 unmodified lines
864	errSend = mapCodexWebsocketWriteError(sess, conn, errSend)	914	errSend = mapCodexWebsocketWriteError(sess, conn, errSend)
865	helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend)	915	helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend)
866	if sess != nil {	916	if sess != nil {
		917	if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
		918	e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "send_error", errSend)
		919	sess.clearActive(conn, readCh)
		920	sess.reqMu.Unlock()
		921	if !shouldRetryCodexWebsocketSend(errSend) {
		922	return nil, errSend
		923	}
		924	return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
		925	}
867	e.invalidateUpstreamConn(sess, conn, "send_error", errSend)	926	e.invalidateUpstreamConn(sess, conn, "send_error", errSend)
868	if !shouldRetryCodexWebsocketSend(errSend) {	927	if !shouldRetryCodexWebsocketSend(errSend) {
869	sess.clearActive(conn, readCh)	928	sess.clearActive(conn, readCh)
@@ -1964,6 +2023,14 @@ func (e CodexWebsocketsExecutor) readUpstreamLoop(sess codexWebsocketSession, · 1094 unmodified lines
1964	}	2023	}
1965		2024
1966	func (e CodexWebsocketsExecutor) invalidateUpstreamConn(sess codexWebsocketSession, conn *websocket.Conn, reason string, err error) {	2025	func (e CodexWebsocketsExecutor) invalidateUpstreamConn(sess codexWebsocketSession, conn *websocket.Conn, reason string, err error) {
		2026	e.invalidateUpstreamConnWithNotify(sess, conn, reason, err, true)
		2027	}
		2028
		2029	func (e CodexWebsocketsExecutor) invalidateUpstreamConnWithoutDisconnectNotify(sess codexWebsocketSession, conn *websocket.Conn, reason string, err error) {
		2030	e.invalidateUpstreamConnWithNotify(sess, conn, reason, err, false)
		2031	}
		2032
		2033	func (e CodexWebsocketsExecutor) invalidateUpstreamConnWithNotify(sess codexWebsocketSession, conn *websocket.Conn, reason string, err error, notify bool) {
1967	if sess == nil \|\| conn == nil {	2034	if sess == nil \|\| conn == nil {
1968	return	2035	return
1969	}	2036	}
@@ -1989,7 +2056,9 @@ func (e CodexWebsocketsExecutor) invalidateUpstreamConn(sess codexWebsocketSes · 19 unmodified lines
1989	sess.connMu.Unlock()	2056	sess.connMu.Unlock()
1990		2057
1991	logCodexWebsocketDisconnected(sessionID, authID, wsURL, reason, err)	2058	logCodexWebsocketDisconnected(sessionID, authID, wsURL, reason, err)
1992	sess.notifyUpstreamDisconnect(err)	2059	if notify {
		2060	sess.notifyUpstreamDisconnect(err)
		2061	}
1993	if closer != nil {	2062	if closer != nil {
1994	if errClose := closer.Close(); errClose != nil {	2063	if errClose := closer.Close(); errClose != nil {
1995	log.Errorf("codex websockets executor: close websocket error: %v", errClose)	2064	log.Errorf("codex websockets executor: close websocket error: %v", errClose)
@@ -2206,6 +2275,9 @@ func (e CodexAutoExecutor) Execute(ctx context.Context, auth cliproxyauth.Auth · 210 unmodified lines
2206	if cliproxyexecutor.DownstreamWebsocket(ctx) && codexWebsocketsEnabled(auth) {	2275	if cliproxyexecutor.DownstreamWebsocket(ctx) && codexWebsocketsEnabled(auth) {
2207	return e.wsExec.Execute(ctx, auth, req, opts)	2276	return e.wsExec.Execute(ctx, auth, req, opts)
2208	}	2277	}
		2278	if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
		2279	return cliproxyexecutor.Response{}, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
		2280	}
2209	return e.httpExec.Execute(ctx, auth, req, opts)	2281	return e.httpExec.Execute(ctx, auth, req, opts)
2210	}	2282	}
2211		2283
@@ -2216,6 +2288,9 @@ func (e CodexAutoExecutor) ExecuteStream(ctx context.Context, auth cliproxyaut · 4 unmodified lines
2216	if cliproxyexecutor.DownstreamWebsocket(ctx) && codexWebsocketsEnabled(auth) {	2288	if cliproxyexecutor.DownstreamWebsocket(ctx) && codexWebsocketsEnabled(auth) {
2217	return e.wsExec.ExecuteStream(ctx, auth, req, opts)	2289	return e.wsExec.ExecuteStream(ctx, auth, req, opts)
2218	}	2290	}
		2291	if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
		2292	return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
		2293	}
2219	return e.httpExec.ExecuteStream(ctx, auth, req, opts)	2294	return e.httpExec.ExecuteStream(ctx, auth, req, opts)
2220	}	2295	}
2221		2296

internal/runtime/executor/codex_websockets_executor_test.go+170−0

@@ -361,6 +361,176 @@ func TestCodexWebsocketsExecutePreservesPreviousResponseIDUpstream(t *testing.T) · 360 unmodified lines
361	}	361	}
362	}	362	}
363		363
		364	func TestCodexWebsocketsExecuteStreamUpgradeRequiredReturnsWithoutLockingSession(t *testing.T) {
		365	upgradeAttempts := make(chan struct{}, 2)
		366	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		367	if !strings.EqualFold(r.Header.Get("Upgrade"), "websocket") {
		368	t.Errorf("unexpected HTTP fallback request: %s %s", r.Method, r.URL.Path)
		369	w.WriteHeader(http.StatusInternalServerError)
		370	return
		371	}
		372	upgradeAttempts <- struct{}{}
		373	w.WriteHeader(http.StatusUpgradeRequired)
		374	_, _ = w.Write([]byte(`{"error":{"message":"websocket unavailable"}}`))
		375	}))
		376	defer server.Close()
		377
		378	exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}})
		379	const executionSessionID = "ws-upgrade-required-session"
		380	t.Cleanup(func() { exec.CloseExecutionSession(executionSessionID) })
		381	auth := &cliproxyauth.Auth{
		382	ID: "codex-test",
		383	Provider: "codex",
		384	Attributes: map[string]string{
		385	"api_key": "sk-test",
		386	"base_url": server.URL,
		387	},
		388	}
		389	opts := cliproxyexecutor.Options{
		390	SourceFormat: sdktranslator.FromString("openai-response"),
		391	ResponseFormat: sdktranslator.FromString("openai-response"),
		392	Metadata: map[string]any{
		393	cliproxyexecutor.ExecutionSessionMetadataKey: executionSessionID,
		394	},
		395	}
		396	ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
		397
		398	execute := func(payload string) {
		399	t.Helper()
		400	done := make(chan error, 1)
		401	go func() {
		402	_, errExecute := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{
		403	Model: "gpt-5.4",
		404	Payload: []byte(payload),
		405	}, opts)
		406	done <- errExecute
		407	}()
		408
		409	select {
		410	case errExecute := <-done:
		411	if errExecute == nil {
		412	t.Fatal("upgrade-required error = nil")
		413	}
		414	statusErr, ok := errExecute.(interface{ StatusCode() int })
		415	if !ok \|\| statusErr.StatusCode() != http.StatusUpgradeRequired {
		416	t.Fatalf("upgrade-required error = %T %v, want status 426", errExecute, errExecute)
		417	}
		418	case <-time.After(5 * time.Second):
		419	t.Fatal("timed out waiting for upgrade-required error; execution session may still be locked")
		420	}
		421	}
		422
		423	execute(`{"model":"gpt-5.4","generate":false,"input":[]}`)
		424	execute(`{"model":"gpt-5.4","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`)
		425
		426	if got := len(upgradeAttempts); got != 2 {
		427	t.Fatalf("websocket upgrade attempts = %d, want 2", got)
		428	}
		429	}
		430
		431	func TestCodexWebsocketsExecuteStreamHandshakeErrorReturnsWithoutLockingSession(t *testing.T) {
		432	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		433	w.WriteHeader(http.StatusUnauthorized)
		434	_, _ = w.Write([]byte(`{"error":{"message":"unauthorized"}}`))
		435	}))
		436	defer server.Close()
		437
		438	exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}})
		439	const executionSessionID = "ws-handshake-error-session"
		440	t.Cleanup(func() { exec.CloseExecutionSession(executionSessionID) })
		441	auth := &cliproxyauth.Auth{
		442	ID: "codex-test",
		443	Provider: "codex",
		444	Attributes: map[string]string{
		445	"api_key": "sk-test",
		446	"base_url": server.URL,
		447	},
		448	}
		449	opts := cliproxyexecutor.Options{
		450	SourceFormat: sdktranslator.FromString("openai-response"),
		451	Metadata: map[string]any{
		452	cliproxyexecutor.ExecutionSessionMetadataKey: executionSessionID,
		453	},
		454	}
		455
		456	for i := 0; i < 2; i++ {
		457	done := make(chan error, 1)
		458	go func() {
		459	_, errExecute := exec.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{
		460	Model: "gpt-5.4",
		461	Payload: []byte(`{"model":"gpt-5.4","input":[{"type":"message","id":"msg-1"}]}`),
		462	}, opts)
		463	done <- errExecute
		464	}()
		465	select {
		466	case errExecute := <-done:
		467	statusErr, ok := errExecute.(interface{ StatusCode() int })
		468	if !ok \|\| statusErr.StatusCode() != http.StatusUnauthorized {
		469	t.Fatalf("attempt %d error = %T %v, want status 401", i+1, errExecute, errExecute)
		470	}
		471	case <-time.After(5 * time.Second):
		472	t.Fatalf("attempt %d timed out; execution session remained locked", i+1)
		473	}
		474	}
		475	}
		476
		477	func TestExistingWebsocketSessionConnRequiresMatchingHealthyConnection(t *testing.T) {
		478	conn := &websocket.Conn{}
		479	closer := newWebsocketConnectionCloser(conn)
		480	sess := &codexWebsocketSession{
		481	conn: conn,
		482	connCloser: closer,
		483	authID: "auth-a",
		484	wsURL: "ws://example.test/responses",
		485	}
		486	sess.resetUpstreamDisconnectError(conn)
		487	if gotConn, gotCloser := existingWebsocketSessionConn(sess, "auth-a", "ws://example.test/responses"); gotConn != conn \|\| gotCloser != closer {
		488	t.Fatal("matching healthy websocket session was not reusable")
		489	}
		490	if got, _ := existingWebsocketSessionConn(sess, "auth-b", "ws://example.test/responses"); got != nil {
		491	t.Fatal("websocket session matched a different auth")
		492	}
		493	if got, _ := existingWebsocketSessionConn(sess, "auth-a", "ws://other.test/responses"); got != nil {
		494	t.Fatal("websocket session matched a different URL")
		495	}
		496	sess.setUpstreamDisconnectError(conn, errors.New("upstream disconnected"))
		497	if got, _ := existingWebsocketSessionConn(sess, "auth-a", "ws://example.test/responses"); got != nil {
		498	t.Fatal("disconnected websocket session remained reusable")
		499	}
		500	}
		501
		502	func TestCodexAutoExecutorRequiredUpstreamWebsocketRejectsHTTPFallback(t *testing.T) {
		503	exec := NewCodexAutoExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}})
		504	auth := &cliproxyauth.Auth{
		505	ID: "codex-http-only",
		506	Provider: "codex",
		507	Attributes: map[string]string{
		508	"api_key": "sk-test",
		509	},
		510	}
		511	ctx := cliproxyexecutor.WithRequiredUpstreamWebsocket(
		512	cliproxyexecutor.WithDownstreamWebsocket(context.Background()),
		513	)
		514	_, errExecute := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{
		515	Model: "gpt-5.4",
		516	Payload: []byte(`{"model":"gpt-5.4","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`),
		517	}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("openai-response")})
		518	if errExecute == nil {
		519	t.Fatal("ExecuteStream() error = nil, want replay-required error")
		520	}
		521	statusErr, ok := errExecute.(interface{ StatusCode() int })
		522	if !ok \|\| statusErr.StatusCode() != http.StatusUpgradeRequired {
		523	t.Fatalf("ExecuteStream() error = %T %v, want status 426", errExecute, errExecute)
		524	}
		525	if got := gjson.Get(errExecute.Error(), "error.code").String(); got != "upstream_http_replay_required" {
		526	t.Fatalf("ExecuteStream() error code = %q, want upstream_http_replay_required", got)
		527	}
		528	requestScoped, ok := errExecute.(cliproxyexecutor.RequestScopedError)
		529	if !ok \|\| !requestScoped.IsRequestScoped() {
		530	t.Fatalf("ExecuteStream() error = %T, want request-scoped replay signal", errExecute)
		531	}
		532	}
		533
364	func TestCodexWebsocketsExecuteStreamPassesThroughUpstreamWebsocketPayloadForDownstreamWebsocket(t *testing.T) {	534	func TestCodexWebsocketsExecuteStreamPassesThroughUpstreamWebsocketPayloadForDownstreamWebsocket(t *testing.T) {
365	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}	535	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}
366	capturedPayload := make(chan []byte, 1)	536	capturedPayload := make(chan []byte, 1)

internal/runtime/executor/xai_websockets_executor.go+149−21

 }
 type xaiWebsocketIDState struct {
-	mu                   sync.Mutex
+	requestMu                        sync.Mutex
-	downstreamToUpstream map[string]string
+	mu                               sync.Mutex
-	sequence             int
+	downstreamToUpstream             map[string]string
-	transcriptInput      []json.RawMessage
+	sequence                         int
+	transcriptInput                  []json.RawMessage
+	replayCompactedTranscriptOnReset bool
 }
 type xaiWebsocketRequestIDMapper struct {
-	state                *xaiWebsocketIDState
+	state                       *xaiWebsocketIDState
-	downstreamPreviousID string
+	downstreamPreviousID        string
-	upstreamPreviousID   string
+	upstreamPreviousID          string
-	upstreamResponseID   string
+	upstreamResponseID          string
-	downstreamResponseID string
+	downstreamResponseID        string
+	replayedCompactedTranscript bool
 }
 func NewXAIWebsocketsExecutor(cfg *config.Config) *XAIWebsocketsExecutor {
 	return out
 }
-func (s *xaiWebsocketIDState) recordTranscriptTurn(requestPayload []byte, completedPayload []byte) {
+func (s *xaiWebsocketIDState) recordTranscriptTurn(requestPayload []byte, completedPayload []byte, reset bool) {
 	if s == nil || len(requestPayload) == 0 || len(completedPayload) == 0 {
 		return
 	}
 	inputItems := xaiJSONRawMessages(gjson.GetBytes(requestPayload, "input"))
 	outputItems := xaiJSONRawMessages(gjson.GetBytes(completedPayload, "response.output"))
-	if len(inputItems) == 0 && len(outputItems) == 0 {
-		return
-	}
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	if strings.TrimSpace(gjson.GetBytes(requestPayload, "previous_response_id").String()) == "" {
+	if reset {
 		s.transcriptInput = nil
+		s.replayCompactedTranscriptOnReset = false
+	}
+	if len(inputItems) == 0 && len(outputItems) == 0 {
+		return
 	}
 	s.transcriptInput = append(s.transcriptInput, inputItems...)
 	s.transcriptInput = append(s.transcriptInput, outputItems...)
 	}
 	s.mu.Lock()
 	s.transcriptInput = next
+	s.replayCompactedTranscriptOnReset = len(next) > 0
 	s.mu.Unlock()
 }
+func (s *xaiWebsocketIDState) prependCompactedTranscriptOnReset(payload []byte) ([]byte, bool) {
+	if s == nil || len(payload) == 0 {
+		return payload, false
+	}
+	s.mu.Lock()
+	if !s.replayCompactedTranscriptOnReset || len(s.transcriptInput) == 0 {
+		s.mu.Unlock()
+		return payload, false
+	}
+	prefix := make([]json.RawMessage, 0, len(s.transcriptInput))
+	for _, item := range s.transcriptInput {
+		prefix = append(prefix, bytes.Clone(item))
+	}
+	s.mu.Unlock()
+	current := xaiJSONRawMessages(gjson.GetBytes(payload, "input"))
+	merged := append(prefix, current...)
+	out, errSet := sjson.SetRawBytes(payload, "input", xaiMarshalRawMessages(merged))
+	if errSet != nil {
+		return payload, false
+	}
+	return out, true
+}
 func xaiJSONRawMessages(result gjson.Result) []json.RawMessage {
 	if !result.Exists() || !result.IsArray() {
 		return nil
 }
 func (m *xaiWebsocketRequestIDMapper) upstreamRequestPayload(payload []byte) []byte {
-	if m == nil || len(payload) == 0 || m.downstreamPreviousID == m.upstreamPreviousID {
+	if m == nil || len(payload) == 0 {
+		return payload
+	}
+	if m.downstreamPreviousID == m.upstreamPreviousID {
+		requestType := strings.TrimSpace(gjson.GetBytes(payload, "type").String())
+		if m.downstreamPreviousID == "" && requestType == "response.append" && m.state != nil {
+			out, replayed := m.state.prependCompactedTranscriptOnReset(payload)
+			m.replayedCompactedTranscript = replayed
+			return out
+		}
 		return payload
 	}
 	if m.upstreamPreviousID == "" {
 		if errDelete == nil {
 			if m.downstreamPreviousID != "" && m.state != nil {
 				out = m.state.prependTranscriptInput(out)
+				m.replayedCompactedTranscript = true
 			}
 			return out
 		}
 	if stateSessionID == "" {
 		stateSessionID = executionSessionID
 	}
-	idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload)
+	state := getXAIWebsocketIDState(e.idStore, stateSessionID)
+	stateRequestLocked := false
+	stateRequestLockTransferred := false
+	if executionSessionID == "" && state != nil {
+		state.requestMu.Lock()
+		stateRequestLocked = true
+	}
+	defer func() {
+		if stateRequestLocked && !stateRequestLockTransferred {
+			state.requestMu.Unlock()
+		}
+	}()
 	if xaiInputHasItemType(req.Payload, "compaction_trigger") {
+		if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
+			return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
+		}
+		if executionSessionID != "" {
+			sess := e.getOrCreateSession(executionSessionID)
+			if sess != nil {
+				sess.reqMu.Lock()
+				defer sess.reqMu.Unlock()
+			}
+		}
+		idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload)
 		return e.executeCompactionTriggerFromWebsocketContext(ctx, auth, req, opts, idMapper)
 	}
 			sess.reqMu.Lock()
 		}
 	}
+	idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload)
 	if idMapper != nil {
 		if websocketSessionTargetChanged(sess, authID, wsURL) {
 			idMapper.upstreamPreviousID = ""
 	wsHeaders := applyXAIWebsocketHeaders(http.Header{}, auth, token, prepared.sessionID)
 	wsReqBody := buildXAIWebsocketRequestBody(prepared.body)
+	requestType := strings.TrimSpace(gjson.GetBytes(req.Payload, "type").String())
+	transcriptReset := strings.TrimSpace(gjson.GetBytes(wsReqBody, "previous_response_id").String()) == "" &&
+		(requestType != "response.append" || (idMapper != nil && idMapper.replayedCompactedTranscript))
 	warmupRequest := xaiWebsocketGenerateFalse(wsReqBody)
 	wsReqLog := helps.UpstreamRequestLog{
 	helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog)
 	logXAIWebsocketRequest(executionSessionID, authID, wsURL, wsReqBody)
-	conn, closer, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
+	var conn *websocket.Conn
+	var closer *websocketConnectionCloser
+	var respHS *http.Response
+	var errDial error
+	if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
+		conn, closer = existingWebsocketSessionConn(sess, authID, wsURL)
+		if conn == nil {
+			if sess != nil {
+				sess.reqMu.Unlock()
+			}
+			return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
+		}
+	} else {
+		conn, closer, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
+	}
 	var upstreamHeaders http.Header
 	if respHS != nil {
 		upstreamHeaders = respHS.Header.Clone()
 		errSend = mapXAIWebsocketWriteError(sess, conn, errSend)
 		helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend)
 		if sess != nil {
+			if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
+				e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "send_error", errSend)
+				sess.clearActive(conn, readCh)
+				sess.reqMu.Unlock()
+				if !shouldRetryXAIWebsocketSend(errSend) {
+					return nil, errSend
+				}
+				return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
+			}
 			e.invalidateUpstreamConn(sess, conn, "send_error", errSend)
 			if !shouldRetryXAIWebsocketSend(errSend) {
 				sess.clearActive(conn, readCh)
 	}
 	out := make(chan cliproxyexecutor.StreamChunk)
+	if stateRequestLocked {
+		stateRequestLockTransferred = true
+	}
 	go func() {
+		if stateRequestLocked {
+			defer state.requestMu.Unlock()
+		}
 		terminateReason := "completed"
 		var terminateErr error
 				case "response.created":
 					if warmupRequest {
 						warmupCompletedPayload = buildXAIWebsocketWarmupCompletedPayload(payload)
+						if idMapper != nil && idMapper.state != nil && !recordedTranscript {
+							idMapper.state.recordTranscriptTurn(wsReqBody, warmupCompletedPayload, transcriptReset)
+							recordedTranscript = true
+						}
 						logXAIWebsocketWarmupCompleted(executionSessionID, authID, wsURL, payload)
 					}
 				case "response.output_item.done":
 					payload = xaiNormalizeReasoningSummaryData(payload)
 					cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, payload)
 					if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript {
-						idMapper.state.recordTranscriptTurn(wsReqBody, payload)
+						idMapper.state.recordTranscriptTurn(wsReqBody, payload, transcriptReset)
 						recordedTranscript = true
 					}
 				case "response.done":
 						reporter.Publish(ctx, detail)
 					}
 					if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript {
-						idMapper.state.recordTranscriptTurn(wsReqBody, payload)
+						idMapper.state.recordTranscriptTurn(wsReqBody, payload, transcriptReset)
 						recordedTranscript = true
 					}
 				}
 		return nil, err
 	}
-	responseID := xaiCompactionResponseID(data)
+	responseID, compactionItem, errValidate := validateXAIWebsocketCompactionResponse(data)
-	idMapper.state.replaceTranscriptWithItems(xaiCompactionOutputItem(data, responseID))
+	if errValidate != nil {
+		return nil, errValidate
+	}
+	idMapper.state.replaceTranscriptWithItems(compactionItem)
 	idMapper.state.mapDownstreamToUpstream(responseID, "")
 	headers = headers.Clone()
 	return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out}, nil
 }
+func validateXAIWebsocketCompactionResponse(data []byte) (string, []byte, error) {
+	if len(data) == 0 || !json.Valid(data) {
+		return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction returned invalid JSON"}
+	}
+	responseIDResult := gjson.GetBytes(data, "id")
+	output := gjson.GetBytes(data, "output")
+	if responseIDResult.Type != gjson.String || strings.TrimSpace(responseIDResult.String()) == "" || !output.Exists() || !output.IsArray() {
+		return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction response is missing compacted state"}
+	}
+	items := output.Array()
+	if len(items) == 0 {
+		return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction response is missing compacted state"}
+	}
+	item := items[0]
+	itemType := item.Get("type")
+	encryptedContent := item.Get("encrypted_content")
+	if item.Type != gjson.JSON || itemType.Type != gjson.String || strings.TrimSpace(itemType.String()) != "compaction" ||
+		encryptedContent.Type != gjson.String || strings.TrimSpace(encryptedContent.String()) == "" {
+		return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction response is missing compacted state"}
+	}
+	normalizedResponseID := xaiCompactionResponseID(data)
+	return normalizedResponseID, xaiCompactionOutputItem(data, normalizedResponseID), nil
+}
 func buildXAIWebsocketCompactionPayload(payload []byte, transcriptInput []byte) ([]byte, error) {
 	if len(payload) == 0 {
 		payload = []byte(`{}`)
 	if cliproxyexecutor.DownstreamWebsocket(ctx) && xaiWebsocketsEnabled(auth) {
 		return e.wsExec.ExecuteStream(ctx, auth, req, opts)
 	}
+	if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
+		return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
+	}
 	return e.httpExec.ExecuteStream(ctx, auth, req, opts)
 }

internal/runtime/executor/xai_websockets_executor_test.go+165−0

@@ -34,6 +34,50 @@ func TestXAIWebsocketsEnabledForConfigAPIKey(t *testing.T) { · 33 unmodified lines
34	}	34	}
35	}	35	}
36		36
		37	func TestXAIAutoExecutorRequiredUpstreamWebsocketRejectsHTTPFallback(t *testing.T) {
		38	exec := NewXAIAutoExecutor(&config.Config{})
		39	auth := &cliproxyauth.Auth{
		40	ID: "xai-http-only",
		41	Provider: "xai",
		42	Attributes: map[string]string{
		43	"api_key": "xai-key",
		44	},
		45	}
		46	ctx := cliproxyexecutor.WithRequiredUpstreamWebsocket(
		47	cliproxyexecutor.WithDownstreamWebsocket(context.Background()),
		48	)
		49	_, errExecute := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{
		50	Model: "grok-4",
		51	Payload: []byte(`{"model":"grok-4","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`),
		52	}, cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("openai-response")})
		53	if errExecute == nil {
		54	t.Fatal("ExecuteStream() error = nil, want replay-required error")
		55	}
		56	statusErr, ok := errExecute.(interface{ StatusCode() int })
		57	if !ok \|\| statusErr.StatusCode() != http.StatusUpgradeRequired {
		58	t.Fatalf("ExecuteStream() error = %T %v, want status 426", errExecute, errExecute)
		59	}
		60	if got := gjson.Get(errExecute.Error(), "error.code").String(); got != "upstream_http_replay_required" {
		61	t.Fatalf("ExecuteStream() error code = %q, want upstream_http_replay_required", got)
		62	}
		63	requestScoped, ok := errExecute.(cliproxyexecutor.RequestScopedError)
		64	if !ok \|\| !requestScoped.IsRequestScoped() {
		65	t.Fatalf("ExecuteStream() error = %T, want request-scoped replay signal", errExecute)
		66	}
		67	}
		68
		69	func TestXAIWebsocketsRequiredUpstreamRejectsCompactionHTTPFallback(t *testing.T) {
		70	exec := NewXAIWebsocketsExecutor(&config.Config{})
		71	ctx := cliproxyexecutor.WithRequiredUpstreamWebsocket(context.Background())
		72	_, errExecute := exec.ExecuteStream(ctx, &cliproxyauth.Auth{}, cliproxyexecutor.Request{
		73	Model: "grok-4",
		74	Payload: []byte(`{"model":"grok-4","input":[{"type":"compaction_trigger"}]}`),
		75	}, cliproxyexecutor.Options{})
		76	if !cliproxyexecutor.IsUpstreamWebsocketReplayRequired(errExecute) {
		77	t.Fatalf("ExecuteStream() error = %T %v, want replay-required", errExecute, errExecute)
		78	}
		79	}
		80
37	func TestMapXAIWebsocketWriteErrorStopsRetryForMessageTooBig(t *testing.T) {	81	func TestMapXAIWebsocketWriteErrorStopsRetryForMessageTooBig(t *testing.T) {
38	networkWriteErr := errors.New("write: broken pipe")	82	networkWriteErr := errors.New("write: broken pipe")
39	tests := []struct {	83	tests := []struct {
@@ -1246,6 +1290,127 @@ func TestXAIWebsocketsExecuteStreamCompactionTriggerUsesHTTPCompactWithRecordedC · 1206 unmodified lines
1246	}	1290	}
1247	}	1291	}
1248		1292
		1293	func TestXAIWebsocketPostCompactionAppendWithoutPreviousReplaysCompactedTranscript(t *testing.T) {
		1294	store := &xaiWebsocketIDStateStore{sessions: make(map[string]*xaiWebsocketIDState)}
		1295	state := getXAIWebsocketIDState(store, "post-compaction-append-session")
		1296	state.replaceTranscriptWithItems([]byte(`{"type":"compaction","encrypted_content":"compact-state"}`))
		1297	state.mapDownstreamToUpstream("resp-compact", "")
		1298
		1299	fullReset := []byte(`{"type":"response.create","model":"grok-4.3","input":[{"type":"message","id":"msg-full"}]}`)
		1300	fullMapper := newXAIWebsocketRequestIDMapper(store, "post-compaction-append-session", fullReset)
		1301	if full := fullMapper.upstreamRequestPayload(fullReset); len(gjson.GetBytes(full, "input").Array()) != 1 {
		1302	t.Fatalf("self-contained response.create unexpectedly replayed compacted transcript: %s", full)
		1303	}
		1304
		1305	payload := []byte(`{"type":"response.append","model":"grok-4.3","input":[{"type":"message","id":"msg-2","role":"user","content":"second"}]}`)
		1306	mapper := newXAIWebsocketRequestIDMapper(store, "post-compaction-append-session", payload)
		1307	got := mapper.upstreamRequestPayload(payload)
		1308	input := gjson.GetBytes(got, "input").Array()
		1309	if len(input) != 2 {
		1310	t.Fatalf("post-compaction append input len = %d, want 2: %s", len(input), got)
		1311	}
		1312	if gotType := input[0].Get("type").String(); gotType != "compaction" {
		1313	t.Fatalf("post-compaction append input[0].type = %q, want compaction: %s", gotType, got)
		1314	}
		1315	if gotID := input[1].Get("id").String(); gotID != "msg-2" {
		1316	t.Fatalf("post-compaction append input[1].id = %q, want msg-2: %s", gotID, got)
		1317	}
		1318
		1319	state.recordTranscriptTurn(got, []byte(`{"type":"response.completed","response":{"id":"resp-after-compact","output":[{"type":"message","id":"out-2"}]}}`), true)
		1320	nextPayload := []byte(`{"type":"response.create","model":"grok-4.3","input":[{"type":"message","id":"msg-3"}]}`)
		1321	nextMapper := newXAIWebsocketRequestIDMapper(store, "post-compaction-append-session", nextPayload)
		1322	next := nextMapper.upstreamRequestPayload(nextPayload)
		1323	if nextInput := gjson.GetBytes(next, "input").Array(); len(nextInput) != 1 \|\| nextInput[0].Get("id").String() != "msg-3" {
		1324	t.Fatalf("compacted transcript replay was not cleared after success: %s", next)
		1325	}
		1326	}
		1327
		1328	func TestXAIWebsocketPostCompactionWarmupPreservesTranscriptForLaterCompaction(t *testing.T) {
		1329	store := &xaiWebsocketIDStateStore{sessions: make(map[string]*xaiWebsocketIDState)}
		1330	state := getXAIWebsocketIDState(store, "warmup-reset-session")
		1331	state.replaceTranscriptWithItems([]byte(`{"type":"compaction","encrypted_content":"compact-state"}`))
		1332
		1333	warmupPayload := []byte(`{"type":"response.append","model":"grok-4.3","generate":false,"input":[{"type":"message","id":"warmup-context"}]}`)
		1334	warmupMapper := newXAIWebsocketRequestIDMapper(store, "warmup-reset-session", warmupPayload)
		1335	warmupUpstream := warmupMapper.upstreamRequestPayload(warmupPayload)
		1336	if !warmupMapper.replayedCompactedTranscript {
		1337	t.Fatal("post-compaction warmup did not mark full transcript replay")
		1338	}
		1339	state.recordTranscriptTurn(
		1340	warmupUpstream,
		1341	[]byte(`{"type":"response.completed","response":{"id":"resp-warmup","output":[]}}`),
		1342	true,
		1343	)
		1344
		1345	appendPayload := []byte(`{"type":"response.append","model":"grok-4.3","input":[{"type":"message","id":"msg-after-warmup"}]}`)
		1346	appendMapper := newXAIWebsocketRequestIDMapper(store, "warmup-reset-session", appendPayload)
		1347	appendUpstream := appendMapper.upstreamRequestPayload(appendPayload)
		1348	input := gjson.GetBytes(appendUpstream, "input").Array()
		1349	if len(input) != 1 \|\| input[0].Get("id").String() != "msg-after-warmup" {
		1350	t.Fatalf("warmup retained pending replay instead of native append: %s", appendUpstream)
		1351	}
		1352	state.recordTranscriptTurn(
		1353	appendUpstream,
		1354	[]byte(`{"type":"response.completed","response":{"id":"resp-after-warmup","output":[{"type":"message","id":"out-after-warmup"}]}}`),
		1355	false,
		1356	)
		1357
		1358	transcript := gjson.ParseBytes(state.snapshotTranscriptInput()).Array()
		1359	wantTypes := []string{"compaction", "message", "message", "message"}
		1360	if len(transcript) != len(wantTypes) {
		1361	t.Fatalf("post-warmup transcript len = %d, want %d: %s", len(transcript), len(wantTypes), state.snapshotTranscriptInput())
		1362	}
		1363	for i, wantType := range wantTypes {
		1364	if gotType := transcript[i].Get("type").String(); gotType != wantType {
		1365	t.Fatalf("post-warmup transcript[%d].type = %q, want %q: %s", i, gotType, wantType, state.snapshotTranscriptInput())
		1366	}
		1367	}
		1368	}
		1369
		1370	func TestXAIWebsocketEmptyFullResetClearsPendingCompactionReplay(t *testing.T) {
		1371	store := &xaiWebsocketIDStateStore{sessions: make(map[string]*xaiWebsocketIDState)}
		1372	state := getXAIWebsocketIDState(store, "empty-reset-session")
		1373	state.replaceTranscriptWithItems([]byte(`{"type":"compaction","encrypted_content":"stale-compact-state"}`))
		1374	state.recordTranscriptTurn(
		1375	[]byte(`{"type":"response.create","model":"grok-4.3","input":[]}`),
		1376	[]byte(`{"type":"response.completed","response":{"id":"resp-empty","output":[]}}`),
		1377	true,
		1378	)
		1379
		1380	appendPayload := []byte(`{"type":"response.append","model":"grok-4.3","input":[{"type":"message","id":"msg-new"}]}`)
		1381	mapper := newXAIWebsocketRequestIDMapper(store, "empty-reset-session", appendPayload)
		1382	got := mapper.upstreamRequestPayload(appendPayload)
		1383	input := gjson.GetBytes(got, "input").Array()
		1384	if len(input) != 1 \|\| input[0].Get("id").String() != "msg-new" {
		1385	t.Fatalf("empty full reset retained stale compaction replay: %s", got)
		1386	}
		1387	}
		1388
		1389	func TestValidateXAIWebsocketCompactionResponse(t *testing.T) {
		1390	valid := []byte(`{"id":"resp_compact","output":[{"type":"compaction","encrypted_content":"opaque-state"}]}`)
		1391	responseID, item, err := validateXAIWebsocketCompactionResponse(valid)
		1392	if err != nil {
		1393	t.Fatalf("valid compaction response error: %v", err)
		1394	}
		1395	if responseID != "resp_compact" \|\| gjson.GetBytes(item, "encrypted_content").String() != "opaque-state" {
		1396	t.Fatalf("validated compaction response = id:%q item:%s", responseID, item)
		1397	}
		1398
		1399	for _, payload := range [][]byte{
		1400	nil,
		1401	[]byte(`{}`),
		1402	[]byte(`{"id":"resp_empty","output":[]}`),
		1403	[]byte(`{"id":123,"output":[{"type":"compaction","encrypted_content":"opaque"}]}`),
		1404	[]byte(`{"id":"resp_object","output":{"0":{"type":"compaction","encrypted_content":"opaque"}}}`),
		1405	[]byte(`{"id":"resp_numeric_state","output":[{"type":"compaction","encrypted_content":123}]}`),
		1406	[]byte(`{"id":"resp_missing_state","output":[{"type":"compaction"}]}`),
		1407	} {
		1408	if _, _, errInvalid := validateXAIWebsocketCompactionResponse(payload); errInvalid == nil {
		1409	t.Fatalf("invalid compaction response accepted: %s", payload)
		1410	}
		1411	}
		1412	}
		1413
1249	func TestBuildXAIWebsocketRequestBodySetsStoreAndKeepsPromptCacheKey(t *testing.T) {	1414	func TestBuildXAIWebsocketRequestBodySetsStoreAndKeepsPromptCacheKey(t *testing.T) {
1250	body := []byte(`{"model":"grok-4.3","stream":true,"stream_options":{"include_usage":true},"background":true,"prompt_cache_key":"cache-1","previous_response_id":"resp-prev","instructions":"system prompt","input":[{"type":"message","role":"user","content":"hello"}]}`)	1415	body := []byte(`{"model":"grok-4.3","stream":true,"stream_options":{"include_usage":true},"background":true,"prompt_cache_key":"cache-1","previous_response_id":"resp-prev","instructions":"system prompt","input":[{"type":"message","role":"user","content":"hello"}]}`)
1251		1416

sdk/api/handlers/handlers.go+25−1

@@ -64,6 +64,7 @@ const ( · 63 unmodified lines
64		64
65	type pinnedAuthContextKey struct{}	65	type pinnedAuthContextKey struct{}
66	type selectedAuthCallbackContextKey struct{}	66	type selectedAuthCallbackContextKey struct{}
		67	type preparedModelRouteContextKey struct{}
67	type executionSessionContextKey struct{}	68	type executionSessionContextKey struct{}
68	type disallowFreeAuthContextKey struct{}	69	type disallowFreeAuthContextKey struct{}
69		70
@@ -142,6 +143,26 @@ func WithSelectedAuthIDCallback(ctx context.Context, callback func(string)) cont · 72 unmodified lines
142	return context.WithValue(ctx, selectedAuthCallbackContextKey{}, callback)	143	return context.WithValue(ctx, selectedAuthCallbackContextKey{}, callback)
143	}	144	}
144		145
		146	// PrepareStreamModelRoute resolves a stream route once and stores it on the returned context for execution.
		147	// The boolean reports whether the route overrides normal model-to-provider resolution.
		148	func (h *BaseAPIHandler) PrepareStreamModelRoute(ctx context.Context, handlerType string, modelName string, rawJSON []byte) (context.Context, bool) {
		149	if ctx == nil {
		150	ctx = context.Background()
		151	}
		152	decision := h.applyModelRouter(ctx, handlerType, modelName, rawJSON, true, modelExecutionOptions{})
		153	ctx = context.WithValue(ctx, preparedModelRouteContextKey{}, decision)
		154	hasOverride := strings.TrimSpace(decision.ExecutorPluginID) != "" \|\| strings.TrimSpace(decision.Provider) != ""
		155	return ctx, hasOverride
		156	}
		157
		158	func preparedModelRouteFromContext(ctx context.Context) (modelRouteDecision, bool) {
		159	if ctx == nil {
		160	return modelRouteDecision{}, false
		161	}
		162	decision, ok := ctx.Value(preparedModelRouteContextKey{}).(modelRouteDecision)
		163	return decision, ok
		164	}
		165
145	// WithExecutionSessionID returns a child context tagged with a long-lived execution session ID.	166	// WithExecutionSessionID returns a child context tagged with a long-lived execution session ID.
146	func WithExecutionSessionID(ctx context.Context, sessionID string) context.Context {	167	func WithExecutionSessionID(ctx context.Context, sessionID string) context.Context {
147	sessionID = strings.TrimSpace(sessionID)	168	sessionID = strings.TrimSpace(sessionID)
@@ -1144,7 +1165,10 @@ func (h *BaseAPIHandler) executeStreamWithAuthManager(ctx context.Context, handl · 996 unmodified lines
1144		1165
1145	func (h BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context, entryProtocol, exitProtocol, modelName string, rawJSON []byte, alt string, allowImageModel bool, execOptions modelExecutionOptions) (<-chan []byte, http.Header, <-chan interfaces.ErrorMessage) {	1166	func (h BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context, entryProtocol, exitProtocol, modelName string, rawJSON []byte, alt string, allowImageModel bool, execOptions modelExecutionOptions) (<-chan []byte, http.Header, <-chan interfaces.ErrorMessage) {
1146	originalRequestedModel := modelName	1167	originalRequestedModel := modelName
1147	routeDecision := h.applyModelRouter(ctx, entryProtocol, modelName, rawJSON, true, execOptions)	1168	routeDecision, preparedRoute := preparedModelRouteFromContext(ctx)
		1169	if !preparedRoute {
		1170	routeDecision = h.applyModelRouter(ctx, entryProtocol, modelName, rawJSON, true, execOptions)
		1171	}
1148	responseProtocol := modelExecutionResponseProtocol(entryProtocol, exitProtocol)	1172	responseProtocol := modelExecutionResponseProtocol(entryProtocol, exitProtocol)
1149	if errMsg := validateNativeInteractionsExecution(entryProtocol, execOptions, routeDecision); errMsg != nil {	1173	if errMsg := validateNativeInteractionsExecution(entryProtocol, execOptions, routeDecision); errMsg != nil {
1150	errChan := make(chan *interfaces.ErrorMessage, 1)	1174	errChan := make(chan *interfaces.ErrorMessage, 1)

sdk/api/handlers/handlers_model_router_test.go+32−0

@@ -464,6 +464,38 @@ func TestHandlerModelRouterRoutesStreamBeforeRequestDetails(t *testing.T) { · 463 unmodified lines
464	}	464	}
465	}	465	}
466		466
		467	func TestPrepareStreamModelRouteReusesDecisionDuringExecution(t *testing.T) {
		468	const model = "prepared-router-model"
		469	const targetPluginID = "prepared-stream-plugin"
		470	routeCalls := 0
		471	host := &handlerDirectExecutorRouteHost{}
		472	host.hasRouters = true
		473	host.route = func(context.Context, pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
		474	routeCalls++
		475	return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
		476	}
		477	handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
		478	handler.SetModelRouterHost(host)
		479	body := []byte(`{"model":"prepared-router-model","stream":true}`)
		480	ctx, routedToPlugin := handler.PrepareStreamModelRoute(context.Background(), "openai", model, body)
		481	if !routedToPlugin {
		482	t.Fatal("PrepareStreamModelRoute() did not detect plugin executor route")
		483	}
		484
		485	dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(ctx, "openai", model, body, "")
		486	for range dataChan {
		487	}
		488	if errMsg := <-errChan; errMsg != nil {
		489	t.Fatalf("ExecuteStreamWithAuthManager() error = %+v", errMsg)
		490	}
		491	if routeCalls != 1 {
		492	t.Fatalf("model router calls = %d, want 1", routeCalls)
		493	}
		494	if host.lastPluginID != targetPluginID {
		495	t.Fatalf("plugin id = %q, want %q", host.lastPluginID, targetPluginID)
		496	}
		497	}
		498
467	func TestExecuteModelPropagatesRouterSkipPluginID(t *testing.T) {	499	func TestExecuteModelPropagatesRouterSkipPluginID(t *testing.T) {
468	model := "model-execution-router-skip-model"	500	model := "model-execution-router-skip-model"
469	requestBody := []byte(fmt.Sprintf(`{"model":%q}`, model))	501	requestBody := []byte(fmt.Sprintf(`{"model":%q}`, model))

sdk/api/handlers/openai/openai_responses_websocket.go+232−75

@@ -33,15 +33,19 @@ import ( · 32 unmodified lines
33	)	33	)
34		34
35	const (	35	const (
36	wsRequestTypeCreate = "response.create"	36	wsRequestTypeCreate = "response.create"
37	wsRequestTypeAppend = "response.append"	37	wsRequestTypeAppend = "response.append"
38	wsEventTypeError = "error"	38	wsEventTypeError = "error"
39	wsEventTypeCompleted = "response.completed"	39	wsEventTypeCompleted = "response.completed"
40	wsEventTypeDone = "response.done"	40	wsEventTypeDone = "response.done"
41	wsDoneMarker = "[DONE]"	41	wsDoneMarker = "[DONE]"
42	wsTurnStateHeader = "x-codex-turn-state"	42	wsTurnStateHeader = "x-codex-turn-state"
43	wsTimelineBodyKey = "WEBSOCKET_TIMELINE_OVERRIDE"	43	wsTimelineBodyKey = "WEBSOCKET_TIMELINE_OVERRIDE"
44	wsCloseReasonMaxBytes = 123	44	wsCloseReasonMaxBytes = 123
		45	wsHTTPReplayRequiredCloseReason = "upstream requires HTTP replay"
		46	responsesWebsocketUpstreamModeUnknown = ""
		47	responsesWebsocketUpstreamModeWS = "websocket"
		48	responsesWebsocketUpstreamModeHTTP = "http"
45		49
46	codexLocalCompactionSummaryPrefix = "Another language model started to solve this problem and produced a summary of its thinking process. You also have access to the state of the tools that were used by that language model. Use this to build on the work that has already been done and avoid duplicating work. Here is the summary produced by the other language model, use the information in this summary to assist with your own analysis:"	50	codexLocalCompactionSummaryPrefix = "Another language model started to solve this problem and produced a summary of its thinking process. You also have access to the state of the tools that were used by that language model. Use this to build on the work that has already been done and avoid duplicating work. Here is the summary produced by the other language model, use the information in this summary to assist with your own analysis:"
47	)	51	)
@@ -74,6 +78,14 @@ func websocketClosePayloadForUpstreamError(err error) (bool, []byte) { · 26 unmodified lines
74	return false, nil	78	return false, nil
75	}	79	}
76		80
		81	errText := err.Error()
		82	if cliproxyexecutor.IsUpstreamWebsocketReplayRequired(err) {
		83	return true, websocket.FormatCloseMessage(
		84	websocket.CloseServiceRestart,
		85	truncateWebsocketCloseReason(wsHTTPReplayRequiredCloseReason, wsCloseReasonMaxBytes),
		86	)
		87	}
		88
77	code := 0	89	code := 0
78	reason := ""	90	reason := ""
79	var closeErr *websocket.CloseError	91	var closeErr *websocket.CloseError
@@ -85,7 +97,6 @@ func websocketClosePayloadForUpstreamError(err error) (bool, []byte) { · 5 unmodified lines
85	StatusCode() int	97	StatusCode() int
86	}	98	}
87	var statusErr statusCoder	99	var statusErr statusCoder
88	errText := err.Error()
89	if !errors.As(err, &statusErr) \|\| statusErr.StatusCode() != http.StatusRequestEntityTooLarge \|\|	100	if !errors.As(err, &statusErr) \|\| statusErr.StatusCode() != http.StatusRequestEntityTooLarge \|\|
90	gjson.Get(errText, "error.code").String() != "message_too_big" {	101	gjson.Get(errText, "error.code").String() != "message_too_big" {
91	return false, nil	102	return false, nil
@@ -419,20 +430,35 @@ func (h OpenAIResponsesAPIHandler) ResponsesWebsocket(c gin.Context) { · 327 unmodified lines
419	// Preserve independent upstream auth affinity when a downstream session switches providers.	430	// Preserve independent upstream auth affinity when a downstream session switches providers.
420	pinnedAuthByProvider := make(map[string]responsesWebsocketPinnedAuthState)	431	pinnedAuthByProvider := make(map[string]responsesWebsocketPinnedAuthState)
421	passthroughModelName := ""	432	passthroughModelName := ""
		433	upstreamMode := responsesWebsocketUpstreamModeUnknown
		434	upstreamWebsocketAuthID := ""
422	sessionAuthByIDWithSource := func(authID string) (*coreauth.Auth, bool, bool) {	435	sessionAuthByIDWithSource := func(authID string) (*coreauth.Auth, bool, bool) {
423	if h == nil \|\| h.AuthManager == nil {	436	if h == nil \|\| h.AuthManager == nil {
424	return nil, false, false	437	return nil, false, false
425	}	438	}
		439	// Prefer the current manager view so hot-reloaded transport eligibility is
		440	// observed even when the execution session still holds an older auth snapshot.
		441	if auth, ok := h.AuthManager.GetByID(authID); ok {
		442	return auth, false, true
		443	}
426	if auth, ok := h.AuthManager.GetExecutionSessionAuthByID(passthroughSessionID, authID); ok {	444	if auth, ok := h.AuthManager.GetExecutionSessionAuthByID(passthroughSessionID, authID); ok {
427	return auth, true, true	445	return auth, true, true
428	}	446	}
429	auth, ok := h.AuthManager.GetByID(authID)	447	return nil, false, false
430	return auth, false, ok
431	}	448	}
432	sessionAuthByID := func(authID string) (*coreauth.Auth, bool) {	449	sessionAuthByID := func(authID string) (*coreauth.Auth, bool) {
433	auth, _, ok := sessionAuthByIDWithSource(authID)	450	auth, _, ok := sessionAuthByIDWithSource(authID)
434	return auth, ok	451	return auth, ok
435	}	452	}
		453	upstreamModeForAuth := func(auth *coreauth.Auth) string {
		454	if auth != nil && websocketUpstreamSupportsIncrementalInput(auth.Attributes, auth.Metadata) {
		455	provider := strings.ToLower(strings.TrimSpace(auth.Provider))
		456	if provider == "codex" \|\| provider == "xai" {
		457	return responsesWebsocketUpstreamModeWS
		458	}
		459	}
		460	return responsesWebsocketUpstreamModeHTTP
		461	}
436	rememberPinnedAuth := func(authID string, modelName string) {	462	rememberPinnedAuth := func(authID string, modelName string) {
437	authID = strings.TrimSpace(authID)	463	authID = strings.TrimSpace(authID)
438	auth, ok := sessionAuthByID(authID)	464	auth, ok := sessionAuthByID(authID)
@@ -454,7 +480,6 @@ func (h OpenAIResponsesAPIHandler) ResponsesWebsocket(c gin.Context) { · 15 unmodified lines
454	}	480	}
455	pinnedAuthID = ""	481	pinnedAuthID = ""
456	}	482	}
457	forceTranscriptReplayNextRequest := false
458		483
459	for {	484	for {
460	msgType, payload, errReadMessage := conn.ReadMessage()	485	msgType, payload, errReadMessage := conn.ReadMessage()
@@ -488,6 +513,13 @@ func (h OpenAIResponsesAPIHandler) ResponsesWebsocket(c gin.Context) { · 27 unmodified lines
488	if requestModelName == "" {	513	if requestModelName == "" {
489	requestModelName = strings.TrimSpace(gjson.GetBytes(lastRequest, "model").String())	514	requestModelName = strings.TrimSpace(gjson.GetBytes(lastRequest, "model").String())
490	}	515	}
		516	executionParent := context.WithValue(c.Request.Context(), "gin", c)
		517	executionParent, routeOverridesModelResolution := h.PrepareStreamModelRoute(
		518	executionParent,
		519	h.HandlerType(),
		520	requestModelName,
		521	payload,
		522	)
491	if pinnedAuthID != "" {	523	if pinnedAuthID != "" {
492	pinnedAuth, homeRuntime, ok := sessionAuthByIDWithSource(pinnedAuthID)	524	pinnedAuth, homeRuntime, ok := sessionAuthByIDWithSource(pinnedAuthID)
493	providerKey := ""	525	providerKey := ""
@@ -514,15 +546,40 @@ func (h OpenAIResponsesAPIHandler) ResponsesWebsocket(c gin.Context) { · 20 unmodified lines
514	}	546	}
515	}	547	}
516	useUpstreamWebsocketPassthrough := h.responsesWebsocketUsesUpstreamWebsocketPassthrough(requestModelName)	548	useUpstreamWebsocketPassthrough := h.responsesWebsocketUsesUpstreamWebsocketPassthrough(requestModelName)
		549	if pinnedAuthID != "" {
		550	if pinnedAuth, ok := sessionAuthByID(pinnedAuthID); ok && responsesWebsocketAuthSupportsIncrementalInput(pinnedAuth) {
		551	provider := strings.ToLower(strings.TrimSpace(pinnedAuth.Provider))
		552	useUpstreamWebsocketPassthrough = provider == "codex" \|\| provider == "xai"
		553	}
		554	}
		555	nativeWebsocketPassthrough := !routeOverridesModelResolution && responsesWebsocketNativePassthroughAllowed(
		556	upstreamMode,
		557	useUpstreamWebsocketPassthrough,
		558	pinnedAuthID,
		559	upstreamWebsocketAuthID,
		560	)
		561	requestRequiresCurrentUpstreamWebsocket := responsesWebsocketRequestRequiresCurrentUpstream(payload)
		562	if upstreamMode == responsesWebsocketUpstreamModeWS && !nativeWebsocketPassthrough {
		563	if requestRequiresCurrentUpstreamWebsocket {
		564	replayErr := responsesWebsocketHTTPReplayRequiredError()
		565	wsTerminateErr = replayErr
		566	matched, errClose := writer.closeForUpstreamError(replayErr)
		567	if !matched {
		568	_ = conn.Close()
		569	} else if errClose != nil && !errors.Is(errClose, websocket.ErrCloseSent) {
		570	log.Debugf("responses websocket: replay close failed id=%s error=%v", passthroughSessionID, errClose)
		571	}
		572	return
		573	}
		574	// A full response.create is already a self-contained reset and can safely
		575	// establish a new upstream transport without another replay.
		576	}
517	if explicitRequestModelName != "" && !useUpstreamWebsocketPassthrough {	577	if explicitRequestModelName != "" && !useUpstreamWebsocketPassthrough {
518	passthroughModelName = ""	578	passthroughModelName = ""
519	}	579	}
520	allowIncrementalInputWithPreviousResponseID := false	580
521	allowCompactionReplayBypass := false	581	allowCompactionReplayBypass := false
522	if !useUpstreamWebsocketPassthrough {	582	if !nativeWebsocketPassthrough {
523	// Downstream websocket with CPA-mediated upstream (HTTP/SSE) always uses merged
524	// transcript replay. Incremental previous_response_id is reserved for end-to-end
525	// upstream websocket passthrough only.
526	if pinnedAuthID != "" {	583	if pinnedAuthID != "" {
527	if pinnedAuth, ok := sessionAuthByID(pinnedAuthID); ok && pinnedAuth != nil {	584	if pinnedAuth, ok := sessionAuthByID(pinnedAuthID); ok && pinnedAuth != nil {
528	allowCompactionReplayBypass = responsesWebsocketAuthSupportsCompactionReplay(pinnedAuth)	585	allowCompactionReplayBypass = responsesWebsocketAuthSupportsCompactionReplay(pinnedAuth)
@@ -535,8 +592,10 @@ func (h OpenAIResponsesAPIHandler) ResponsesWebsocket(c gin.Context) { · 6 unmodified lines
535	var requestJSON []byte	592	var requestJSON []byte
536	var updatedLastRequest []byte	593	var updatedLastRequest []byte
537	var errMsg *interfaces.ErrorMessage	594	var errMsg *interfaces.ErrorMessage
538	if useUpstreamWebsocketPassthrough {	595	if nativeWebsocketPassthrough {
539	requestJSON, errMsg = normalizeResponsesWebsocketPassthroughRequest(payload, requestModelName)	596	requestJSON, errMsg = normalizeResponsesWebsocketPassthroughRequest(payload, requestModelName)
		597	} else if len(lastRequest) == 0 && strings.TrimSpace(gjson.GetBytes(payload, "previous_response_id").String()) != "" {
		598	errMsg = responsesWebsocketPreviousResponseNotFoundError()
540	} else {	599	} else {
541	requestJSON, updatedLastRequest, errMsg = normalizeResponsesWebsocketRequestWithIncrementalState(	600	requestJSON, updatedLastRequest, errMsg = normalizeResponsesWebsocketRequestWithIncrementalState(
542	payload,	601	payload,
@@ -544,7 +603,7 @@ func (h OpenAIResponsesAPIHandler) ResponsesWebsocket(c gin.Context) { · 1 unmodified line
544	lastResponseOutput,	603	lastResponseOutput,
545	lastResponseID,	604	lastResponseID,
546	lastResponsePendingToolCallIDs,	605	lastResponsePendingToolCallIDs,
547	allowIncrementalInputWithPreviousResponseID,	606	false,
548	allowCompactionReplayBypass,	607	allowCompactionReplayBypass,
549	)	608	)
550	}	609	}
@@ -570,7 +629,7 @@ func (h OpenAIResponsesAPIHandler) ResponsesWebsocket(c gin.Context) { · 19 unmodified lines
570	}	629	}
571	continue	630	continue
572	}	631	}
573	if !useUpstreamWebsocketPassthrough && shouldHandleResponsesWebsocketPrewarmLocally(payload, lastRequest, allowIncrementalInputWithPreviousResponseID) {	632	if !useUpstreamWebsocketPassthrough && shouldHandleResponsesWebsocketPrewarmLocally(payload, lastRequest, false) {
574	if updated, errDelete := sjson.DeleteBytes(requestJSON, "generate"); errDelete == nil {	633	if updated, errDelete := sjson.DeleteBytes(requestJSON, "generate"); errDelete == nil {
575	requestJSON = updated	634	requestJSON = updated
576	}	635	}
@@ -588,84 +647,120 @@ func (h OpenAIResponsesAPIHandler) ResponsesWebsocket(c gin.Context) { · 11 unmodified lines
588	continue	647	continue
589	}	648	}
590		649
		650	toolCacheTurn := newResponsesWebsocketToolCacheTurn(downstreamSessionKey)
591	previousLastRequest := bytes.Clone(lastRequest)	651	previousLastRequest := bytes.Clone(lastRequest)
592	previousLastResponseOutput := bytes.Clone(lastResponseOutput)	652	previousLastResponseOutput := bytes.Clone(lastResponseOutput)
593	previousLastResponseID := lastResponseID	653	previousLastResponseID := lastResponseID
594	previousLastResponsePendingToolCallIDs := append([]string(nil), lastResponsePendingToolCallIDs...)	654	previousLastResponsePendingToolCallIDs := append([]string(nil), lastResponsePendingToolCallIDs...)
595	forcedTranscriptReplay := forceTranscriptReplayNextRequest	655	if nativeWebsocketPassthrough {
596	if useUpstreamWebsocketPassthrough {
597	if modelName := strings.TrimSpace(gjson.GetBytes(requestJSON, "model").String()); modelName != "" {	656	if modelName := strings.TrimSpace(gjson.GetBytes(requestJSON, "model").String()); modelName != "" {
598	passthroughModelName = modelName	657	passthroughModelName = modelName
599	}	658	}
600	if forcedTranscriptReplay {
601	forceTranscriptReplayNextRequest = false
602	}
603	} else {	659	} else {
604	requestJSON = repairResponsesWebsocketToolCalls(downstreamSessionKey, requestJSON)	660	toolCacheTurn.recordRequest(requestJSON)
		661	requestJSON = repairResponsesWebsocketToolCallsWithoutRecording(downstreamSessionKey, requestJSON)
605	requestJSON = dedupeResponsesWebsocketInputItemsByID(requestJSON)	662	requestJSON = dedupeResponsesWebsocketInputItemsByID(requestJSON)
606	updatedLastRequest = bytes.Clone(requestJSON)	663	updatedLastRequest = bytes.Clone(requestJSON)
607	lastRequest = updatedLastRequest	664	lastRequest = updatedLastRequest
608	if forcedTranscriptReplay {
609	forceTranscriptReplayNextRequest = false
610	}
611	}	665	}
612		666
613	modelName := gjson.GetBytes(requestJSON, "model").String()	667	modelName := gjson.GetBytes(requestJSON, "model").String()
614	lastAttemptedAuthID := pinnedAuthID	668	lastAttemptedAuthID := pinnedAuthID
615	cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())	669	attemptedUpstreamMode := responsesWebsocketUpstreamModeUnknown
		670	selectedAuthObserved := false
		671	pinnedAuthAttempted := false
		672	cliCtx, cliCancel := h.GetContextWithCancel(h, c, executionParent)
616	cliCtx = cliproxyexecutor.WithDownstreamWebsocket(cliCtx)	673	cliCtx = cliproxyexecutor.WithDownstreamWebsocket(cliCtx)
		674	if nativeWebsocketPassthrough && requestRequiresCurrentUpstreamWebsocket {
		675	cliCtx = cliproxyexecutor.WithRequiredUpstreamWebsocket(cliCtx)
		676	}
617	cliCtx = handlers.WithExecutionSessionID(cliCtx, passthroughSessionID)	677	cliCtx = handlers.WithExecutionSessionID(cliCtx, passthroughSessionID)
618	if pinnedAuthID != "" {	678	cliCtx = handlers.WithSelectedAuthIDCallback(cliCtx, func(authID string) {
		679	authID = strings.TrimSpace(authID)
		680	if authID == "" \|\| h == nil \|\| h.AuthManager == nil {
		681	return
		682	}
		683	lastAttemptedAuthID = authID
		684	selectedAuthObserved = true
		685	pinnedAuthAttempted = pinnedAuthAttempted \|\| (pinnedAuthID != "" && authID == pinnedAuthID)
		686	selectedAuth, ok := sessionAuthByID(authID)
		687	if !ok \|\| selectedAuth == nil {
		688	return
		689	}
		690	attemptedUpstreamMode = upstreamModeForAuth(selectedAuth)
		691	})
		692	if pinnedAuthID != "" && !routeOverridesModelResolution {
619	cliCtx = handlers.WithPinnedAuthID(cliCtx, pinnedAuthID)	693	cliCtx = handlers.WithPinnedAuthID(cliCtx, pinnedAuthID)
620	} else {
621	cliCtx = handlers.WithSelectedAuthIDCallback(cliCtx, func(authID string) {
622	authID = strings.TrimSpace(authID)
623	if authID == "" \|\| h == nil \|\| h.AuthManager == nil {
624	return
625	}
626	lastAttemptedAuthID = authID
627	selectedAuth, ok := sessionAuthByID(authID)
628	if !ok \|\| selectedAuth == nil {
629	return
630	}
631	if websocketUpstreamSupportsIncrementalInput(selectedAuth.Attributes, selectedAuth.Metadata) {
632	rememberPinnedAuth(authID, modelName)
633	}
634	})
635	}	694	}
636	dataChan, _, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, requestJSON, "")	695	dataChan, _, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, requestJSON, "")
637		696	if !selectedAuthObserved {
638	completedOutput, completedResponseID, completedPendingToolCallIDs, forwardErrMsg, errForward := h.forwardResponsesWebsocket(c, writer, cliCancel, dataChan, errChan, wsTimelineLog, passthroughSessionID)	697	// Plugin/alternate routes bypass auth selection. Keep canonical HTTP-mode
		698	// state instead of inheriting the previous pinned websocket mode.
		699	attemptedUpstreamMode = responsesWebsocketUpstreamModeHTTP
		700	}
		701	// A connection-scoped continuation cannot rotate credentials in place. Suppress
		702	// credential errors and make the client replay the full turn on a new socket.
		703	replayPinnedAuthFailure := func(errMsg *interfaces.ErrorMessage) bool {
		704	return nativeWebsocketPassthrough && requestRequiresCurrentUpstreamWebsocket && pinnedAuthAttempted &&
		705	shouldReplayResponsesWebsocketPinnedAuthFailure(errMsg)
		706	}
		707
		708	completedOutput, completedResponseID, completedPendingToolCallIDs, forwardErrMsg, errForward := h.forwardResponsesWebsocket(
		709	c,
		710	writer,
		711	cliCancel,
		712	dataChan,
		713	errChan,
		714	wsTimelineLog,
		715	passthroughSessionID,
		716	responsesWebsocketForwardOptions{
		717	toolCacheTurn: toolCacheTurn,
		718	suppressError: replayPinnedAuthFailure,
		719	},
		720	)
639	if errForward != nil {	721	if errForward != nil {
640	wsTerminateErr = errForward	722	wsTerminateErr = errForward
641	log.Warnf("responses websocket: forward failed id=%s error=%v", passthroughSessionID, errForward)	723	if !errors.Is(errForward, websocket.ErrCloseSent) {
		724	log.Warnf("responses websocket: forward failed id=%s error=%v", passthroughSessionID, errForward)
		725	}
642	return	726	return
643	}	727	}
644	if forwardErrMsg == nil && !useUpstreamWebsocketPassthrough && lastAttemptedAuthID != "" {	728	if forwardErrMsg != nil {
645	if selectedAuth, ok := sessionAuthByID(lastAttemptedAuthID); ok && selectedAuth != nil {	729	lastRequest = previousLastRequest
646	if websocketUpstreamSupportsIncrementalInput(selectedAuth.Attributes, selectedAuth.Metadata) {	730	lastResponseOutput = previousLastResponseOutput
647	rememberPinnedAuth(lastAttemptedAuthID, modelName)	731	lastResponseID = previousLastResponseID
648	} else if pinnedAuthID != "" {	732	lastResponsePendingToolCallIDs = previousLastResponsePendingToolCallIDs
649	if pinnedAuth, ok := sessionAuthByID(pinnedAuthID); ok && pinnedAuth != nil && websocketUpstreamSupportsIncrementalInput(pinnedAuth.Attributes, pinnedAuth.Metadata) {	733	if pinnedAuthAttempted && shouldReleaseResponsesWebsocketPinnedAuth(forwardErrMsg) {
650	rememberPinnedAuth(lastAttemptedAuthID, modelName)	734	forgetPinnedAuth()
651	}
652	}
653	}	735	}
654	}	736	if replayPinnedAuthFailure(forwardErrMsg) {
655	if shouldReleaseResponsesWebsocketPinnedAuth(forwardErrMsg) {	737	replayErr := responsesWebsocketHTTPReplayRequiredError()
656	forgetPinnedAuth()	738	wsTerminateErr = replayErr
657	forceTranscriptReplayNextRequest = true	739	matched, errClose := writer.closeForUpstreamError(replayErr)
658	if useUpstreamWebsocketPassthrough {	740	if !matched {
659	passthroughModelName = ""	741	_ = conn.Close()
660	} else {	742	} else if errClose != nil && !errors.Is(errClose, websocket.ErrCloseSent) {
661	lastRequest = previousLastRequest	743	log.Debugf("responses websocket: credential replay close failed id=%s error=%v", passthroughSessionID, errClose)
662	lastResponseOutput = previousLastResponseOutput	744	}
663	lastResponseID = previousLastResponseID	745	return
664	lastResponsePendingToolCallIDs = previousLastResponsePendingToolCallIDs
665	}	746	}
666	continue	747	continue
667	}	748	}
668	if !useUpstreamWebsocketPassthrough {	749
		750	toolCacheTurn.commit()
		751	upstreamMode = attemptedUpstreamMode
		752	if upstreamMode == responsesWebsocketUpstreamModeWS {
		753	upstreamWebsocketAuthID = lastAttemptedAuthID
		754	if lastAttemptedAuthID != "" {
		755	rememberPinnedAuth(lastAttemptedAuthID, modelName)
		756	}
		757	passthroughModelName = modelName
		758	lastRequest = nil
		759	lastResponseOutput = []byte("[]")
		760	lastResponseID = ""
		761	lastResponsePendingToolCallIDs = nil
		762	} else {
		763	upstreamWebsocketAuthID = ""
669	lastResponseOutput = completedOutput	764	lastResponseOutput = completedOutput
670	lastResponseID = strings.TrimSpace(completedResponseID)	765	lastResponseID = strings.TrimSpace(completedResponseID)
671	lastResponsePendingToolCallIDs = append([]string(nil), completedPendingToolCallIDs...)	766	lastResponsePendingToolCallIDs = append([]string(nil), completedPendingToolCallIDs...)
@@ -673,6 +768,20 @@ func (h OpenAIResponsesAPIHandler) ResponsesWebsocket(c gin.Context) { · 1 unmodified line
673	}	768	}
674	}	769	}
675		770
		771	func responsesWebsocketHTTPReplayRequiredError() error {
		772	return cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
		773	}
		774
		775	func responsesWebsocketRequestRequiresCurrentUpstream(payload []byte) bool {
		776	return strings.TrimSpace(gjson.GetBytes(payload, "previous_response_id").String()) != "" \|\|
		777	strings.TrimSpace(gjson.GetBytes(payload, "type").String()) == wsRequestTypeAppend
		778	}
		779
		780	func responsesWebsocketNativePassthroughAllowed(upstreamMode string, useUpstreamWebsocket bool, pinnedAuthID string, upstreamAuthID string) bool {
		781	return upstreamMode == responsesWebsocketUpstreamModeWS && useUpstreamWebsocket &&
		782	strings.TrimSpace(pinnedAuthID) != "" && strings.TrimSpace(pinnedAuthID) == strings.TrimSpace(upstreamAuthID)
		783	}
		784
676	func websocketClientAddress(c *gin.Context) string {	785	func websocketClientAddress(c *gin.Context) string {
677	if c == nil \|\| c.Request == nil {	786	if c == nil \|\| c.Request == nil {
678	return ""	787	return ""
@@ -694,6 +803,15 @@ func websocketUpgradeHeaders(req *http.Request) http.Header { · 15 unmodified lines
694	return headers	803	return headers
695	}	804	}
696		805
		806	func responsesWebsocketPreviousResponseNotFoundError() *interfaces.ErrorMessage {
		807	return &interfaces.ErrorMessage{
		808	StatusCode: http.StatusConflict,
		809	Error: errors.New(
		810	`{"error":{"message":"Previous response is not available on this websocket; resend the full conversation input without previous_response_id","type":"invalid_request_error","code":"previous_response_not_found","param":"previous_response_id"}}`,
		811	),
		812	}
		813	}
		814
697	func normalizeResponsesWebsocketRequest(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte) ([]byte, []byte, *interfaces.ErrorMessage) {	815	func normalizeResponsesWebsocketRequest(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte) ([]byte, []byte, *interfaces.ErrorMessage) {
698	return normalizeResponsesWebsocketRequestWithMode(rawJSON, lastRequest, lastResponseOutput, true, true)	816	return normalizeResponsesWebsocketRequestWithMode(rawJSON, lastRequest, lastResponseOutput, true, true)
699	}	817	}
@@ -1570,6 +1688,11 @@ func normalizeJSONArrayRaw(raw []byte) string { · 870 unmodified lines
1570	return "[]"	1688	return "[]"
1571	}	1689	}
1572		1690
		1691	type responsesWebsocketForwardOptions struct {
		1692	toolCacheTurn *responsesWebsocketToolCacheTurn
		1693	suppressError func(*interfaces.ErrorMessage) bool
		1694	}
		1695
1573	func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket(	1696	func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket(
1574	c *gin.Context,	1697	c *gin.Context,
1575	writer *responsesWebsocketWriter,	1698	writer *responsesWebsocketWriter,
@@ -1578,7 +1701,13 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( · 2 unmodified lines
1578	errs <-chan *interfaces.ErrorMessage,	1701	errs <-chan *interfaces.ErrorMessage,
1579	wsTimelineLog websocketTimelineAppender,	1702	wsTimelineLog websocketTimelineAppender,
1580	sessionID string,	1703	sessionID string,
		1704	options ...responsesWebsocketForwardOptions,
1581	) ([]byte, string, []string, *interfaces.ErrorMessage, error) {	1705	) ([]byte, string, []string, *interfaces.ErrorMessage, error) {
		1706	var opts responsesWebsocketForwardOptions
		1707	if len(options) > 0 {
		1708	opts = options[0]
		1709	}
		1710	toolCacheTurn := opts.toolCacheTurn
1582	completed := false	1711	completed := false
1583	completedOutput := []byte("[]")	1712	completedOutput := []byte("[]")
1584	completedResponseID := ""	1713	completedResponseID := ""
@@ -1602,6 +1731,10 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( · 17 unmodified lines
1602	}	1731	}
1603	if errMsg != nil {	1732	if errMsg != nil {
1604	h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), errMsg)	1733	h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), errMsg)
		1734	if opts.suppressError != nil && opts.suppressError(errMsg) {
		1735	cancel(errMsg.Error)
		1736	return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, nil
		1737	}
1605	markAPIResponseTimestamp(c)	1738	markAPIResponseTimestamp(c)
1606	if matched, errClose := writer.closeForUpstreamError(errMsg.Error); matched {	1739	if matched, errClose := writer.closeForUpstreamError(errMsg.Error); matched {
1607	cancel(errMsg.Error)	1740	cancel(errMsg.Error)
@@ -1676,7 +1809,11 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( · 68 unmodified lines
1676	if isResponsesWebsocketCompletionEvent(eventType) {	1809	if isResponsesWebsocketCompletionEvent(eventType) {
1677	payloads[i] = restoreResponsesWebsocketCompletionOutput(payloads[i], outputItemsByIndex, outputItemsFallback)	1810	payloads[i] = restoreResponsesWebsocketCompletionOutput(payloads[i], outputItemsByIndex, outputItemsFallback)
1678	}	1811	}
1679	recordResponsesWebsocketToolCallsFromPayload(downstreamSessionKey, payloads[i])	1812	if toolCacheTurn != nil {
		1813	toolCacheTurn.recordResponse(payloads[i])
		1814	} else {
		1815	recordResponsesWebsocketToolCallsFromPayload(downstreamSessionKey, payloads[i])
		1816	}
1680	recordPendingToolCallIDsFromPayload(pendingToolCallIDs, payloads[i])	1817	recordPendingToolCallIDsFromPayload(pendingToolCallIDs, payloads[i])
1681	var payloadErrMsg *interfaces.ErrorMessage	1818	var payloadErrMsg *interfaces.ErrorMessage
1682	if eventType == wsEventTypeError {	1819	if eventType == wsEventTypeError {
@@ -1684,6 +1821,10 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( · 1 unmodified line
1684	if h != nil {	1821	if h != nil {
1685	h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), payloadErrMsg)	1822	h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), payloadErrMsg)
1686	}	1823	}
		1824	if opts.suppressError != nil && opts.suppressError(payloadErrMsg) {
		1825	cancel(payloadErrMsg.Error)
		1826	return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), payloadErrMsg, nil
		1827	}
1687	} else if isResponsesWebsocketCompletionEvent(eventType) {	1828	} else if isResponsesWebsocketCompletionEvent(eventType) {
1688	completed = true	1829	completed = true
1689	completedOutput = responseCompletedOutputFromPayload(payloads[i], outputItemsByIndex, outputItemsFallback)	1830	completedOutput = responseCompletedOutputFromPayload(payloads[i], outputItemsByIndex, outputItemsFallback)
@@ -1716,9 +1857,9 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( · 26 unmodified lines
1716	}	1857	}
1717	}	1858	}
1718		1859
1719	func shouldReleaseResponsesWebsocketPinnedAuth(errMsg *interfaces.ErrorMessage) bool {	1860	func responsesWebsocketErrorStatus(errMsg *interfaces.ErrorMessage) int {
1720	if errMsg == nil {	1861	if errMsg == nil {
1721	return false	1862	return 0
1722	}	1863	}
1723	status := errMsg.StatusCode	1864	status := errMsg.StatusCode
1724	if status <= 0 && errMsg.Error != nil {	1865	if status <= 0 && errMsg.Error != nil {
@@ -1726,7 +1867,23 @@ func shouldReleaseResponsesWebsocketPinnedAuth(errMsg *interfaces.ErrorMessage) · 1 unmodified line
1726	status = se.StatusCode()	1867	status = se.StatusCode()
1727	}	1868	}
1728	}	1869	}
1729	switch status {	1870	return status
		1871	}
		1872
		1873	func shouldReplayResponsesWebsocketPinnedAuthFailure(errMsg *interfaces.ErrorMessage) bool {
		1874	switch responsesWebsocketErrorStatus(errMsg) {
		1875	case http.StatusUnauthorized, http.StatusTooManyRequests:
		1876	return true
		1877	default:
		1878	return false
		1879	}
		1880	}
		1881
		1882	func shouldReleaseResponsesWebsocketPinnedAuth(errMsg *interfaces.ErrorMessage) bool {
		1883	if errMsg == nil {
		1884	return false
		1885	}
		1886	switch responsesWebsocketErrorStatus(errMsg) {
1730	case http.StatusUnauthorized,	1887	case http.StatusUnauthorized,
1731	http.StatusPaymentRequired,	1888	http.StatusPaymentRequired,
1732	http.StatusForbidden,	1889	http.StatusForbidden,

sdk/api/handlers/openai/openai_responses_websocket_test.go+1061−252

@@ -9,6 +9,7 @@ import ( · 8 unmodified lines
9	"maps"	9	"maps"
10	"net/http"	10	"net/http"
11	"net/http/httptest"	11	"net/http/httptest"
		12	"strconv"
12	"strings"	13	"strings"
13	"sync"	14	"sync"
14	"sync/atomic"	15	"sync/atomic"
@@ -27,6 +28,7 @@ import ( · 12 unmodified lines
27	"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry"	28	"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry"
28	coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"	29	coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
29	sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"	30	sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
		31	"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
30	"github.com/tidwall/gjson"	32	"github.com/tidwall/gjson"
31	)	33	)
32		34
@@ -160,6 +162,51 @@ func TestResponsesWebsocketHomeSelectedAuthCallbackPinsAndReusesFirstSelection(t · 127 unmodified lines
160	}	162	}
161	}	163	}
162		164
		165	func TestWebsocketReplayCloseRequiresTypedSignal(t *testing.T) {
		166	matched, payload := websocketClosePayloadForUpstreamError(responsesWebsocketHTTPReplayRequiredError())
		167	if !matched \|\| len(payload) == 0 {
		168	t.Fatalf("typed replay signal matched=%t payload_len=%d, want close payload", matched, len(payload))
		169	}
		170	spoofed := websocketPinnedFailoverStatusError{
		171	status: http.StatusUpgradeRequired,
		172	msg: `{"error":{"code":"upstream_http_replay_required"}}`,
		173	}
		174	if matched, _ := websocketClosePayloadForUpstreamError(spoofed); matched {
		175	t.Fatal("untyped upstream error spoofed replay close")
		176	}
		177	}
		178
		179	func TestResponsesWebsocketRequestRequiresCurrentUpstream(t *testing.T) {
		180	cases := []struct {
		181	name string
		182	payload string
		183	want bool
		184	}{
		185	{name: "incremental create", payload: `{"type":"response.create","previous_response_id":"resp-1","input":[]}`, want: true},
		186	{name: "append", payload: `{"type":"response.append","input":[]}`, want: true},
		187	{name: "full create", payload: `{"type":"response.create","input":[]}`, want: false},
		188	}
		189	for _, tc := range cases {
		190	t.Run(tc.name, func(t *testing.T) {
		191	if got := responsesWebsocketRequestRequiresCurrentUpstream([]byte(tc.payload)); got != tc.want {
		192	t.Fatalf("responsesWebsocketRequestRequiresCurrentUpstream() = %t, want %t", got, tc.want)
		193	}
		194	})
		195	}
		196	}
		197
		198	func TestResponsesWebsocketNativePassthroughRequiresImmediatelyPreviousAuth(t *testing.T) {
		199	if !responsesWebsocketNativePassthroughAllowed(responsesWebsocketUpstreamModeWS, true, "auth-a", "auth-a") {
		200	t.Fatal("matching immediate websocket auth did not allow native passthrough")
		201	}
		202	if responsesWebsocketNativePassthroughAllowed(responsesWebsocketUpstreamModeWS, true, "auth-a", "auth-b") {
		203	t.Fatal("restored auth from an older provider session allowed native passthrough")
		204	}
		205	if responsesWebsocketNativePassthroughAllowed(responsesWebsocketUpstreamModeHTTP, true, "auth-a", "auth-a") {
		206	t.Fatal("HTTP mode allowed native websocket passthrough")
		207	}
		208	}
		209
163	func TestWriteWebsocketCloseForUpstreamErrorMirrorsMessageTooBig(t *testing.T) {	210	func TestWriteWebsocketCloseForUpstreamErrorMirrorsMessageTooBig(t *testing.T) {
164	tests := []struct {	211	tests := []struct {
165	name string	212	name string
@@ -439,6 +486,22 @@ type websocketProviderCaptureExecutor struct { · 273 unmodified lines
439	websocketCaptureExecutor	486	websocketCaptureExecutor
440	}	487	}
441		488
		489	type websocketProviderRouteHost struct{}
		490
		491	func (*websocketProviderRouteHost) HasModelRouters() bool { return true }
		492
		493	func (*websocketProviderRouteHost) RouteModel(_ context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
		494	if !gjson.GetBytes(req.Body, "route_to_claude").Bool() {
		495	return pluginapi.ModelRouteResponse{}, false
		496	}
		497	return pluginapi.ModelRouteResponse{
		498	Handled: true,
		499	TargetKind: pluginapi.ModelRouteTargetProvider,
		500	Target: "claude",
		501	TargetModel: "claude-provider-route-target",
		502	}, true
		503	}
		504
442	type websocketCompactionCaptureExecutor struct {	505	type websocketCompactionCaptureExecutor struct {
443	mu sync.Mutex	506	mu sync.Mutex
444	streamPayloads [][]byte	507	streamPayloads [][]byte
@@ -481,10 +544,11 @@ type websocketAuthCaptureExecutor struct { · 36 unmodified lines
481	}	544	}
482		545
483	type websocketPinnedFailoverExecutor struct {	546	type websocketPinnedFailoverExecutor struct {
484	mu sync.Mutex	547	mu sync.Mutex
485	authIDs []string	548	failStatus int
486	calls map[string]int	549	authIDs []string
487	payloads map[string][][]byte	550	calls map[string]int
		551	payloads map[string][][]byte
488	}	552	}
489		553
490	type websocketBootstrapFallbackExecutor struct {	554	type websocketBootstrapFallbackExecutor struct {
@@ -494,12 +558,21 @@ type websocketBootstrapFallbackExecutor struct { · 3 unmodified lines
494	}	558	}
495		559
496	type websocketDirectCaptureExecutor struct {	560	type websocketDirectCaptureExecutor struct {
		561	mu sync.Mutex
		562	provider string
		563	failStatus int
		564	authIDs []string
		565	models []string
		566	payloads [][]byte
		567	requiredUpstreamWebsocket []bool
		568	done chan struct{}
		569	doneOnce sync.Once
		570	}
		571
		572	type websocketCanonicalRollbackExecutor struct {
497	mu sync.Mutex	573	mu sync.Mutex
498	provider string
499	authIDs []string
500	payloads [][]byte	574	payloads [][]byte
501	done chan struct{}	575	calls int
502	doneOnce sync.Once
503	}	576	}
504		577
505	type websocketPinnedFailoverStatusError struct {	578	type websocketPinnedFailoverStatusError struct {
@@ -534,7 +607,7 @@ func (e *websocketBootstrapFallbackExecutor) ExecuteStream(_ context.Context, au · 28 unmodified lines
534	chunks := make(chan coreexecutor.StreamChunk, 1)	607	chunks := make(chan coreexecutor.StreamChunk, 1)
535	if authID == "auth-ws" {	608	if authID == "auth-ws" {
536	chunks <- coreexecutor.StreamChunk{Err: websocketPinnedFailoverStatusError{	609	chunks <- coreexecutor.StreamChunk{Err: websocketPinnedFailoverStatusError{
537	status: http.StatusServiceUnavailable,	610	status: http.StatusUpgradeRequired,
538	msg: `{"error":{"message":"websocket bootstrap failed","type":"server_error","code":"ws_failed"}}`,	611	msg: `{"error":{"message":"websocket bootstrap failed","type":"server_error","code":"ws_failed"}}`,
539	}}	612	}}
540	close(chunks)	613	close(chunks)
@@ -586,18 +659,29 @@ func (e websocketDirectCaptureExecutor) Execute(context.Context, coreauth.Auth · 45 unmodified lines
586	return coreexecutor.Response{}, errors.New("not implemented")	659	return coreexecutor.Response{}, errors.New("not implemented")
587	}	660	}
588		661
589	func (e websocketDirectCaptureExecutor) ExecuteStream(_ context.Context, auth coreauth.Auth, req coreexecutor.Request, _ coreexecutor.Options) (*coreexecutor.StreamResult, error) {	662	func (e websocketDirectCaptureExecutor) ExecuteStream(ctx context.Context, auth coreauth.Auth, req coreexecutor.Request, _ coreexecutor.Options) (*coreexecutor.StreamResult, error) {
590	authID := ""	663	authID := ""
591	if auth != nil {	664	if auth != nil {
592	authID = auth.ID	665	authID = auth.ID
593	}	666	}
594	e.mu.Lock()	667	e.mu.Lock()
595	e.authIDs = append(e.authIDs, authID)	668	e.authIDs = append(e.authIDs, authID)
		669	e.models = append(e.models, req.Model)
596	e.payloads = append(e.payloads, bytes.Clone(req.Payload))	670	e.payloads = append(e.payloads, bytes.Clone(req.Payload))
		671	e.requiredUpstreamWebsocket = append(e.requiredUpstreamWebsocket, coreexecutor.RequiredUpstreamWebsocket(ctx))
597	count := len(e.payloads)	672	count := len(e.payloads)
		673	failStatus := e.failStatus
598	e.mu.Unlock()	674	e.mu.Unlock()
599		675
600	chunks := make(chan coreexecutor.StreamChunk, 1)	676	chunks := make(chan coreexecutor.StreamChunk, 1)
		677	if failStatus > 0 {
		678	chunks <- coreexecutor.StreamChunk{Err: websocketPinnedFailoverStatusError{
		679	status: failStatus,
		680	msg: `{"error":{"message":"routed provider failed","type":"authentication_error","code":"invalid_api_key"}}`,
		681	}}
		682	close(chunks)
		683	return &coreexecutor.StreamResult{Chunks: chunks}, nil
		684	}
601	responseID := fmt.Sprintf("resp-%d", count)	685	responseID := fmt.Sprintf("resp-%d", count)
602	chunks <- coreexecutor.StreamChunk{Payload: []byte(fmt.Sprintf(`{"type":"response.completed","response":{"id":%q,"output":[{"type":"message","id":"out-%d"}]}}`, responseID, count))}	686	chunks <- coreexecutor.StreamChunk{Payload: []byte(fmt.Sprintf(`{"type":"response.completed","response":{"id":%q,"output":[{"type":"message","id":"out-%d"}]}}`, responseID, count))}
603	close(chunks)	687	close(chunks)
@@ -637,6 +721,67 @@ func (e *websocketDirectCaptureExecutor) AuthIDs() []string { · 33 unmodified lines
637	return append([]string(nil), e.authIDs...)	721	return append([]string(nil), e.authIDs...)
638	}	722	}
639		723
		724	func (e *websocketDirectCaptureExecutor) Models() []string {
		725	e.mu.Lock()
		726	defer e.mu.Unlock()
		727	return append([]string(nil), e.models...)
		728	}
		729
		730	func (e *websocketDirectCaptureExecutor) RequiredUpstreamWebsocketFlags() []bool {
		731	e.mu.Lock()
		732	defer e.mu.Unlock()
		733	return append([]bool(nil), e.requiredUpstreamWebsocket...)
		734	}
		735
		736	func (e *websocketCanonicalRollbackExecutor) Identifier() string { return "xai" }
		737
		738	func (e websocketCanonicalRollbackExecutor) Execute(context.Context, coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
		739	return coreexecutor.Response{}, errors.New("not implemented")
		740	}
		741
		742	func (e websocketCanonicalRollbackExecutor) ExecuteStream(_ context.Context, _ coreauth.Auth, req coreexecutor.Request, _ coreexecutor.Options) (*coreexecutor.StreamResult, error) {
		743	e.mu.Lock()
		744	e.calls++
		745	call := e.calls
		746	e.payloads = append(e.payloads, bytes.Clone(req.Payload))
		747	e.mu.Unlock()
		748
		749	chunks := make(chan coreexecutor.StreamChunk, 1)
		750	if call == 2 {
		751	chunks <- coreexecutor.StreamChunk{Err: websocketPinnedFailoverStatusError{
		752	status: http.StatusBadRequest,
		753	msg: `{"error":{"message":"bad turn","type":"invalid_request_error","code":"invalid_request"}}`,
		754	}}
		755	close(chunks)
		756	return &coreexecutor.StreamResult{Chunks: chunks}, nil
		757	}
		758	chunks <- coreexecutor.StreamChunk{Payload: []byte(fmt.Sprintf(`{"type":"response.completed","response":{"id":"resp-%d","output":[{"type":"message","id":"out-%d"}]}}`, call, call))}
		759	close(chunks)
		760	return &coreexecutor.StreamResult{Chunks: chunks}, nil
		761	}
		762
		763	func (e websocketCanonicalRollbackExecutor) Refresh(_ context.Context, auth coreauth.Auth) (*coreauth.Auth, error) {
		764	return auth, nil
		765	}
		766
		767	func (e websocketCanonicalRollbackExecutor) CountTokens(context.Context, coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
		768	return coreexecutor.Response{}, errors.New("not implemented")
		769	}
		770
		771	func (e websocketCanonicalRollbackExecutor) HttpRequest(context.Context, coreauth.Auth, http.Request) (http.Response, error) {
		772	return nil, errors.New("not implemented")
		773	}
		774
		775	func (e *websocketCanonicalRollbackExecutor) Payloads() [][]byte {
		776	e.mu.Lock()
		777	defer e.mu.Unlock()
		778	out := make([][]byte, len(e.payloads))
		779	for i := range e.payloads {
		780	out[i] = bytes.Clone(e.payloads[i])
		781	}
		782	return out
		783	}
		784
640	type websocketUpstreamDisconnectExecutor struct {	785	type websocketUpstreamDisconnectExecutor struct {
641	mu sync.Mutex	786	mu sync.Mutex
642	provider string	787	provider string
@@ -753,7 +898,7 @@ func (e *websocketAuthCaptureExecutor) AuthIDs() []string { · 110 unmodified lines
753	return append([]string(nil), e.authIDs...)	898	return append([]string(nil), e.authIDs...)
754	}	899	}
755		900
756	func (e *websocketPinnedFailoverExecutor) Identifier() string { return "test-provider" }	901	func (e *websocketPinnedFailoverExecutor) Identifier() string { return "xai" }
757		902
758	func (e websocketPinnedFailoverExecutor) Execute(context.Context, coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {	903	func (e websocketPinnedFailoverExecutor) Execute(context.Context, coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
759	return coreexecutor.Response{}, errors.New("not implemented")	904	return coreexecutor.Response{}, errors.New("not implemented")
@@ -781,8 +926,8 @@ func (e *websocketPinnedFailoverExecutor) ExecuteStream(_ context.Context, auth · 21 unmodified lines
781	if authID == "auth-a" && call == 2 {	926	if authID == "auth-a" && call == 2 {
782	chunks := make(chan coreexecutor.StreamChunk, 1)	927	chunks := make(chan coreexecutor.StreamChunk, 1)
783	chunks <- coreexecutor.StreamChunk{Err: websocketPinnedFailoverStatusError{	928	chunks <- coreexecutor.StreamChunk{Err: websocketPinnedFailoverStatusError{
784	status: http.StatusTooManyRequests,	929	status: e.failStatus,
785	msg: `{"error":{"message":"quota exhausted","type":"rate_limit_error","code":"rate_limit_exceeded"}}`,	930	msg: fmt.Sprintf(`{"error":{"message":"credential failed","status":%d}}`, e.failStatus),
786	}}	931	}}
787	close(chunks)	932	close(chunks)
788	return &coreexecutor.StreamResult{Chunks: chunks}, nil	933	return &coreexecutor.StreamResult{Chunks: chunks}, nil
@@ -1598,6 +1743,44 @@ func TestRepairResponsesWebsocketToolCallsInsertsCachedOutput(t *testing.T) { · 809 unmodified lines
1598	}	1743	}
1599	}	1744	}
1600		1745
		1746	func TestResponsesWebsocketToolCacheTurnCommitsOnlyOnSuccess(t *testing.T) {
		1747	const sessionKey = "tool-cache-turn-commit-session"
		1748	defer defaultWebsocketToolOutputCache.deleteSession(sessionKey)
		1749	defer defaultWebsocketToolCallCache.deleteSession(sessionKey)
		1750
		1751	turn := newResponsesWebsocketToolCacheTurn(sessionKey)
		1752	turn.recordRequest([]byte(`{"input":[{"type":"function_call_output","id":"fco-1","call_id":"call-1","output":"cached result"}]}`))
		1753	beforeCommit := repairResponsesWebsocketToolCallsWithoutRecording(sessionKey, []byte(`{"input":[{"type":"function_call","id":"fc-next","call_id":"call-1","name":"lookup","arguments":"{}"}]}`))
		1754	if gjson.GetBytes(beforeCommit, "input.#").Int() != 0 {
		1755	t.Fatalf("uncommitted turn populated global cache: %s", beforeCommit)
		1756	}
		1757
		1758	turn.commit()
		1759	afterCommit := repairResponsesWebsocketToolCallsWithoutRecording(sessionKey, []byte(`{"input":[{"type":"function_call","id":"fc-next","call_id":"call-1","name":"lookup","arguments":"{}"}]}`))
		1760	input := gjson.GetBytes(afterCommit, "input").Array()
		1761	if len(input) != 2 \|\| input[1].Get("output").String() != "cached result" {
		1762	t.Fatalf("committed turn was not available to tool repair: %s", afterCommit)
		1763	}
		1764	}
		1765
		1766	func TestResponsesWebsocketToolCacheRetainPreventsOverlappingReleaseDeletion(t *testing.T) {
		1767	const sessionKey = "tool-cache-overlapping-retain-session"
		1768	retainResponsesWebsocketToolCaches(sessionKey)
		1769	retainResponsesWebsocketToolCaches(sessionKey)
		1770	turn := newResponsesWebsocketToolCacheTurn(sessionKey)
		1771	turn.recordRequest([]byte(`{"input":[{"type":"function_call_output","id":"fco-1","call_id":"call-1","output":"kept"}]}`))
		1772	turn.commit()
		1773
		1774	releaseResponsesWebsocketToolCaches(sessionKey)
		1775	if _, ok := defaultWebsocketToolOutputCache.get(sessionKey, "call-1"); !ok {
		1776	t.Fatal("first overlapping release deleted active session cache")
		1777	}
		1778	releaseResponsesWebsocketToolCaches(sessionKey)
		1779	if _, ok := defaultWebsocketToolOutputCache.get(sessionKey, "call-1"); ok {
		1780	t.Fatal("final release did not delete session cache")
		1781	}
		1782	}
		1783
1601	func TestRepairResponsesWebsocketToolCallsDropsOrphanFunctionCall(t *testing.T) {	1784	func TestRepairResponsesWebsocketToolCallsDropsOrphanFunctionCall(t *testing.T) {
1602	cache := newWebsocketToolOutputCache(time.Minute, 10)	1785	cache := newWebsocketToolOutputCache(time.Minute, 10)
1603	sessionKey := "session-1"	1786	sessionKey := "session-1"
@@ -2419,7 +2602,7 @@ func TestResponsesWebsocketCodexWebsocketPassthroughPassesCompactedRequestWithou · 815 unmodified lines
2419	}	2602	}
2420	}	2603	}
2421		2604
2422	func TestResponsesWebsocketXAIWebsocketPassthroughCarriesPreviousResponseID(t *testing.T) {	2605	func TestResponsesWebsocketXAIWebsocketPassthroughKeepsNativeIncrementalRequest(t *testing.T) {
2423	gin.SetMode(gin.TestMode)	2606	gin.SetMode(gin.TestMode)
2424		2607
2425	modelName := "xai-websocket-passthrough-model"	2608	modelName := "xai-websocket-passthrough-model"
@@ -2483,165 +2666,660 @@ func TestResponsesWebsocketXAIWebsocketPassthroughCarriesPreviousResponseID(t *t · 57 unmodified lines
2483	}	2666	}
2484	secondPayload := payloads[1]	2667	secondPayload := payloads[1]
2485	if got := gjson.GetBytes(secondPayload, "type").String(); got != wsRequestTypeCreate {	2668	if got := gjson.GetBytes(secondPayload, "type").String(); got != wsRequestTypeCreate {
2486	t.Fatalf("second xai passthrough type = %s, want %s: %s", got, wsRequestTypeCreate, secondPayload)	2669	t.Fatalf("incremental xai payload type = %q, want %q: %s", got, wsRequestTypeCreate, secondPayload)
2487	}	2670	}
2488	if got := gjson.GetBytes(secondPayload, "model").String(); got != modelName {	2671	if got := gjson.GetBytes(secondPayload, "model").String(); got != modelName {
2489	t.Fatalf("second xai payload model = %s, want %s", got, modelName)	2672	t.Fatalf("second xai payload model = %s, want %s", got, modelName)
2490	}	2673	}
2491	if got := gjson.GetBytes(secondPayload, "previous_response_id").String(); got != "resp-1" {	2674	if got := gjson.GetBytes(secondPayload, "previous_response_id").String(); got != "resp-1" {
2492	t.Fatalf("second xai previous_response_id = %s, want resp-1: %s", got, secondPayload)	2675	t.Fatalf("second xai previous_response_id = %q, want resp-1: %s", got, secondPayload)
2493	}	2676	}
2494	input := gjson.GetBytes(secondPayload, "input").Array()	2677	input := gjson.GetBytes(secondPayload, "input").Array()
2495	if len(input) != 1 {	2678	if len(input) != 1 \|\| input[0].Get("id").String() != "msg-2" {
2496	t.Fatalf("second xai passthrough input len = %d, want 1: %s", len(input), secondPayload)	2679	t.Fatalf("second xai incremental input is not the client delta: %s", secondPayload)
2497	}
2498	if input[0].Get("id").String() != "msg-2" {
2499	t.Fatalf("second xai passthrough input must contain only the new turn: %s", secondPayload)
2500	}
2501	if bytes.Contains(secondPayload, []byte(`"id":"msg-1"`)) \|\| bytes.Contains(secondPayload, []byte(`"id":"out-1"`)) {
2502	t.Fatalf("second xai passthrough payload contains stale transcript state: %s", secondPayload)
2503	}	2680	}
2504	authIDs := executor.AuthIDs()	2681	authIDs := executor.AuthIDs()
2505	if len(authIDs) != 2 \|\| authIDs[0] != "auth-xai-ws" \|\| authIDs[1] != "auth-xai-ws" {	2682	if len(authIDs) != 2 \|\| authIDs[0] != "auth-xai-ws" \|\| authIDs[1] != "auth-xai-ws" {
2506	t.Fatalf("xai websocket auth IDs = %v, want [auth-xai-ws auth-xai-ws]", authIDs)	2683	t.Fatalf("xai websocket auth IDs = %v, want [auth-xai-ws auth-xai-ws]", authIDs)
2507	}	2684	}
		2685	if got := executor.RequiredUpstreamWebsocketFlags(); len(got) != 2 \|\| got[0] \|\| !got[1] {
		2686	t.Fatalf("required upstream websocket flags = %v, want [false true]", got)
		2687	}
2508	}	2688	}
2509		2689
2510	func TestResponsesWebsocketSwitchesPinnedAuthAcrossProviders(t *testing.T) {	2690	func TestResponsesWebsocketFullRequestCanRouteFromNativeWebsocketToBuiltInProvider(t *testing.T) {
2511	for _, testCase := range []struct {	2691	gin.SetMode(gin.TestMode)
2512	name string
2513	xaiWebsockets bool
2514	returnToDifferentXAIModel bool
2515	}{
2516	{name: "xai SSE", xaiWebsockets: false},
2517	{name: "xai websocket", xaiWebsockets: true},
2518	{name: "xai websocket different model", xaiWebsockets: true, returnToDifferentXAIModel: true},
2519	} {
2520	t.Run(testCase.name, func(t *testing.T) {
2521	gin.SetMode(gin.TestMode)
2522
2523	xaiModel := "xai-provider-switch-" + strings.ReplaceAll(testCase.name, " ", "-")
2524	returnXAIModel := xaiModel
2525	if testCase.returnToDifferentXAIModel {
2526	returnXAIModel += "-return"
2527	}
2528	codexModel := "codex-provider-switch-" + strings.ReplaceAll(testCase.name, " ", "-")
2529	xaiExecutor := &websocketDirectCaptureExecutor{provider: "xai"}
2530	codexExecutor := &websocketDirectCaptureExecutor{provider: "codex"}
2531
2532	xaiAuth := &coreauth.Auth{
2533	ID: "auth-" + xaiModel,
2534	Provider: "xai",
2535	Status: coreauth.StatusActive,
2536	}
2537	if testCase.xaiWebsockets {
2538	xaiAuth.Attributes = map[string]string{"websockets": "true"}
2539	}
2540	codexAuth := &coreauth.Auth{
2541	ID: "auth-" + codexModel,
2542	Provider: "codex",
2543	Status: coreauth.StatusActive,
2544	Attributes: map[string]string{"websockets": "true"},
2545	}
2546	selector := &orderedWebsocketSelector{order: []string{xaiAuth.ID, codexAuth.ID}}
2547	manager := coreauth.NewManager(nil, selector, nil)
2548	manager.RegisterExecutor(xaiExecutor)
2549	manager.RegisterExecutor(codexExecutor)
2550	if _, errRegister := manager.Register(context.Background(), xaiAuth); errRegister != nil {
2551	t.Fatalf("Register xAI auth: %v", errRegister)
2552	}
2553	if _, errRegister := manager.Register(context.Background(), codexAuth); errRegister != nil {
2554	t.Fatalf("Register Codex auth: %v", errRegister)
2555	}
2556
2557	registry.GetGlobalRegistry().RegisterClient(xaiAuth.ID, xaiAuth.Provider, []*registry.ModelInfo{{ID: xaiModel}})
2558	registry.GetGlobalRegistry().RegisterClient(codexAuth.ID, codexAuth.Provider, []*registry.ModelInfo{{ID: codexModel}})
2559	registeredAuthIDs := []string{xaiAuth.ID, codexAuth.ID}
2560	xaiAlternateAuthID := ""
2561	if testCase.xaiWebsockets {
2562	xaiAlternateAuth := &coreauth.Auth{
2563	ID: "auth-alternate-" + xaiModel,
2564	Provider: "xai",
2565	Status: coreauth.StatusActive,
2566	Attributes: map[string]string{"websockets": "true"},
2567	}
2568	xaiAlternateAuthID = xaiAlternateAuth.ID
2569	selector.order = append(selector.order, xaiAlternateAuth.ID)
2570	if _, errRegister := manager.Register(context.Background(), xaiAlternateAuth); errRegister != nil {
2571	t.Fatalf("Register alternate xAI auth: %v", errRegister)
2572	}
2573	alternateModels := []*registry.ModelInfo{{ID: xaiModel}}
2574	if testCase.returnToDifferentXAIModel {
2575	alternateModels = []*registry.ModelInfo{{ID: returnXAIModel}}
2576	}
2577	registry.GetGlobalRegistry().RegisterClient(xaiAlternateAuth.ID, xaiAlternateAuth.Provider, alternateModels)
2578	registeredAuthIDs = append(registeredAuthIDs, xaiAlternateAuth.ID)
2579	}
2580	t.Cleanup(func() {
2581	for _, authID := range registeredAuthIDs {
2582	registry.GetGlobalRegistry().UnregisterClient(authID)
2583	}
2584	})
2585		2692
2586	base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)	2693	const sourceModel = "codex-provider-route-source"
2587	h := NewOpenAIResponsesAPIHandler(base)	2694	const targetModel = "claude-provider-route-target"
2588	router := gin.New()	2695	codexExecutor := &websocketDirectCaptureExecutor{provider: "codex"}
2589	router.GET("/v1/responses/ws", h.ResponsesWebsocket)	2696	claudeExecutor := &websocketDirectCaptureExecutor{provider: "claude"}
		2697	manager := coreauth.NewManager(nil, nil, nil)
		2698	manager.RegisterExecutor(codexExecutor)
		2699	manager.RegisterExecutor(claudeExecutor)
		2700	codexAuth := &coreauth.Auth{
		2701	ID: "auth-codex-provider-route",
		2702	Provider: "codex",
		2703	Status: coreauth.StatusActive,
		2704	Attributes: map[string]string{"websockets": "true"},
		2705	}
		2706	claudeAuth := &coreauth.Auth{
		2707	ID: "auth-claude-provider-route",
		2708	Provider: "claude",
		2709	Status: coreauth.StatusActive,
		2710	}
		2711	for _, auth := range []*coreauth.Auth{codexAuth, claudeAuth} {
		2712	if _, err := manager.Register(context.Background(), auth); err != nil {
		2713	t.Fatalf("Register auth %s: %v", auth.ID, err)
		2714	}
		2715	}
		2716	registry.GetGlobalRegistry().RegisterClient(codexAuth.ID, codexAuth.Provider, []*registry.ModelInfo{{ID: sourceModel}})
		2717	registry.GetGlobalRegistry().RegisterClient(claudeAuth.ID, claudeAuth.Provider, []*registry.ModelInfo{{ID: targetModel}})
		2718	t.Cleanup(func() {
		2719	registry.GetGlobalRegistry().UnregisterClient(codexAuth.ID)
		2720	registry.GetGlobalRegistry().UnregisterClient(claudeAuth.ID)
		2721	})
2590		2722
2591	server := httptest.NewServer(router)	2723	base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
2592	defer server.Close()	2724	base.SetModelRouterHost(&websocketProviderRouteHost{})
		2725	h := NewOpenAIResponsesAPIHandler(base)
		2726	router := gin.New()
		2727	router.GET("/v1/responses/ws", h.ResponsesWebsocket)
		2728	server := httptest.NewServer(router)
		2729	defer server.Close()
2593		2730
2594	wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses/ws"	2731	wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses/ws"
2595	conn, _, errDial := websocket.DefaultDialer.Dial(wsURL, nil)	2732	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
2596	if errDial != nil {	2733	if err != nil {
2597	t.Fatalf("dial websocket: %v", errDial)	2734	t.Fatalf("dial websocket: %v", err)
2598	}	2735	}
2599	defer func() {	2736	defer func() { _ = conn.Close() }()
2600	if errClose := conn.Close(); errClose != nil {
2601	t.Errorf("close websocket: %v", errClose)
2602	}
2603	}()
2604		2737
2605	requests := []string{	2738	firstRequest := []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"}]}`, sourceModel))
2606	fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-xai-1"}]}`, xaiModel),	2739	if errWrite := conn.WriteMessage(websocket.TextMessage, firstRequest); errWrite != nil {
2607	fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-codex-1"}]}`, codexModel),	2740	t.Fatalf("write first websocket message: %v", errWrite)
2608	fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-xai-2"}]}`, returnXAIModel),	2741	}
2609	`{"type":"response.create","input":[{"type":"message","id":"msg-xai-3"}]}`,	2742	if _, _, errRead := conn.ReadMessage(); errRead != nil {
2610	}	2743	t.Fatalf("read first websocket response: %v", errRead)
2611	for index, request := range requests {	2744	}
2612	if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(request)); errWrite != nil {
2613	t.Fatalf("write websocket message %d: %v", index+1, errWrite)
2614	}
2615	_, payload, errRead := conn.ReadMessage()
2616	if errRead != nil {
2617	t.Fatalf("read websocket response %d: %v", index+1, errRead)
2618	}
2619	if got := gjson.GetBytes(payload, "type").String(); got != wsEventTypeCompleted {
2620	t.Fatalf("response %d type = %s, want %s: %s", index+1, got, wsEventTypeCompleted, payload)
2621	}
2622	}
2623		2745
2624	wantReturnAuthID := xaiAuth.ID	2746	routedRequest := []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"route_to_claude":true,"input":[{"type":"message","id":"msg-routed"}]}`, sourceModel))
2625	if testCase.returnToDifferentXAIModel {	2747	if errWrite := conn.WriteMessage(websocket.TextMessage, routedRequest); errWrite != nil {
2626	wantReturnAuthID = xaiAlternateAuthID	2748	t.Fatalf("write routed websocket message: %v", errWrite)
2627	}	2749	}
2628	if got := xaiExecutor.AuthIDs(); len(got) != 3 \|\| got[0] != xaiAuth.ID \|\| got[1] != wantReturnAuthID \|\| got[2] != wantReturnAuthID {	2750	_, response, errRead := conn.ReadMessage()
2629	t.Fatalf("xAI auth IDs = %v, want [%s %s %s]", got, xaiAuth.ID, wantReturnAuthID, wantReturnAuthID)	2751	if errRead != nil {
2630	}	2752	t.Fatalf("read routed websocket response: %v", errRead)
2631	if got := codexExecutor.AuthIDs(); len(got) != 1 \|\| got[0] != codexAuth.ID {	2753	}
2632	t.Fatalf("Codex auth IDs = %v, want [%s]", got, codexAuth.ID)	2754	if got := gjson.GetBytes(response, "type").String(); got != wsEventTypeCompleted {
2633	}	2755	t.Fatalf("routed response type = %q, want %q: %s", got, wsEventTypeCompleted, response)
2634	})	2756	}
		2757	if got := len(codexExecutor.Payloads()); got != 1 {
		2758	t.Fatalf("codex payload count = %d, want 1", got)
		2759	}
		2760	claudePayloads := claudeExecutor.Payloads()
		2761	if len(claudePayloads) != 1 {
		2762	t.Fatalf("claude payload count = %d, want 1", len(claudePayloads))
		2763	}
		2764	if got := claudeExecutor.Models(); len(got) != 1 \|\| got[0] != targetModel {
		2765	t.Fatalf("routed models = %v, want [%s]", got, targetModel)
2635	}	2766	}
2636	}	2767	}
2637		2768
2638	func TestResponsesWebsocketPinnedAuthMatchesModel(t *testing.T) {	2769	func TestResponsesWebsocketFailedProviderRoutePreservesNativeWebsocketPin(t *testing.T) {
2639	modelA := "xai-pinned-auth-model-a"	2770	gin.SetMode(gin.TestMode)
2640	modelB := "xai-pinned-auth-model-b"	2771
2641	auth := &coreauth.Auth{ID: "xai-pinned-auth", Provider: "xai", Status: coreauth.StatusActive}	2772	const sourceModel = "codex-provider-route-failure-source"
2642	otherAuthID := "xai-pinned-auth-other"	2773	const targetModel = "claude-provider-route-target"
2643	registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: modelA}})	2774	codexExecutor := &websocketDirectCaptureExecutor{provider: "codex"}
2644	registry.GetGlobalRegistry().RegisterClient(otherAuthID, auth.Provider, []*registry.ModelInfo{{ID: modelB}})	2775	claudeExecutor := &websocketDirectCaptureExecutor{provider: "claude", failStatus: http.StatusUnauthorized}
		2776	manager := coreauth.NewManager(nil, nil, nil)
		2777	manager.RegisterExecutor(codexExecutor)
		2778	manager.RegisterExecutor(claudeExecutor)
		2779	codexAuth := &coreauth.Auth{ID: "auth-codex-provider-route-failure", Provider: "codex", Status: coreauth.StatusActive, Attributes: map[string]string{"websockets": "true"}}
		2780	claudeAuth := &coreauth.Auth{ID: "auth-claude-provider-route-failure", Provider: "claude", Status: coreauth.StatusActive}
		2781	for _, auth := range []*coreauth.Auth{codexAuth, claudeAuth} {
		2782	if _, err := manager.Register(context.Background(), auth); err != nil {
		2783	t.Fatalf("Register auth %s: %v", auth.ID, err)
		2784	}
		2785	}
		2786	registry.GetGlobalRegistry().RegisterClient(codexAuth.ID, codexAuth.Provider, []*registry.ModelInfo{{ID: sourceModel}})
		2787	registry.GetGlobalRegistry().RegisterClient(claudeAuth.ID, claudeAuth.Provider, []*registry.ModelInfo{{ID: targetModel}})
		2788	t.Cleanup(func() {
		2789	registry.GetGlobalRegistry().UnregisterClient(codexAuth.ID)
		2790	registry.GetGlobalRegistry().UnregisterClient(claudeAuth.ID)
		2791	})
		2792
		2793	base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
		2794	base.SetModelRouterHost(&websocketProviderRouteHost{})
		2795	h := NewOpenAIResponsesAPIHandler(base)
		2796	router := gin.New()
		2797	router.GET("/v1/responses/ws", h.ResponsesWebsocket)
		2798	server := httptest.NewServer(router)
		2799	defer server.Close()
		2800
		2801	wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses/ws"
		2802	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
		2803	if err != nil {
		2804	t.Fatalf("dial websocket: %v", err)
		2805	}
		2806	defer func() { _ = conn.Close() }()
		2807
		2808	requests := []string{
		2809	fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"}]}`, sourceModel),
		2810	fmt.Sprintf(`{"type":"response.create","model":%q,"route_to_claude":true,"input":[{"type":"message","id":"msg-routed"}]}`, sourceModel),
		2811	`{"type":"response.create","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`,
		2812	}
		2813	wantTypes := []string{wsEventTypeCompleted, wsEventTypeError, wsEventTypeCompleted}
		2814	for i, request := range requests {
		2815	if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(request)); errWrite != nil {
		2816	t.Fatalf("write request %d: %v", i+1, errWrite)
		2817	}
		2818	_, response, errRead := conn.ReadMessage()
		2819	if errRead != nil {
		2820	t.Fatalf("read response %d: %v", i+1, errRead)
		2821	}
		2822	if got := gjson.GetBytes(response, "type").String(); got != wantTypes[i] {
		2823	t.Fatalf("response %d type = %q, want %q: %s", i+1, got, wantTypes[i], response)
		2824	}
		2825	}
		2826
		2827	codexPayloads := codexExecutor.Payloads()
		2828	if len(codexPayloads) != 2 {
		2829	t.Fatalf("codex payload count = %d, want 2", len(codexPayloads))
		2830	}
		2831	if got := gjson.GetBytes(codexPayloads[1], "previous_response_id").String(); got != "resp-1" {
		2832	t.Fatalf("resumed codex previous_response_id = %q, want resp-1: %s", got, codexPayloads[1])
		2833	}
		2834	if got := len(claudeExecutor.Payloads()); got != 1 {
		2835	t.Fatalf("claude payload count = %d, want 1", got)
		2836	}
		2837	}
		2838
		2839	func TestResponsesWebsocketDeltaRouteToBuiltInProviderRequiresFullReplay(t *testing.T) {
		2840	gin.SetMode(gin.TestMode)
		2841
		2842	const sourceModel = "codex-provider-route-delta-source"
		2843	const targetModel = "claude-provider-route-target"
		2844	codexExecutor := &websocketDirectCaptureExecutor{provider: "codex"}
		2845	claudeExecutor := &websocketDirectCaptureExecutor{provider: "claude"}
		2846	manager := coreauth.NewManager(nil, nil, nil)
		2847	manager.RegisterExecutor(codexExecutor)
		2848	manager.RegisterExecutor(claudeExecutor)
		2849	codexAuth := &coreauth.Auth{ID: "auth-codex-provider-route-delta", Provider: "codex", Status: coreauth.StatusActive, Attributes: map[string]string{"websockets": "true"}}
		2850	claudeAuth := &coreauth.Auth{ID: "auth-claude-provider-route-delta", Provider: "claude", Status: coreauth.StatusActive}
		2851	for _, auth := range []*coreauth.Auth{codexAuth, claudeAuth} {
		2852	if _, err := manager.Register(context.Background(), auth); err != nil {
		2853	t.Fatalf("Register auth %s: %v", auth.ID, err)
		2854	}
		2855	}
		2856	registry.GetGlobalRegistry().RegisterClient(codexAuth.ID, codexAuth.Provider, []*registry.ModelInfo{{ID: sourceModel}})
		2857	registry.GetGlobalRegistry().RegisterClient(claudeAuth.ID, claudeAuth.Provider, []*registry.ModelInfo{{ID: targetModel}})
		2858	t.Cleanup(func() {
		2859	registry.GetGlobalRegistry().UnregisterClient(codexAuth.ID)
		2860	registry.GetGlobalRegistry().UnregisterClient(claudeAuth.ID)
		2861	})
		2862
		2863	base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
		2864	base.SetModelRouterHost(&websocketProviderRouteHost{})
		2865	h := NewOpenAIResponsesAPIHandler(base)
		2866	router := gin.New()
		2867	router.GET("/v1/responses/ws", h.ResponsesWebsocket)
		2868	server := httptest.NewServer(router)
		2869	defer server.Close()
		2870
		2871	wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses/ws"
		2872	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
		2873	if err != nil {
		2874	t.Fatalf("dial websocket: %v", err)
		2875	}
		2876	defer func() { _ = conn.Close() }()
		2877
		2878	firstRequest := []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"}]}`, sourceModel))
		2879	if errWrite := conn.WriteMessage(websocket.TextMessage, firstRequest); errWrite != nil {
		2880	t.Fatalf("write first websocket message: %v", errWrite)
		2881	}
		2882	if _, _, errRead := conn.ReadMessage(); errRead != nil {
		2883	t.Fatalf("read first websocket response: %v", errRead)
		2884	}
		2885
		2886	routedDelta := []byte(`{"type":"response.create","route_to_claude":true,"previous_response_id":"resp-1","input":[{"type":"message","id":"msg-routed"}]}`)
		2887	if errWrite := conn.WriteMessage(websocket.TextMessage, routedDelta); errWrite != nil {
		2888	t.Fatalf("write routed delta: %v", errWrite)
		2889	}
		2890	_, _, errRead := conn.ReadMessage()
		2891	var closeErr *websocket.CloseError
		2892	if !errors.As(errRead, &closeErr) {
		2893	t.Fatalf("routed delta error = %v, want websocket close", errRead)
		2894	}
		2895	if closeErr.Code != websocket.CloseServiceRestart \|\| closeErr.Text != wsHTTPReplayRequiredCloseReason {
		2896	t.Fatalf("routed delta close = %d %q, want %d %q", closeErr.Code, closeErr.Text, websocket.CloseServiceRestart, wsHTTPReplayRequiredCloseReason)
		2897	}
		2898	if got := len(codexExecutor.Payloads()); got != 1 {
		2899	t.Fatalf("codex payload count = %d, want 1", got)
		2900	}
		2901	if got := len(claudeExecutor.Payloads()); got != 0 {
		2902	t.Fatalf("claude payload count = %d, want 0 before full replay", got)
		2903	}
		2904	}
		2905
		2906	func TestResponsesWebsocketClosesForHTTPReplayWhenWebsocketEligibilityChanges(t *testing.T) {
		2907	gin.SetMode(gin.TestMode)
		2908
		2909	modelName := "xai-websocket-mode-change-model"
		2910	executor := &websocketDirectCaptureExecutor{provider: "xai", done: make(chan struct{})}
		2911	manager := coreauth.NewManager(nil, nil, nil)
		2912	manager.RegisterExecutor(executor)
		2913	auth := &coreauth.Auth{
		2914	ID: "auth-xai-mode-change",
		2915	Provider: "xai",
		2916	Status: coreauth.StatusActive,
		2917	Attributes: map[string]string{"websockets": "true"},
		2918	}
		2919	if _, err := manager.Register(context.Background(), auth); err != nil {
		2920	t.Fatalf("Register auth: %v", err)
		2921	}
		2922	registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: modelName}})
		2923	t.Cleanup(func() {
		2924	registry.GetGlobalRegistry().UnregisterClient(auth.ID)
		2925	})
		2926
		2927	base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
		2928	h := NewOpenAIResponsesAPIHandler(base)
		2929	router := gin.New()
		2930	router.GET("/v1/responses/ws", h.ResponsesWebsocket)
		2931
		2932	server := httptest.NewServer(router)
		2933	defer server.Close()
		2934
		2935	wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses/ws"
		2936	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
		2937	if err != nil {
		2938	t.Fatalf("dial websocket: %v", err)
		2939	}
		2940	defer func() { _ = conn.Close() }()
		2941
		2942	firstRequest := []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"}]}`, modelName))
		2943	if errWrite := conn.WriteMessage(websocket.TextMessage, firstRequest); errWrite != nil {
		2944	t.Fatalf("write first websocket message: %v", errWrite)
		2945	}
		2946	if _, _, errRead := conn.ReadMessage(); errRead != nil {
		2947	t.Fatalf("read first websocket response: %v", errRead)
		2948	}
		2949
		2950	secondRequest := []byte(`{"type":"response.create","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`)
		2951	if errWrite := conn.WriteMessage(websocket.TextMessage, secondRequest); errWrite != nil {
		2952	t.Fatalf("write second websocket message: %v", errWrite)
		2953	}
		2954	if _, _, errRead := conn.ReadMessage(); errRead != nil {
		2955	t.Fatalf("read second websocket response: %v", errRead)
		2956	}
		2957
		2958	updatedAuth := &coreauth.Auth{
		2959	ID: auth.ID,
		2960	Provider: auth.Provider,
		2961	Status: coreauth.StatusActive,
		2962	}
		2963	if _, errUpdate := manager.Update(context.Background(), updatedAuth); errUpdate != nil {
		2964	t.Fatalf("Update auth: %v", errUpdate)
		2965	}
		2966
		2967	thirdRequest := []byte(`{"type":"response.create","previous_response_id":"resp-2","input":[{"type":"message","id":"msg-3"}]}`)
		2968	if errWrite := conn.WriteMessage(websocket.TextMessage, thirdRequest); errWrite != nil {
		2969	t.Fatalf("write third websocket message: %v", errWrite)
		2970	}
		2971	_, _, errRead := conn.ReadMessage()
		2972	var closeErr *websocket.CloseError
		2973	if !errors.As(errRead, &closeErr) {
		2974	t.Fatalf("third response error = %v, want websocket close", errRead)
		2975	}
		2976	if closeErr.Code != websocket.CloseServiceRestart \|\| closeErr.Text != wsHTTPReplayRequiredCloseReason {
		2977	t.Fatalf("third response close = %d %q, want %d %q", closeErr.Code, closeErr.Text, websocket.CloseServiceRestart, wsHTTPReplayRequiredCloseReason)
		2978	}
		2979
		2980	payloads := executor.Payloads()
		2981	if len(payloads) != 2 {
		2982	t.Fatalf("executor payload count = %d, want 2; transport switch must not call HTTP upstream", len(payloads))
		2983	}
		2984	second := payloads[1]
		2985	if got := gjson.GetBytes(second, "previous_response_id").String(); got != "resp-1" {
		2986	t.Fatalf("stable websocket previous_response_id = %q, want resp-1: %s", got, second)
		2987	}
		2988	if input := gjson.GetBytes(second, "input").Array(); len(input) != 1 \|\| input[0].Get("id").String() != "msg-2" {
		2989	t.Fatalf("stable websocket payload is not incremental: %s", second)
		2990	}
		2991
		2992	replayConn, _, errDialReplay := websocket.DefaultDialer.Dial(wsURL, nil)
		2993	if errDialReplay != nil {
		2994	t.Fatalf("dial replay websocket: %v", errDialReplay)
		2995	}
		2996	defer func() { _ = replayConn.Close() }()
		2997	fullReplay := []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"},{"type":"message","id":"out-1"},{"type":"message","id":"msg-2"},{"type":"message","id":"out-2"},{"type":"message","id":"msg-3"}]}`, modelName))
		2998	if errWrite := replayConn.WriteMessage(websocket.TextMessage, fullReplay); errWrite != nil {
		2999	t.Fatalf("write full replay: %v", errWrite)
		3000	}
		3001	if _, _, errReadReplay := replayConn.ReadMessage(); errReadReplay != nil {
		3002	t.Fatalf("read full replay response: %v", errReadReplay)
		3003	}
		3004	deltaAfterReplay := []byte(`{"type":"response.create","previous_response_id":"resp-3","input":[{"type":"message","id":"msg-4"}]}`)
		3005	if errWrite := replayConn.WriteMessage(websocket.TextMessage, deltaAfterReplay); errWrite != nil {
		3006	t.Fatalf("write delta after replay: %v", errWrite)
		3007	}
		3008	if _, _, errReadReplay := replayConn.ReadMessage(); errReadReplay != nil {
		3009	t.Fatalf("read delta after replay response: %v", errReadReplay)
		3010	}
		3011
		3012	payloads = executor.Payloads()
		3013	if len(payloads) != 4 {
		3014	t.Fatalf("executor payload count after replay = %d, want 4", len(payloads))
		3015	}
		3016	httpDelta := payloads[3]
		3017	if gjson.GetBytes(httpDelta, "previous_response_id").Exists() {
		3018	t.Fatalf("HTTP-mode delta retained previous_response_id: %s", httpDelta)
		3019	}
		3020	input := gjson.GetBytes(httpDelta, "input").Array()
		3021	wantIDs := []string{"msg-1", "out-1", "msg-2", "out-2", "msg-3", "out-3", "msg-4"}
		3022	if len(input) != len(wantIDs) {
		3023	t.Fatalf("HTTP-mode canonical input len = %d, want %d: %s", len(input), len(wantIDs), httpDelta)
		3024	}
		3025	for i, wantID := range wantIDs {
		3026	if got := input[i].Get("id").String(); got != wantID {
		3027	t.Fatalf("HTTP-mode canonical input[%d].id = %q, want %q: %s", i, got, wantID, httpDelta)
		3028	}
		3029	}
		3030	}
		3031
		3032	func TestResponsesWebsocketRejectsUnknownPreviousResponseOnNewSocket(t *testing.T) {
		3033	gin.SetMode(gin.TestMode)
		3034
		3035	modelName := "xai-websocket-reconnect-model"
		3036	executor := &websocketDirectCaptureExecutor{provider: "xai"}
		3037	manager := coreauth.NewManager(nil, nil, nil)
		3038	manager.RegisterExecutor(executor)
		3039	auth := &coreauth.Auth{
		3040	ID: "auth-xai-reconnect",
		3041	Provider: "xai",
		3042	Status: coreauth.StatusActive,
		3043	Attributes: map[string]string{"websockets": "true"},
		3044	}
		3045	if _, err := manager.Register(context.Background(), auth); err != nil {
		3046	t.Fatalf("Register auth: %v", err)
		3047	}
		3048	registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: modelName}})
		3049	t.Cleanup(func() {
		3050	registry.GetGlobalRegistry().UnregisterClient(auth.ID)
		3051	})
		3052
		3053	base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
		3054	h := NewOpenAIResponsesAPIHandler(base)
		3055	router := gin.New()
		3056	router.GET("/v1/responses/ws", h.ResponsesWebsocket)
		3057
		3058	server := httptest.NewServer(router)
		3059	defer server.Close()
		3060
		3061	wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses/ws"
		3062	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
		3063	if err != nil {
		3064	t.Fatalf("dial websocket: %v", err)
		3065	}
		3066	defer func() { _ = conn.Close() }()
		3067
		3068	request := []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"previous_response_id":"resp-old","input":[{"type":"message","id":"msg-2","role":"user","content":"second"}]}`, modelName))
		3069	if errWrite := conn.WriteMessage(websocket.TextMessage, request); errWrite != nil {
		3070	t.Fatalf("write websocket message: %v", errWrite)
		3071	}
		3072	_, payload, errRead := conn.ReadMessage()
		3073	if errRead != nil {
		3074	t.Fatalf("read websocket response: %v", errRead)
		3075	}
		3076	if got := gjson.GetBytes(payload, "type").String(); got != wsEventTypeError {
		3077	t.Fatalf("response type = %q, want %q: %s", got, wsEventTypeError, payload)
		3078	}
		3079	if got := int(gjson.GetBytes(payload, "status").Int()); got != http.StatusConflict {
		3080	t.Fatalf("response status = %d, want %d: %s", got, http.StatusConflict, payload)
		3081	}
		3082	if got := gjson.GetBytes(payload, "error.code").String(); got != "previous_response_not_found" {
		3083	t.Fatalf("response error code = %q, want previous_response_not_found: %s", got, payload)
		3084	}
		3085	if got := len(executor.Payloads()); got != 0 {
		3086	t.Fatalf("executor payload count = %d, want 0", got)
		3087	}
		3088
		3089	recoveryRequest := []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"},{"type":"message","id":"out-1","role":"assistant"},{"type":"message","id":"msg-2"}]}`, modelName))
		3090	if errWrite := conn.WriteMessage(websocket.TextMessage, recoveryRequest); errWrite != nil {
		3091	t.Fatalf("write full recovery message: %v", errWrite)
		3092	}
		3093	_, recoveryPayload, errRead := conn.ReadMessage()
		3094	if errRead != nil {
		3095	t.Fatalf("read full recovery response: %v", errRead)
		3096	}
		3097	if got := gjson.GetBytes(recoveryPayload, "type").String(); got != wsEventTypeCompleted {
		3098	t.Fatalf("recovery response type = %q, want %q: %s", got, wsEventTypeCompleted, recoveryPayload)
		3099	}
		3100	payloads := executor.Payloads()
		3101	if len(payloads) != 1 {
		3102	t.Fatalf("executor payload count after recovery = %d, want 1", len(payloads))
		3103	}
		3104	if got := len(gjson.GetBytes(payloads[0], "input").Array()); got != 3 {
		3105	t.Fatalf("full recovery input len = %d, want 3: %s", got, payloads[0])
		3106	}
		3107	}
		3108
		3109	func TestResponsesWebsocketRollsBackCanonicalTranscriptAfterNonRetryableError(t *testing.T) {
		3110	gin.SetMode(gin.TestMode)
		3111
		3112	modelName := "xai-websocket-rollback-model"
		3113	executor := &websocketCanonicalRollbackExecutor{}
		3114	manager := coreauth.NewManager(nil, nil, nil)
		3115	manager.RegisterExecutor(executor)
		3116	auth := &coreauth.Auth{
		3117	ID: "auth-xai-rollback",
		3118	Provider: "xai",
		3119	Status: coreauth.StatusActive,
		3120	}
		3121	if _, err := manager.Register(context.Background(), auth); err != nil {
		3122	t.Fatalf("Register auth: %v", err)
		3123	}
		3124	registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: modelName}})
		3125	t.Cleanup(func() {
		3126	registry.GetGlobalRegistry().UnregisterClient(auth.ID)
		3127	})
		3128
		3129	base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
		3130	h := NewOpenAIResponsesAPIHandler(base)
		3131	router := gin.New()
		3132	router.GET("/v1/responses/ws", h.ResponsesWebsocket)
		3133
		3134	server := httptest.NewServer(router)
		3135	defer server.Close()
		3136
		3137	wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses/ws"
		3138	conn, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{"Session-Id": []string{"rollback-tool-cache-session"}})
		3139	if err != nil {
		3140	t.Fatalf("dial websocket: %v", err)
		3141	}
		3142	defer func() { _ = conn.Close() }()
		3143
		3144	requests := []string{
		3145	fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"}]}`, modelName),
		3146	`{"type":"response.create","previous_response_id":"resp-1","input":[{"type":"function_call","id":"fc-failed","call_id":"failed-call","name":"failed_tool","arguments":"{}"},{"type":"function_call_output","id":"fco-failed","call_id":"failed-call","output":"must-not-survive"}]}`,
		3147	`{"type":"response.create","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-3"},{"type":"function_call","id":"fc-retry","call_id":"failed-call","name":"failed_tool","arguments":"{}"}]}`,
		3148	}
		3149	wantTypes := []string{wsEventTypeCompleted, wsEventTypeError, wsEventTypeCompleted}
		3150	for i := range requests {
		3151	if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(requests[i])); errWrite != nil {
		3152	t.Fatalf("write websocket message %d: %v", i+1, errWrite)
		3153	}
		3154	_, payload, errRead := conn.ReadMessage()
		3155	if errRead != nil {
		3156	t.Fatalf("read websocket response %d: %v", i+1, errRead)
		3157	}
		3158	if got := gjson.GetBytes(payload, "type").String(); got != wantTypes[i] {
		3159	t.Fatalf("response %d type = %q, want %q: %s", i+1, got, wantTypes[i], payload)
		3160	}
		3161	}
		3162
		3163	payloads := executor.Payloads()
		3164	if len(payloads) != 3 {
		3165	t.Fatalf("executor payload count = %d, want 3", len(payloads))
		3166	}
		3167	third := payloads[2]
		3168	if gjson.GetBytes(third, "previous_response_id").Exists() {
		3169	t.Fatalf("retry payload must not depend on previous_response_id: %s", third)
		3170	}
		3171	input := gjson.GetBytes(third, "input").Array()
		3172	if len(input) != 3 {
		3173	t.Fatalf("retry canonical input len = %d, want 3: %s", len(input), third)
		3174	}
		3175	wantIDs := []string{"msg-1", "out-1", "msg-3"}
		3176	for i, wantID := range wantIDs {
		3177	if got := input[i].Get("id").String(); got != wantID {
		3178	t.Fatalf("retry canonical input[%d].id = %q, want %q: %s", i, got, wantID, third)
		3179	}
		3180	}
		3181	if bytes.Contains(third, []byte(`"id":"fc-failed"`)) \|\| bytes.Contains(third, []byte(`"id":"fco-failed"`)) {
		3182	t.Fatalf("failed turn leaked into retry transcript: %s", third)
		3183	}
		3184	if bytes.Contains(third, []byte(`"call_id":"failed-call"`)) \|\| bytes.Contains(third, []byte("must-not-survive")) {
		3185	t.Fatalf("failed turn contaminated tool repair cache: %s", third)
		3186	}
		3187	}
		3188
		3189	func TestResponsesWebsocketSwitchesPinnedAuthAcrossProviders(t *testing.T) {
		3190	for _, testCase := range []struct {
		3191	name string
		3192	xaiWebsockets bool
		3193	returnToDifferentXAIModel bool
		3194	}{
		3195	{name: "xai SSE", xaiWebsockets: false},
		3196	{name: "xai websocket", xaiWebsockets: true},
		3197	{name: "xai websocket different model", xaiWebsockets: true, returnToDifferentXAIModel: true},
		3198	} {
		3199	t.Run(testCase.name, func(t *testing.T) {
		3200	gin.SetMode(gin.TestMode)
		3201
		3202	xaiModel := "xai-provider-switch-" + strings.ReplaceAll(testCase.name, " ", "-")
		3203	returnXAIModel := xaiModel
		3204	if testCase.returnToDifferentXAIModel {
		3205	returnXAIModel += "-return"
		3206	}
		3207	codexModel := "codex-provider-switch-" + strings.ReplaceAll(testCase.name, " ", "-")
		3208	xaiExecutor := &websocketDirectCaptureExecutor{provider: "xai"}
		3209	codexExecutor := &websocketDirectCaptureExecutor{provider: "codex"}
		3210
		3211	xaiAuth := &coreauth.Auth{
		3212	ID: "auth-" + xaiModel,
		3213	Provider: "xai",
		3214	Status: coreauth.StatusActive,
		3215	}
		3216	if testCase.xaiWebsockets {
		3217	xaiAuth.Attributes = map[string]string{"websockets": "true"}
		3218	}
		3219	codexAuth := &coreauth.Auth{
		3220	ID: "auth-" + codexModel,
		3221	Provider: "codex",
		3222	Status: coreauth.StatusActive,
		3223	Attributes: map[string]string{"websockets": "true"},
		3224	}
		3225	selector := &orderedWebsocketSelector{order: []string{xaiAuth.ID, codexAuth.ID}}
		3226	manager := coreauth.NewManager(nil, selector, nil)
		3227	manager.RegisterExecutor(xaiExecutor)
		3228	manager.RegisterExecutor(codexExecutor)
		3229	if _, errRegister := manager.Register(context.Background(), xaiAuth); errRegister != nil {
		3230	t.Fatalf("Register xAI auth: %v", errRegister)
		3231	}
		3232	if _, errRegister := manager.Register(context.Background(), codexAuth); errRegister != nil {
		3233	t.Fatalf("Register Codex auth: %v", errRegister)
		3234	}
		3235
		3236	registry.GetGlobalRegistry().RegisterClient(xaiAuth.ID, xaiAuth.Provider, []*registry.ModelInfo{{ID: xaiModel}})
		3237	registry.GetGlobalRegistry().RegisterClient(codexAuth.ID, codexAuth.Provider, []*registry.ModelInfo{{ID: codexModel}})
		3238	registeredAuthIDs := []string{xaiAuth.ID, codexAuth.ID}
		3239	if testCase.xaiWebsockets {
		3240	xaiAlternateAuth := &coreauth.Auth{
		3241	ID: "auth-alternate-" + xaiModel,
		3242	Provider: "xai",
		3243	Status: coreauth.StatusActive,
		3244	Attributes: map[string]string{"websockets": "true"},
		3245	}
		3246	selector.order = append(selector.order, xaiAlternateAuth.ID)
		3247	if _, errRegister := manager.Register(context.Background(), xaiAlternateAuth); errRegister != nil {
		3248	t.Fatalf("Register alternate xAI auth: %v", errRegister)
		3249	}
		3250	alternateModels := []*registry.ModelInfo{{ID: xaiModel}}
		3251	if testCase.returnToDifferentXAIModel {
		3252	alternateModels = []*registry.ModelInfo{{ID: returnXAIModel}}
		3253	}
		3254	registry.GetGlobalRegistry().RegisterClient(xaiAlternateAuth.ID, xaiAlternateAuth.Provider, alternateModels)
		3255	registeredAuthIDs = append(registeredAuthIDs, xaiAlternateAuth.ID)
		3256	}
		3257	t.Cleanup(func() {
		3258	for _, authID := range registeredAuthIDs {
		3259	registry.GetGlobalRegistry().UnregisterClient(authID)
		3260	}
		3261	})
		3262
		3263	base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
		3264	h := NewOpenAIResponsesAPIHandler(base)
		3265	router := gin.New()
		3266	router.GET("/v1/responses/ws", h.ResponsesWebsocket)
		3267
		3268	server := httptest.NewServer(router)
		3269	defer server.Close()
		3270
		3271	wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses/ws"
		3272	conn, _, errDial := websocket.DefaultDialer.Dial(wsURL, nil)
		3273	if errDial != nil {
		3274	t.Fatalf("dial websocket: %v", errDial)
		3275	}
		3276	defer func() {
		3277	if errClose := conn.Close(); errClose != nil {
		3278	t.Errorf("close websocket: %v", errClose)
		3279	}
		3280	}()
		3281
		3282	requests := []string{
		3283	fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-xai-1"}]}`, xaiModel),
		3284	fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-codex-1"}]}`, codexModel),
		3285	fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-xai-2"}]}`, returnXAIModel),
		3286	`{"type":"response.create","input":[{"type":"message","id":"msg-xai-3"}]}`,
		3287	}
		3288	for index, request := range requests {
		3289	turn := index + 1
		3290	if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(request)); errWrite != nil {
		3291	t.Fatalf("write websocket message %d: %v", turn, errWrite)
		3292	}
		3293	_, payload, errRead := conn.ReadMessage()
		3294	if errRead != nil {
		3295	t.Fatalf("read websocket response %d: %v", turn, errRead)
		3296	}
		3297	if got := gjson.GetBytes(payload, "type").String(); got != wsEventTypeCompleted {
		3298	t.Fatalf("response %d type = %s, want %s: %s", turn, got, wsEventTypeCompleted, payload)
		3299	}
		3300	}
		3301
		3302	wantReturnAuthID := xaiAuth.ID
		3303	if testCase.returnToDifferentXAIModel {
		3304	wantReturnAuthID = "auth-alternate-" + xaiModel
		3305	}
		3306	if got := xaiExecutor.AuthIDs(); len(got) != 3 \|\| got[0] != xaiAuth.ID \|\| got[1] != wantReturnAuthID \|\| got[2] != wantReturnAuthID {
		3307	t.Fatalf("xAI auth IDs = %v, want [%s %s %s]", got, xaiAuth.ID, wantReturnAuthID, wantReturnAuthID)
		3308	}
		3309	if got := codexExecutor.AuthIDs(); len(got) != 1 \|\| got[0] != codexAuth.ID {
		3310	t.Fatalf("Codex auth IDs = %v, want [%s]", got, codexAuth.ID)
		3311	}
		3312	})
		3313	}
		3314	}
		3315
		3316	func TestResponsesWebsocketPinnedAuthMatchesModel(t *testing.T) {
		3317	modelA := "xai-pinned-auth-model-a"
		3318	modelB := "xai-pinned-auth-model-b"
		3319	auth := &coreauth.Auth{ID: "xai-pinned-auth", Provider: "xai", Status: coreauth.StatusActive}
		3320	otherAuthID := "xai-pinned-auth-other"
		3321	registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: modelA}})
		3322	registry.GetGlobalRegistry().RegisterClient(otherAuthID, auth.Provider, []*registry.ModelInfo{{ID: modelB}})
2645	t.Cleanup(func() {	3323	t.Cleanup(func() {
2646	registry.GetGlobalRegistry().UnregisterClient(auth.ID)	3324	registry.GetGlobalRegistry().UnregisterClient(auth.ID)
2647	registry.GetGlobalRegistry().UnregisterClient(otherAuthID)	3325	registry.GetGlobalRegistry().UnregisterClient(otherAuthID)
@@ -3055,7 +3733,7 @@ func TestResponsesWebsocketDoesNotInjectPreviousResponseIDWhenPendingToolOutputM · 407 unmodified lines
3055	func TestResponsesWebsocketStripsGenerateWhenWebsocketAttemptFallsBackToHTTP(t *testing.T) {	3733	func TestResponsesWebsocketStripsGenerateWhenWebsocketAttemptFallsBackToHTTP(t *testing.T) {
3056	gin.SetMode(gin.TestMode)	3734	gin.SetMode(gin.TestMode)
3057		3735
3058	selector := &orderedWebsocketSelector{order: []string{"auth-ws", "auth-http"}}	3736	selector := &orderedWebsocketSelector{order: []string{"auth-ws", "auth-http", "auth-http"}}
3059	executor := &websocketBootstrapFallbackExecutor{}	3737	executor := &websocketBootstrapFallbackExecutor{}
3060	manager := coreauth.NewManager(nil, selector, nil)	3738	manager := coreauth.NewManager(nil, selector, nil)
3061	manager.RegisterExecutor(executor)	3739	manager.RegisterExecutor(executor)
@@ -3131,6 +3809,21 @@ func TestResponsesWebsocketStripsGenerateWhenWebsocketAttemptFallsBackToHTTP(t * · 69 unmodified lines
3131	if gjson.GetBytes(httpPayloads[0], "generate").Exists() {	3809	if gjson.GetBytes(httpPayloads[0], "generate").Exists() {
3132	t.Fatalf("generate leaked after HTTP fallback: %s", httpPayloads[0])	3810	t.Fatalf("generate leaked after HTTP fallback: %s", httpPayloads[0])
3133	}	3811	}
		3812
		3813	secondRequest := `{"type":"response.create","previous_response_id":"resp-http","input":[{"type":"message","id":"msg-2"}]}`
		3814	if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(secondRequest)); errWrite != nil {
		3815	t.Fatalf("write second websocket message: %v", errWrite)
		3816	}
		3817	_, secondPayload, errReadSecond := conn.ReadMessage()
		3818	if errReadSecond != nil {
		3819	t.Fatalf("read second websocket message: %v", errReadSecond)
		3820	}
		3821	if got := gjson.GetBytes(secondPayload, "type").String(); got != wsEventTypeCompleted {
		3822	t.Fatalf("second payload type = %s, want %s: %s", got, wsEventTypeCompleted, secondPayload)
		3823	}
		3824	if got := executor.AuthIDs(); len(got) != 3 \|\| got[2] != "auth-http" {
		3825	t.Fatalf("selected auth IDs after HTTP retry = %v, want [auth-ws auth-http auth-http]", got)
		3826	}
3134	}	3827	}
3135		3828
3136	func TestWebsocketClientAddressUsesGinClientIP(t *testing.T) {	3829	func TestWebsocketClientAddressUsesGinClientIP(t *testing.T) {
@@ -3228,45 +3921,38 @@ func TestResponsesWebsocketPinsOnlyWebsocketCapableAuth(t *testing.T) { · 91 unmodified lines
3228	}	3921	}
3229	}	3922	}
3230		3923
3231	func TestResponsesWebsocketReleasesPinnedAuthAfterQuotaError(t *testing.T) {	3924	func TestResponsesWebsocketUsesNativeIncrementalAfterPinningWebsocketAuthFromMixedPool(t *testing.T) {
3232	gin.SetMode(gin.TestMode)	3925	gin.SetMode(gin.TestMode)
3233		3926
3234	selector := &orderedWebsocketSelector{order: []string{"auth-a", "auth-b"}}	3927	modelName := "xai-mixed-pool-model"
3235	executor := &websocketPinnedFailoverExecutor{}	3928	selector := &orderedWebsocketSelector{order: []string{"auth-http", "auth-ws"}}
		3929	executor := &websocketDirectCaptureExecutor{provider: "xai"}
3236	manager := coreauth.NewManager(nil, selector, nil)	3930	manager := coreauth.NewManager(nil, selector, nil)
3237	manager.RegisterExecutor(executor)	3931	manager.RegisterExecutor(executor)
3238		3932	authHTTP := &coreauth.Auth{ID: "auth-http", Provider: "xai", Status: coreauth.StatusActive}
3239	authA := &coreauth.Auth{	3933	if _, err := manager.Register(context.Background(), authHTTP); err != nil {
3240	ID: "auth-a",	3934	t.Fatalf("Register HTTP auth: %v", err)
3241	Provider: executor.Identifier(),
3242	Status: coreauth.StatusActive,
3243	Attributes: map[string]string{"websockets": "true"},
3244	}
3245	if _, err := manager.Register(context.Background(), authA); err != nil {
3246	t.Fatalf("Register auth A: %v", err)
3247	}	3935	}
3248	authB := &coreauth.Auth{	3936	authWS := &coreauth.Auth{
3249	ID: "auth-b",	3937	ID: "auth-ws",
3250	Provider: executor.Identifier(),	3938	Provider: "xai",
3251	Status: coreauth.StatusActive,	3939	Status: coreauth.StatusActive,
3252	Attributes: map[string]string{"websockets": "true"},	3940	Attributes: map[string]string{"websockets": "true"},
3253	}	3941	}
3254	if _, err := manager.Register(context.Background(), authB); err != nil {	3942	if _, err := manager.Register(context.Background(), authWS); err != nil {
3255	t.Fatalf("Register auth B: %v", err)	3943	t.Fatalf("Register websocket auth: %v", err)
3256	}	3944	}
3257		3945	registry.GetGlobalRegistry().RegisterClient(authHTTP.ID, authHTTP.Provider, []*registry.ModelInfo{{ID: modelName}})
3258	registry.GetGlobalRegistry().RegisterClient(authA.ID, authA.Provider, []*registry.ModelInfo{{ID: "quota-model"}})	3946	registry.GetGlobalRegistry().RegisterClient(authWS.ID, authWS.Provider, []*registry.ModelInfo{{ID: modelName}})
3259	registry.GetGlobalRegistry().RegisterClient(authB.ID, authB.Provider, []*registry.ModelInfo{{ID: "quota-model"}})
3260	t.Cleanup(func() {	3947	t.Cleanup(func() {
3261	registry.GetGlobalRegistry().UnregisterClient(authA.ID)	3948	registry.GetGlobalRegistry().UnregisterClient(authHTTP.ID)
3262	registry.GetGlobalRegistry().UnregisterClient(authB.ID)	3949	registry.GetGlobalRegistry().UnregisterClient(authWS.ID)
3263	})	3950	})
3264		3951
3265	base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)	3952	base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
3266	h := NewOpenAIResponsesAPIHandler(base)	3953	h := NewOpenAIResponsesAPIHandler(base)
3267	router := gin.New()	3954	router := gin.New()
3268	router.GET("/v1/responses/ws", h.ResponsesWebsocket)	3955	router.GET("/v1/responses/ws", h.ResponsesWebsocket)
3269
3270	server := httptest.NewServer(router)	3956	server := httptest.NewServer(router)
3271	defer server.Close()	3957	defer server.Close()
3272		3958
@@ -3275,49 +3961,169 @@ func TestResponsesWebsocketReleasesPinnedAuthAfterQuotaError(t *testing.T) { · 2 unmodified lines
3275	if err != nil {	3961	if err != nil {
3276	t.Fatalf("dial websocket: %v", err)	3962	t.Fatalf("dial websocket: %v", err)
3277	}	3963	}
3278	defer func() {	3964	defer func() { _ = conn.Close() }()
3279	if errClose := conn.Close(); errClose != nil {
3280	t.Fatalf("close websocket: %v", errClose)
3281	}
3282	}()
3283		3965
3284	requests := []string{	3966	requests := []string{
3285	`{"type":"response.create","model":"quota-model","input":[{"type":"message","id":"msg-1"}]}`,	3967	fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"}]}`, modelName),
3286	`{"type":"response.create","previous_response_id":"resp-auth-a-1","input":[{"type":"message","id":"msg-2"}]}`,	3968	`{"type":"response.create","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`,
3287	`{"type":"response.create","previous_response_id":"resp-auth-a-1","input":[{"type":"message","id":"msg-3"}]}`,	3969	`{"type":"response.create","previous_response_id":"resp-2","input":[{"type":"message","id":"msg-3"}]}`,
3288	}	3970	}
3289	wantTypes := []string{wsEventTypeCompleted, wsEventTypeError, wsEventTypeCompleted}
3290	for i := range requests {	3971	for i := range requests {
3291	if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(requests[i])); errWrite != nil {	3972	if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(requests[i])); errWrite != nil {
3292	t.Fatalf("write websocket message %d: %v", i+1, errWrite)	3973	t.Fatalf("write websocket message %d: %v", i+1, errWrite)
3293	}	3974	}
3294	_, payload, errReadMessage := conn.ReadMessage()	3975	if _, _, errRead := conn.ReadMessage(); errRead != nil {
3295	if errReadMessage != nil {	3976	t.Fatalf("read websocket response %d: %v", i+1, errRead)
3296	t.Fatalf("read websocket message %d: %v", i+1, errReadMessage)
3297	}
3298	if got := gjson.GetBytes(payload, "type").String(); got != wantTypes[i] {
3299	t.Fatalf("message %d payload type = %s, want %s: %s", i+1, got, wantTypes[i], payload)
3300	}
3301	if i == 1 && int(gjson.GetBytes(payload, "status").Int()) != http.StatusTooManyRequests {
3302	t.Fatalf("quota payload status = %d, want %d: %s", gjson.GetBytes(payload, "status").Int(), http.StatusTooManyRequests, payload)
3303	}	3977	}
3304	}	3978	}
3305		3979
3306	if got := executor.AuthIDs(); len(got) != 3 \|\| got[0] != "auth-a" \|\| got[1] != "auth-a" \|\| got[2] != "auth-b" {	3980	if got := executor.AuthIDs(); len(got) != 3 \|\| got[0] != "auth-http" \|\| got[1] != "auth-ws" \|\| got[2] != "auth-ws" {
3307	t.Fatalf("selected auth IDs = %v, want [auth-a auth-a auth-b]", got)	3981	t.Fatalf("selected auth IDs = %v, want [auth-http auth-ws auth-ws]", got)
		3982	}
		3983	payloads := executor.Payloads()
		3984	if len(payloads) != 3 {
		3985	t.Fatalf("payload count = %d, want 3", len(payloads))
		3986	}
		3987	if gjson.GetBytes(payloads[1], "previous_response_id").Exists() \|\| len(gjson.GetBytes(payloads[1], "input").Array()) != 3 {
		3988	t.Fatalf("first request on newly selected websocket auth must be canonical: %s", payloads[1])
		3989	}
		3990	if got := gjson.GetBytes(payloads[2], "previous_response_id").String(); got != "resp-2" {
		3991	t.Fatalf("stable pinned websocket previous_response_id = %q, want resp-2: %s", got, payloads[2])
		3992	}
		3993	input := gjson.GetBytes(payloads[2], "input").Array()
		3994	if len(input) != 1 \|\| input[0].Get("id").String() != "msg-3" {
		3995	t.Fatalf("stable pinned websocket request is not incremental: %s", payloads[2])
		3996	}
		3997	}
		3998
		3999	func TestResponsesWebsocketReplaysImmediatelyAfterPinnedAuthFailure(t *testing.T) {
		4000	gin.SetMode(gin.TestMode)
		4001
		4002	tests := []struct {
		4003	name string
		4004	status int
		4005	backupWebsocket bool
		4006	}{
		4007	{name: "unauthorized to websocket", status: http.StatusUnauthorized, backupWebsocket: true},
		4008	{name: "unauthorized to http", status: http.StatusUnauthorized, backupWebsocket: false},
		4009	{name: "rate limit to websocket", status: http.StatusTooManyRequests, backupWebsocket: true},
		4010	{name: "rate limit to http", status: http.StatusTooManyRequests, backupWebsocket: false},
3308	}	4011	}
		4012	for _, tc := range tests {
		4013	t.Run(tc.name, func(t *testing.T) {
		4014	modelName := fmt.Sprintf("credential-failure-%d-%t-model", tc.status, tc.backupWebsocket)
		4015	selector := &orderedWebsocketSelector{order: []string{"auth-a", "auth-b"}}
		4016	executor := &websocketPinnedFailoverExecutor{failStatus: tc.status}
		4017	manager := coreauth.NewManager(nil, selector, nil)
		4018	manager.RegisterExecutor(executor)
		4019
		4020	authA := &coreauth.Auth{
		4021	ID: "auth-a",
		4022	Provider: executor.Identifier(),
		4023	Status: coreauth.StatusActive,
		4024	Attributes: map[string]string{"websockets": "true"},
		4025	}
		4026	if _, err := manager.Register(context.Background(), authA); err != nil {
		4027	t.Fatalf("Register auth A: %v", err)
		4028	}
		4029	authB := &coreauth.Auth{
		4030	ID: "auth-b",
		4031	Provider: executor.Identifier(),
		4032	Status: coreauth.StatusActive,
		4033	Attributes: map[string]string{"websockets": strconv.FormatBool(tc.backupWebsocket)},
		4034	}
		4035	if _, err := manager.Register(context.Background(), authB); err != nil {
		4036	t.Fatalf("Register auth B: %v", err)
		4037	}
		4038
		4039	registry.GetGlobalRegistry().RegisterClient(authA.ID, authA.Provider, []*registry.ModelInfo{{ID: modelName}})
		4040	registry.GetGlobalRegistry().RegisterClient(authB.ID, authB.Provider, []*registry.ModelInfo{{ID: modelName}})
		4041	t.Cleanup(func() {
		4042	registry.GetGlobalRegistry().UnregisterClient(authA.ID)
		4043	registry.GetGlobalRegistry().UnregisterClient(authB.ID)
		4044	})
		4045
		4046	base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
		4047	h := NewOpenAIResponsesAPIHandler(base)
		4048	router := gin.New()
		4049	router.GET("/v1/responses/ws", h.ResponsesWebsocket)
		4050
		4051	server := httptest.NewServer(router)
		4052	defer server.Close()
3309		4053
3310	authBPayloads := executor.Payloads("auth-b")	4054	wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses/ws"
3311	if len(authBPayloads) != 1 {	4055	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
3312	t.Fatalf("auth-b payload count = %d, want 1", len(authBPayloads))	4056	if err != nil {
		4057	t.Fatalf("dial websocket: %v", err)
		4058	}
		4059	defer func() { _ = conn.Close() }()
		4060
		4061	firstRequest := fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"}]}`, modelName)
		4062	if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(firstRequest)); errWrite != nil {
		4063	t.Fatalf("write first websocket message: %v", errWrite)
		4064	}
		4065	if _, payload, errRead := conn.ReadMessage(); errRead != nil \|\| gjson.GetBytes(payload, "type").String() != wsEventTypeCompleted {
		4066	t.Fatalf("first websocket response = %s, err=%v", payload, errRead)
		4067	}
		4068
		4069	secondRequest := `{"type":"response.create","previous_response_id":"resp-auth-a-1","input":[{"type":"message","id":"msg-2"}]}`
		4070	if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(secondRequest)); errWrite != nil {
		4071	t.Fatalf("write second websocket message: %v", errWrite)
		4072	}
		4073	_, _, errReadClose := conn.ReadMessage()
		4074	var replayClose *websocket.CloseError
		4075	if !errors.As(errReadClose, &replayClose) \|\| replayClose.Code != websocket.CloseServiceRestart \|\| replayClose.Text != wsHTTPReplayRequiredCloseReason {
		4076	t.Fatalf("credential failure response = %v, want replay close %d %q", errReadClose, websocket.CloseServiceRestart, wsHTTPReplayRequiredCloseReason)
		4077	}
		4078	if got := executor.AuthIDs(); len(got) != 2 \|\| got[0] != "auth-a" \|\| got[1] != "auth-a" {
		4079	t.Fatalf("selected auth IDs before replay = %v, want [auth-a auth-a]", got)
		4080	}
		4081
		4082	replayConn, _, errDialReplay := websocket.DefaultDialer.Dial(wsURL, nil)
		4083	if errDialReplay != nil {
		4084	t.Fatalf("dial replay websocket: %v", errDialReplay)
		4085	}
		4086	defer func() { _ = replayConn.Close() }()
		4087	fullReplay := fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"},{"type":"message","id":"out-auth-a-1"},{"type":"message","id":"msg-2"}]}`, modelName)
		4088	if errWrite := replayConn.WriteMessage(websocket.TextMessage, []byte(fullReplay)); errWrite != nil {
		4089	t.Fatalf("write full replay: %v", errWrite)
		4090	}
		4091	if _, replayPayload, errReadReplay := replayConn.ReadMessage(); errReadReplay != nil \|\| gjson.GetBytes(replayPayload, "type").String() != wsEventTypeCompleted {
		4092	t.Fatalf("full replay response = %s, err=%v", replayPayload, errReadReplay)
		4093	}
		4094	if got := executor.AuthIDs(); len(got) != 3 \|\| got[2] != "auth-b" {
		4095	t.Fatalf("selected auth IDs after replay = %v, want [auth-a auth-a auth-b]", got)
		4096	}
		4097	authBPayloads := executor.Payloads("auth-b")
		4098	if len(authBPayloads) != 1 {
		4099	t.Fatalf("auth-b payloads = %d, want 1", len(authBPayloads))
		4100	}
		4101	authBPayload := authBPayloads[0]
		4102	if gjson.GetBytes(authBPayload, "previous_response_id").Exists() \|\| len(gjson.GetBytes(authBPayload, "input").Array()) != 3 {
		4103	t.Fatalf("auth-b did not receive full replay: %s", authBPayload)
		4104	}
		4105	})
3313	}	4106	}
3314	authBPayload := authBPayloads[0]	4107	}
3315	if gjson.GetBytes(authBPayload, "previous_response_id").Exists() {	4108
3316	t.Fatalf("previous_response_id leaked after auth failover: %s", authBPayload)	4109	func TestShouldReplayResponsesWebsocketPinnedAuthFailure(t *testing.T) {
		4110	cases := []struct {
		4111	name string
		4112	err *interfaces.ErrorMessage
		4113	want bool
		4114	}{
		4115	{name: "nil", err: nil, want: false},
		4116	{name: "unauthorized", err: &interfaces.ErrorMessage{StatusCode: http.StatusUnauthorized}, want: true},
		4117	{name: "rate limit", err: &interfaces.ErrorMessage{StatusCode: http.StatusTooManyRequests}, want: true},
		4118	{name: "forbidden", err: &interfaces.ErrorMessage{StatusCode: http.StatusForbidden}, want: false},
		4119	{name: "service unavailable", err: &interfaces.ErrorMessage{StatusCode: http.StatusServiceUnavailable}, want: false},
3317	}	4120	}
3318	authBInput := gjson.GetBytes(authBPayload, "input").Raw	4121	for _, tc := range cases {
3319	if !strings.Contains(authBInput, `"id":"msg-1"`) \|\| !strings.Contains(authBInput, `"id":"msg-3"`) {	4122	t.Run(tc.name, func(t *testing.T) {
3320	t.Fatalf("auth-b replay input missing expected transcript items: %s", authBInput)	4123	if got := shouldReplayResponsesWebsocketPinnedAuthFailure(tc.err); got != tc.want {
		4124	t.Fatalf("shouldReplayResponsesWebsocketPinnedAuthFailure() = %v, want %v", got, tc.want)
		4125	}
		4126	})
3321	}	4127	}
3322	}	4128	}
3323		4129
@@ -3350,7 +4156,7 @@ type websocketPinnedPrematureCloseExecutor struct { · 26 unmodified lines
3350	payloads map[string][][]byte	4156	payloads map[string][][]byte
3351	}	4157	}
3352		4158
3353	func (e *websocketPinnedPrematureCloseExecutor) Identifier() string { return "test-provider" }	4159	func (e *websocketPinnedPrematureCloseExecutor) Identifier() string { return "xai" }
3354		4160
3355	func (e websocketPinnedPrematureCloseExecutor) Execute(context.Context, coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {	4161	func (e websocketPinnedPrematureCloseExecutor) Execute(context.Context, coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
3356	return coreexecutor.Response{}, errors.New("not implemented")	4162	return coreexecutor.Response{}, errors.New("not implemented")
@@ -3470,66 +4276,69 @@ func TestResponsesWebsocketReleasesPinnedAuthAfterStreamClosed408(t *testing.T) · 113 unmodified lines
3470	}	4276	}
3471	}()	4277	}()
3472		4278
3473	requests := []string{	4279	firstRequest := `{"type":"response.create","model":"stream-model","input":[{"type":"message","id":"msg-1"}]}`
3474	`{"type":"response.create","model":"stream-model","input":[{"type":"message","id":"msg-1"}]}`,	4280	if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(firstRequest)); errWrite != nil {
3475	`{"type":"response.create","previous_response_id":"resp-auth-a-1","input":[{"type":"message","id":"msg-2"}]}`,	4281	t.Fatalf("write first websocket message: %v", errWrite)
3476	`{"type":"response.create","previous_response_id":"resp-auth-a-1","input":[{"type":"message","id":"msg-3"}]}`,
3477	}	4282	}
3478	wantTypes := []string{wsEventTypeCompleted, wsEventTypeError, wsEventTypeCompleted}	4283	if _, payload, errRead := conn.ReadMessage(); errRead != nil \|\| gjson.GetBytes(payload, "type").String() != wsEventTypeCompleted {
3479	for i := range requests {	4284	t.Fatalf("first websocket response = %s, err=%v", payload, errRead)
3480	if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(requests[i])); errWrite != nil {	4285	}
3481	t.Fatalf("write websocket message %d: %v", i+1, errWrite)	4286
		4287	secondRequest := `{"type":"response.create","previous_response_id":"resp-auth-a-1","input":[{"type":"message","id":"msg-2"}]}`
		4288	if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(secondRequest)); errWrite != nil {
		4289	t.Fatalf("write second websocket message: %v", errWrite)
		4290	}
		4291	for {
		4292	_, payload, errRead := conn.ReadMessage()
		4293	if errRead != nil {
		4294	t.Fatalf("read stream-closed response: %v", errRead)
3482	}	4295	}
3483	if i == 1 {	4296	eventType := gjson.GetBytes(payload, "type").String()
3484	gotError := false	4297	if eventType == wsEventTypeError {
3485	for {	4298	if got := int(gjson.GetBytes(payload, "status").Int()); got != http.StatusRequestTimeout {
3486	_, payload, errReadMessage := conn.ReadMessage()	4299	t.Fatalf("stream-closed status = %d, want %d: %s", got, http.StatusRequestTimeout, payload)
3487	if errReadMessage != nil {
3488	t.Fatalf("read websocket message %d: %v", i+1, errReadMessage)
3489	}
3490	got := gjson.GetBytes(payload, "type").String()
3491	if got == wsEventTypeError {
3492	if int(gjson.GetBytes(payload, "status").Int()) != http.StatusRequestTimeout {
3493	t.Fatalf("stream-closed payload status = %d, want %d: %s", gjson.GetBytes(payload, "status").Int(), http.StatusRequestTimeout, payload)
3494	}
3495	gotError = true
3496	break
3497	}
3498	if got == wsEventTypeCompleted {
3499	t.Fatalf("message %d unexpectedly completed: %s", i+1, payload)
3500	}
3501	}	4300	}
3502	if !gotError {	4301	break
3503	t.Fatalf("message %d did not return stream-closed error", i+1)
3504	}
3505	continue
3506	}
3507	_, payload, errReadMessage := conn.ReadMessage()
3508	if errReadMessage != nil {
3509	t.Fatalf("read websocket message %d: %v", i+1, errReadMessage)
3510	}	4302	}
3511	if got := gjson.GetBytes(payload, "type").String(); got != wantTypes[i] {	4303	if eventType == wsEventTypeCompleted {
3512	t.Fatalf("message %d payload type = %s, want %s: %s", i+1, got, wantTypes[i], payload)	4304	t.Fatalf("stream-closed turn unexpectedly completed: %s", payload)
3513	}	4305	}
3514	}	4306	}
3515		4307
		4308	thirdDelta := `{"type":"response.create","previous_response_id":"resp-auth-a-1","input":[{"type":"message","id":"msg-3"}]}`
		4309	if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(thirdDelta)); errWrite != nil {
		4310	t.Fatalf("write third websocket message: %v", errWrite)
		4311	}
		4312	_, _, errReadClose := conn.ReadMessage()
		4313	var replayClose *websocket.CloseError
		4314	if !errors.As(errReadClose, &replayClose) \|\| replayClose.Code != websocket.CloseServiceRestart {
		4315	t.Fatalf("third websocket response error = %v, want replay close", errReadClose)
		4316	}
		4317	if got := executor.AuthIDs(); len(got) != 2 \|\| got[0] != "auth-a" \|\| got[1] != "auth-a" {
		4318	t.Fatalf("selected auth IDs before replay = %v, want [auth-a auth-a]", got)
		4319	}
		4320
		4321	replayConn, _, errDialReplay := websocket.DefaultDialer.Dial(wsURL, nil)
		4322	if errDialReplay != nil {
		4323	t.Fatalf("dial replay websocket: %v", errDialReplay)
		4324	}
		4325	defer func() { _ = replayConn.Close() }()
		4326	fullReplay := `{"type":"response.create","model":"stream-model","input":[{"type":"message","id":"msg-1"},{"type":"message","id":"out-auth-a-1"},{"type":"message","id":"msg-3"}]}`
		4327	if errWrite := replayConn.WriteMessage(websocket.TextMessage, []byte(fullReplay)); errWrite != nil {
		4328	t.Fatalf("write full replay: %v", errWrite)
		4329	}
		4330	if _, replayResponse, errReadReplay := replayConn.ReadMessage(); errReadReplay != nil \|\| gjson.GetBytes(replayResponse, "type").String() != wsEventTypeCompleted {
		4331	t.Fatalf("full replay response = %s, err=%v", replayResponse, errReadReplay)
		4332	}
3516	authIDs := executor.AuthIDs()	4333	authIDs := executor.AuthIDs()
3517	if len(authIDs) != 3 \|\| authIDs[0] != "auth-a" \|\| authIDs[1] != "auth-a" {	4334	if len(authIDs) != 3 \|\| authIDs[0] != "auth-a" \|\| authIDs[1] != "auth-a" {
3518	t.Fatalf("selected auth IDs = %v, want auth-a for first two turns", authIDs)	4335	t.Fatalf("selected auth IDs after replay = %v, want auth-a for the first two turns", authIDs)
3519	}	4336	}
3520
3521	replayAuthID := authIDs[2]	4337	replayAuthID := authIDs[2]
3522	replayPayloads := executor.Payloads(replayAuthID)	4338	replayPayloads := executor.Payloads(replayAuthID)
3523	if len(replayPayloads) == 0 {
3524	t.Fatalf("replay auth %s has no payloads", replayAuthID)
3525	}
3526	replayPayload := replayPayloads[len(replayPayloads)-1]	4339	replayPayload := replayPayloads[len(replayPayloads)-1]
3527	if gjson.GetBytes(replayPayload, "previous_response_id").Exists() {	4340	if gjson.GetBytes(replayPayload, "previous_response_id").Exists() \|\| len(gjson.GetBytes(replayPayload, "input").Array()) != 3 {
3528	t.Fatalf("previous_response_id leaked after stream-closed replay: %s", replayPayload)	4341	t.Fatalf("replay auth %s did not receive full replay: %s", replayAuthID, replayPayload)
3529	}
3530	replayInput := gjson.GetBytes(replayPayload, "input").Raw
3531	if !strings.Contains(replayInput, `"id":"msg-1"`) \|\| !strings.Contains(replayInput, `"id":"msg-3"`) {
3532	t.Fatalf("replay input missing expected transcript items: %s", replayInput)
3533	}	4342	}
3534	}	4343	}
3535		4344

sdk/api/handlers/openai/openai_responses_websocket_toolcall_repair.go+118−5

@@ -19,6 +19,7 @@ const ( · 18 unmodified lines
19	var defaultWebsocketToolOutputCache = newWebsocketToolOutputCache(0, websocketToolOutputCacheMaxPerSession)	19	var defaultWebsocketToolOutputCache = newWebsocketToolOutputCache(0, websocketToolOutputCacheMaxPerSession)
20	var defaultWebsocketToolCallCache = newWebsocketToolOutputCache(0, websocketToolOutputCacheMaxPerSession)	20	var defaultWebsocketToolCallCache = newWebsocketToolOutputCache(0, websocketToolOutputCacheMaxPerSession)
21	var defaultWebsocketToolSessionRefs = newWebsocketToolSessionRefCounter()	21	var defaultWebsocketToolSessionRefs = newWebsocketToolSessionRefCounter()
		22	var defaultWebsocketToolCacheTransactionMu sync.RWMutex
22		23
23	type websocketToolOutputCache struct {	24	type websocketToolOutputCache struct {
24	mu sync.Mutex	25	mu sync.Mutex
@@ -33,6 +34,14 @@ type websocketToolOutputSession struct { · 8 unmodified lines
33	order []string	34	order []string
34	}	35	}
35		36
		37	type responsesWebsocketToolCacheTurn struct {
		38	sessionKey string
		39	outputs map[string]json.RawMessage
		40	outputOrder []string
		41	calls map[string]json.RawMessage
		42	callOrder []string
		43	}
		44
36	func newWebsocketToolOutputCache(ttl time.Duration, maxPerSession int) *websocketToolOutputCache {	45	func newWebsocketToolOutputCache(ttl time.Duration, maxPerSession int) *websocketToolOutputCache {
37	if ttl < 0 {	46	if ttl < 0 {
38	ttl = websocketToolOutputCacheTTL	47	ttl = websocketToolOutputCacheTTL
@@ -196,6 +205,8 @@ func (c *websocketToolSessionRefCounter) release(sessionKey string) bool { · 157 unmodified lines
196	}	205	}
197		206
198	func retainResponsesWebsocketToolCaches(sessionKey string) {	207	func retainResponsesWebsocketToolCaches(sessionKey string) {
		208	defaultWebsocketToolCacheTransactionMu.Lock()
		209	defer defaultWebsocketToolCacheTransactionMu.Unlock()
199	if defaultWebsocketToolSessionRefs == nil {	210	if defaultWebsocketToolSessionRefs == nil {
200	return	211	return
201	}	212	}
@@ -203,13 +214,14 @@ func retainResponsesWebsocketToolCaches(sessionKey string) { · 1 unmodified line
203	}	214	}
204		215
205	func releaseResponsesWebsocketToolCaches(sessionKey string) {	216	func releaseResponsesWebsocketToolCaches(sessionKey string) {
		217	defaultWebsocketToolCacheTransactionMu.Lock()
		218	defer defaultWebsocketToolCacheTransactionMu.Unlock()
206	if defaultWebsocketToolSessionRefs == nil {	219	if defaultWebsocketToolSessionRefs == nil {
207	return	220	return
208	}	221	}
209	if !defaultWebsocketToolSessionRefs.release(sessionKey) {	222	if !defaultWebsocketToolSessionRefs.release(sessionKey) {
210	return	223	return
211	}	224	}
212
213	if defaultWebsocketToolOutputCache != nil {	225	if defaultWebsocketToolOutputCache != nil {
214	defaultWebsocketToolOutputCache.deleteSession(sessionKey)	226	defaultWebsocketToolOutputCache.deleteSession(sessionKey)
215	}	227	}
@@ -218,15 +230,114 @@ func releaseResponsesWebsocketToolCaches(sessionKey string) { · 2 unmodified lines
218	}	230	}
219	}	231	}
220		232
		233	func newResponsesWebsocketToolCacheTurn(sessionKey string) *responsesWebsocketToolCacheTurn {
		234	sessionKey = strings.TrimSpace(sessionKey)
		235	if sessionKey == "" {
		236	return nil
		237	}
		238	return &responsesWebsocketToolCacheTurn{
		239	sessionKey: sessionKey,
		240	outputs: make(map[string]json.RawMessage),
		241	calls: make(map[string]json.RawMessage),
		242	}
		243	}
		244
		245	func (t *responsesWebsocketToolCacheTurn) recordRequest(payload []byte) {
		246	if t == nil \|\| len(payload) == 0 {
		247	return
		248	}
		249	input := gjson.GetBytes(payload, "input")
		250	if !input.Exists() \|\| !input.IsArray() {
		251	return
		252	}
		253	for _, item := range input.Array() {
		254	t.recordItem(item)
		255	}
		256	}
		257
		258	func (t *responsesWebsocketToolCacheTurn) recordResponse(payload []byte) {
		259	if t == nil \|\| len(payload) == 0 {
		260	return
		261	}
		262	switch strings.TrimSpace(gjson.GetBytes(payload, "type").String()) {
		263	case "response.completed":
		264	output := gjson.GetBytes(payload, "response.output")
		265	if !output.Exists() \|\| !output.IsArray() {
		266	return
		267	}
		268	for _, item := range output.Array() {
		269	if isCompleteResponsesWebsocketToolCall(item) {
		270	t.recordItem(item)
		271	}
		272	}
		273	case "response.output_item.added", "response.output_item.done":
		274	item := gjson.GetBytes(payload, "item")
		275	if isCompleteResponsesWebsocketToolCall(item) {
		276	t.recordItem(item)
		277	}
		278	}
		279	}
		280
		281	func (t *responsesWebsocketToolCacheTurn) recordItem(item gjson.Result) {
		282	if t == nil \|\| !item.Exists() {
		283	return
		284	}
		285	callID := strings.TrimSpace(item.Get("call_id").String())
		286	if callID == "" \|\| strings.TrimSpace(item.Raw) == "" {
		287	return
		288	}
		289	raw := append(json.RawMessage(nil), item.Raw...)
		290	switch {
		291	case isResponsesToolCallOutputType(item.Get("type").String()):
		292	if _, exists := t.outputs[callID]; !exists {
		293	t.outputOrder = append(t.outputOrder, callID)
		294	}
		295	t.outputs[callID] = raw
		296	case isResponsesToolCallType(item.Get("type").String()):
		297	if _, exists := t.calls[callID]; !exists {
		298	t.callOrder = append(t.callOrder, callID)
		299	}
		300	t.calls[callID] = raw
		301	}
		302	}
		303
		304	func (t *responsesWebsocketToolCacheTurn) commit() {
		305	if t == nil \|\| t.sessionKey == "" {
		306	return
		307	}
		308	defaultWebsocketToolCacheTransactionMu.Lock()
		309	defer defaultWebsocketToolCacheTransactionMu.Unlock()
		310	if defaultWebsocketToolOutputCache != nil {
		311	for _, callID := range t.outputOrder {
		312	defaultWebsocketToolOutputCache.record(t.sessionKey, callID, t.outputs[callID])
		313	}
		314	}
		315	if defaultWebsocketToolCallCache != nil {
		316	for _, callID := range t.callOrder {
		317	defaultWebsocketToolCallCache.record(t.sessionKey, callID, t.calls[callID])
		318	}
		319	}
		320	}
		321
221	func repairResponsesWebsocketToolCalls(sessionKey string, payload []byte) []byte {	322	func repairResponsesWebsocketToolCalls(sessionKey string, payload []byte) []byte {
222	return repairResponsesWebsocketToolCallsWithCaches(defaultWebsocketToolOutputCache, defaultWebsocketToolCallCache, sessionKey, payload)	323	return repairResponsesWebsocketToolCallsWithCaches(defaultWebsocketToolOutputCache, defaultWebsocketToolCallCache, sessionKey, payload)
223	}	324	}
224		325
		326	func repairResponsesWebsocketToolCallsWithoutRecording(sessionKey string, payload []byte) []byte {
		327	defaultWebsocketToolCacheTransactionMu.RLock()
		328	defer defaultWebsocketToolCacheTransactionMu.RUnlock()
		329	return repairResponsesWebsocketToolCallsWithCachesMode(defaultWebsocketToolOutputCache, defaultWebsocketToolCallCache, sessionKey, payload, false)
		330	}
		331
225	func repairResponsesWebsocketToolCallsWithCache(cache *websocketToolOutputCache, sessionKey string, payload []byte) []byte {	332	func repairResponsesWebsocketToolCallsWithCache(cache *websocketToolOutputCache, sessionKey string, payload []byte) []byte {
226	return repairResponsesWebsocketToolCallsWithCaches(cache, nil, sessionKey, payload)	333	return repairResponsesWebsocketToolCallsWithCaches(cache, nil, sessionKey, payload)
227	}	334	}
228		335
229	func repairResponsesWebsocketToolCallsWithCaches(outputCache, callCache *websocketToolOutputCache, sessionKey string, payload []byte) []byte {	336	func repairResponsesWebsocketToolCallsWithCaches(outputCache, callCache *websocketToolOutputCache, sessionKey string, payload []byte) []byte {
		337	return repairResponsesWebsocketToolCallsWithCachesMode(outputCache, callCache, sessionKey, payload, true)
		338	}
		339
		340	func repairResponsesWebsocketToolCallsWithCachesMode(outputCache, callCache *websocketToolOutputCache, sessionKey string, payload []byte, record bool) []byte {
230	sessionKey = strings.TrimSpace(sessionKey)	341	sessionKey = strings.TrimSpace(sessionKey)
231	if sessionKey == "" \|\| outputCache == nil \|\| len(payload) == 0 {	342	if sessionKey == "" \|\| outputCache == nil \|\| len(payload) == 0 {
232	return payload	343	return payload
@@ -238,7 +349,7 @@ func repairResponsesWebsocketToolCallsWithCaches(outputCache, callCache *websock · 5 unmodified lines
238	}	349	}
239		350
240	allowOrphanOutputs := strings.TrimSpace(gjson.GetBytes(payload, "previous_response_id").String()) != ""	351	allowOrphanOutputs := strings.TrimSpace(gjson.GetBytes(payload, "previous_response_id").String()) != ""
241	updatedRaw, errRepair := repairResponsesToolCallsArray(outputCache, callCache, sessionKey, input.Raw, allowOrphanOutputs)	352	updatedRaw, errRepair := repairResponsesToolCallsArray(outputCache, callCache, sessionKey, input.Raw, allowOrphanOutputs, record)
242	if errRepair != nil \|\| updatedRaw == "" \|\| updatedRaw == input.Raw {	353	if errRepair != nil \|\| updatedRaw == "" \|\| updatedRaw == input.Raw {
243	return payload	354	return payload
244	}	355	}
@@ -250,7 +361,7 @@ func repairResponsesWebsocketToolCallsWithCaches(outputCache, callCache *websock · 5 unmodified lines
250	return updated	361	return updated
251	}	362	}
252		363
253	func repairResponsesToolCallsArray(outputCache, callCache *websocketToolOutputCache, sessionKey string, rawArray string, allowOrphanOutputs bool) (string, error) {	364	func repairResponsesToolCallsArray(outputCache, callCache *websocketToolOutputCache, sessionKey string, rawArray string, allowOrphanOutputs bool, record bool) (string, error) {
254	rawArray = strings.TrimSpace(rawArray)	365	rawArray = strings.TrimSpace(rawArray)
255	if rawArray == "" {	366	if rawArray == "" {
256	return "[]", nil	367	return "[]", nil
@@ -276,14 +387,16 @@ func repairResponsesToolCallsArray(outputCache, callCache *websocketToolOutputCa · 19 unmodified lines
276	continue	387	continue
277	}	388	}
278	outputPresent[callID] = struct{}{}	389	outputPresent[callID] = struct{}{}
279	outputCache.record(sessionKey, callID, item)	390	if record {
		391	outputCache.record(sessionKey, callID, item)
		392	}
280	case isResponsesToolCallType(itemType):	393	case isResponsesToolCallType(itemType):
281	callID := strings.TrimSpace(gjson.GetBytes(item, "call_id").String())	394	callID := strings.TrimSpace(gjson.GetBytes(item, "call_id").String())
282	if callID == "" {	395	if callID == "" {
283	continue	396	continue
284	}	397	}
285	callPresent[callID] = struct{}{}	398	callPresent[callID] = struct{}{}
286	if callCache != nil {	399	if record && callCache != nil {
287	callCache.record(sessionKey, callID, item)	400	callCache.record(sessionKey, callID, item)
288	}	401	}
289	}	402	}

sdk/cliproxy/executor/context.go+19−0

@@ -3,6 +3,7 @@ package executor · 2 unmodified lines
3	import "context"	3	import "context"
4		4
5	type downstreamWebsocketContextKey struct{}	5	type downstreamWebsocketContextKey struct{}
		6	type requireUpstreamWebsocketContextKey struct{}
6		7
7	// WithDownstreamWebsocket marks the current request as coming from a downstream websocket connection.	8	// WithDownstreamWebsocket marks the current request as coming from a downstream websocket connection.
8	func WithDownstreamWebsocket(ctx context.Context) context.Context {	9	func WithDownstreamWebsocket(ctx context.Context) context.Context {
@@ -21,3 +22,21 @@ func DownstreamWebsocket(ctx context.Context) bool { · 12 unmodified lines
21	enabled, ok := raw.(bool)	22	enabled, ok := raw.(bool)
22	return ok && enabled	23	return ok && enabled
23	}	24	}
		25
		26	// WithRequiredUpstreamWebsocket marks a request whose incremental context is valid only on the current upstream websocket.
		27	func WithRequiredUpstreamWebsocket(ctx context.Context) context.Context {
		28	if ctx == nil {
		29	ctx = context.Background()
		30	}
		31	return context.WithValue(ctx, requireUpstreamWebsocketContextKey{}, true)
		32	}
		33
		34	// RequiredUpstreamWebsocket reports whether falling back to an HTTP upstream would lose request context.
		35	func RequiredUpstreamWebsocket(ctx context.Context) bool {
		36	if ctx == nil {
		37	return false
		38	}
		39	raw := ctx.Value(requireUpstreamWebsocketContextKey{})
		40	enabled, ok := raw.(bool)
		41	return ok && enabled
		42	}

sdk/cliproxy/executor/websocket.go+29−0

@@ -0,0 +1,29 @@
	1	package executor
	2
	3	import (
	4	"errors"
	5	"net/http"
	6	)
	7
	8	// UpstreamWebsocketReplayRequiredError indicates that an incremental request
	9	// cannot safely continue because its upstream websocket is no longer reusable.
	10	type UpstreamWebsocketReplayRequiredError struct{}
	11
	12	func (*UpstreamWebsocketReplayRequiredError) Error() string {
	13	return `{"error":{"message":"upstream transport requires full HTTP replay","type":"server_error","code":"upstream_http_replay_required","status":426}}`
	14	}
	15
	16	func (*UpstreamWebsocketReplayRequiredError) StatusCode() int { return http.StatusUpgradeRequired }
	17
	18	func (*UpstreamWebsocketReplayRequiredError) IsRequestScoped() bool { return true }
	19
	20	// NewUpstreamWebsocketReplayRequiredError creates a request-scoped replay signal.
	21	func NewUpstreamWebsocketReplayRequiredError() error {
	22	return &UpstreamWebsocketReplayRequiredError{}
	23	}
	24
	25	// IsUpstreamWebsocketReplayRequired reports whether err is the internal replay signal.
	26	func IsUpstreamWebsocketReplayRequired(err error) bool {
	27	var replayErr *UpstreamWebsocketReplayRequiredError
	28	return errors.As(err, &replayErr)
	29	}

sdk/cliproxy/executor/websocket_test.go+25−0

@@ -0,0 +1,25 @@
	1	package executor
	2
	3	import (
	4	"fmt"
	5	"net/http"
	6	"testing"
	7	)
	8
	9	func TestUpstreamWebsocketReplayRequiredError(t *testing.T) {
	10	err := NewUpstreamWebsocketReplayRequiredError()
	11	if !IsUpstreamWebsocketReplayRequired(err) {
	12	t.Fatal("replay error was not recognized")
	13	}
	14	if !IsUpstreamWebsocketReplayRequired(fmt.Errorf("wrapped: %w", err)) {
	15	t.Fatal("wrapped replay error was not recognized")
	16	}
	17	statusErr, ok := err.(interface{ StatusCode() int })
	18	if !ok \|\| statusErr.StatusCode() != http.StatusUpgradeRequired {
	19	t.Fatalf("replay error = %T %v, want status 426", err, err)
	20	}
	21	requestErr, ok := err.(RequestScopedError)
	22	if !ok \|\| !requestErr.IsRequestScoped() {
	23	t.Fatalf("replay error = %T, want request scoped", err)
	24	}
	25	}