diff --git a/sdk/api/handlers/openai/openai_responses_websocket.go b/sdk/api/handlers/openai/openai_responses_websocket.go index a0aa52c9..2d017d36 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket.go +++ b/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 { diff --git a/sdk/api/handlers/openai/openai_responses_websocket_test.go b/sdk/api/handlers/openai/openai_responses_websocket_test.go index f50f4469..99f42623 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket_test.go +++ b/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 diff --git a/sdk/api/handlers/openai/openai_responses_websocket_toolcall_repair.go b/sdk/api/handlers/openai/openai_responses_websocket_toolcall_repair.go index 0c3ab58d..85260ad4 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket_toolcall_repair.go +++ b/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) } }