diff --git a/internal/runtime/executor/xai_websockets_executor.go b/internal/runtime/executor/xai_websockets_executor.go index 1f04538e..1e2e7c82 100644 --- a/internal/runtime/executor/xai_websockets_executor.go +++ b/internal/runtime/executor/xai_websockets_executor.go @@ -48,18 +48,21 @@ type xaiWebsocketIDStateStore struct { } type xaiWebsocketIDState struct { - mu sync.Mutex - downstreamToUpstream map[string]string - sequence int - transcriptInput []json.RawMessage + requestMu sync.Mutex + mu sync.Mutex + downstreamToUpstream map[string]string + sequence int + transcriptInput []json.RawMessage + replayCompactedTranscriptOnReset bool } type xaiWebsocketRequestIDMapper struct { - state *xaiWebsocketIDState - downstreamPreviousID string - upstreamPreviousID string - upstreamResponseID string - downstreamResponseID string + state *xaiWebsocketIDState + downstreamPreviousID string + upstreamPreviousID string + upstreamResponseID string + downstreamResponseID string + replayedCompactedTranscript bool } func NewXAIWebsocketsExecutor(cfg *config.Config) *XAIWebsocketsExecutor { @@ -177,20 +180,21 @@ func (s *xaiWebsocketIDState) prependTranscriptInput(payload []byte) []byte { return out } -func (s *xaiWebsocketIDState) recordTranscriptTurn(requestPayload []byte, completedPayload []byte) { +func (s *xaiWebsocketIDState) recordTranscriptTurn(requestPayload []byte, completedPayload []byte, reset bool) { if s == nil || len(requestPayload) == 0 || len(completedPayload) == 0 { return } inputItems := xaiJSONRawMessages(gjson.GetBytes(requestPayload, "input")) outputItems := xaiJSONRawMessages(gjson.GetBytes(completedPayload, "response.output")) - if len(inputItems) == 0 && len(outputItems) == 0 { - return - } s.mu.Lock() defer s.mu.Unlock() - if strings.TrimSpace(gjson.GetBytes(requestPayload, "previous_response_id").String()) == "" { + if reset { s.transcriptInput = nil + s.replayCompactedTranscriptOnReset = false + } + if len(inputItems) == 0 && len(outputItems) == 0 { + return } s.transcriptInput = append(s.transcriptInput, inputItems...) s.transcriptInput = append(s.transcriptInput, outputItems...) @@ -210,9 +214,34 @@ func (s *xaiWebsocketIDState) replaceTranscriptWithItems(items ...[]byte) { } s.mu.Lock() s.transcriptInput = next + s.replayCompactedTranscriptOnReset = len(next) > 0 s.mu.Unlock() } +func (s *xaiWebsocketIDState) prependCompactedTranscriptOnReset(payload []byte) ([]byte, bool) { + if s == nil || len(payload) == 0 { + return payload, false + } + s.mu.Lock() + if !s.replayCompactedTranscriptOnReset || len(s.transcriptInput) == 0 { + s.mu.Unlock() + return payload, false + } + prefix := make([]json.RawMessage, 0, len(s.transcriptInput)) + for _, item := range s.transcriptInput { + prefix = append(prefix, bytes.Clone(item)) + } + s.mu.Unlock() + + current := xaiJSONRawMessages(gjson.GetBytes(payload, "input")) + merged := append(prefix, current...) + out, errSet := sjson.SetRawBytes(payload, "input", xaiMarshalRawMessages(merged)) + if errSet != nil { + return payload, false + } + return out, true +} + func xaiJSONRawMessages(result gjson.Result) []json.RawMessage { if !result.Exists() || !result.IsArray() { return nil @@ -243,7 +272,16 @@ func xaiMarshalRawMessages(items []json.RawMessage) []byte { } func (m *xaiWebsocketRequestIDMapper) upstreamRequestPayload(payload []byte) []byte { - if m == nil || len(payload) == 0 || m.downstreamPreviousID == m.upstreamPreviousID { + if m == nil || len(payload) == 0 { + return payload + } + if m.downstreamPreviousID == m.upstreamPreviousID { + requestType := strings.TrimSpace(gjson.GetBytes(payload, "type").String()) + if m.downstreamPreviousID == "" && requestType == "response.append" && m.state != nil { + out, replayed := m.state.prependCompactedTranscriptOnReset(payload) + m.replayedCompactedTranscript = replayed + return out + } return payload } if m.upstreamPreviousID == "" { @@ -251,6 +289,7 @@ func (m *xaiWebsocketRequestIDMapper) upstreamRequestPayload(payload []byte) []b if errDelete == nil { if m.downstreamPreviousID != "" && m.state != nil { out = m.state.prependTranscriptInput(out) + m.replayedCompactedTranscript = true } return out } @@ -396,11 +435,30 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox if stateSessionID == "" { stateSessionID = executionSessionID } - idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload) + state := getXAIWebsocketIDState(e.idStore, stateSessionID) + stateRequestLocked := false + stateRequestLockTransferred := false + if executionSessionID == "" && state != nil { + state.requestMu.Lock() + stateRequestLocked = true + } + defer func() { + if stateRequestLocked && !stateRequestLockTransferred { + state.requestMu.Unlock() + } + }() if xaiInputHasItemType(req.Payload, "compaction_trigger") { if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() } + if executionSessionID != "" { + sess := e.getOrCreateSession(executionSessionID) + if sess != nil { + sess.reqMu.Lock() + defer sess.reqMu.Unlock() + } + } + idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload) return e.executeCompactionTriggerFromWebsocketContext(ctx, auth, req, opts, idMapper) } @@ -440,6 +498,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox sess.reqMu.Lock() } } + idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload) if idMapper != nil { if websocketSessionTargetChanged(sess, authID, wsURL) { idMapper.upstreamPreviousID = "" @@ -450,6 +509,9 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox wsHeaders := applyXAIWebsocketHeaders(http.Header{}, auth, token, prepared.sessionID) wsReqBody := buildXAIWebsocketRequestBody(prepared.body) + requestType := strings.TrimSpace(gjson.GetBytes(req.Payload, "type").String()) + transcriptReset := strings.TrimSpace(gjson.GetBytes(wsReqBody, "previous_response_id").String()) == "" && + (requestType != "response.append" || (idMapper != nil && idMapper.replayedCompactedTranscript)) warmupRequest := xaiWebsocketGenerateFalse(wsReqBody) wsReqLog := helps.UpstreamRequestLog{ @@ -580,7 +642,13 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox } out := make(chan cliproxyexecutor.StreamChunk) + if stateRequestLocked { + stateRequestLockTransferred = true + } go func() { + if stateRequestLocked { + defer state.requestMu.Unlock() + } terminateReason := "completed" var terminateErr error @@ -687,6 +755,10 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox case "response.created": if warmupRequest { warmupCompletedPayload = buildXAIWebsocketWarmupCompletedPayload(payload) + if idMapper != nil && idMapper.state != nil && !recordedTranscript { + idMapper.state.recordTranscriptTurn(wsReqBody, warmupCompletedPayload, transcriptReset) + recordedTranscript = true + } logXAIWebsocketWarmupCompleted(executionSessionID, authID, wsURL, payload) } case "response.output_item.done": @@ -700,7 +772,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox payload = xaiNormalizeReasoningSummaryData(payload) cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, payload) if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript { - idMapper.state.recordTranscriptTurn(wsReqBody, payload) + idMapper.state.recordTranscriptTurn(wsReqBody, payload, transcriptReset) recordedTranscript = true } case "response.done": @@ -709,7 +781,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox reporter.Publish(ctx, detail) } if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript { - idMapper.state.recordTranscriptTurn(wsReqBody, payload) + idMapper.state.recordTranscriptTurn(wsReqBody, payload, transcriptReset) recordedTranscript = true } } @@ -803,8 +875,11 @@ func (e *XAIWebsocketsExecutor) executeCompactionTriggerFromWebsocketContext(ctx return nil, err } - responseID := xaiCompactionResponseID(data) - idMapper.state.replaceTranscriptWithItems(xaiCompactionOutputItem(data, responseID)) + responseID, compactionItem, errValidate := validateXAIWebsocketCompactionResponse(data) + if errValidate != nil { + return nil, errValidate + } + idMapper.state.replaceTranscriptWithItems(compactionItem) idMapper.state.mapDownstreamToUpstream(responseID, "") headers = headers.Clone() @@ -822,6 +897,30 @@ func (e *XAIWebsocketsExecutor) executeCompactionTriggerFromWebsocketContext(ctx return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out}, nil } +func validateXAIWebsocketCompactionResponse(data []byte) (string, []byte, error) { + if len(data) == 0 || !json.Valid(data) { + return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction returned invalid JSON"} + } + responseIDResult := gjson.GetBytes(data, "id") + output := gjson.GetBytes(data, "output") + if responseIDResult.Type != gjson.String || strings.TrimSpace(responseIDResult.String()) == "" || !output.Exists() || !output.IsArray() { + return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction response is missing compacted state"} + } + items := output.Array() + if len(items) == 0 { + return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction response is missing compacted state"} + } + item := items[0] + itemType := item.Get("type") + encryptedContent := item.Get("encrypted_content") + if item.Type != gjson.JSON || itemType.Type != gjson.String || strings.TrimSpace(itemType.String()) != "compaction" || + encryptedContent.Type != gjson.String || strings.TrimSpace(encryptedContent.String()) == "" { + return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction response is missing compacted state"} + } + normalizedResponseID := xaiCompactionResponseID(data) + return normalizedResponseID, xaiCompactionOutputItem(data, normalizedResponseID), nil +} + func buildXAIWebsocketCompactionPayload(payload []byte, transcriptInput []byte) ([]byte, error) { if len(payload) == 0 { payload = []byte(`{}`) diff --git a/internal/runtime/executor/xai_websockets_executor_test.go b/internal/runtime/executor/xai_websockets_executor_test.go index 75ceae3c..065e73f5 100644 --- a/internal/runtime/executor/xai_websockets_executor_test.go +++ b/internal/runtime/executor/xai_websockets_executor_test.go @@ -1290,6 +1290,127 @@ func TestXAIWebsocketsExecuteStreamCompactionTriggerUsesHTTPCompactWithRecordedC } } +func TestXAIWebsocketPostCompactionAppendWithoutPreviousReplaysCompactedTranscript(t *testing.T) { + store := &xaiWebsocketIDStateStore{sessions: make(map[string]*xaiWebsocketIDState)} + state := getXAIWebsocketIDState(store, "post-compaction-append-session") + state.replaceTranscriptWithItems([]byte(`{"type":"compaction","encrypted_content":"compact-state"}`)) + state.mapDownstreamToUpstream("resp-compact", "") + + fullReset := []byte(`{"type":"response.create","model":"grok-4.3","input":[{"type":"message","id":"msg-full"}]}`) + fullMapper := newXAIWebsocketRequestIDMapper(store, "post-compaction-append-session", fullReset) + if full := fullMapper.upstreamRequestPayload(fullReset); len(gjson.GetBytes(full, "input").Array()) != 1 { + t.Fatalf("self-contained response.create unexpectedly replayed compacted transcript: %s", full) + } + + payload := []byte(`{"type":"response.append","model":"grok-4.3","input":[{"type":"message","id":"msg-2","role":"user","content":"second"}]}`) + mapper := newXAIWebsocketRequestIDMapper(store, "post-compaction-append-session", payload) + got := mapper.upstreamRequestPayload(payload) + input := gjson.GetBytes(got, "input").Array() + if len(input) != 2 { + t.Fatalf("post-compaction append input len = %d, want 2: %s", len(input), got) + } + if gotType := input[0].Get("type").String(); gotType != "compaction" { + t.Fatalf("post-compaction append input[0].type = %q, want compaction: %s", gotType, got) + } + if gotID := input[1].Get("id").String(); gotID != "msg-2" { + t.Fatalf("post-compaction append input[1].id = %q, want msg-2: %s", gotID, got) + } + + state.recordTranscriptTurn(got, []byte(`{"type":"response.completed","response":{"id":"resp-after-compact","output":[{"type":"message","id":"out-2"}]}}`), true) + nextPayload := []byte(`{"type":"response.create","model":"grok-4.3","input":[{"type":"message","id":"msg-3"}]}`) + nextMapper := newXAIWebsocketRequestIDMapper(store, "post-compaction-append-session", nextPayload) + next := nextMapper.upstreamRequestPayload(nextPayload) + if nextInput := gjson.GetBytes(next, "input").Array(); len(nextInput) != 1 || nextInput[0].Get("id").String() != "msg-3" { + t.Fatalf("compacted transcript replay was not cleared after success: %s", next) + } +} + +func TestXAIWebsocketPostCompactionWarmupPreservesTranscriptForLaterCompaction(t *testing.T) { + store := &xaiWebsocketIDStateStore{sessions: make(map[string]*xaiWebsocketIDState)} + state := getXAIWebsocketIDState(store, "warmup-reset-session") + state.replaceTranscriptWithItems([]byte(`{"type":"compaction","encrypted_content":"compact-state"}`)) + + warmupPayload := []byte(`{"type":"response.append","model":"grok-4.3","generate":false,"input":[{"type":"message","id":"warmup-context"}]}`) + warmupMapper := newXAIWebsocketRequestIDMapper(store, "warmup-reset-session", warmupPayload) + warmupUpstream := warmupMapper.upstreamRequestPayload(warmupPayload) + if !warmupMapper.replayedCompactedTranscript { + t.Fatal("post-compaction warmup did not mark full transcript replay") + } + state.recordTranscriptTurn( + warmupUpstream, + []byte(`{"type":"response.completed","response":{"id":"resp-warmup","output":[]}}`), + true, + ) + + appendPayload := []byte(`{"type":"response.append","model":"grok-4.3","input":[{"type":"message","id":"msg-after-warmup"}]}`) + appendMapper := newXAIWebsocketRequestIDMapper(store, "warmup-reset-session", appendPayload) + appendUpstream := appendMapper.upstreamRequestPayload(appendPayload) + input := gjson.GetBytes(appendUpstream, "input").Array() + if len(input) != 1 || input[0].Get("id").String() != "msg-after-warmup" { + t.Fatalf("warmup retained pending replay instead of native append: %s", appendUpstream) + } + state.recordTranscriptTurn( + appendUpstream, + []byte(`{"type":"response.completed","response":{"id":"resp-after-warmup","output":[{"type":"message","id":"out-after-warmup"}]}}`), + false, + ) + + transcript := gjson.ParseBytes(state.snapshotTranscriptInput()).Array() + wantTypes := []string{"compaction", "message", "message", "message"} + if len(transcript) != len(wantTypes) { + t.Fatalf("post-warmup transcript len = %d, want %d: %s", len(transcript), len(wantTypes), state.snapshotTranscriptInput()) + } + for i, wantType := range wantTypes { + if gotType := transcript[i].Get("type").String(); gotType != wantType { + t.Fatalf("post-warmup transcript[%d].type = %q, want %q: %s", i, gotType, wantType, state.snapshotTranscriptInput()) + } + } +} + +func TestXAIWebsocketEmptyFullResetClearsPendingCompactionReplay(t *testing.T) { + store := &xaiWebsocketIDStateStore{sessions: make(map[string]*xaiWebsocketIDState)} + state := getXAIWebsocketIDState(store, "empty-reset-session") + state.replaceTranscriptWithItems([]byte(`{"type":"compaction","encrypted_content":"stale-compact-state"}`)) + state.recordTranscriptTurn( + []byte(`{"type":"response.create","model":"grok-4.3","input":[]}`), + []byte(`{"type":"response.completed","response":{"id":"resp-empty","output":[]}}`), + true, + ) + + appendPayload := []byte(`{"type":"response.append","model":"grok-4.3","input":[{"type":"message","id":"msg-new"}]}`) + mapper := newXAIWebsocketRequestIDMapper(store, "empty-reset-session", appendPayload) + got := mapper.upstreamRequestPayload(appendPayload) + input := gjson.GetBytes(got, "input").Array() + if len(input) != 1 || input[0].Get("id").String() != "msg-new" { + t.Fatalf("empty full reset retained stale compaction replay: %s", got) + } +} + +func TestValidateXAIWebsocketCompactionResponse(t *testing.T) { + valid := []byte(`{"id":"resp_compact","output":[{"type":"compaction","encrypted_content":"opaque-state"}]}`) + responseID, item, err := validateXAIWebsocketCompactionResponse(valid) + if err != nil { + t.Fatalf("valid compaction response error: %v", err) + } + if responseID != "resp_compact" || gjson.GetBytes(item, "encrypted_content").String() != "opaque-state" { + t.Fatalf("validated compaction response = id:%q item:%s", responseID, item) + } + + for _, payload := range [][]byte{ + nil, + []byte(`{}`), + []byte(`{"id":"resp_empty","output":[]}`), + []byte(`{"id":123,"output":[{"type":"compaction","encrypted_content":"opaque"}]}`), + []byte(`{"id":"resp_object","output":{"0":{"type":"compaction","encrypted_content":"opaque"}}}`), + []byte(`{"id":"resp_numeric_state","output":[{"type":"compaction","encrypted_content":123}]}`), + []byte(`{"id":"resp_missing_state","output":[{"type":"compaction"}]}`), + } { + if _, _, errInvalid := validateXAIWebsocketCompactionResponse(payload); errInvalid == nil { + t.Fatalf("invalid compaction response accepted: %s", payload) + } + } +} + func TestBuildXAIWebsocketRequestBodySetsStoreAndKeepsPromptCacheKey(t *testing.T) { 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"}]}`)