fix(xai): preserve compacted websocket transcript state

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(`{}`)

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"}]}`)