fix(responses): commit websocket tool cache atomically

sdk/api/handlers/openai/openai_responses_websocket.go

@@ -647,6 +647,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
 			continue
 		}
 
+		toolCacheTurn := newResponsesWebsocketToolCacheTurn(downstreamSessionKey)
 		previousLastRequest := bytes.Clone(lastRequest)
 		previousLastResponseOutput := bytes.Clone(lastResponseOutput)
 		previousLastResponseID := lastResponseID
@@ -656,7 +657,8 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
 				passthroughModelName = modelName
 			}
 		} else {
-			requestJSON = repairResponsesWebsocketToolCalls(downstreamSessionKey, requestJSON)
+			toolCacheTurn.recordRequest(requestJSON)
+			requestJSON = repairResponsesWebsocketToolCallsWithoutRecording(downstreamSessionKey, requestJSON)
 			requestJSON = dedupeResponsesWebsocketInputItemsByID(requestJSON)
 			updatedLastRequest = bytes.Clone(requestJSON)
 			lastRequest = updatedLastRequest
@@ -712,6 +714,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
 			wsTimelineLog,
 			passthroughSessionID,
 			responsesWebsocketForwardOptions{
+				toolCacheTurn: toolCacheTurn,
 				suppressError: replayPinnedAuthFailure,
 			},
 		)
@@ -744,6 +747,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) {
 			continue
 		}
 
