// Package executor provides runtime execution capabilities for various AI service providers.// This file implements an xAI executor that uses the Responses API WebSocket transport.package executor import (	"bytes"	"context"	"encoding/json"	"fmt"	"net/http"	"net/url"	"strconv"	"strings"	"sync"	"time" 	"github.com/gorilla/websocket"	xaiauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/xai"	"github.com/router-for-me/CLIProxyAPI/v7/internal/config"	"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"	"github.com/router-for-me/CLIProxyAPI/v7/internal/util"	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"	log "github.com/sirupsen/logrus"	"github.com/tidwall/gjson"	"github.com/tidwall/sjson") // XAIWebsocketsExecutor executes xAI Responses requests using a WebSocket transport.type XAIWebsocketsExecutor struct {	*XAIExecutor 	store   *codexWebsocketSessionStore	idStore *xaiWebsocketIDStateStore} var globalXAIWebsocketSessionStore = &codexWebsocketSessionStore{	sessions: make(map[string]*codexWebsocketSession),} var globalXAIWebsocketIDStates = &xaiWebsocketIDStateStore{	sessions: make(map[string]*xaiWebsocketIDState),} type xaiWebsocketIDStateStore struct {	mu       sync.Mutex	sessions map[string]*xaiWebsocketIDState} type xaiWebsocketIDState struct {	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	replayedCompactedTranscript bool} func NewXAIWebsocketsExecutor(cfg *config.Config) *XAIWebsocketsExecutor {	return &XAIWebsocketsExecutor{		XAIExecutor: NewXAIExecutor(cfg),		store:       globalXAIWebsocketSessionStore,		idStore:     globalXAIWebsocketIDStates,	}} func getXAIWebsocketIDState(store *xaiWebsocketIDStateStore, sessionID string) *xaiWebsocketIDState {	sessionID = strings.TrimSpace(sessionID)	if sessionID == "" || store == nil {		return nil	}	store.mu.Lock()	defer store.mu.Unlock()	if store.sessions == nil {		store.sessions = make(map[string]*xaiWebsocketIDState)	}	if state := store.sessions[sessionID]; state != nil {		return state	}	state := &xaiWebsocketIDState{		downstreamToUpstream: make(map[string]string),	}	store.sessions[sessionID] = state	return state} func deleteXAIWebsocketIDState(store *xaiWebsocketIDStateStore, sessionID string) {	sessionID = strings.TrimSpace(sessionID)	if sessionID == "" || store == nil {		return	}	store.mu.Lock()	delete(store.sessions, sessionID)	store.mu.Unlock()} func newXAIWebsocketRequestIDMapper(store *xaiWebsocketIDStateStore, sessionID string, downstreamRequest []byte) *xaiWebsocketRequestIDMapper {	state := getXAIWebsocketIDState(store, sessionID)	if state == nil {		return nil	}	downstreamPreviousID := strings.TrimSpace(gjson.GetBytes(downstreamRequest, "previous_response_id").String())	upstreamPreviousID := downstreamPreviousID	if downstreamPreviousID != "" {		upstreamPreviousID = state.upstreamIDForDownstream(downstreamPreviousID)	}	return &xaiWebsocketRequestIDMapper{		state:                state,		downstreamPreviousID: downstreamPreviousID,		upstreamPreviousID:   upstreamPreviousID,	}} func (s *xaiWebsocketIDState) upstreamIDForDownstream(downstreamID string) string {	downstreamID = strings.TrimSpace(downstreamID)	if s == nil || downstreamID == "" {		return downstreamID	}	s.mu.Lock()	defer s.mu.Unlock()	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, 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")) 	s.mu.Lock()	defer s.mu.Unlock()	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...)} 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.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	}	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 {	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 == "" {		out, errDelete := sjson.DeleteBytes(payload, "previous_response_id")		if errDelete == nil {			if m.downstreamPreviousID != "" && m.state != nil {				out = m.state.prependTranscriptInput(out)				m.replayedCompactedTranscript = true			}			return out		}		return payload	}	out, errSet := sjson.SetBytes(payload, "previous_response_id", m.upstreamPreviousID)	if errSet != nil {		return payload	}	return out} func (m *xaiWebsocketRequestIDMapper) downstreamResponsePayload(payload []byte) []byte {	if m == nil || len(payload) == 0 {		return payload	}	upstreamResponseID := strings.TrimSpace(gjson.GetBytes(payload, "response.id").String())	downstreamResponseID := m.downstreamIDForUpstreamResponse(upstreamResponseID)	if downstreamResponseID == "" {		return payload	}	return rewriteXAIWebsocketDownstreamIDs(payload, m.upstreamResponseID, downstreamResponseID, m.upstreamPreviousID, m.downstreamPreviousID)} func (m *xaiWebsocketRequestIDMapper) downstreamIDForUpstreamResponse(upstreamResponseID string) string {	upstreamResponseID = strings.TrimSpace(upstreamResponseID)	if m == nil || m.state == nil {		return upstreamResponseID	}	if m.upstreamResponseID != "" {		return m.downstreamResponseID	}	if upstreamResponseID == "" {		return ""	} 	m.state.mu.Lock()	defer m.state.mu.Unlock()	m.upstreamResponseID = upstreamResponseID	m.downstreamResponseID = upstreamResponseID	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} func rewriteXAIWebsocketDownstreamIDs(payload []byte, upstreamResponseID string, downstreamResponseID string, upstreamPreviousID string, downstreamPreviousID string) []byte {	upstreamResponseID = strings.TrimSpace(upstreamResponseID)	downstreamResponseID = strings.TrimSpace(downstreamResponseID)	upstreamPreviousID = strings.TrimSpace(upstreamPreviousID)	downstreamPreviousID = strings.TrimSpace(downstreamPreviousID)	if len(payload) == 0 || (upstreamResponseID == downstreamResponseID && upstreamPreviousID == downstreamPreviousID) {		return payload	} 	var value any	decoder := json.NewDecoder(bytes.NewReader(payload))	decoder.UseNumber()	if errDecode := decoder.Decode(&value); errDecode != nil {		return payload	}	if !rewriteXAIWebsocketDownstreamIDValue(value, upstreamResponseID, downstreamResponseID, upstreamPreviousID, downstreamPreviousID, "") {		return payload	}	out, errMarshal := json.Marshal(value)	if errMarshal != nil {		return payload	}	return out} func rewriteXAIWebsocketDownstreamIDValue(value any, upstreamResponseID string, downstreamResponseID string, upstreamPreviousID string, downstreamPreviousID string, key string) bool {	switch typed := value.(type) {	case map[string]any:		changed := false		for childKey, childValue := range typed {			if childString, ok := childValue.(string); ok {				replaced := rewriteXAIWebsocketDownstreamIDString(childString, childKey, upstreamResponseID, downstreamResponseID, upstreamPreviousID, downstreamPreviousID)				if replaced != childString {					typed[childKey] = replaced					changed = true				}				continue			}			if rewriteXAIWebsocketDownstreamIDValue(childValue, upstreamResponseID, downstreamResponseID, upstreamPreviousID, downstreamPreviousID, childKey) {				changed = true			}		}		return changed	case []any:		changed := false		for i := range typed {			if rewriteXAIWebsocketDownstreamIDValue(typed[i], upstreamResponseID, downstreamResponseID, upstreamPreviousID, downstreamPreviousID, key) {				changed = true			}		}		return changed	default:		return false	}} func rewriteXAIWebsocketDownstreamIDString(value string, key string, upstreamResponseID string, downstreamResponseID string, upstreamPreviousID string, downstreamPreviousID string) string {	switch key {	case "id", "item_id":		if upstreamResponseID != "" && downstreamResponseID != "" && downstreamResponseID != upstreamResponseID && strings.Contains(value, upstreamResponseID) {			return strings.ReplaceAll(value, upstreamResponseID, downstreamResponseID)		}	case "previous_response_id":		if upstreamPreviousID != "" && downstreamPreviousID != "" && value == upstreamPreviousID {			return downstreamPreviousID		}	}	return value} func (e *XAIWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {	if e == nil || e.XAIExecutor == nil {		return cliproxyexecutor.Response{}, fmt.Errorf("xai websockets executor: executor is nil")	}	return e.XAIExecutor.Execute(ctx, auth, req, opts)} func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (_ *cliproxyexecutor.StreamResult, err error) {	if e == nil || e.XAIExecutor == nil {		return nil, fmt.Errorf("xai websockets executor: executor is nil")	}	if ctx == nil {		ctx = context.Background()	}	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	}	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)	} 	// Keep websocket on the official API base URL (or an explicit non-default	// base_url). Do not reuse xaiChatBaseURL: cli-chat-proxy only accepts HTTP	// POST and returns 405 for websocket upgrades.	token, baseURL := xaiCreds(auth)	if baseURL == "" {		baseURL = xaiauth.DefaultAPIBaseURL	} 	prepared, err := e.prepareResponsesWebsocketRequest(ctx, req, opts)	if err != nil {		return nil, err	} 	reporter := helps.NewExecutorUsageReporter(ctx, e, prepared.baseModel, auth)	defer reporter.TrackFailure(ctx, &err) 	httpURL := strings.TrimSuffix(baseURL, "/") + "/responses"	wsURL, err := buildXAIResponsesWebsocketURL(httpURL)	if err != nil {		return nil, err	} 	var authID, authLabel, authType, authValue string	if auth != nil {		authID = auth.ID		authLabel = auth.Label		authType, authValue = auth.AccountInfo()	} 	var sess *codexWebsocketSession	isEphemeralSession := false	if executionSessionID != "" {		sess = e.getOrCreateSession(executionSessionID)		if sess != nil {			sess.reqMu.Lock()		}

Lines 1–500 of 1807. Load more