From f33bc56bb947134e4fc10bdd3212404c56dcbaf8 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Mon, 15 Jun 2026 02:41:35 +0000 Subject: [PATCH] feat(websockets): add transcript state tracking and compaction trigger support - Added methods for managing and tracking WebSocket transcript state, including recording, prepending, and replacing transcript inputs. - Implemented `executeCompactionTriggerFromWebsocketContext` to support compaction triggers using recorded transcript context. - Enhanced upstream-downstream ID mapping with additional utilities and state synchronization. - Expanded test coverage to validate transcript state management, compaction payload generation, and WebSocket response handling. --- internal/runtime/executor/xai_websockets_executor.go | 205 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++----- internal/runtime/executor/xai_websockets_executor_test.go | 165 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 file(s) changed, 365 insertion(s)(+), 5 deletion(s)(-) diff --git a/internal/runtime/executor/xai_websockets_executor.go b/internal/runtime/executor/xai_websockets_executor.go --- a/internal/runtime/executor/xai_websockets_executor.go +++ b/internal/runtime/executor/xai_websockets_executor.go @@ -52,6 +52,7 @@ mu sync.Mutex downstreamToUpstream map[string]string sequence int + transcriptInput []json.RawMessage } type xaiWebsocketRequestIDMapper struct { @@ -124,10 +125,122 @@ } s.mu.Lock() defer s.mu.Unlock() - if upstreamID := strings.TrimSpace(s.downstreamToUpstream[downstreamID]); upstreamID != "" { - return upstreamID + if upstreamID, ok := s.downstreamToUpstream[downstreamID]; ok { + return strings.TrimSpace(upstreamID) } return downstreamID +} + +func (s *xaiWebsocketIDState) mapDownstreamToUpstream(downstreamID string, upstreamID string) { + downstreamID = strings.TrimSpace(downstreamID) + if s == nil || downstreamID == "" { + return + } + s.mu.Lock() + if s.downstreamToUpstream == nil { + s.downstreamToUpstream = make(map[string]string) + } + s.downstreamToUpstream[downstreamID] = strings.TrimSpace(upstreamID) + s.mu.Unlock() +} + +func (s *xaiWebsocketIDState) snapshotTranscriptInput() []byte { + if s == nil { + return nil + } + s.mu.Lock() + defer s.mu.Unlock() + if len(s.transcriptInput) == 0 { + return nil + } + return xaiMarshalRawMessages(s.transcriptInput) +} + +func (s *xaiWebsocketIDState) prependTranscriptInput(payload []byte) []byte { + if s == nil || len(payload) == 0 { + return payload + } + s.mu.Lock() + prefix := make([]json.RawMessage, 0, len(s.transcriptInput)) + for _, item := range s.transcriptInput { + prefix = append(prefix, bytes.Clone(item)) + } + s.mu.Unlock() + if len(prefix) == 0 { + return payload + } + current := xaiJSONRawMessages(gjson.GetBytes(payload, "input")) + merged := append(prefix, current...) + out, errSet := sjson.SetRawBytes(payload, "input", xaiMarshalRawMessages(merged)) + if errSet != nil { + return payload + } + return out +} + +func (s *xaiWebsocketIDState) recordTranscriptTurn(requestPayload []byte, completedPayload []byte) { + 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()) == "" { + s.transcriptInput = nil + } + s.transcriptInput = append(s.transcriptInput, inputItems...) + s.transcriptInput = append(s.transcriptInput, outputItems...) +} + +func (s *xaiWebsocketIDState) replaceTranscriptWithItems(items ...[]byte) { + if s == nil { + return + } + next := make([]json.RawMessage, 0, len(items)) + for _, item := range items { + item = bytes.TrimSpace(item) + if len(item) == 0 || !json.Valid(item) { + continue + } + next = append(next, bytes.Clone(item)) + } + s.mu.Lock() + s.transcriptInput = next + s.mu.Unlock() +} + +func xaiJSONRawMessages(result gjson.Result) []json.RawMessage { + if !result.Exists() || !result.IsArray() { + return nil + } + items := result.Array() + out := make([]json.RawMessage, 0, len(items)) + for _, item := range items { + raw := bytes.TrimSpace([]byte(item.Raw)) + if len(raw) == 0 || !json.Valid(raw) { + continue + } + out = append(out, bytes.Clone(raw)) + } + return out +} + +func xaiMarshalRawMessages(items []json.RawMessage) []byte { + var buf bytes.Buffer + buf.WriteByte('[') + for i, item := range items { + if i > 0 { + buf.WriteByte(',') + } + buf.Write(bytes.TrimSpace(item)) + } + buf.WriteByte(']') + return buf.Bytes() } func (m *xaiWebsocketRequestIDMapper) upstreamRequestPayload(payload []byte) []byte { @@ -137,6 +250,9 @@ if m.upstreamPreviousID == "" { out, errDelete := sjson.DeleteBytes(payload, "previous_response_id") if errDelete == nil { + if m.downstreamPreviousID != "" && m.state != nil { + out = m.state.prependTranscriptInput(out) + } return out } return payload @@ -275,12 +391,16 @@ if opts.Alt == "responses/compact" { return nil, statusErr{code: http.StatusBadRequest, msg: "streaming not supported for /responses/compact"} } + executionSessionID := executionSessionIDFromOptions(opts) + stateSessionID := xaiExecutionSessionID(req, opts) + if stateSessionID == "" { + stateSessionID = executionSessionID + } + idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload) if xaiInputHasItemType(req.Payload, "compaction_trigger") { - return e.XAIExecutor.ExecuteStream(ctx, auth, req, opts) + return e.executeCompactionTriggerFromWebsocketContext(ctx, auth, req, opts, idMapper) } - executionSessionID := executionSessionIDFromOptions(opts) - idMapper := newXAIWebsocketRequestIDMapper(e.idStore, executionSessionID, req.Payload) token, baseURL := xaiCreds(auth) if baseURL == "" { baseURL = xaiauth.DefaultAPIBaseURL @@ -450,6 +570,7 @@ var param any outputItemsByIndex := make(map[int64][]byte) var outputItemsFallback [][]byte + recordedTranscript := false for { if ctx != nil && ctx.Err() != nil { terminateReason = "context_done" @@ -524,10 +645,18 @@ reporter.Publish(ctx, detail) } payload = xaiPatchCompletedOutput(payload, outputItemsByIndex, outputItemsFallback) + if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript { + idMapper.state.recordTranscriptTurn(wsReqBody, payload) + recordedTranscript = true + } case "response.done": logXAIWebsocketTerminalResponse(executionSessionID, authID, wsURL, eventType, payload) if detail, ok := helps.ParseCodexUsage(payload); ok { reporter.Publish(ctx, detail) + } + if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript { + idMapper.state.recordTranscriptTurn(wsReqBody, payload) + recordedTranscript = true } } @@ -587,6 +716,72 @@ } }() return &cliproxyexecutor.StreamResult{Headers: upstreamHeaders, Chunks: out}, nil +} + +func (e *XAIWebsocketsExecutor) executeCompactionTriggerFromWebsocketContext(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, idMapper *xaiWebsocketRequestIDMapper) (*cliproxyexecutor.StreamResult, error) { + if idMapper == nil || idMapper.state == nil { + return nil, statusErr{code: http.StatusBadRequest, msg: "xai websocket compaction context is unavailable"} + } + transcriptInput := idMapper.state.snapshotTranscriptInput() + if len(transcriptInput) == 0 { + return nil, statusErr{code: http.StatusBadRequest, msg: "xai websocket compaction context is empty"} + } + authID := "" + if auth != nil { + authID = auth.ID + } + log.Infof( + "xai websockets: compact fallback session=%s auth=%s input_items=%d", + xaiExecutionSessionID(req, opts), + strings.TrimSpace(authID), + len(gjson.ParseBytes(transcriptInput).Array()), + ) + compactPayload, err := buildXAIWebsocketCompactionPayload(req.Payload, transcriptInput) + if err != nil { + return nil, err + } + compactReq := req + compactReq.Payload = compactPayload + + prepared, data, headers, err := e.XAIExecutor.executeCompactRequest(ctx, auth, compactReq, opts) + if err != nil { + return nil, err + } + + responseID := xaiCompactionResponseID(data) + idMapper.state.replaceTranscriptWithItems(xaiCompactionOutputItem(data, responseID)) + idMapper.state.mapDownstreamToUpstream(responseID, "") + + headers = headers.Clone() + if headers == nil { + headers = make(http.Header) + } + headers.Set("Content-Type", "text/event-stream") + + chunks := xaiBuildCompactionTriggerStreamChunks(prepared, data) + out := make(chan cliproxyexecutor.StreamChunk, len(chunks)) + for _, chunk := range chunks { + out <- cliproxyexecutor.StreamChunk{Payload: chunk} + } + close(out) + return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out}, nil +} + +func buildXAIWebsocketCompactionPayload(payload []byte, transcriptInput []byte) ([]byte, error) { + if len(payload) == 0 { + payload = []byte(`{}`) + } + if len(transcriptInput) == 0 { + transcriptInput = []byte("[]") + } + out := bytes.Clone(payload) + var err error + out, err = sjson.SetRawBytes(out, "input", transcriptInput) + if err != nil { + return nil, err + } + out, _ = sjson.DeleteBytes(out, "previous_response_id") + return out, nil } func xaiWebsocketGenerateFalse(payload []byte) bool { diff --git a/internal/runtime/executor/xai_websockets_executor_test.go b/internal/runtime/executor/xai_websockets_executor_test.go --- a/internal/runtime/executor/xai_websockets_executor_test.go +++ b/internal/runtime/executor/xai_websockets_executor_test.go @@ -4,6 +4,7 @@ "bytes" "context" "fmt" + "io" "net/http" "net/http/httptest" "strings" @@ -242,6 +243,170 @@ thirdUpstreamPrevious := <-capturedPreviousIDs if thirdUpstreamPrevious != "resp-real" { t.Fatalf("third upstream previous_response_id = %q, want resp-real", thirdUpstreamPrevious) + } +} + +func TestXAIWebsocketsExecuteStreamCompactionTriggerUsesHTTPCompactWithRecordedContext(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + capturedWebsocketPayload := make(chan []byte, 1) + capturedCompactPayload := make(chan []byte, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/responses": + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade websocket: %v", err) + return + } + defer func() { _ = conn.Close() }() + + for i := 0; i < 2; i++ { + _ = conn.SetReadDeadline(time.Now().Add(5 * time.Second)) + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + t.Errorf("read upstream websocket message: %v", errRead) + return + } + capturedWebsocketPayload <- bytes.Clone(payload) + completed := []byte(`{"type":"response.completed","response":{"id":"resp-real","output":[{"type":"message","id":"out-1","role":"assistant","content":"first answer"}],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) + if i == 1 { + completed = []byte(`{"type":"response.completed","response":{"id":"resp-after-compact","output":[{"type":"message","id":"out-2","role":"assistant","content":"second answer"}],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) + } + if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { + t.Errorf("write completed websocket message: %v", errWrite) + return + } + } + case "/responses/compact": + body, errRead := io.ReadAll(r.Body) + if errRead != nil { + t.Errorf("read compact body: %v", errRead) + return + } + capturedCompactPayload <- bytes.Clone(body) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"resp_compact","model":"grok-4.3","output":[{"type":"compaction","encrypted_content":"opaque"}],"usage":{"input_tokens":1,"output_tokens":2,"total_tokens":3}}`)) + default: + t.Errorf("path = %q, want /responses", r.URL.Path) + http.Error(w, "unexpected path", http.StatusNotFound) + } + })) + defer server.Close() + + exec := NewXAIWebsocketsExecutor(&config.Config{}) + exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + exec.idStore = &xaiWebsocketIDStateStore{sessions: make(map[string]*xaiWebsocketIDState)} + auth := &cliproxyauth.Auth{ + ID: "xai-auth-compaction", + Provider: "xai", + Attributes: map[string]string{ + "base_url": server.URL, + "websockets": "true", + }, + Metadata: map[string]any{"access_token": "xai-token"}, + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Stream: true, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "xai-compaction-session", + }, + } + + result, err := exec.ExecuteStream(cliproxyexecutor.WithDownstreamWebsocket(context.Background()), auth, cliproxyexecutor.Request{ + Model: "grok-4.3", + Payload: []byte(`{"model":"grok-4.3","stream":true,"input":[{"type":"message","id":"msg-1","role":"user","content":"first"}]}`), + }, opts) + if err != nil { + t.Fatalf("ExecuteStream first turn error: %v", err) + } + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatalf("stream chunk error = %v", chunk.Err) + } + } + + select { + case payload := <-capturedWebsocketPayload: + if got := gjson.GetBytes(payload, "type").String(); got != "response.create" { + t.Fatalf("type = %q, want response.create; payload=%s", got, payload) + } + input := gjson.GetBytes(payload, "input") + if !input.IsArray() || len(input.Array()) != 1 { + t.Fatalf("input = %s, want one first-turn item", input.Raw) + } + if gjson.GetBytes(payload, "stream").Exists() { + t.Fatalf("stream must be omitted for xAI websocket payload: %s", payload) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for upstream websocket payload") + } + + compactResult, err := exec.ExecuteStream(cliproxyexecutor.WithDownstreamWebsocket(context.Background()), auth, cliproxyexecutor.Request{ + Model: "grok-4.3", + Payload: []byte(`{"model":"grok-4.3","stream":true,"previous_response_id":"resp-real-xai-1","input":[{"type":"compaction_trigger"}]}`), + }, opts) + if err != nil { + t.Fatalf("ExecuteStream compaction trigger error: %v", err) + } + for chunk := range compactResult.Chunks { + if chunk.Err != nil { + t.Fatalf("compact stream chunk error = %v", chunk.Err) + } + } + + select { + case payload := <-capturedCompactPayload: + if xaiInputHasItemType(payload, "compaction_trigger") { + t.Fatalf("compaction_trigger reached xai compact body: %s", payload) + } + input := gjson.GetBytes(payload, "input") + if !input.IsArray() || len(input.Array()) != 2 { + t.Fatalf("compact input = %s, want first request input plus response output", input.Raw) + } + if got := input.Array()[0].Get("id").String(); got != "msg-1" { + t.Fatalf("compact input[0].id = %q, want msg-1; payload=%s", got, payload) + } + if got := input.Array()[1].Get("id").String(); got != "out-1" { + t.Fatalf("compact input[1].id = %q, want out-1; payload=%s", got, payload) + } + if got := gjson.GetBytes(payload, "previous_response_id").String(); got != "" { + t.Fatalf("compact previous_response_id = %q, want empty; payload=%s", got, payload) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for compact HTTP payload") + } + + nextResult, err := exec.ExecuteStream(cliproxyexecutor.WithDownstreamWebsocket(context.Background()), auth, cliproxyexecutor.Request{ + Model: "grok-4.3", + Payload: []byte(`{"model":"grok-4.3","stream":true,"previous_response_id":"resp_compact","input":[{"type":"message","id":"msg-2","role":"user","content":"second"}]}`), + }, opts) + if err != nil { + t.Fatalf("ExecuteStream post-compaction turn error: %v", err) + } + for chunk := range nextResult.Chunks { + if chunk.Err != nil { + t.Fatalf("post-compaction stream chunk error = %v", chunk.Err) + } + } + select { + case payload := <-capturedWebsocketPayload: + if got := gjson.GetBytes(payload, "previous_response_id").String(); got != "" { + t.Fatalf("post-compaction previous_response_id = %q, want empty; payload=%s", got, payload) + } + input := gjson.GetBytes(payload, "input") + if !input.IsArray() || len(input.Array()) != 2 { + t.Fatalf("post-compaction input = %s, want compaction item plus new message", input.Raw) + } + if got := input.Array()[0].Get("type").String(); got != "compaction" { + t.Fatalf("post-compaction input[0].type = %q, want compaction; payload=%s", got, payload) + } + if got := input.Array()[1].Get("id").String(); got != "msg-2" { + t.Fatalf("post-compaction input[1].id = %q, want msg-2; payload=%s", got, payload) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for post-compaction websocket payload") } } -- tangled.sh