+		toolCacheTurn.commit()
 		upstreamMode = attemptedUpstreamMode
 		if upstreamMode == responsesWebsocketUpstreamModeWS {
 			upstreamWebsocketAuthID = lastAttemptedAuthID
@@ -1685,6 +1689,7 @@ func normalizeJSONArrayRaw(raw []byte) string {
 }
 
 type responsesWebsocketForwardOptions struct {
+	toolCacheTurn *responsesWebsocketToolCacheTurn
 	suppressError func(*interfaces.ErrorMessage) bool
 }
 
@@ -1702,6 +1707,7 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket(
 	if len(options) > 0 {
 		opts = options[0]
 	}
+	toolCacheTurn := opts.toolCacheTurn
 	completed := false
 	completedOutput := []byte("[]")
 	completedResponseID := ""
@@ -1803,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 {

sdk/api/handlers/openai/openai_responses_websocket_test.go

@@ -1608,6 +1608,44 @@ func TestRepairResponsesWebsocketToolCallsInsertsCachedOutput(t *testing.T) {
 	}
 }
 
+func TestResponsesWebsocketToolCacheTurnCommitsOnlyOnSuccess(t *testing.T) {
+	const sessionKey = "tool-cache-turn-commit-session"
+	defer defaultWebsocketToolOutputCache.deleteSession(sessionKey)
+	defer defaultWebsocketToolCallCache.deleteSession(sessionKey)
+
+	turn := newResponsesWebsocketToolCacheTurn(sessionKey)
+	turn.recordRequest([]byte(`{"input":[{"type":"function_call_output","id":"fco-1","call_id":"call-1","output":"cached result"}]}`))
+	beforeCommit := repairResponsesWebsocketToolCallsWithoutRecording(sessionKey, []byte(`{"input":[{"type":"function_call","id":"fc-next","call_id":"call-1","name":"lookup","arguments":"{}"}]}`))
+	if gjson.GetBytes(beforeCommit, "input.#").Int() != 0 {
+		t.Fatalf("uncommitted turn populated global cache: %s", beforeCommit)
+	}
+
+	turn.commit()
+	afterCommit := repairResponsesWebsocketToolCallsWithoutRecording(sessionKey, []byte(`{"input":[{"type":"function_call","id":"fc-next","call_id":"call-1","name":"lookup","arguments":"{}"}]}`))
+	input := gjson.GetBytes(afterCommit, "input").Array()
+	if len(input) != 2 || input[1].Get("output").String() != "cached result" {
+		t.Fatalf("committed turn was not available to tool repair: %s", afterCommit)
+	}
+}
+
+func TestResponsesWebsocketToolCacheRetainPreventsOverlappingReleaseDeletion(t *testing.T) {
+	const sessionKey = "tool-cache-overlapping-retain-session"
+	retainResponsesWebsocketToolCaches(sessionKey)
+	retainResponsesWebsocketToolCaches(sessionKey)
+	turn := newResponsesWebsocketToolCacheTurn(sessionKey)
+	turn.recordRequest([]byte(`{"input":[{"type":"function_call_output","id":"fco-1","call_id":"call-1","output":"kept"}]}`))
+	turn.commit()
+
+	releaseResponsesWebsocketToolCaches(sessionKey)
+	if _, ok := defaultWebsocketToolOutputCache.get(sessionKey, "call-1"); !ok {
+		t.Fatal("first overlapping release deleted active session cache")
+	}
+	releaseResponsesWebsocketToolCaches(sessionKey)
+	if _, ok := defaultWebsocketToolOutputCache.get(sessionKey, "call-1"); ok {
+		t.Fatal("final release did not delete session cache")
+	}
+}
+
 func TestRepairResponsesWebsocketToolCallsDropsOrphanFunctionCall(t *testing.T) {
 	cache := newWebsocketToolOutputCache(time.Minute, 10)
 	sessionKey := "session-1"
@@ -2933,6 +2971,86 @@ func TestResponsesWebsocketRejectsUnknownPreviousResponseOnNewSocket(t *testing.
 	}
 }
 
+func TestResponsesWebsocketRollsBackCanonicalTranscriptAfterNonRetryableError(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	modelName := "xai-websocket-rollback-model"
+	executor := &websocketCanonicalRollbackExecutor{}
+	manager := coreauth.NewManager(nil, nil, nil)
+	manager.RegisterExecutor(executor)
+	auth := &coreauth.Auth{
+		ID:       "auth-xai-rollback",
+		Provider: "xai",
+		Status:   coreauth.StatusActive,
+	}
+	if _, err := manager.Register(context.Background(), auth); err != nil {
+		t.Fatalf("Register auth: %v", err)
+	}
+	registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: modelName}})
+	t.Cleanup(func() {
+		registry.GetGlobalRegistry().UnregisterClient(auth.ID)
+	})
+
+	base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
+	h := NewOpenAIResponsesAPIHandler(base)
+	router := gin.New()
+	router.GET("/v1/responses/ws", h.ResponsesWebsocket)
+
+	server := httptest.NewServer(router)
+	defer server.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses/ws"
+	conn, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{"Session-Id": []string{"rollback-tool-cache-session"}})
+	if err != nil {
+		t.Fatalf("dial websocket: %v", err)
+	}
+	defer func() { _ = conn.Close() }()
+
+	requests := []string{
+		fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"}]}`, modelName),
+		`{"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"}]}`,
+		`{"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":"{}"}]}`,
+	}
+	wantTypes := []string{wsEventTypeCompleted, wsEventTypeError, wsEventTypeCompleted}
+	for i := range requests {
+		if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(requests[i])); errWrite != nil {
+			t.Fatalf("write websocket message %d: %v", i+1, errWrite)
+		}
+		_, payload, errRead := conn.ReadMessage()
+		if errRead != nil {
+			t.Fatalf("read websocket response %d: %v", i+1, errRead)
+		}
+		if got := gjson.GetBytes(payload, "type").String(); got != wantTypes[i] {
+			t.Fatalf("response %d type = %q, want %q: %s", i+1, got, wantTypes[i], payload)
+		}
+	}
+
+	payloads := executor.Payloads()
+	if len(payloads) != 3 {
+		t.Fatalf("executor payload count = %d, want 3", len(payloads))
+	}
+	third := payloads[2]
+	if gjson.GetBytes(third, "previous_response_id").Exists() {
+		t.Fatalf("retry payload must not depend on previous_response_id: %s", third)
+	}
+	input := gjson.GetBytes(third, "input").Array()
+	if len(input) != 3 {
+		t.Fatalf("retry canonical input len = %d, want 3: %s", len(input), third)
+	}
+	wantIDs := []string{"msg-1", "out-1", "msg-3"}
+	for i, wantID := range wantIDs {
+		if got := input[i].Get("id").String(); got != wantID {
+			t.Fatalf("retry canonical input[%d].id = %q, want %q: %s", i, got, wantID, third)
+		}
+	}
+	if bytes.Contains(third, []byte(`"id":"fc-failed"`)) || bytes.Contains(third, []byte(`"id":"fco-failed"`)) {
+		t.Fatalf("failed turn leaked into retry transcript: %s", third)
+	}
+	if bytes.Contains(third, []byte(`"call_id":"failed-call"`)) || bytes.Contains(third, []byte("must-not-survive")) {
+		t.Fatalf("failed turn contaminated tool repair cache: %s", third)
+	}
+}
+
 func TestResponsesWebsocketSwitchesPinnedAuthAcrossProviders(t *testing.T) {
 	for _, testCase := range []struct {
 		name                      string

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