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

sdk/api/handlers/handlers.go

@@ -64,6 +64,7 @@ const (
 
 type pinnedAuthContextKey struct{}
 type selectedAuthCallbackContextKey struct{}
+type preparedModelRouteContextKey struct{}
 type executionSessionContextKey struct{}
 type disallowFreeAuthContextKey struct{}
 
@@ -142,6 +143,26 @@ func WithSelectedAuthIDCallback(ctx context.Context, callback func(string)) cont
 	return context.WithValue(ctx, selectedAuthCallbackContextKey{}, callback)
 }
 
+// PrepareStreamModelRoute resolves a stream route once and stores it on the returned context for execution.
+// The boolean reports whether the route overrides normal model-to-provider resolution.
+func (h *BaseAPIHandler) PrepareStreamModelRoute(ctx context.Context, handlerType string, modelName string, rawJSON []byte) (context.Context, bool) {
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	decision := h.applyModelRouter(ctx, handlerType, modelName, rawJSON, true, modelExecutionOptions{})
+	ctx = context.WithValue(ctx, preparedModelRouteContextKey{}, decision)
+	hasOverride := strings.TrimSpace(decision.ExecutorPluginID) != "" || strings.TrimSpace(decision.Provider) != ""
+	return ctx, hasOverride
+}
+
+func preparedModelRouteFromContext(ctx context.Context) (modelRouteDecision, bool) {
+	if ctx == nil {
+		return modelRouteDecision{}, false
+	}
+	decision, ok := ctx.Value(preparedModelRouteContextKey{}).(modelRouteDecision)
+	return decision, ok
+}
+
 // WithExecutionSessionID returns a child context tagged with a long-lived execution session ID.
 func WithExecutionSessionID(ctx context.Context, sessionID string) context.Context {
 	sessionID = strings.TrimSpace(sessionID)
@@ -1144,7 +1165,10 @@ func (h *BaseAPIHandler) executeStreamWithAuthManager(ctx context.Context, handl
 
 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) {
 	originalRequestedModel := modelName
-	routeDecision := h.applyModelRouter(ctx, entryProtocol, modelName, rawJSON, true, execOptions)
+	routeDecision, preparedRoute := preparedModelRouteFromContext(ctx)
+	if !preparedRoute {
+		routeDecision = h.applyModelRouter(ctx, entryProtocol, modelName, rawJSON, true, execOptions)
+	}
 	responseProtocol := modelExecutionResponseProtocol(entryProtocol, exitProtocol)
 	if errMsg := validateNativeInteractionsExecution(entryProtocol, execOptions, routeDecision); errMsg != nil {
 		errChan := make(chan *interfaces.ErrorMessage, 1)

sdk/api/handlers/handlers_model_router_test.go

@@ -464,6 +464,38 @@ func TestHandlerModelRouterRoutesStreamBeforeRequestDetails(t *testing.T) {
 	}
 }
 
+func TestPrepareStreamModelRouteReusesDecisionDuringExecution(t *testing.T) {
+	const model = "prepared-router-model"
+	const targetPluginID = "prepared-stream-plugin"
+	routeCalls := 0
+	host := &handlerDirectExecutorRouteHost{}
+	host.hasRouters = true
+	host.route = func(context.Context, pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {
+		routeCalls++
+		return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true
+	}
+	handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)
+	handler.SetModelRouterHost(host)
+	body := []byte(`{"model":"prepared-router-model","stream":true}`)
+	ctx, routedToPlugin := handler.PrepareStreamModelRoute(context.Background(), "openai", model, body)
+	if !routedToPlugin {
+		t.Fatal("PrepareStreamModelRoute() did not detect plugin executor route")
+	}
+
+	dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(ctx, "openai", model, body, "")
+	for range dataChan {
+	}
+	if errMsg := <-errChan; errMsg != nil {
+		t.Fatalf("ExecuteStreamWithAuthManager() error = %+v", errMsg)
+	}
+	if routeCalls != 1 {
+		t.Fatalf("model router calls = %d, want 1", routeCalls)
+	}
+	if host.lastPluginID != targetPluginID {
+		t.Fatalf("plugin id = %q, want %q", host.lastPluginID, targetPluginID)
+	}
+}
+
 func TestExecuteModelPropagatesRouterSkipPluginID(t *testing.T) {
 	model := "model-execution-router-skip-model"
 	requestBody := []byte(fmt.Sprintf(`{"model":%q}`, model))

sdk/api/handlers/openai/openai_responses_websocket.go

@@ -33,15 +33,19 @@ import (
 )
 
 const (
-	wsRequestTypeCreate   = "response.create"
-	wsRequestTypeAppend   = "response.append"
-	wsEventTypeError      = "error"
-	wsEventTypeCompleted  = "response.completed"
-	wsEventTypeDone       = "response.done"
-	wsDoneMarker          = "[DONE]"
-	wsTurnStateHeader     = "x-codex-turn-state"
-	wsTimelineBodyKey     = "WEBSOCKET_TIMELINE_OVERRIDE"
-	wsCloseReasonMaxBytes = 123
+	wsRequestTypeCreate                   = "response.create"
+	wsRequestTypeAppend                   = "response.append"
+	wsEventTypeError                      = "error"
+	wsEventTypeCompleted                  = "response.completed"
+	wsEventTypeDone                       = "response.done"
+	wsDoneMarker                          = "[DONE]"
+	wsTurnStateHeader                     = "x-codex-turn-state"
+	wsTimelineBodyKey                     = "WEBSOCKET_TIMELINE_OVERRIDE"
+	wsCloseReasonMaxBytes                 = 123
+	wsHTTPReplayRequiredCloseReason       = "upstream requires HTTP replay"
+	responsesWebsocketUpstreamModeUnknown = ""
+	responsesWebsocketUpstreamModeWS      = "websocket"
+	responsesWebsocketUpstreamModeHTTP    = "http"
 
 	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:"
 )
@@ -74,6 +78,14 @@ func websocketClosePayloadForUpstreamError(err error) (bool, []byte) {
 		return false, nil
 	}
 
+	errText := err.Error()
+	if cliproxyexecutor.IsUpstreamWebsocketReplayRequired(err) {
+		return true, websocket.FormatCloseMessage(
+			websocket.CloseServiceRestart,
+			truncateWebsocketCloseReason(wsHTTPReplayRequiredCloseReason, wsCloseReasonMaxBytes),
+		)
+	}
+
 	code := 0
 	reason := ""
 	var closeErr *websocket.CloseError
@@ -85,7 +97,6 @@ func websocketClosePayloadForUpstreamError(err error) (bool, []byte) {
 			StatusCode() int
 		}
 		var statusErr statusCoder
-		errText := err.Error()
 		if !errors.As(err, &statusErr) || statusErr.StatusCode() != http.StatusRequestEntityTooLarge ||
 			gjson.Get(errText, "error.code").String() != "message_too_big" {
 			return false, nil
@@ -419,20 +430,35 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
 	// Preserve independent upstream auth affinity when a downstream session switches providers.
 	pinnedAuthByProvider := make(map[string]responsesWebsocketPinnedAuthState)
 	passthroughModelName := ""
+	upstreamMode := responsesWebsocketUpstreamModeUnknown
+	upstreamWebsocketAuthID := ""
 	sessionAuthByIDWithSource := func(authID string) (*coreauth.Auth, bool, bool) {
 		if h == nil || h.AuthManager == nil {
 			return nil, false, false
 		}
+		// Prefer the current manager view so hot-reloaded transport eligibility is
+		// observed even when the execution session still holds an older auth snapshot.
+		if auth, ok := h.AuthManager.GetByID(authID); ok {
+			return auth, false, true
+		}
 		if auth, ok := h.AuthManager.GetExecutionSessionAuthByID(passthroughSessionID, authID); ok {
 			return auth, true, true
 		}
-		auth, ok := h.AuthManager.GetByID(authID)
-		return auth, false, ok
+		return nil, false, false
 	}
 	sessionAuthByID := func(authID string) (*coreauth.Auth, bool) {
 		auth, _, ok := sessionAuthByIDWithSource(authID)
 		return auth, ok
 	}
+	upstreamModeForAuth := func(auth *coreauth.Auth) string {
+		if auth != nil && websocketUpstreamSupportsIncrementalInput(auth.Attributes, auth.Metadata) {
+			provider := strings.ToLower(strings.TrimSpace(auth.Provider))
+			if provider == "codex" || provider == "xai" {
+				return responsesWebsocketUpstreamModeWS
+			}
+		}
+		return responsesWebsocketUpstreamModeHTTP
+	}
 	rememberPinnedAuth := func(authID string, modelName string) {
 		authID = strings.TrimSpace(authID)
 		auth, ok := sessionAuthByID(authID)
@@ -454,7 +480,6 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
 		}
 		pinnedAuthID = ""
 	}
-	forceTranscriptReplayNextRequest := false
 
 	for {
 		msgType, payload, errReadMessage := conn.ReadMessage()
@@ -488,6 +513,13 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
 		if requestModelName == "" {
 			requestModelName = strings.TrimSpace(gjson.GetBytes(lastRequest, "model").String())
 		}
+		executionParent := context.WithValue(c.Request.Context(), "gin", c)
+		executionParent, routeOverridesModelResolution := h.PrepareStreamModelRoute(
+			executionParent,
+			h.HandlerType(),
+			requestModelName,
+			payload,
+		)
 		if pinnedAuthID != "" {
 			pinnedAuth, homeRuntime, ok := sessionAuthByIDWithSource(pinnedAuthID)
 			providerKey := ""
@@ -514,15 +546,40 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
 			}
 		}
 		useUpstreamWebsocketPassthrough := h.responsesWebsocketUsesUpstreamWebsocketPassthrough(requestModelName)
+		if pinnedAuthID != "" {
+			if pinnedAuth, ok := sessionAuthByID(pinnedAuthID); ok && responsesWebsocketAuthSupportsIncrementalInput(pinnedAuth) {
+				provider := strings.ToLower(strings.TrimSpace(pinnedAuth.Provider))
+				useUpstreamWebsocketPassthrough = provider == "codex" || provider == "xai"
+			}
+		}
+		nativeWebsocketPassthrough := !routeOverridesModelResolution && responsesWebsocketNativePassthroughAllowed(
+			upstreamMode,
+			useUpstreamWebsocketPassthrough,
+			pinnedAuthID,
+			upstreamWebsocketAuthID,
+		)
+		requestRequiresCurrentUpstreamWebsocket := responsesWebsocketRequestRequiresCurrentUpstream(payload)
+		if upstreamMode == responsesWebsocketUpstreamModeWS && !nativeWebsocketPassthrough {
+			if requestRequiresCurrentUpstreamWebsocket {
+				replayErr := responsesWebsocketHTTPReplayRequiredError()
+				wsTerminateErr = replayErr
+				matched, errClose := writer.closeForUpstreamError(replayErr)
+				if !matched {
+					_ = conn.Close()
+				} else if errClose != nil && !errors.Is(errClose, websocket.ErrCloseSent) {
+					log.Debugf("responses websocket: replay close failed id=%s error=%v", passthroughSessionID, errClose)
+				}
+				return
+			}
+			// A full response.create is already a self-contained reset and can safely
+			// establish a new upstream transport without another replay.
+		}
 		if explicitRequestModelName != "" && !useUpstreamWebsocketPassthrough {
 			passthroughModelName = ""
 		}
-		allowIncrementalInputWithPreviousResponseID := false
+
 		allowCompactionReplayBypass := false
-		if !useUpstreamWebsocketPassthrough {
-			// Downstream websocket with CPA-mediated upstream (HTTP/SSE) always uses merged
-			// transcript replay. Incremental previous_response_id is reserved for end-to-end
-			// upstream websocket passthrough only.
+		if !nativeWebsocketPassthrough {
 			if pinnedAuthID != "" {
 				if pinnedAuth, ok := sessionAuthByID(pinnedAuthID); ok && pinnedAuth != nil {
 					allowCompactionReplayBypass = responsesWebsocketAuthSupportsCompactionReplay(pinnedAuth)
@@ -535,8 +592,10 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
 		var requestJSON []byte
 		var updatedLastRequest []byte
 		var errMsg *interfaces.ErrorMessage
-		if useUpstreamWebsocketPassthrough {
+		if nativeWebsocketPassthrough {
 			requestJSON, errMsg = normalizeResponsesWebsocketPassthroughRequest(payload, requestModelName)
+		} else if len(lastRequest) == 0 && strings.TrimSpace(gjson.GetBytes(payload, "previous_response_id").String()) != "" {
+			errMsg = responsesWebsocketPreviousResponseNotFoundError()
 		} else {
 			requestJSON, updatedLastRequest, errMsg = normalizeResponsesWebsocketRequestWithIncrementalState(
 				payload,
@@ -544,7 +603,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
 				lastResponseOutput,
 				lastResponseID,
 				lastResponsePendingToolCallIDs,
-				allowIncrementalInputWithPreviousResponseID,
+				false,
 				allowCompactionReplayBypass,
 			)
 		}
@@ -570,7 +629,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
 			}
 			continue
 		}
-		if !useUpstreamWebsocketPassthrough && shouldHandleResponsesWebsocketPrewarmLocally(payload, lastRequest, allowIncrementalInputWithPreviousResponseID) {
+		if !useUpstreamWebsocketPassthrough && shouldHandleResponsesWebsocketPrewarmLocally(payload, lastRequest, false) {
 			if updated, errDelete := sjson.DeleteBytes(requestJSON, "generate"); errDelete == nil {
 				requestJSON = updated
 			}
@@ -588,84 +647,120 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
 			continue
 		}
 
+		toolCacheTurn := newResponsesWebsocketToolCacheTurn(downstreamSessionKey)
 		previousLastRequest := bytes.Clone(lastRequest)
 		previousLastResponseOutput := bytes.Clone(lastResponseOutput)
 		previousLastResponseID := lastResponseID
 		previousLastResponsePendingToolCallIDs := append([]string(nil), lastResponsePendingToolCallIDs...)
-		forcedTranscriptReplay := forceTranscriptReplayNextRequest
-		if useUpstreamWebsocketPassthrough {
+		if nativeWebsocketPassthrough {
 			if modelName := strings.TrimSpace(gjson.GetBytes(requestJSON, "model").String()); modelName != "" {
 				passthroughModelName = modelName
 			}
-			if forcedTranscriptReplay {
-				forceTranscriptReplayNextRequest = false
-			}
 		} else {
-			requestJSON = repairResponsesWebsocketToolCalls(downstreamSessionKey, requestJSON)
+			toolCacheTurn.recordRequest(requestJSON)
+			requestJSON = repairResponsesWebsocketToolCallsWithoutRecording(downstreamSessionKey, requestJSON)
 			requestJSON = dedupeResponsesWebsocketInputItemsByID(requestJSON)
 			updatedLastRequest = bytes.Clone(requestJSON)
 			lastRequest = updatedLastRequest
-			if forcedTranscriptReplay {
-				forceTranscriptReplayNextRequest = false
-			}
 		}
 
 		modelName := gjson.GetBytes(requestJSON, "model").String()
 		lastAttemptedAuthID := pinnedAuthID
-		cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
+		attemptedUpstreamMode := responsesWebsocketUpstreamModeUnknown
+		selectedAuthObserved := false
+		pinnedAuthAttempted := false
+		cliCtx, cliCancel := h.GetContextWithCancel(h, c, executionParent)
 		cliCtx = cliproxyexecutor.WithDownstreamWebsocket(cliCtx)
+		if nativeWebsocketPassthrough && requestRequiresCurrentUpstreamWebsocket {
+			cliCtx = cliproxyexecutor.WithRequiredUpstreamWebsocket(cliCtx)
+		}
 		cliCtx = handlers.WithExecutionSessionID(cliCtx, passthroughSessionID)
-		if pinnedAuthID != "" {
+		cliCtx = handlers.WithSelectedAuthIDCallback(cliCtx, func(authID string) {
+			authID = strings.TrimSpace(authID)
+			if authID == "" || h == nil || h.AuthManager == nil {
+				return
+			}
+			lastAttemptedAuthID = authID
+			selectedAuthObserved = true
+			pinnedAuthAttempted = pinnedAuthAttempted || (pinnedAuthID != "" && authID == pinnedAuthID)
+			selectedAuth, ok := sessionAuthByID(authID)
+			if !ok || selectedAuth == nil {
+				return
+			}
+			attemptedUpstreamMode = upstreamModeForAuth(selectedAuth)
+		})
+		if pinnedAuthID != "" && !routeOverridesModelResolution {
 			cliCtx = handlers.WithPinnedAuthID(cliCtx, pinnedAuthID)
-		} else {
-			cliCtx = handlers.WithSelectedAuthIDCallback(cliCtx, func(authID string) {
-				authID = strings.TrimSpace(authID)
-				if authID == "" || h == nil || h.AuthManager == nil {
-					return
-				}
-				lastAttemptedAuthID = authID
-				selectedAuth, ok := sessionAuthByID(authID)
-				if !ok || selectedAuth == nil {
-					return
-				}
-				if websocketUpstreamSupportsIncrementalInput(selectedAuth.Attributes, selectedAuth.Metadata) {
-					rememberPinnedAuth(authID, modelName)
-				}
-			})
 		}
 		dataChan, _, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, requestJSON, "")
+		if !selectedAuthObserved {
+			// Plugin/alternate routes bypass auth selection. Keep canonical HTTP-mode
+			// state instead of inheriting the previous pinned websocket mode.
+			attemptedUpstreamMode = responsesWebsocketUpstreamModeHTTP
+		}
+		// A connection-scoped continuation cannot rotate credentials in place. Suppress
+		// credential errors and make the client replay the full turn on a new socket.
+		replayPinnedAuthFailure := func(errMsg *interfaces.ErrorMessage) bool {
+			return nativeWebsocketPassthrough && requestRequiresCurrentUpstreamWebsocket && pinnedAuthAttempted &&
+				shouldReplayResponsesWebsocketPinnedAuthFailure(errMsg)
+		}
 
-		completedOutput, completedResponseID, completedPendingToolCallIDs, forwardErrMsg, errForward := h.forwardResponsesWebsocket(c, writer, cliCancel, dataChan, errChan, wsTimelineLog, passthroughSessionID)
+		completedOutput, completedResponseID, completedPendingToolCallIDs, forwardErrMsg, errForward := h.forwardResponsesWebsocket(
+			c,
+			writer,
+			cliCancel,
+			dataChan,
+			errChan,
+			wsTimelineLog,
+			passthroughSessionID,
+			responsesWebsocketForwardOptions{
+				toolCacheTurn: toolCacheTurn,
+				suppressError: replayPinnedAuthFailure,
+			},
+		)
 		if errForward != nil {
 			wsTerminateErr = errForward
-			log.Warnf("responses websocket: forward failed id=%s error=%v", passthroughSessionID, errForward)
+			if !errors.Is(errForward, websocket.ErrCloseSent) {
+				log.Warnf("responses websocket: forward failed id=%s error=%v", passthroughSessionID, errForward)
+			}
 			return
 		}
-		if forwardErrMsg == nil && !useUpstreamWebsocketPassthrough && lastAttemptedAuthID != "" {
-			if selectedAuth, ok := sessionAuthByID(lastAttemptedAuthID); ok && selectedAuth != nil {
-				if websocketUpstreamSupportsIncrementalInput(selectedAuth.Attributes, selectedAuth.Metadata) {
-					rememberPinnedAuth(lastAttemptedAuthID, modelName)
-				} else if pinnedAuthID != "" {
-					if pinnedAuth, ok := sessionAuthByID(pinnedAuthID); ok && pinnedAuth != nil && websocketUpstreamSupportsIncrementalInput(pinnedAuth.Attributes, pinnedAuth.Metadata) {
-						rememberPinnedAuth(lastAttemptedAuthID, modelName)
-					}
-				}
+		if forwardErrMsg != nil {
+			lastRequest = previousLastRequest
+			lastResponseOutput = previousLastResponseOutput
+			lastResponseID = previousLastResponseID
+			lastResponsePendingToolCallIDs = previousLastResponsePendingToolCallIDs
+			if pinnedAuthAttempted && shouldReleaseResponsesWebsocketPinnedAuth(forwardErrMsg) {
+				forgetPinnedAuth()
 			}
-		}
-		if shouldReleaseResponsesWebsocketPinnedAuth(forwardErrMsg) {
-			forgetPinnedAuth()
-			forceTranscriptReplayNextRequest = true
-			if useUpstreamWebsocketPassthrough {
-				passthroughModelName = ""
-			} else {
-				lastRequest = previousLastRequest
-				lastResponseOutput = previousLastResponseOutput
-				lastResponseID = previousLastResponseID
-				lastResponsePendingToolCallIDs = previousLastResponsePendingToolCallIDs
+			if replayPinnedAuthFailure(forwardErrMsg) {
+				replayErr := responsesWebsocketHTTPReplayRequiredError()
+				wsTerminateErr = replayErr
+				matched, errClose := writer.closeForUpstreamError(replayErr)
+				if !matched {
+					_ = conn.Close()
+				} else if errClose != nil && !errors.Is(errClose, websocket.ErrCloseSent) {
+					log.Debugf("responses websocket: credential replay close failed id=%s error=%v", passthroughSessionID, errClose)
+				}
+				return
 			}
 			continue
 		}
-		if !useUpstreamWebsocketPassthrough {
+
+		toolCacheTurn.commit()
+		upstreamMode = attemptedUpstreamMode
+		if upstreamMode == responsesWebsocketUpstreamModeWS {
+			upstreamWebsocketAuthID = lastAttemptedAuthID
+			if lastAttemptedAuthID != "" {
+				rememberPinnedAuth(lastAttemptedAuthID, modelName)
+			}
+			passthroughModelName = modelName
+			lastRequest = nil
+			lastResponseOutput = []byte("[]")
+			lastResponseID = ""
+			lastResponsePendingToolCallIDs = nil
+		} else {
+			upstreamWebsocketAuthID = ""
 			lastResponseOutput = completedOutput
 			lastResponseID = strings.TrimSpace(completedResponseID)
 			lastResponsePendingToolCallIDs = append([]string(nil), completedPendingToolCallIDs...)
@@ -673,6 +768,20 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
 	}
 }
 
+func responsesWebsocketHTTPReplayRequiredError() error {
+	return cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError()
+}
+
+func responsesWebsocketRequestRequiresCurrentUpstream(payload []byte) bool {
+	return strings.TrimSpace(gjson.GetBytes(payload, "previous_response_id").String()) != "" ||
+		strings.TrimSpace(gjson.GetBytes(payload, "type").String()) == wsRequestTypeAppend
+}
+
+func responsesWebsocketNativePassthroughAllowed(upstreamMode string, useUpstreamWebsocket bool, pinnedAuthID string, upstreamAuthID string) bool {
+	return upstreamMode == responsesWebsocketUpstreamModeWS && useUpstreamWebsocket &&
+		strings.TrimSpace(pinnedAuthID) != "" && strings.TrimSpace(pinnedAuthID) == strings.TrimSpace(upstreamAuthID)
+}
+
 func websocketClientAddress(c *gin.Context) string {
 	if c == nil || c.Request == nil {
 		return ""
@@ -694,6 +803,15 @@ func websocketUpgradeHeaders(req *http.Request) http.Header {
 	return headers
 }
 
+func responsesWebsocketPreviousResponseNotFoundError() *interfaces.ErrorMessage {
+	return &interfaces.ErrorMessage{
+		StatusCode: http.StatusConflict,
+		Error: errors.New(
+			`{"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"}}`,
+		),
+	}
+}
+
 func normalizeResponsesWebsocketRequest(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte) ([]byte, []byte, *interfaces.ErrorMessage) {
 	return normalizeResponsesWebsocketRequestWithMode(rawJSON, lastRequest, lastResponseOutput, true, true)
 }
@@ -1570,6 +1688,11 @@ func normalizeJSONArrayRaw(raw []byte) string {
 	return "[]"
 }
 
+type responsesWebsocketForwardOptions struct {
+	toolCacheTurn *responsesWebsocketToolCacheTurn
+	suppressError func(*interfaces.ErrorMessage) bool
+}
+
 func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket(
 	c *gin.Context,
 	writer *responsesWebsocketWriter,
@@ -1578,7 +1701,13 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket(
 	errs <-chan *interfaces.ErrorMessage,
 	wsTimelineLog websocketTimelineAppender,
 	sessionID string,
+	options ...responsesWebsocketForwardOptions,
 ) ([]byte, string, []string, *interfaces.ErrorMessage, error) {
+	var opts responsesWebsocketForwardOptions
+	if len(options) > 0 {
+		opts = options[0]
+	}
+	toolCacheTurn := opts.toolCacheTurn
 	completed := false
 	completedOutput := []byte("[]")
 	completedResponseID := ""
@@ -1602,6 +1731,10 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket(
 			}
 			if errMsg != nil {
 				h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), errMsg)
+				if opts.suppressError != nil && opts.suppressError(errMsg) {
+					cancel(errMsg.Error)
+					return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, nil
+				}
 				markAPIResponseTimestamp(c)
 				if matched, errClose := writer.closeForUpstreamError(errMsg.Error); matched {
 					cancel(errMsg.Error)
@@ -1676,7 +1809,11 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket(
 				if isResponsesWebsocketCompletionEvent(eventType) {
 					payloads[i] = restoreResponsesWebsocketCompletionOutput(payloads[i], outputItemsByIndex, outputItemsFallback)
 				}
-				recordResponsesWebsocketToolCallsFromPayload(downstreamSessionKey, payloads[i])
+				if toolCacheTurn != nil {
+					toolCacheTurn.recordResponse(payloads[i])
+				} else {
+					recordResponsesWebsocketToolCallsFromPayload(downstreamSessionKey, payloads[i])
+				}
 				recordPendingToolCallIDsFromPayload(pendingToolCallIDs, payloads[i])
 				var payloadErrMsg *interfaces.ErrorMessage
 				if eventType == wsEventTypeError {
@@ -1684,6 +1821,10 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket(
 					if h != nil {
 						h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), payloadErrMsg)
 					}
+					if opts.suppressError != nil && opts.suppressError(payloadErrMsg) {
+						cancel(payloadErrMsg.Error)
+						return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), payloadErrMsg, nil
+					}
 				} else if isResponsesWebsocketCompletionEvent(eventType) {
 					completed = true
 					completedOutput = responseCompletedOutputFromPayload(payloads[i], outputItemsByIndex, outputItemsFallback)
@@ -1716,9 +1857,9 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket(
 	}
 }
 
-func shouldReleaseResponsesWebsocketPinnedAuth(errMsg *interfaces.ErrorMessage) bool {
+func responsesWebsocketErrorStatus(errMsg *interfaces.ErrorMessage) int {
 	if errMsg == nil {
-		return false
+		return 0
 	}
 	status := errMsg.StatusCode
 	if status <= 0 && errMsg.Error != nil {
@@ -1726,7 +1867,23 @@ func shouldReleaseResponsesWebsocketPinnedAuth(errMsg *interfaces.ErrorMessage)
 			status = se.StatusCode()
 		}
 	}
-	switch status {
+	return status
+}
+
+func shouldReplayResponsesWebsocketPinnedAuthFailure(errMsg *interfaces.ErrorMessage) bool {
+	switch responsesWebsocketErrorStatus(errMsg) {
+	case http.StatusUnauthorized, http.StatusTooManyRequests:
+		return true
+	default:
+		return false
+	}
+}
+
+func shouldReleaseResponsesWebsocketPinnedAuth(errMsg *interfaces.ErrorMessage) bool {
+	if errMsg == nil {
+		return false
+	}
+	switch responsesWebsocketErrorStatus(errMsg) {
 	case http.StatusUnauthorized,
 		http.StatusPaymentRequired,
 		http.StatusForbidden,

sdk/api/handlers/openai/openai_responses_websocket_toolcall_repair.go

@@ -19,6 +19,7 @@ const (
 var defaultWebsocketToolOutputCache = newWebsocketToolOutputCache(0, websocketToolOutputCacheMaxPerSession)
 var defaultWebsocketToolCallCache = newWebsocketToolOutputCache(0, websocketToolOutputCacheMaxPerSession)
 var defaultWebsocketToolSessionRefs = newWebsocketToolSessionRefCounter()
+var defaultWebsocketToolCacheTransactionMu sync.RWMutex
 
 type websocketToolOutputCache struct {
 	mu            sync.Mutex
@@ -33,6 +34,14 @@ type websocketToolOutputSession struct {
 	order    []string
 }
 
+type responsesWebsocketToolCacheTurn struct {
+	sessionKey  string
+	outputs     map[string]json.RawMessage
+	outputOrder []string
+	calls       map[string]json.RawMessage
+	callOrder   []string
+}
+
 func newWebsocketToolOutputCache(ttl time.Duration, maxPerSession int) *websocketToolOutputCache {
 	if ttl < 0 {
 		ttl = websocketToolOutputCacheTTL
@@ -196,6 +205,8 @@ func (c *websocketToolSessionRefCounter) release(sessionKey string) bool {
 }
 
 func retainResponsesWebsocketToolCaches(sessionKey string) {
+	defaultWebsocketToolCacheTransactionMu.Lock()
+	defer defaultWebsocketToolCacheTransactionMu.Unlock()
 	if defaultWebsocketToolSessionRefs == nil {
 		return
 	}
@@ -203,13 +214,14 @@ func retainResponsesWebsocketToolCaches(sessionKey string) {
 }
 
 func releaseResponsesWebsocketToolCaches(sessionKey string) {
+	defaultWebsocketToolCacheTransactionMu.Lock()
+	defer defaultWebsocketToolCacheTransactionMu.Unlock()
 	if defaultWebsocketToolSessionRefs == nil {
 		return
 	}
 	if !defaultWebsocketToolSessionRefs.release(sessionKey) {
 		return
 	}
-
 	if defaultWebsocketToolOutputCache != nil {
 		defaultWebsocketToolOutputCache.deleteSession(sessionKey)
 	}
@@ -218,15 +230,114 @@ func releaseResponsesWebsocketToolCaches(sessionKey string) {
 	}
 }
 
+func newResponsesWebsocketToolCacheTurn(sessionKey string) *responsesWebsocketToolCacheTurn {
+	sessionKey = strings.TrimSpace(sessionKey)
+	if sessionKey == "" {
+		return nil
+	}
+	return &responsesWebsocketToolCacheTurn{
+		sessionKey: sessionKey,
+		outputs:    make(map[string]json.RawMessage),
+		calls:      make(map[string]json.RawMessage),
+	}
+}
+
+func (t *responsesWebsocketToolCacheTurn) recordRequest(payload []byte) {
+	if t == nil || len(payload) == 0 {
+		return
+	}
+	input := gjson.GetBytes(payload, "input")
+	if !input.Exists() || !input.IsArray() {
+		return
+	}
+	for _, item := range input.Array() {
+		t.recordItem(item)
+	}
+}
+
+func (t *responsesWebsocketToolCacheTurn) recordResponse(payload []byte) {
+	if t == nil || len(payload) == 0 {
+		return
+	}
+	switch strings.TrimSpace(gjson.GetBytes(payload, "type").String()) {
+	case "response.completed":
+		output := gjson.GetBytes(payload, "response.output")
+		if !output.Exists() || !output.IsArray() {
+			return
+		}
+		for _, item := range output.Array() {
+			if isCompleteResponsesWebsocketToolCall(item) {
+				t.recordItem(item)
+			}
+		}
+	case "response.output_item.added", "response.output_item.done":
+		item := gjson.GetBytes(payload, "item")
+		if isCompleteResponsesWebsocketToolCall(item) {
+			t.recordItem(item)
+		}
+	}
+}
+
+func (t *responsesWebsocketToolCacheTurn) recordItem(item gjson.Result) {
+	if t == nil || !item.Exists() {
+		return
+	}
+	callID := strings.TrimSpace(item.Get("call_id").String())
+	if callID == "" || strings.TrimSpace(item.Raw) == "" {
+		return
+	}
+	raw := append(json.RawMessage(nil), item.Raw...)
+	switch {
+	case isResponsesToolCallOutputType(item.Get("type").String()):
+		if _, exists := t.outputs[callID]; !exists {
+			t.outputOrder = append(t.outputOrder, callID)
+		}
+		t.outputs[callID] = raw
+	case isResponsesToolCallType(item.Get("type").String()):
+		if _, exists := t.calls[callID]; !exists {
+			t.callOrder = append(t.callOrder, callID)
+		}
+		t.calls[callID] = raw
+	}
+}
+
+func (t *responsesWebsocketToolCacheTurn) commit() {
+	if t == nil || t.sessionKey == "" {
+		return
+	}
+	defaultWebsocketToolCacheTransactionMu.Lock()
+	defer defaultWebsocketToolCacheTransactionMu.Unlock()
+	if defaultWebsocketToolOutputCache != nil {
+		for _, callID := range t.outputOrder {
+			defaultWebsocketToolOutputCache.record(t.sessionKey, callID, t.outputs[callID])
+		}
+	}
+	if defaultWebsocketToolCallCache != nil {
+		for _, callID := range t.callOrder {
+			defaultWebsocketToolCallCache.record(t.sessionKey, callID, t.calls[callID])
+		}
+	}
+}
+
 func repairResponsesWebsocketToolCalls(sessionKey string, payload []byte) []byte {
 	return repairResponsesWebsocketToolCallsWithCaches(defaultWebsocketToolOutputCache, defaultWebsocketToolCallCache, sessionKey, payload)
 }
 
+func repairResponsesWebsocketToolCallsWithoutRecording(sessionKey string, payload []byte) []byte {
+	defaultWebsocketToolCacheTransactionMu.RLock()
+	defer defaultWebsocketToolCacheTransactionMu.RUnlock()
+	return repairResponsesWebsocketToolCallsWithCachesMode(defaultWebsocketToolOutputCache, defaultWebsocketToolCallCache, sessionKey, payload, false)
+}
+
 func repairResponsesWebsocketToolCallsWithCache(cache *websocketToolOutputCache, sessionKey string, payload []byte) []byte {
 	return repairResponsesWebsocketToolCallsWithCaches(cache, nil, sessionKey, payload)
 }
 
 func repairResponsesWebsocketToolCallsWithCaches(outputCache, callCache *websocketToolOutputCache, sessionKey string, payload []byte) []byte {
+	return repairResponsesWebsocketToolCallsWithCachesMode(outputCache, callCache, sessionKey, payload, true)
+}
+
+func repairResponsesWebsocketToolCallsWithCachesMode(outputCache, callCache *websocketToolOutputCache, sessionKey string, payload []byte, record bool) []byte {
 	sessionKey = strings.TrimSpace(sessionKey)
 	if sessionKey == "" || outputCache == nil || len(payload) == 0 {
 		return payload
@@ -238,7 +349,7 @@ func repairResponsesWebsocketToolCallsWithCaches(outputCache, callCache *websock
 	}
 
 	allowOrphanOutputs := strings.TrimSpace(gjson.GetBytes(payload, "previous_response_id").String()) != ""
-	updatedRaw, errRepair := repairResponsesToolCallsArray(outputCache, callCache, sessionKey, input.Raw, allowOrphanOutputs)
+	updatedRaw, errRepair := repairResponsesToolCallsArray(outputCache, callCache, sessionKey, input.Raw, allowOrphanOutputs, record)
 	if errRepair != nil || updatedRaw == "" || updatedRaw == input.Raw {
 		return payload
 	}
@@ -250,7 +361,7 @@ func repairResponsesWebsocketToolCallsWithCaches(outputCache, callCache *websock
 	return updated
 }
 
-func repairResponsesToolCallsArray(outputCache, callCache *websocketToolOutputCache, sessionKey string, rawArray string, allowOrphanOutputs bool) (string, error) {
+func repairResponsesToolCallsArray(outputCache, callCache *websocketToolOutputCache, sessionKey string, rawArray string, allowOrphanOutputs bool, record bool) (string, error) {
 	rawArray = strings.TrimSpace(rawArray)
 	if rawArray == "" {
 		return "[]", nil
@@ -276,14 +387,16 @@ func repairResponsesToolCallsArray(outputCache, callCache *websocketToolOutputCa
 				continue
 			}
 			outputPresent[callID] = struct{}{}
-			outputCache.record(sessionKey, callID, item)
+			if record {
+				outputCache.record(sessionKey, callID, item)
+			}
 		case isResponsesToolCallType(itemType):
 			callID := strings.TrimSpace(gjson.GetBytes(item, "call_id").String())
 			if callID == "" {
 				continue
 			}
 			callPresent[callID] = struct{}{}
-			if callCache != nil {
+			if record && callCache != nil {
 				callCache.record(sessionKey, callID, item)
 			}
 		}