feat(executor): improve downstream response ID rewrite and add test for repeated response scenarios

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

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 }}