From 2fa4dabe93a24fd8c6fdb56e85bb6860ef6c46f4 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Fri, 26 Jun 2026 16:34:02 +0800 Subject: [PATCH] feat(executor): improve downstream response ID rewrite and add test for repeated response scenarios --- .../executor/xai_websockets_executor.go | 9 +- .../executor/xai_websockets_executor_test.go | 97 +++++++++++++++++++ 2 files changed, 102 insertions(+), 4 deletions(-) diff --git a/internal/runtime/executor/xai_websockets_executor.go b/internal/runtime/executor/xai_websockets_executor.go index c9333a748..d7bed502f 100644 --- a/internal/runtime/executor/xai_websockets_executor.go +++ b/internal/runtime/executor/xai_websockets_executor.go @@ -292,13 +292,14 @@ func (m *xaiWebsocketRequestIDMapper) downstreamIDForUpstreamResponse(upstreamRe defer m.state.mu.Unlock() m.upstreamResponseID = upstreamResponseID m.downstreamResponseID = upstreamResponseID - if m.downstreamPreviousID != "" && m.upstreamPreviousID != "" && upstreamResponseID == m.upstreamPreviousID { - m.state.sequence++ - m.downstreamResponseID = fmt.Sprintf("%s-xai-%d", upstreamResponseID, m.state.sequence) - } if m.state.downstreamToUpstream == nil { m.state.downstreamToUpstream = make(map[string]string) } + _, upstreamResponseIDSeen := m.state.downstreamToUpstream[upstreamResponseID] + if (m.downstreamPreviousID != "" && m.upstreamPreviousID != "" && upstreamResponseID == m.upstreamPreviousID) || upstreamResponseIDSeen { + m.state.sequence++ + m.downstreamResponseID = fmt.Sprintf("%s-xai-%d", upstreamResponseID, m.state.sequence) + } m.state.downstreamToUpstream[upstreamResponseID] = upstreamResponseID m.state.downstreamToUpstream[m.downstreamResponseID] = upstreamResponseID return m.downstreamResponseID diff --git a/internal/runtime/executor/xai_websockets_executor_test.go b/internal/runtime/executor/xai_websockets_executor_test.go index a3a717e6a..237e8ab5c 100644 --- a/internal/runtime/executor/xai_websockets_executor_test.go +++ b/internal/runtime/executor/xai_websockets_executor_test.go @@ -330,6 +330,103 @@ func TestXAIWebsocketsExecuteStreamRewritesRepeatedResponseIDForDownstream(t *te } } +func TestXAIWebsocketsExecuteStreamRewritesRepeatedResponseIDWithoutPreviousResponseID(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + capturedPreviousIDs := make(chan string, 2) + releaseServer := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + 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++ { + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + t.Errorf("read upstream websocket message: %v", errRead) + return + } + capturedPreviousIDs <- gjson.GetBytes(payload, "previous_response_id").String() + completed := []byte(`{"type":"response.completed","response":{"id":"resp-real","output":[{"id":"msg_resp-real","type":"message","status":"completed","role":"assistant","content":[{"type":"output_text","text":"ok"}]}],"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 + } + } + <-releaseServer + })) + defer server.Close() + defer close(releaseServer) + + 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-id-map-no-prev", + 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, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "xai-id-map-no-prev-session", + }, + } + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + + runRequest := func(content string) (string, string) { + body := []byte(fmt.Sprintf(`{"model":"grok-4.3","input":[{"type":"message","role":"user","content":%q}]}`, content)) + result, err := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{Model: "grok-4.3", Payload: body}, opts) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + select { + case chunk, ok := <-result.Chunks: + if !ok { + t.Fatal("stream closed before completed chunk") + } + if chunk.Err != nil { + t.Fatalf("chunk error = %v", chunk.Err) + } + payload := bytes.TrimSpace(chunk.Payload) + return gjson.GetBytes(payload, "response.id").String(), + gjson.GetBytes(payload, "response.output.0.id").String() + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for completed chunk") + } + return "", "" + } + + firstDownstreamID, firstOutputID := runRequest("first") + if firstDownstreamID != "resp-real" { + t.Fatalf("first downstream id = %q, want resp-real", firstDownstreamID) + } + if firstOutputID != "msg_resp-real" { + t.Fatalf("first output item id = %q, want msg_resp-real", firstOutputID) + } + if firstUpstreamPrevious := <-capturedPreviousIDs; firstUpstreamPrevious != "" { + t.Fatalf("first upstream previous_response_id = %q, want empty", firstUpstreamPrevious) + } + + secondDownstreamID, secondOutputID := runRequest("second") + if secondDownstreamID == "" || secondDownstreamID == "resp-real" { + t.Fatalf("second downstream id = %q, want synthetic id different from resp-real", secondDownstreamID) + } + if secondOutputID == "msg_resp-real" || !strings.Contains(secondOutputID, secondDownstreamID) { + t.Fatalf("second output item id = %q, want rewritten id containing %q", secondOutputID, secondDownstreamID) + } + if secondUpstreamPrevious := <-capturedPreviousIDs; secondUpstreamPrevious != "" { + t.Fatalf("second upstream previous_response_id = %q, want empty", secondUpstreamPrevious) + } +} + func TestXAIWebsocketsExecuteStreamCompactionTriggerUsesHTTPCompactWithRecordedContext(t *testing.T) { nativeEncryptedContent := testValidGrokEncryptedContent() upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}