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

@@ -1139,6 +1139,7 @@ func (e CodexExecutor) Execute(ctx context.Context, auth cliproxyauth.Auth, re · 1138 unmodified lines
1139	1139	body = helps.SetStringIfDifferent(body, "model", baseModel)
1140	1140	body = helps.SetBoolIfDifferent(body, "stream", true)
1141	1141	body, _ = sjson.DeleteBytes(body, "previous_response_id")
	1142	+ body, _ = sjson.DeleteBytes(body, "generate")
1142	1143	body, _ = sjson.DeleteBytes(body, "prompt_cache_retention")
1143	1144	body, _ = sjson.DeleteBytes(body, "safety_identifier")
1144	1145	body, _ = sjson.DeleteBytes(body, "stream_options")
@@ -1408,6 +1409,7 @@ func (e CodexExecutor) ExecuteStream(ctx context.Context, auth cliproxyauth.Au · 263 unmodified lines
1408	1409	requestPath := helps.PayloadRequestPath(opts)
1409	1410	body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers)
1410	1411	body, _ = sjson.DeleteBytes(body, "previous_response_id")
	1412	+ body, _ = sjson.DeleteBytes(body, "generate")
1411	1413	body, _ = sjson.DeleteBytes(body, "prompt_cache_retention")
1412	1414	body, _ = sjson.DeleteBytes(body, "safety_identifier")
1413	1415	body, _ = sjson.DeleteBytes(body, "stream_options")
@@ -1580,6 +1582,7 @@ func (e CodexExecutor) CountTokens(ctx context.Context, auth cliproxyauth.Auth · 166 unmodified lines
1580	1582
1581	1583	body = helps.SetStringIfDifferent(body, "model", baseModel)
1582	1584	body, _ = sjson.DeleteBytes(body, "previous_response_id")
	1585	+ body, _ = sjson.DeleteBytes(body, "generate")
1583	1586	body, _ = sjson.DeleteBytes(body, "prompt_cache_retention")
1584	1587	body, _ = sjson.DeleteBytes(body, "safety_identifier")
1585	1588	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	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	372	func detachMismatchedWebsocketSessionConn(sess codexWebsocketSession, authID string, wsURL string) (websocket.Conn, *websocketConnectionCloser, string, string, cliproxyexecutor.ExecutionLifecycle) {
356	373	if sess == nil {
357	374	return nil, nil, "", "", nil
@@ -532,14 +549,25 @@ func (e CodexWebsocketsExecutor) Execute(ctx context.Context, auth cliproxyaut · 174 unmodified lines
532	549	}
533	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	564	if errDial != nil {
537	565	bodyErr := websocketHandshakeBody(respHS)
538	566	if respHS != nil {
539	567	helps.RecordAPIWebsocketUpgradeRejection(ctx, e.cfg, websocketUpgradeRequestLog(wsReqLog), respHS.StatusCode, respHS.Header.Clone(), bodyErr)
540	568	}
541	569	if respHS != nil && respHS.StatusCode == http.StatusUpgradeRequired {
542		− if opts.ExecutionLifecycle != nil {
	570	+ if opts.ExecutionLifecycle != nil \|\| cliproxyexecutor.DownstreamWebsocket(ctx) {
543	571	return resp, statusErr{code: respHS.StatusCode, msg: string(bodyErr)}
544	572	}
545	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	610	if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil {
583	611	errSend = mapCodexWebsocketWriteError(sess, conn, errSend)
584	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	621	e.invalidateUpstreamConn(sess, conn, "send_error", errSend)
586	622	if !shouldRetryCodexWebsocketSend(errSend) {
587	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	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	864	var upstreamHeaders http.Header
815	865	if respHS != nil {
816	866	upstreamHeaders = respHS.Header.Clone()
@@ -824,7 +874,7 @@ func (e CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth clipr · 7 unmodified lines
824	874	if sess != nil {
825	875	sess.reqMu.Unlock()
826	876	}
827		− if opts.ExecutionLifecycle != nil {
	877	+ if opts.ExecutionLifecycle != nil \|\| cliproxyexecutor.DownstreamWebsocket(ctx) {
828	878	return nil, statusErr{code: respHS.StatusCode, msg: string(bodyErr)}
829	879	}
830	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	914	errSend = mapCodexWebsocketWriteError(sess, conn, errSend)
865	915	helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend)
866	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	926	e.invalidateUpstreamConn(sess, conn, "send_error", errSend)
868	927	if !shouldRetryCodexWebsocketSend(errSend) {
869	928	sess.clearActive(conn, readCh)
@@ -1964,6 +2023,14 @@ func (e CodexWebsocketsExecutor) readUpstreamLoop(sess codexWebsocketSession, · 1094 unmodified lines
1964	2023	}
1965	2024
1966	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	2034	if sess == nil \|\| conn == nil {
1968	2035	return
1969	2036	}
@@ -1989,7 +2056,9 @@ func (e CodexWebsocketsExecutor) invalidateUpstreamConn(sess codexWebsocketSes · 19 unmodified lines
1989	2056	sess.connMu.Unlock()
1990	2057
1991	2058	logCodexWebsocketDisconnected(sessionID, authID, wsURL, reason, err)
1992		− sess.notifyUpstreamDisconnect(err)
	2059	+ if notify {
	2060	+ sess.notifyUpstreamDisconnect(err)
	2061	+ }
1993	2062	if closer != nil {
1994	2063	if errClose := closer.Close(); errClose != nil {
1995	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	2275	if cliproxyexecutor.DownstreamWebsocket(ctx) && codexWebsocketsEnabled(auth) {
2207	2276	return e.wsExec.Execute(ctx, auth, req, opts)
2208	2277	}
	2278	+ if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
	2279	+ return cliproxyexecutor.Response{}, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
	2280	+ }
2209	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	2288	if cliproxyexecutor.DownstreamWebsocket(ctx) && codexWebsocketsEnabled(auth) {
2217	2289	return e.wsExec.ExecuteStream(ctx, auth, req, opts)
2218	2290	}
	2291	+ if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
	2292	+ return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
	2293	+ }
2219	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	534	func TestCodexWebsocketsExecuteStreamPassesThroughUpstreamWebsocketPayloadForDownstreamWebsocket(t *testing.T) {
365	535	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}
366	536	capturedPayload := make(chan []byte, 1)

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

@@ -48,18 +48,21 @@ type xaiWebsocketIDStateStore struct { · 47 unmodified lines
48	48	}
49	49
50	50	type xaiWebsocketIDState struct {
51		− mu sync.Mutex
52		− downstreamToUpstream map[string]string
53		− sequence int
54		− transcriptInput []json.RawMessage
	51	+ requestMu sync.Mutex
	52	+ mu sync.Mutex
	53	+ downstreamToUpstream map[string]string
	54	+ sequence int
	55	+ transcriptInput []json.RawMessage
	56	+ replayCompactedTranscriptOnReset bool
55	57	}
56	58
57	59	type xaiWebsocketRequestIDMapper struct {
58		− state *xaiWebsocketIDState
59		− downstreamPreviousID string
60		− upstreamPreviousID string
61		− upstreamResponseID string
62		− downstreamResponseID string
	60	+ state *xaiWebsocketIDState
	61	+ downstreamPreviousID string
	62	+ upstreamPreviousID string
	63	+ upstreamResponseID string
	64	+ downstreamResponseID string
	65	+ replayedCompactedTranscript bool
63	66	}
64	67
65	68	func NewXAIWebsocketsExecutor(cfg config.Config) XAIWebsocketsExecutor {
@@ -177,20 +180,21 @@ func (s *xaiWebsocketIDState) prependTranscriptInput(payload []byte) []byte { · 111 unmodified lines
177	180	return out
178	181	}
179	182
180		−func (s *xaiWebsocketIDState) recordTranscriptTurn(requestPayload []byte, completedPayload []byte) {
	183	+func (s *xaiWebsocketIDState) recordTranscriptTurn(requestPayload []byte, completedPayload []byte, reset bool) {
181	184	if s == nil \|\| len(requestPayload) == 0 \|\| len(completedPayload) == 0 {
182	185	return
183	186	}
184	187	inputItems := xaiJSONRawMessages(gjson.GetBytes(requestPayload, "input"))
185	188	outputItems := xaiJSONRawMessages(gjson.GetBytes(completedPayload, "response.output"))
186		− if len(inputItems) == 0 && len(outputItems) == 0 {
187		− return
188		− }
189	189
190	190	s.mu.Lock()
191	191	defer s.mu.Unlock()
192		− if strings.TrimSpace(gjson.GetBytes(requestPayload, "previous_response_id").String()) == "" {
	192	+ if reset {
193	193	s.transcriptInput = nil
	194	+ s.replayCompactedTranscriptOnReset = false
	195	+ }
	196	+ if len(inputItems) == 0 && len(outputItems) == 0 {
	197	+ return
194	198	}
195	199	s.transcriptInput = append(s.transcriptInput, inputItems...)
196	200	s.transcriptInput = append(s.transcriptInput, outputItems...)
@@ -210,9 +214,34 @@ func (s *xaiWebsocketIDState) replaceTranscriptWithItems(items ...[]byte) { · 13 unmodified lines
210	214	}
211	215	s.mu.Lock()
212	216	s.transcriptInput = next
	217	+ s.replayCompactedTranscriptOnReset = len(next) > 0
213	218	s.mu.Unlock()
214	219	}
215	220
	221	+func (s *xaiWebsocketIDState) prependCompactedTranscriptOnReset(payload []byte) ([]byte, bool) {
	222	+ if s == nil \|\| len(payload) == 0 {
	223	+ return payload, false
	224	+ }
	225	+ s.mu.Lock()
	226	+ if !s.replayCompactedTranscriptOnReset \|\| len(s.transcriptInput) == 0 {
	227	+ s.mu.Unlock()
	228	+ return payload, false
	229	+ }
	230	+ prefix := make([]json.RawMessage, 0, len(s.transcriptInput))
	231	+ for _, item := range s.transcriptInput {
	232	+ prefix = append(prefix, bytes.Clone(item))
	233	+ }
	234	+ s.mu.Unlock()
	235	+
	236	+ current := xaiJSONRawMessages(gjson.GetBytes(payload, "input"))
	237	+ merged := append(prefix, current...)
	238	+ out, errSet := sjson.SetRawBytes(payload, "input", xaiMarshalRawMessages(merged))
	239	+ if errSet != nil {
	240	+ return payload, false
	241	+ }
	242	+ return out, true
	243	+}
	244	+
216	245	func xaiJSONRawMessages(result gjson.Result) []json.RawMessage {
217	246	if !result.Exists() \|\| !result.IsArray() {
218	247	return nil
@@ -243,7 +272,16 @@ func xaiMarshalRawMessages(items []json.RawMessage) []byte { · 24 unmodified lines
243	272	}
244	273
245	274	func (m *xaiWebsocketRequestIDMapper) upstreamRequestPayload(payload []byte) []byte {
246		− if m == nil \|\| len(payload) == 0 \|\| m.downstreamPreviousID == m.upstreamPreviousID {
	275	+ if m == nil \|\| len(payload) == 0 {
	276	+ return payload
	277	+ }
	278	+ if m.downstreamPreviousID == m.upstreamPreviousID {
	279	+ requestType := strings.TrimSpace(gjson.GetBytes(payload, "type").String())
	280	+ if m.downstreamPreviousID == "" && requestType == "response.append" && m.state != nil {
	281	+ out, replayed := m.state.prependCompactedTranscriptOnReset(payload)
	282	+ m.replayedCompactedTranscript = replayed
	283	+ return out
	284	+ }
247	285	return payload
248	286	}
249	287	if m.upstreamPreviousID == "" {
@@ -251,6 +289,7 @@ func (m *xaiWebsocketRequestIDMapper) upstreamRequestPayload(payload []byte) []b · 1 unmodified line
251	289	if errDelete == nil {
252	290	if m.downstreamPreviousID != "" && m.state != nil {
253	291	out = m.state.prependTranscriptInput(out)
	292	+ m.replayedCompactedTranscript = true
254	293	}
255	294	return out
256	295	}
@@ -396,8 +435,30 @@ func (e XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth cliprox · 139 unmodified lines
396	435	if stateSessionID == "" {
397	436	stateSessionID = executionSessionID
398	437	}
399		− idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload)
	438	+ state := getXAIWebsocketIDState(e.idStore, stateSessionID)
	439	+ stateRequestLocked := false
	440	+ stateRequestLockTransferred := false
	441	+ if executionSessionID == "" && state != nil {
	442	+ state.requestMu.Lock()
	443	+ stateRequestLocked = true
	444	+ }
	445	+ defer func() {
	446	+ if stateRequestLocked && !stateRequestLockTransferred {
	447	+ state.requestMu.Unlock()
	448	+ }
	449	+ }()
400	450	if xaiInputHasItemType(req.Payload, "compaction_trigger") {
	451	+ if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
	452	+ return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
	453	+ }
	454	+ if executionSessionID != "" {
	455	+ sess := e.getOrCreateSession(executionSessionID)
	456	+ if sess != nil {
	457	+ sess.reqMu.Lock()
	458	+ defer sess.reqMu.Unlock()
	459	+ }
	460	+ }
	461	+ idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload)
401	462	return e.executeCompactionTriggerFromWebsocketContext(ctx, auth, req, opts, idMapper)
402	463	}
403	464
@@ -437,6 +498,7 @@ func (e XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth cliprox · 33 unmodified lines
437	498	sess.reqMu.Lock()
438	499	}
439	500	}
	501	+ idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload)
440	502	if idMapper != nil {
441	503	if websocketSessionTargetChanged(sess, authID, wsURL) {
442	504	idMapper.upstreamPreviousID = ""
@@ -447,6 +509,9 @@ func (e XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth cliprox · 4 unmodified lines
447	509
448	510	wsHeaders := applyXAIWebsocketHeaders(http.Header{}, auth, token, prepared.sessionID)
449	511	wsReqBody := buildXAIWebsocketRequestBody(prepared.body)
	512	+ requestType := strings.TrimSpace(gjson.GetBytes(req.Payload, "type").String())
	513	+ transcriptReset := strings.TrimSpace(gjson.GetBytes(wsReqBody, "previous_response_id").String()) == "" &&
	514	+ (requestType != "response.append" \|\| (idMapper != nil && idMapper.replayedCompactedTranscript))
450	515	warmupRequest := xaiWebsocketGenerateFalse(wsReqBody)
451	516
452	517	wsReqLog := helps.UpstreamRequestLog{
@@ -463,7 +528,21 @@ func (e XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth cliprox · 10 unmodified lines
463	528	helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog)
464	529	logXAIWebsocketRequest(executionSessionID, authID, wsURL, wsReqBody)
465	530
466		− conn, closer, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
	531	+ var conn *websocket.Conn
	532	+ var closer *websocketConnectionCloser
	533	+ var respHS *http.Response
	534	+ var errDial error
	535	+ if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
	536	+ conn, closer = existingWebsocketSessionConn(sess, authID, wsURL)
	537	+ if conn == nil {
	538	+ if sess != nil {
	539	+ sess.reqMu.Unlock()
	540	+ }
	541	+ return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
	542	+ }
	543	+ } else {
	544	+ conn, closer, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders)
	545	+ }
467	546	var upstreamHeaders http.Header
468	547	if respHS != nil {
469	548	upstreamHeaders = respHS.Header.Clone()
@@ -508,6 +587,15 @@ func (e XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth cliprox · 38 unmodified lines
508	587	errSend = mapXAIWebsocketWriteError(sess, conn, errSend)
509	588	helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend)
510	589	if sess != nil {
	590	+ if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
	591	+ e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "send_error", errSend)
	592	+ sess.clearActive(conn, readCh)
	593	+ sess.reqMu.Unlock()
	594	+ if !shouldRetryXAIWebsocketSend(errSend) {
	595	+ return nil, errSend
	596	+ }
	597	+ return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
	598	+ }
511	599	e.invalidateUpstreamConn(sess, conn, "send_error", errSend)
512	600	if !shouldRetryXAIWebsocketSend(errSend) {
513	601	sess.clearActive(conn, readCh)
@@ -570,7 +658,13 @@ func (e XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth cliprox · 56 unmodified lines
570	658	}
571	659
572	660	out := make(chan cliproxyexecutor.StreamChunk)
	661	+ if stateRequestLocked {
	662	+ stateRequestLockTransferred = true
	663	+ }
573	664	go func() {
	665	+ if stateRequestLocked {
	666	+ defer state.requestMu.Unlock()
	667	+ }
574	668	terminateReason := "completed"
575	669	var terminateErr error
576	670
@@ -677,6 +771,10 @@ func (e XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth cliprox · 100 unmodified lines
677	771	case "response.created":
678	772	if warmupRequest {
679	773	warmupCompletedPayload = buildXAIWebsocketWarmupCompletedPayload(payload)
	774	+ if idMapper != nil && idMapper.state != nil && !recordedTranscript {
	775	+ idMapper.state.recordTranscriptTurn(wsReqBody, warmupCompletedPayload, transcriptReset)
	776	+ recordedTranscript = true
	777	+ }
680	778	logXAIWebsocketWarmupCompleted(executionSessionID, authID, wsURL, payload)
681	779	}
682	780	case "response.output_item.done":
@@ -690,7 +788,7 @@ func (e XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth cliprox · 7 unmodified lines
690	788	payload = xaiNormalizeReasoningSummaryData(payload)
691	789	cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, payload)
692	790	if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript {
693		− idMapper.state.recordTranscriptTurn(wsReqBody, payload)
	791	+ idMapper.state.recordTranscriptTurn(wsReqBody, payload, transcriptReset)
694	792	recordedTranscript = true
695	793	}
696	794	case "response.done":
@@ -699,7 +797,7 @@ func (e XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth cliprox · 2 unmodified lines
699	797	reporter.Publish(ctx, detail)
700	798	}
701	799	if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript {
702		− idMapper.state.recordTranscriptTurn(wsReqBody, payload)
	800	+ idMapper.state.recordTranscriptTurn(wsReqBody, payload, transcriptReset)
703	801	recordedTranscript = true
704	802	}
705	803	}
@@ -793,8 +891,11 @@ func (e *XAIWebsocketsExecutor) executeCompactionTriggerFromWebsocketContext(ctx · 87 unmodified lines
793	891	return nil, err
794	892	}
795	893
796		− responseID := xaiCompactionResponseID(data)
797		− idMapper.state.replaceTranscriptWithItems(xaiCompactionOutputItem(data, responseID))
	894	+ responseID, compactionItem, errValidate := validateXAIWebsocketCompactionResponse(data)
	895	+ if errValidate != nil {
	896	+ return nil, errValidate
	897	+ }
	898	+ idMapper.state.replaceTranscriptWithItems(compactionItem)
798	899	idMapper.state.mapDownstreamToUpstream(responseID, "")
799	900
800	901	headers = headers.Clone()
@@ -812,6 +913,30 @@ func (e *XAIWebsocketsExecutor) executeCompactionTriggerFromWebsocketContext(ctx · 11 unmodified lines
812	913	return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out}, nil
813	914	}
814	915
	916	+func validateXAIWebsocketCompactionResponse(data []byte) (string, []byte, error) {
	917	+ if len(data) == 0 \|\| !json.Valid(data) {
	918	+ return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction returned invalid JSON"}
	919	+ }
	920	+ responseIDResult := gjson.GetBytes(data, "id")
	921	+ output := gjson.GetBytes(data, "output")
	922	+ if responseIDResult.Type != gjson.String \|\| strings.TrimSpace(responseIDResult.String()) == "" \|\| !output.Exists() \|\| !output.IsArray() {
	923	+ return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction response is missing compacted state"}
	924	+ }
	925	+ items := output.Array()
	926	+ if len(items) == 0 {
	927	+ return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction response is missing compacted state"}
	928	+ }
	929	+ item := items[0]
	930	+ itemType := item.Get("type")
	931	+ encryptedContent := item.Get("encrypted_content")
	932	+ if item.Type != gjson.JSON \|\| itemType.Type != gjson.String \|\| strings.TrimSpace(itemType.String()) != "compaction" \|\|
	933	+ encryptedContent.Type != gjson.String \|\| strings.TrimSpace(encryptedContent.String()) == "" {
	934	+ return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction response is missing compacted state"}
	935	+ }
	936	+ normalizedResponseID := xaiCompactionResponseID(data)
	937	+ return normalizedResponseID, xaiCompactionOutputItem(data, normalizedResponseID), nil
	938	+}
	939	+
815	940	func buildXAIWebsocketCompactionPayload(payload []byte, transcriptInput []byte) ([]byte, error) {
816	941	if len(payload) == 0 {
817	942	payload = []byte(`{}`)
@@ -1497,6 +1622,9 @@ func (e XAIAutoExecutor) ExecuteStream(ctx context.Context, auth cliproxyauth. · 679 unmodified lines
1497	1622	if cliproxyexecutor.DownstreamWebsocket(ctx) && xaiWebsocketsEnabled(auth) {
1498	1623	return e.wsExec.ExecuteStream(ctx, auth, req, opts)
1499	1624	}
	1625	+ if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) {
	1626	+ return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
	1627	+ }
1500	1628	return e.httpExec.ExecuteStream(ctx, auth, req, opts)
1501	1629	}
1502	1630

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	81	func TestMapXAIWebsocketWriteErrorStopsRetryForMessageTooBig(t *testing.T) {
38	82	networkWriteErr := errors.New("write: broken pipe")
39	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	1414	func TestBuildXAIWebsocketRequestBodySetsStoreAndKeepsPromptCacheKey(t *testing.T) {
1250	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	65	type pinnedAuthContextKey struct{}
66	66	type selectedAuthCallbackContextKey struct{}
	67	+type preparedModelRouteContextKey struct{}
67	68	type executionSessionContextKey struct{}
68	69	type disallowFreeAuthContextKey struct{}
69	70
@@ -142,6 +143,26 @@ func WithSelectedAuthIDCallback(ctx context.Context, callback func(string)) cont · 72 unmodified lines
142	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	166	// WithExecutionSessionID returns a child context tagged with a long-lived execution session ID.
146	167	func WithExecutionSessionID(ctx context.Context, sessionID string) context.Context {
147	168	sessionID = strings.TrimSpace(sessionID)
@@ -1144,7 +1165,10 @@ func (h *BaseAPIHandler) executeStreamWithAuthManager(ctx context.Context, handl · 996 unmodified lines
1144	1165
1145	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	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	1172	responseProtocol := modelExecutionResponseProtocol(entryProtocol, exitProtocol)
1149	1173	if errMsg := validateNativeInteractionsExecution(entryProtocol, execOptions, routeDecision); errMsg != nil {
1150	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	499	func TestExecuteModelPropagatesRouterSkipPluginID(t *testing.T) {
468	500	model := "model-execution-router-skip-model"
469	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	35	const (
36		− wsRequestTypeCreate = "response.create"
37		− wsRequestTypeAppend = "response.append"
38		− wsEventTypeError = "error"
39		− wsEventTypeCompleted = "response.completed"
40		− wsEventTypeDone = "response.done"
41		− wsDoneMarker = "[DONE]"
42		− wsTurnStateHeader = "x-codex-turn-state"
43		− wsTimelineBodyKey = "WEBSOCKET_TIMELINE_OVERRIDE"
44		− wsCloseReasonMaxBytes = 123
	36	+ wsRequestTypeCreate = "response.create"
	37	+ wsRequestTypeAppend = "response.append"
	38	+ wsEventTypeError = "error"
	39	+ wsEventTypeCompleted = "response.completed"
	40	+ wsEventTypeDone = "response.done"
	41	+ wsDoneMarker = "[DONE]"
	42	+ wsTurnStateHeader = "x-codex-turn-state"
	43	+ wsTimelineBodyKey = "WEBSOCKET_TIMELINE_OVERRIDE"
	44	+ wsCloseReasonMaxBytes = 123
	45	+ wsHTTPReplayRequiredCloseReason = "upstream requires HTTP replay"
	46	+ responsesWebsocketUpstreamModeUnknown = ""
	47	+ responsesWebsocketUpstreamModeWS = "websocket"
	48	+ responsesWebsocketUpstreamModeHTTP = "http"
45	49
46	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	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	89	code := 0
78	90	reason := ""
79	91	var closeErr *websocket.CloseError
@@ -85,7 +97,6 @@ func websocketClosePayloadForUpstreamError(err error) (bool, []byte) { · 5 unmodified lines
85	97	StatusCode() int
86	98	}
87	99	var statusErr statusCoder
88		− errText := err.Error()
89	100	if !errors.As(err, &statusErr) \|\| statusErr.StatusCode() != http.StatusRequestEntityTooLarge \|\|
90	101	gjson.Get(errText, "error.code").String() != "message_too_big" {
91	102	return false, nil
@@ -419,20 +430,35 @@ func (h OpenAIResponsesAPIHandler) ResponsesWebsocket(c gin.Context) { · 327 unmodified lines
419	430	// Preserve independent upstream auth affinity when a downstream session switches providers.
420	431	pinnedAuthByProvider := make(map[string]responsesWebsocketPinnedAuthState)
421	432	passthroughModelName := ""
	433	+ upstreamMode := responsesWebsocketUpstreamModeUnknown
	434	+ upstreamWebsocketAuthID := ""
422	435	sessionAuthByIDWithSource := func(authID string) (*coreauth.Auth, bool, bool) {
423	436	if h == nil \|\| h.AuthManager == nil {
424	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	444	if auth, ok := h.AuthManager.GetExecutionSessionAuthByID(passthroughSessionID, authID); ok {
427	445	return auth, true, true
428	446	}
429		− auth, ok := h.AuthManager.GetByID(authID)
430		− return auth, false, ok
	447	+ return nil, false, false
431	448	}
432	449	sessionAuthByID := func(authID string) (*coreauth.Auth, bool) {
433	450	auth, _, ok := sessionAuthByIDWithSource(authID)
434	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	462	rememberPinnedAuth := func(authID string, modelName string) {
437	463	authID = strings.TrimSpace(authID)
438	464	auth, ok := sessionAuthByID(authID)
@@ -454,7 +480,6 @@ func (h OpenAIResponsesAPIHandler) ResponsesWebsocket(c gin.Context) { · 15 unmodified lines
454	480	}
455	481	pinnedAuthID = ""
456	482	}
457		− forceTranscriptReplayNextRequest := false
458	483
459	484	for {
460	485	msgType, payload, errReadMessage := conn.ReadMessage()
@@ -488,6 +513,13 @@ func (h OpenAIResponsesAPIHandler) ResponsesWebsocket(c gin.Context) { · 27 unmodified lines
488	513	if requestModelName == "" {
489	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	523	if pinnedAuthID != "" {
492	524	pinnedAuth, homeRuntime, ok := sessionAuthByIDWithSource(pinnedAuthID)
493	525	providerKey := ""
@@ -514,15 +546,40 @@ func (h OpenAIResponsesAPIHandler) ResponsesWebsocket(c gin.Context) { · 20 unmodified lines
514	546	}
515	547	}
516	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	577	if explicitRequestModelName != "" && !useUpstreamWebsocketPassthrough {
518	578	passthroughModelName = ""
519	579	}
520		− allowIncrementalInputWithPreviousResponseID := false
	580	+
521	581	allowCompactionReplayBypass := false
522		− if !useUpstreamWebsocketPassthrough {
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.
	582	+ if !nativeWebsocketPassthrough {
526	583	if pinnedAuthID != "" {
527	584	if pinnedAuth, ok := sessionAuthByID(pinnedAuthID); ok && pinnedAuth != nil {
528	585	allowCompactionReplayBypass = responsesWebsocketAuthSupportsCompactionReplay(pinnedAuth)
@@ -535,8 +592,10 @@ func (h OpenAIResponsesAPIHandler) ResponsesWebsocket(c gin.Context) { · 6 unmodified lines
535	592	var requestJSON []byte
536	593	var updatedLastRequest []byte
537	594	var errMsg *interfaces.ErrorMessage
538		− if useUpstreamWebsocketPassthrough {
	595	+ if nativeWebsocketPassthrough {
539	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	599	} else {
541	600	requestJSON, updatedLastRequest, errMsg = normalizeResponsesWebsocketRequestWithIncrementalState(
542	601	payload,
@@ -544,7 +603,7 @@ func (h OpenAIResponsesAPIHandler) ResponsesWebsocket(c gin.Context) { · 1 unmodified line
544	603	lastResponseOutput,
545	604	lastResponseID,
546	605	lastResponsePendingToolCallIDs,
547		− allowIncrementalInputWithPreviousResponseID,
	606	+ false,
548	607	allowCompactionReplayBypass,
549	608	)
550	609	}
@@ -570,7 +629,7 @@ func (h OpenAIResponsesAPIHandler) ResponsesWebsocket(c gin.Context) { · 19 unmodified lines
570	629	}
571	630	continue
572	631	}
573		− if !useUpstreamWebsocketPassthrough && shouldHandleResponsesWebsocketPrewarmLocally(payload, lastRequest, allowIncrementalInputWithPreviousResponseID) {
	632	+ if !useUpstreamWebsocketPassthrough && shouldHandleResponsesWebsocketPrewarmLocally(payload, lastRequest, false) {
574	633	if updated, errDelete := sjson.DeleteBytes(requestJSON, "generate"); errDelete == nil {
575	634	requestJSON = updated
576	635	}
@@ -588,84 +647,120 @@ func (h OpenAIResponsesAPIHandler) ResponsesWebsocket(c gin.Context) { · 11 unmodified lines
588	647	continue
589	648	}
590	649
	650	+ toolCacheTurn := newResponsesWebsocketToolCacheTurn(downstreamSessionKey)
591	651	previousLastRequest := bytes.Clone(lastRequest)
592	652	previousLastResponseOutput := bytes.Clone(lastResponseOutput)
593	653	previousLastResponseID := lastResponseID
594	654	previousLastResponsePendingToolCallIDs := append([]string(nil), lastResponsePendingToolCallIDs...)
595		− forcedTranscriptReplay := forceTranscriptReplayNextRequest
596		− if useUpstreamWebsocketPassthrough {
	655	+ if nativeWebsocketPassthrough {
597	656	if modelName := strings.TrimSpace(gjson.GetBytes(requestJSON, "model").String()); modelName != "" {
598	657	passthroughModelName = modelName
599	658	}
600		− if forcedTranscriptReplay {
601		− forceTranscriptReplayNextRequest = false
602		− }
603	659	} else {
604		− requestJSON = repairResponsesWebsocketToolCalls(downstreamSessionKey, requestJSON)
	660	+ toolCacheTurn.recordRequest(requestJSON)
	661	+ requestJSON = repairResponsesWebsocketToolCallsWithoutRecording(downstreamSessionKey, requestJSON)
605	662	requestJSON = dedupeResponsesWebsocketInputItemsByID(requestJSON)
606	663	updatedLastRequest = bytes.Clone(requestJSON)
607	664	lastRequest = updatedLastRequest
608		− if forcedTranscriptReplay {
609		− forceTranscriptReplayNextRequest = false
610		− }
611	665	}
612	666
613	667	modelName := gjson.GetBytes(requestJSON, "model").String()
614	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	673	cliCtx = cliproxyexecutor.WithDownstreamWebsocket(cliCtx)
	674	+ if nativeWebsocketPassthrough && requestRequiresCurrentUpstreamWebsocket {
	675	+ cliCtx = cliproxyexecutor.WithRequiredUpstreamWebsocket(cliCtx)
	676	+ }
617	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	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	695	dataChan, _, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, requestJSON, "")
637		−
638		− completedOutput, completedResponseID, completedPendingToolCallIDs, forwardErrMsg, errForward := h.forwardResponsesWebsocket(c, writer, cliCancel, dataChan, errChan, wsTimelineLog, passthroughSessionID)
	696	+ if !selectedAuthObserved {
	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	721	if errForward != nil {
640	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	726	return
643	727	}
644		− if forwardErrMsg == nil && !useUpstreamWebsocketPassthrough && lastAttemptedAuthID != "" {
645		− if selectedAuth, ok := sessionAuthByID(lastAttemptedAuthID); ok && selectedAuth != nil {
646		− if websocketUpstreamSupportsIncrementalInput(selectedAuth.Attributes, selectedAuth.Metadata) {
647		− rememberPinnedAuth(lastAttemptedAuthID, modelName)
648		− } else if pinnedAuthID != "" {
649		− if pinnedAuth, ok := sessionAuthByID(pinnedAuthID); ok && pinnedAuth != nil && websocketUpstreamSupportsIncrementalInput(pinnedAuth.Attributes, pinnedAuth.Metadata) {
650		− rememberPinnedAuth(lastAttemptedAuthID, modelName)
651		− }
652		− }
	728	+ if forwardErrMsg != nil {
	729	+ lastRequest = previousLastRequest
	730	+ lastResponseOutput = previousLastResponseOutput
	731	+ lastResponseID = previousLastResponseID
	732	+ lastResponsePendingToolCallIDs = previousLastResponsePendingToolCallIDs
	733	+ if pinnedAuthAttempted && shouldReleaseResponsesWebsocketPinnedAuth(forwardErrMsg) {
	734	+ forgetPinnedAuth()
653	735	}
654		− }
655		− if shouldReleaseResponsesWebsocketPinnedAuth(forwardErrMsg) {
656		− forgetPinnedAuth()
657		− forceTranscriptReplayNextRequest = true
658		− if useUpstreamWebsocketPassthrough {
659		− passthroughModelName = ""
660		− } else {
661		− lastRequest = previousLastRequest
662		− lastResponseOutput = previousLastResponseOutput
663		− lastResponseID = previousLastResponseID
664		− lastResponsePendingToolCallIDs = previousLastResponsePendingToolCallIDs
	736	+ if replayPinnedAuthFailure(forwardErrMsg) {
	737	+ replayErr := responsesWebsocketHTTPReplayRequiredError()
	738	+ wsTerminateErr = replayErr
	739	+ matched, errClose := writer.closeForUpstreamError(replayErr)
	740	+ if !matched {
	741	+ _ = conn.Close()
	742	+ } else if errClose != nil && !errors.Is(errClose, websocket.ErrCloseSent) {
	743	+ log.Debugf("responses websocket: credential replay close failed id=%s error=%v", passthroughSessionID, errClose)
	744	+ }
	745	+ return
665	746	}
666	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	764	lastResponseOutput = completedOutput
670	765	lastResponseID = strings.TrimSpace(completedResponseID)
671	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	785	func websocketClientAddress(c *gin.Context) string {
677	786	if c == nil \|\| c.Request == nil {
678	787	return ""
@@ -694,6 +803,15 @@ func websocketUpgradeHeaders(req *http.Request) http.Header { · 15 unmodified lines
694	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	815	func normalizeResponsesWebsocketRequest(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte) ([]byte, []byte, *interfaces.ErrorMessage) {
698	816	return normalizeResponsesWebsocketRequestWithMode(rawJSON, lastRequest, lastResponseOutput, true, true)
699	817	}
@@ -1570,6 +1688,11 @@ func normalizeJSONArrayRaw(raw []byte) string { · 870 unmodified lines
1570	1688	return "[]"
1571	1689	}
1572	1690
	1691	+type responsesWebsocketForwardOptions struct {
	1692	+ toolCacheTurn *responsesWebsocketToolCacheTurn
	1693	+ suppressError func(*interfaces.ErrorMessage) bool
	1694	+}
	1695	+
1573	1696	func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket(
1574	1697	c *gin.Context,
1575	1698	writer *responsesWebsocketWriter,
@@ -1578,7 +1701,13 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( · 2 unmodified lines
1578	1701	errs <-chan *interfaces.ErrorMessage,
1579	1702	wsTimelineLog websocketTimelineAppender,
1580	1703	sessionID string,
	1704	+ options ...responsesWebsocketForwardOptions,
1581	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	1711	completed := false
1583	1712	completedOutput := []byte("[]")
1584	1713	completedResponseID := ""
@@ -1602,6 +1731,10 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( · 17 unmodified lines
1602	1731	}
1603	1732	if errMsg != nil {
1604	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	1738	markAPIResponseTimestamp(c)
1606	1739	if matched, errClose := writer.closeForUpstreamError(errMsg.Error); matched {
1607	1740	cancel(errMsg.Error)
@@ -1676,7 +1809,11 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( · 68 unmodified lines
1676	1809	if isResponsesWebsocketCompletionEvent(eventType) {
1677	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	1817	recordPendingToolCallIDsFromPayload(pendingToolCallIDs, payloads[i])
1681	1818	var payloadErrMsg *interfaces.ErrorMessage
1682	1819	if eventType == wsEventTypeError {
@@ -1684,6 +1821,10 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( · 1 unmodified line
1684	1821	if h != nil {
1685	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	1828	} else if isResponsesWebsocketCompletionEvent(eventType) {
1688	1829	completed = true
1689	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	1861	if errMsg == nil {
1721		− return false
	1862	+ return 0
1722	1863	}
1723	1864	status := errMsg.StatusCode
1724	1865	if status <= 0 && errMsg.Error != nil {
@@ -1726,7 +1867,23 @@ func shouldReleaseResponsesWebsocketPinnedAuth(errMsg *interfaces.ErrorMessage) · 1 unmodified line
1726	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	1887	case http.StatusUnauthorized,
1731	1888	http.StatusPaymentRequired,
1732	1889	http.StatusForbidden,

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

@@ -9,6 +9,7 @@ import ( · 8 unmodified lines
9	9	"maps"
10	10	"net/http"
11	11	"net/http/httptest"
	12	+ "strconv"
12	13	"strings"
13	14	"sync"
14	15	"sync/atomic"
@@ -27,6 +28,7 @@ import ( · 12 unmodified lines
27	28	"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry"
28	29	coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
29	30	sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
	31	+ "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
30	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	210	func TestWriteWebsocketCloseForUpstreamErrorMirrorsMessageTooBig(t *testing.T) {
164	211	tests := []struct {
165	212	name string
@@ -439,6 +486,22 @@ type websocketProviderCaptureExecutor struct { · 273 unmodified lines
439	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	505	type websocketCompactionCaptureExecutor struct {
443	506	mu sync.Mutex
444	507	streamPayloads [][]byte
@@ -481,10 +544,11 @@ type websocketAuthCaptureExecutor struct { · 36 unmodified lines
481	544	}
482	545
483	546	type websocketPinnedFailoverExecutor struct {
484		− mu sync.Mutex
485		− authIDs []string
486		− calls map[string]int
487		− payloads map[string][][]byte
	547	+ mu sync.Mutex
	548	+ failStatus int
	549	+ authIDs []string
	550	+ calls map[string]int
	551	+ payloads map[string][][]byte
488	552	}
489	553
490	554	type websocketBootstrapFallbackExecutor struct {
@@ -494,12 +558,21 @@ type websocketBootstrapFallbackExecutor struct { · 3 unmodified lines
494	558	}
495	559
496	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	573	mu sync.Mutex
498		− provider string
499		− authIDs []string
500	574	payloads [][]byte
501		− done chan struct{}
502		− doneOnce sync.Once
	575	+ calls int
503	576	}
504	577
505	578	type websocketPinnedFailoverStatusError struct {
@@ -534,7 +607,7 @@ func (e *websocketBootstrapFallbackExecutor) ExecuteStream(_ context.Context, au · 28 unmodified lines
534	607	chunks := make(chan coreexecutor.StreamChunk, 1)
535	608	if authID == "auth-ws" {
536	609	chunks <- coreexecutor.StreamChunk{Err: websocketPinnedFailoverStatusError{
537		− status: http.StatusServiceUnavailable,
	610	+ status: http.StatusUpgradeRequired,
538	611	msg: `{"error":{"message":"websocket bootstrap failed","type":"server_error","code":"ws_failed"}}`,
539	612	}}
540	613	close(chunks)
@@ -586,18 +659,29 @@ func (e websocketDirectCaptureExecutor) Execute(context.Context, coreauth.Auth · 45 unmodified lines
586	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	663	authID := ""
591	664	if auth != nil {
592	665	authID = auth.ID
593	666	}
594	667	e.mu.Lock()
595	668	e.authIDs = append(e.authIDs, authID)
	669	+ e.models = append(e.models, req.Model)
596	670	e.payloads = append(e.payloads, bytes.Clone(req.Payload))
	671	+ e.requiredUpstreamWebsocket = append(e.requiredUpstreamWebsocket, coreexecutor.RequiredUpstreamWebsocket(ctx))
597	672	count := len(e.payloads)
	673	+ failStatus := e.failStatus
598	674	e.mu.Unlock()
599	675
600	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	685	responseID := fmt.Sprintf("resp-%d", count)
602	686	chunks <- coreexecutor.StreamChunk{Payload: []byte(fmt.Sprintf(`{"type":"response.completed","response":{"id":%q,"output":[{"type":"message","id":"out-%d"}]}}`, responseID, count))}
603	687	close(chunks)
@@ -637,6 +721,67 @@ func (e *websocketDirectCaptureExecutor) AuthIDs() []string { · 33 unmodified lines
637	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	785	type websocketUpstreamDisconnectExecutor struct {
641	786	mu sync.Mutex
642	787	provider string
@@ -753,7 +898,7 @@ func (e *websocketAuthCaptureExecutor) AuthIDs() []string { · 110 unmodified lines
753	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	903	func (e websocketPinnedFailoverExecutor) Execute(context.Context, coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
759	904	return coreexecutor.Response{}, errors.New("not implemented")
@@ -781,8 +926,8 @@ func (e *websocketPinnedFailoverExecutor) ExecuteStream(_ context.Context, auth · 21 unmodified lines
781	926	if authID == "auth-a" && call == 2 {
782	927	chunks := make(chan coreexecutor.StreamChunk, 1)
783	928	chunks <- coreexecutor.StreamChunk{Err: websocketPinnedFailoverStatusError{
784		− status: http.StatusTooManyRequests,
785		− msg: `{"error":{"message":"quota exhausted","type":"rate_limit_error","code":"rate_limit_exceeded"}}`,
	929	+ status: e.failStatus,
	930	+ msg: fmt.Sprintf(`{"error":{"message":"credential failed","status":%d}}`, e.failStatus),
786	931	}}
787	932	close(chunks)
788	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	1784	func TestRepairResponsesWebsocketToolCallsDropsOrphanFunctionCall(t *testing.T) {
1602	1785	cache := newWebsocketToolOutputCache(time.Minute, 10)
1603	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	2606	gin.SetMode(gin.TestMode)
2424	2607
2425	2608	modelName := "xai-websocket-passthrough-model"
@@ -2483,165 +2666,660 @@ func TestResponsesWebsocketXAIWebsocketPassthroughCarriesPreviousResponseID(t *t · 57 unmodified lines
2483	2666	}
2484	2667	secondPayload := payloads[1]
2485	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	2671	if got := gjson.GetBytes(secondPayload, "model").String(); got != modelName {
2489	2672	t.Fatalf("second xai payload model = %s, want %s", got, modelName)
2490	2673	}
2491	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	2677	input := gjson.GetBytes(secondPayload, "input").Array()
2495		− if len(input) != 1 {
2496		− t.Fatalf("second xai passthrough input len = %d, want 1: %s", len(input), 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)
	2678	+ if len(input) != 1 \|\| input[0].Get("id").String() != "msg-2" {
	2679	+ t.Fatalf("second xai incremental input is not the client delta: %s", secondPayload)
2503	2680	}
2504	2681	authIDs := executor.AuthIDs()
2505	2682	if len(authIDs) != 2 \|\| authIDs[0] != "auth-xai-ws" \|\| authIDs[1] != "auth-xai-ws" {
2506	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) {
2511		− for _, testCase := range []struct {
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		− })
	2690	+func TestResponsesWebsocketFullRequestCanRouteFromNativeWebsocketToBuiltInProvider(t *testing.T) {
	2691	+ gin.SetMode(gin.TestMode)
2585	2692
2586		− base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
2587		− h := NewOpenAIResponsesAPIHandler(base)
2588		− router := gin.New()
2589		− router.GET("/v1/responses/ws", h.ResponsesWebsocket)
	2693	+ const sourceModel = "codex-provider-route-source"
	2694	+ const targetModel = "claude-provider-route-target"
	2695	+ codexExecutor := &websocketDirectCaptureExecutor{provider: "codex"}
	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)
2592		− defer server.Close()
	2723	+ base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
	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"
2595		− conn, _, errDial := websocket.DefaultDialer.Dial(wsURL, nil)
2596		− if errDial != nil {
2597		− t.Fatalf("dial websocket: %v", errDial)
2598		− }
2599		− defer func() {
2600		− if errClose := conn.Close(); errClose != nil {
2601		− t.Errorf("close websocket: %v", errClose)
2602		− }
2603		− }()
	2731	+ wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses/ws"
	2732	+ conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
	2733	+ if err != nil {
	2734	+ t.Fatalf("dial websocket: %v", err)
	2735	+ }
	2736	+ defer func() { _ = conn.Close() }()
2604	2737
2605		− requests := []string{
2606		− fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-xai-1"}]}`, xaiModel),
2607		− fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-codex-1"}]}`, codexModel),
2608		− fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-xai-2"}]}`, returnXAIModel),
2609		− `{"type":"response.create","input":[{"type":"message","id":"msg-xai-3"}]}`,
2610		− }
2611		− for index, request := range requests {
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		− }
	2738	+ firstRequest := []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"}]}`, sourceModel))
	2739	+ if errWrite := conn.WriteMessage(websocket.TextMessage, firstRequest); errWrite != nil {
	2740	+ t.Fatalf("write first websocket message: %v", errWrite)
	2741	+ }
	2742	+ if _, _, errRead := conn.ReadMessage(); errRead != nil {
	2743	+ t.Fatalf("read first websocket response: %v", errRead)
	2744	+ }
2623	2745
2624		− wantReturnAuthID := xaiAuth.ID
2625		− if testCase.returnToDifferentXAIModel {
2626		− wantReturnAuthID = xaiAlternateAuthID
2627		− }
2628		− if got := xaiExecutor.AuthIDs(); len(got) != 3 \|\| got[0] != xaiAuth.ID \|\| got[1] != wantReturnAuthID \|\| got[2] != wantReturnAuthID {
2629		− t.Fatalf("xAI auth IDs = %v, want [%s %s %s]", got, xaiAuth.ID, wantReturnAuthID, wantReturnAuthID)
2630		− }
2631		− if got := codexExecutor.AuthIDs(); len(got) != 1 \|\| got[0] != codexAuth.ID {
2632		− t.Fatalf("Codex auth IDs = %v, want [%s]", got, codexAuth.ID)
2633		− }
2634		− })
	2746	+ routedRequest := []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"route_to_claude":true,"input":[{"type":"message","id":"msg-routed"}]}`, sourceModel))
	2747	+ if errWrite := conn.WriteMessage(websocket.TextMessage, routedRequest); errWrite != nil {
	2748	+ t.Fatalf("write routed websocket message: %v", errWrite)
	2749	+ }
	2750	+ _, response, errRead := conn.ReadMessage()
	2751	+ if errRead != nil {
	2752	+ t.Fatalf("read routed websocket response: %v", errRead)
	2753	+ }
	2754	+ if got := gjson.GetBytes(response, "type").String(); got != wsEventTypeCompleted {
	2755	+ t.Fatalf("routed response type = %q, want %q: %s", got, wsEventTypeCompleted, response)
	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) {
2639		− modelA := "xai-pinned-auth-model-a"
2640		− modelB := "xai-pinned-auth-model-b"
2641		− auth := &coreauth.Auth{ID: "xai-pinned-auth", Provider: "xai", Status: coreauth.StatusActive}
2642		− otherAuthID := "xai-pinned-auth-other"
2643		− registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: modelA}})
2644		− registry.GetGlobalRegistry().RegisterClient(otherAuthID, auth.Provider, []*registry.ModelInfo{{ID: modelB}})
	2769	+func TestResponsesWebsocketFailedProviderRoutePreservesNativeWebsocketPin(t *testing.T) {
	2770	+ gin.SetMode(gin.TestMode)
	2771	+
	2772	+ const sourceModel = "codex-provider-route-failure-source"
	2773	+ const targetModel = "claude-provider-route-target"
	2774	+ codexExecutor := &websocketDirectCaptureExecutor{provider: "codex"}
	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	3323	t.Cleanup(func() {
2646	3324	registry.GetGlobalRegistry().UnregisterClient(auth.ID)
2647	3325	registry.GetGlobalRegistry().UnregisterClient(otherAuthID)
@@ -3055,7 +3733,7 @@ func TestResponsesWebsocketDoesNotInjectPreviousResponseIDWhenPendingToolOutputM · 407 unmodified lines
3055	3733	func TestResponsesWebsocketStripsGenerateWhenWebsocketAttemptFallsBackToHTTP(t *testing.T) {
3056	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	3737	executor := &websocketBootstrapFallbackExecutor{}
3060	3738	manager := coreauth.NewManager(nil, selector, nil)
3061	3739	manager.RegisterExecutor(executor)
@@ -3131,6 +3809,21 @@ func TestResponsesWebsocketStripsGenerateWhenWebsocketAttemptFallsBackToHTTP(t * · 69 unmodified lines
3131	3809	if gjson.GetBytes(httpPayloads[0], "generate").Exists() {
3132	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	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	3925	gin.SetMode(gin.TestMode)
3233	3926
3234		− selector := &orderedWebsocketSelector{order: []string{"auth-a", "auth-b"}}
3235		− executor := &websocketPinnedFailoverExecutor{}
	3927	+ modelName := "xai-mixed-pool-model"
	3928	+ selector := &orderedWebsocketSelector{order: []string{"auth-http", "auth-ws"}}
	3929	+ executor := &websocketDirectCaptureExecutor{provider: "xai"}
3236	3930	manager := coreauth.NewManager(nil, selector, nil)
3237	3931	manager.RegisterExecutor(executor)
3238		−
3239		− authA := &coreauth.Auth{
3240		− ID: "auth-a",
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)
	3932	+ authHTTP := &coreauth.Auth{ID: "auth-http", Provider: "xai", Status: coreauth.StatusActive}
	3933	+ if _, err := manager.Register(context.Background(), authHTTP); err != nil {
	3934	+ t.Fatalf("Register HTTP auth: %v", err)
3247	3935	}
3248		− authB := &coreauth.Auth{
3249		− ID: "auth-b",
3250		− Provider: executor.Identifier(),
	3936	+ authWS := &coreauth.Auth{
	3937	+ ID: "auth-ws",
	3938	+ Provider: "xai",
3251	3939	Status: coreauth.StatusActive,
3252	3940	Attributes: map[string]string{"websockets": "true"},
3253	3941	}
3254		− if _, err := manager.Register(context.Background(), authB); err != nil {
3255		− t.Fatalf("Register auth B: %v", err)
	3942	+ if _, err := manager.Register(context.Background(), authWS); err != nil {
	3943	+ t.Fatalf("Register websocket auth: %v", err)
3256	3944	}
3257		−
3258		− registry.GetGlobalRegistry().RegisterClient(authA.ID, authA.Provider, []*registry.ModelInfo{{ID: "quota-model"}})
3259		− registry.GetGlobalRegistry().RegisterClient(authB.ID, authB.Provider, []*registry.ModelInfo{{ID: "quota-model"}})
	3945	+ registry.GetGlobalRegistry().RegisterClient(authHTTP.ID, authHTTP.Provider, []*registry.ModelInfo{{ID: modelName}})
	3946	+ registry.GetGlobalRegistry().RegisterClient(authWS.ID, authWS.Provider, []*registry.ModelInfo{{ID: modelName}})
3260	3947	t.Cleanup(func() {
3261		− registry.GetGlobalRegistry().UnregisterClient(authA.ID)
3262		− registry.GetGlobalRegistry().UnregisterClient(authB.ID)
	3948	+ registry.GetGlobalRegistry().UnregisterClient(authHTTP.ID)
	3949	+ registry.GetGlobalRegistry().UnregisterClient(authWS.ID)
3263	3950	})
3264	3951
3265	3952	base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
3266	3953	h := NewOpenAIResponsesAPIHandler(base)
3267	3954	router := gin.New()
3268	3955	router.GET("/v1/responses/ws", h.ResponsesWebsocket)
3269		−
3270	3956	server := httptest.NewServer(router)
3271	3957	defer server.Close()
3272	3958
@@ -3275,49 +3961,169 @@ func TestResponsesWebsocketReleasesPinnedAuthAfterQuotaError(t *testing.T) { · 2 unmodified lines
3275	3961	if err != nil {
3276	3962	t.Fatalf("dial websocket: %v", err)
3277	3963	}
3278		− defer func() {
3279		− if errClose := conn.Close(); errClose != nil {
3280		− t.Fatalf("close websocket: %v", errClose)
3281		− }
3282		− }()
	3964	+ defer func() { _ = conn.Close() }()
3283	3965
3284	3966	requests := []string{
3285		− `{"type":"response.create","model":"quota-model","input":[{"type":"message","id":"msg-1"}]}`,
3286		− `{"type":"response.create","previous_response_id":"resp-auth-a-1","input":[{"type":"message","id":"msg-2"}]}`,
3287		− `{"type":"response.create","previous_response_id":"resp-auth-a-1","input":[{"type":"message","id":"msg-3"}]}`,
	3967	+ fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"}]}`, modelName),
	3968	+ `{"type":"response.create","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`,
	3969	+ `{"type":"response.create","previous_response_id":"resp-2","input":[{"type":"message","id":"msg-3"}]}`,
3288	3970	}
3289		− wantTypes := []string{wsEventTypeCompleted, wsEventTypeError, wsEventTypeCompleted}
3290	3971	for i := range requests {
3291	3972	if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(requests[i])); errWrite != nil {
3292	3973	t.Fatalf("write websocket message %d: %v", i+1, errWrite)
3293	3974	}
3294		− _, payload, errReadMessage := conn.ReadMessage()
3295		− if errReadMessage != nil {
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)
	3975	+ if _, _, errRead := conn.ReadMessage(); errRead != nil {
	3976	+ t.Fatalf("read websocket response %d: %v", i+1, errRead)
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" {
3307		− t.Fatalf("selected auth IDs = %v, want [auth-a auth-a auth-b]", got)
	3980	+ if got := executor.AuthIDs(); len(got) != 3 \|\| got[0] != "auth-http" \|\| got[1] != "auth-ws" \|\| got[2] != "auth-ws" {
	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")
3311		− if len(authBPayloads) != 1 {
3312		− t.Fatalf("auth-b payload count = %d, want 1", len(authBPayloads))
	4054	+ wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses/ws"
	4055	+ conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
	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]
3315		− if gjson.GetBytes(authBPayload, "previous_response_id").Exists() {
3316		− t.Fatalf("previous_response_id leaked after auth failover: %s", authBPayload)
	4107	+}
	4108	+
	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
3319		− if !strings.Contains(authBInput, `"id":"msg-1"`) \|\| !strings.Contains(authBInput, `"id":"msg-3"`) {
3320		− t.Fatalf("auth-b replay input missing expected transcript items: %s", authBInput)
	4121	+ for _, tc := range cases {
	4122	+ t.Run(tc.name, func(t *testing.T) {
	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	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	4161	func (e websocketPinnedPrematureCloseExecutor) Execute(context.Context, coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
3356	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{
3474		− `{"type":"response.create","model":"stream-model","input":[{"type":"message","id":"msg-1"}]}`,
3475		− `{"type":"response.create","previous_response_id":"resp-auth-a-1","input":[{"type":"message","id":"msg-2"}]}`,
3476		− `{"type":"response.create","previous_response_id":"resp-auth-a-1","input":[{"type":"message","id":"msg-3"}]}`,
	4279	+ firstRequest := `{"type":"response.create","model":"stream-model","input":[{"type":"message","id":"msg-1"}]}`
	4280	+ if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(firstRequest)); errWrite != nil {
	4281	+ t.Fatalf("write first websocket message: %v", errWrite)
3477	4282	}
3478		− wantTypes := []string{wsEventTypeCompleted, wsEventTypeError, wsEventTypeCompleted}
3479		− for i := range requests {
3480		− if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(requests[i])); errWrite != nil {
3481		− t.Fatalf("write websocket message %d: %v", i+1, errWrite)
	4283	+ if _, payload, errRead := conn.ReadMessage(); errRead != nil \|\| gjson.GetBytes(payload, "type").String() != wsEventTypeCompleted {
	4284	+ t.Fatalf("first websocket response = %s, err=%v", payload, errRead)
	4285	+ }
	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 {
3484		− gotError := false
3485		− for {
3486		− _, payload, errReadMessage := conn.ReadMessage()
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		− }
	4296	+ eventType := gjson.GetBytes(payload, "type").String()
	4297	+ if eventType == wsEventTypeError {
	4298	+ if got := int(gjson.GetBytes(payload, "status").Int()); got != http.StatusRequestTimeout {
	4299	+ t.Fatalf("stream-closed status = %d, want %d: %s", got, http.StatusRequestTimeout, payload)
3501	4300	}
3502		− if !gotError {
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)
	4301	+ break
3510	4302	}
3511		− if got := gjson.GetBytes(payload, "type").String(); got != wantTypes[i] {
3512		− t.Fatalf("message %d payload type = %s, want %s: %s", i+1, got, wantTypes[i], payload)
	4303	+ if eventType == wsEventTypeCompleted {
	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	4333	authIDs := executor.AuthIDs()
3517	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	4337	replayAuthID := authIDs[2]
3522	4338	replayPayloads := executor.Payloads(replayAuthID)
3523		− if len(replayPayloads) == 0 {
3524		− t.Fatalf("replay auth %s has no payloads", replayAuthID)
3525		− }
3526	4339	replayPayload := replayPayloads[len(replayPayloads)-1]
3527		− if gjson.GetBytes(replayPayload, "previous_response_id").Exists() {
3528		− t.Fatalf("previous_response_id leaked after stream-closed replay: %s", 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)
	4340	+ if gjson.GetBytes(replayPayload, "previous_response_id").Exists() \|\| len(gjson.GetBytes(replayPayload, "input").Array()) != 3 {
	4341	+ t.Fatalf("replay auth %s did not receive full replay: %s", replayAuthID, replayPayload)
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	19	var defaultWebsocketToolOutputCache = newWebsocketToolOutputCache(0, websocketToolOutputCacheMaxPerSession)
20	20	var defaultWebsocketToolCallCache = newWebsocketToolOutputCache(0, websocketToolOutputCacheMaxPerSession)
21	21	var defaultWebsocketToolSessionRefs = newWebsocketToolSessionRefCounter()
	22	+var defaultWebsocketToolCacheTransactionMu sync.RWMutex
22	23
23	24	type websocketToolOutputCache struct {
24	25	mu sync.Mutex
@@ -33,6 +34,14 @@ type websocketToolOutputSession struct { · 8 unmodified lines
33	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	45	func newWebsocketToolOutputCache(ttl time.Duration, maxPerSession int) *websocketToolOutputCache {
37	46	if ttl < 0 {
38	47	ttl = websocketToolOutputCacheTTL
@@ -196,6 +205,8 @@ func (c *websocketToolSessionRefCounter) release(sessionKey string) bool { · 157 unmodified lines
196	205	}
197	206
198	207	func retainResponsesWebsocketToolCaches(sessionKey string) {
	208	+ defaultWebsocketToolCacheTransactionMu.Lock()
	209	+ defer defaultWebsocketToolCacheTransactionMu.Unlock()
199	210	if defaultWebsocketToolSessionRefs == nil {
200	211	return
201	212	}
@@ -203,13 +214,14 @@ func retainResponsesWebsocketToolCaches(sessionKey string) { · 1 unmodified line
203	214	}
204	215
205	216	func releaseResponsesWebsocketToolCaches(sessionKey string) {
	217	+ defaultWebsocketToolCacheTransactionMu.Lock()
	218	+ defer defaultWebsocketToolCacheTransactionMu.Unlock()
206	219	if defaultWebsocketToolSessionRefs == nil {
207	220	return
208	221	}
209	222	if !defaultWebsocketToolSessionRefs.release(sessionKey) {
210	223	return
211	224	}
212		−
213	225	if defaultWebsocketToolOutputCache != nil {
214	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	322	func repairResponsesWebsocketToolCalls(sessionKey string, payload []byte) []byte {
222	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	332	func repairResponsesWebsocketToolCallsWithCache(cache *websocketToolOutputCache, sessionKey string, payload []byte) []byte {
226	333	return repairResponsesWebsocketToolCallsWithCaches(cache, nil, sessionKey, payload)
227	334	}
228	335
229	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	341	sessionKey = strings.TrimSpace(sessionKey)
231	342	if sessionKey == "" \|\| outputCache == nil \|\| len(payload) == 0 {
232	343	return payload
@@ -238,7 +349,7 @@ func repairResponsesWebsocketToolCallsWithCaches(outputCache, callCache *websock · 5 unmodified lines
238	349	}
239	350
240	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	353	if errRepair != nil \|\| updatedRaw == "" \|\| updatedRaw == input.Raw {
243	354	return payload
244	355	}
@@ -250,7 +361,7 @@ func repairResponsesWebsocketToolCallsWithCaches(outputCache, callCache *websock · 5 unmodified lines
250	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	365	rawArray = strings.TrimSpace(rawArray)
255	366	if rawArray == "" {
256	367	return "[]", nil
@@ -276,14 +387,16 @@ func repairResponsesToolCallsArray(outputCache, callCache *websocketToolOutputCa · 19 unmodified lines
276	387	continue
277	388	}
278	389	outputPresent[callID] = struct{}{}
279		− outputCache.record(sessionKey, callID, item)
	390	+ if record {
	391	+ outputCache.record(sessionKey, callID, item)
	392	+ }
280	393	case isResponsesToolCallType(itemType):
281	394	callID := strings.TrimSpace(gjson.GetBytes(item, "call_id").String())
282	395	if callID == "" {
283	396	continue
284	397	}
285	398	callPresent[callID] = struct{}{}
286		− if callCache != nil {
	399	+ if record && callCache != nil {
287	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	3	import "context"
4	4
5	5	type downstreamWebsocketContextKey struct{}
	6	+type requireUpstreamWebsocketContextKey struct{}
6	7
7	8	// WithDownstreamWebsocket marks the current request as coming from a downstream websocket connection.
8	9	func WithDownstreamWebsocket(ctx context.Context) context.Context {
@@ -21,3 +22,21 @@ func DownstreamWebsocket(ctx context.Context) bool { · 12 unmodified lines
21	22	enabled, ok := raw.(bool)
22	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	+}