1644 lines · 50.6 KB · Go
RawDownload
package executor import (	"bytes"	"context"	"encoding/json"	"fmt"	"net/http"	"strconv"	"strings" 	"github.com/google/uuid"	xaiauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/xai"	"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"	"github.com/router-for-me/CLIProxyAPI/v7/internal/thinking"	"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"	sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"	log "github.com/sirupsen/logrus"	"github.com/tidwall/gjson"	"github.com/tidwall/sjson") type xaiPreparedRequest struct {	baseModel             string	from                  sdktranslator.Format	responseFormat        sdktranslator.Format	to                    sdktranslator.Format	originalPayload       []byte	body                  []byte	namespaceTools        map[string]xaiNamespaceToolRef	clientDeclaredTools   map[xaiClientToolKey]struct{}	sessionID             string	replayScope           xaiReasoningReplayScope	filterInternalXSearch bool} type xaiNamespaceToolRef struct {	namespace    string	name         string	isDispatcher bool} // xaiClientToolKey identifies a client-declared callable tool using the// post-restore Responses shape (short name + optional namespace) and the// effective upstream tool type after normalizeXAITool (client custom tools are// sent as function). Response call types are matched against this effective// kind so internal custom_tool_call traces are not exempted merely because a// client declared an ordinary function/custom tool with the same short name,// while legitimate function_call responses for normalized custom tools are kept.type xaiClientToolKey struct {	namespace string	name      string	toolType  string} func (e *XAIExecutor) prepareResponsesRequest(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, stream bool) (*xaiPreparedRequest, error) {	return e.prepareResponsesRequestTo(ctx, req, opts, stream, sdktranslator.FormatCodex)} func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, stream bool, to sdktranslator.Format) (*xaiPreparedRequest, error) {	baseModel := thinking.ParseSuffix(req.Model).ModelName	from := opts.SourceFormat	responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts)	originalPayloadSource := req.Payload	if len(opts.OriginalRequest) > 0 {		originalPayloadSource = opts.OriginalRequest	}	originalPayload := bytes.Clone(originalPayloadSource)	originalTranslated := helps.TranslateRequestWithAPIKeyModelCompatibility(ctx, opts.Headers, e.cfg, from, to, baseModel, originalPayload, stream, helps.APIKeyModelIsCompat(req))	originalTranslated = preserveXAIResponsesOutputControls(originalTranslated, originalPayload, from)	body := helps.TranslateRequestWithAPIKeyModelCompatibility(ctx, opts.Headers, e.cfg, from, to, baseModel, bytes.Clone(req.Payload), stream, helps.APIKeyModelIsCompat(req))	body = preserveXAIResponsesOutputControls(body, req.Payload, from) 	var err error	body, err = helps.ApplyRequestThinking(body, req, opts, from.String(), e.Identifier(), e.Identifier())	if err != nil {		return nil, err	} 	requestedModel := helps.PayloadRequestedModel(opts, req.Model)	requestPath := helps.PayloadRequestPath(opts)	body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers)	body = helps.SetStringIfDifferent(body, "model", baseModel)	body = helps.SetBoolIfDifferent(body, "stream", stream)	body, _ = sjson.DeleteBytes(body, "previous_response_id")	body, _ = sjson.DeleteBytes(body, "prompt_cache_retention")	body, _ = sjson.DeleteBytes(body, "safety_identifier")	body, _ = sjson.DeleteBytes(body, "stream_options")	body = helps.RewriteCodexMultiAgentV2Input(ctx, opts.Headers, body, e.cfg)	willInjectXSearch := e.cfg != nil && e.cfg.XAI.InjectXSearch	shouldFold := xaiShouldFoldNamespaceTools(body, willInjectXSearch)	namespaceTools := collectXAINamespaceToolRefsWithFold(body, shouldFold)	// Collect before normalizeXAITools flattens namespace wrappers so keys match	// the post-restore (namespace, short-name) shape used by the response filter.	clientDeclaredTools := collectXAIClientDeclaredToolKeys(body)	body = normalizeXAIToolsWithFold(body, shouldFold)	body = promoteXAIAdditionalTools(body)	// Drop choices that point at tools removed by normalizeXAITools before any	// configured x_search injection, so no surviving choice references a deleted tool.	body = normalizeXAINamespaceToolChoiceWithFold(body, shouldFold)	body = normalizeXAIForcedWebSearchToolChoice(body)	// Prune before rewriting image_generation choices so older models that still	// strip the tool do not keep a leftover "required" selection.	body = pruneXAIOrphanedToolChoice(body)	body = normalizeXAIForcedImageGenerationToolChoice(body)	body = normalizeXAIToolChoiceForTools(body)	// Skip x_search injection when the request was forced to image_generation and	// the remaining tools list is only that hosted tool. "required" plus extra	// tools would let Grok call x_search instead of Imagine.	if e.cfg != nil && e.cfg.XAI.InjectXSearch && !xaiToolChoiceRequiresImageGenerationOnly(body) {		body = ensureXAINativeXSearchTool(body)	}	body = clampXAIToolsLimit(body, xaiMaxTools, namespaceTools)	var replayScope xaiReasoningReplayScope	body, replayScope, err = applyXAIReasoningReplayCacheRequired(ctx, from, req, opts, body)	if err != nil {		return nil, err	}	body = normalizeXAIInputCustomToolCalls(body)	body = normalizeXAIInputNamespaceToolCallsWithFold(body, shouldFold)	body = normalizeXAIInputReasoningItems(body)	body = sanitizeXAIInputEncryptedContent(body)	body = normalizeCodexInstructions(body)	body = sanitizeXAIResponsesBody(body, baseModel)	body = normalizeXAIImageRefs(body) 	sessionID, errSession := xaiResolveComposerSessionID(ctx, req, opts, baseModel)	if errSession != nil {		return nil, errSession	}	if sessionID != "" {		body = helps.SetStringIfDifferent(body, "prompt_cache_key", sessionID)	} 	return &xaiPreparedRequest{		baseModel:             baseModel,		from:                  from,		responseFormat:        responseFormat,		to:                    to,		originalPayload:       originalPayload,		body:                  body,		namespaceTools:        namespaceTools,		clientDeclaredTools:   clientDeclaredTools,		sessionID:             sessionID,		replayScope:           replayScope,		filterInternalXSearch: xaiRequestHasNativeXSearch(body),	}, nil} func (e *XAIExecutor) recordXAIRequest(ctx context.Context, auth *cliproxyauth.Auth, url string, headers http.Header, body []byte) {	var authID, authLabel, authType, authValue string	if auth != nil {		authID = auth.ID		authLabel = auth.Label		authType, authValue = auth.AccountInfo()	}	helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{		URL:       url,		Method:    http.MethodPost,		Headers:   headers,		Body:      body,		Provider:  e.Identifier(),		AuthID:    authID,		AuthLabel: authLabel,		AuthType:  authType,		AuthValue: authValue,	})} func xaiCreds(auth *cliproxyauth.Auth) (token, baseURL string) {	if auth == nil {		return "", ""	}	if auth.Attributes != nil {		token = strings.TrimSpace(auth.Attributes["api_key"])		baseURL = strings.TrimSpace(auth.Attributes["base_url"])	}	if auth.Metadata != nil {		if token == "" {			token = xaiMetadataString(auth.Metadata, "access_token")		}		if baseURL == "" {			baseURL = xaiMetadataString(auth.Metadata, "base_url")		}	}	return token, baseURL} // xaiUsingAPI reports whether this xAI auth should use the official API path// for HTTP chat and media. OAuth defaults to false to use Grok Build.func xaiUsingAPI(auth *cliproxyauth.Auth) bool {	if auth == nil {		return true	}	if len(auth.Attributes) > 0 {		if raw := strings.TrimSpace(auth.Attributes[xaiUsingAPIAttr]); raw != "" {			parsed, errParse := strconv.ParseBool(raw)			if errParse == nil {				return parsed			}		}	}	if len(auth.Metadata) > 0 {		raw, ok := auth.Metadata[xaiUsingAPIAttr]		if ok && raw != nil {			switch v := raw.(type) {			case bool:				return v			case string:				parsed, errParse := strconv.ParseBool(strings.TrimSpace(v))				if errParse == nil {					return parsed				}			default:			}		}	}	if raw := strings.TrimSpace(auth.Attributes["auth_kind"]); raw != "" {		return !strings.EqualFold(raw, "oauth")	}	return !strings.EqualFold(xaiMetadataString(auth.Metadata, "auth_kind"), "oauth")} // xaiChatBaseURL returns the base URL for xAI HTTP chat and media (image/video) requests.// When auth using_api is true, the official API base URL logic is used. When it// is false (including its OAuth default), empty or official default base_url is// rewritten to the CLI chat-proxy endpoint; an explicit non-default base_url is// still honored.// Websocket and compact transports intentionally do not use this helper:// cli-chat-proxy does not implement /responses/compact (404) or websocket// upgrades (405).func xaiChatBaseURL(auth *cliproxyauth.Auth) string {	_, baseURL := xaiCreds(auth)	if xaiUsingAPI(auth) {		if baseURL == "" {			return xaiauth.DefaultAPIBaseURL		}		return baseURL	}	if baseURL != "" && !xaiIsDefaultAPIBaseURL(baseURL) {		return baseURL	}	return xaiauth.CLIChatProxyBaseURL} // xaiCompactBaseURL returns the base URL for xAI /responses/compact requests.// Compact must stay on the official API (or an explicit non-CLI-proxy base_url).// Reusing xaiChatBaseURL would pin OAuth traffic to cli-chat-proxy, which returns// 404 for /responses/compact and then cools down the auth pool as not_found.func xaiCompactBaseURL(auth *cliproxyauth.Auth) string {	_, baseURL := xaiCreds(auth)	if baseURL == "" || xaiIsCLIChatProxyBaseURL(baseURL) {		return xaiauth.DefaultAPIBaseURL	}	return baseURL} func xaiNormalizeBaseURL(baseURL string) string {	return strings.TrimRight(strings.TrimSpace(baseURL), "/")} func xaiIsDefaultAPIBaseURL(baseURL string) bool {	return xaiNormalizeBaseURL(baseURL) == xaiNormalizeBaseURL(xaiauth.DefaultAPIBaseURL)} func xaiIsCLIChatProxyBaseURL(baseURL string) bool {	return xaiNormalizeBaseURL(baseURL) == xaiNormalizeBaseURL(xaiauth.CLIChatProxyBaseURL)} // xaiBaseURLSource classifies a resolved xAI base URL for logging.func xaiBaseURLSource(baseURL string) string {	switch {	case xaiIsDefaultAPIBaseURL(baseURL):		return "DefaultAPIBaseURL"	case xaiIsCLIChatProxyBaseURL(baseURL):		return "CLIChatProxyBaseURL"	default:		return "custom"	}} // logXAIResolvedBaseURL emits a console log for the resolved upstream base URL.func logXAIResolvedBaseURL(ctx context.Context, baseURL string) {	helps.LogWithRequestID(ctx).Infof("xai: using base_url=%s source=%s", baseURL, xaiBaseURLSource(baseURL))} func applyXAIHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, sessionID string, clientHeaders ...http.Header) {	applyXAIDefaultHeaders(r, token, stream, sessionID)	applyXAICustomHeaders(r, auth, clientHeaders...)} func applyXAIDefaultHeaders(r *http.Request, token string, stream bool, sessionID string) {	r.Header.Set("Content-Type", "application/json")	if strings.TrimSpace(token) != "" {		r.Header.Set("Authorization", "Bearer "+token)	} else {		r.Header.Del("Authorization")	}	if stream {		r.Header.Set("Accept", "text/event-stream")	} else {		r.Header.Set("Accept", "application/json")	}	r.Header.Set("Connection", "Keep-Alive")	if sessionID != "" {		r.Header.Set("x-grok-conv-id", sessionID)	}} func applyXAICustomHeaders(r *http.Request, auth *cliproxyauth.Auth, clientHeaders ...http.Header) {	var attrs map[string]string	if auth != nil {		attrs = auth.Attributes	}	util.ApplyCustomHeadersFromAttrs(r, attrs, clientHeaders...)} // applyXAIChatHeaders applies standard xAI headers for non-image/video chat// requests. When using_api is true, this matches the standard// applyXAIHeaders behavior. CLI chat-proxy identity headers are only attached// when using_api is false and the resolved chat base URL is the official CLI// chat-proxy endpoint.func applyXAIChatHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, sessionID string, clientHeaders ...http.Header) {	if xaiUsingAPI(auth) {		applyXAIHeaders(r, auth, token, stream, sessionID, clientHeaders...)		return	}	applyXAIDefaultHeaders(r, token, stream, sessionID)	if xaiIsCLIChatProxyBaseURL(xaiChatBaseURL(auth)) {		r.Header.Set(xaiTokenAuthHeader, xaiTokenAuthValue)		r.Header.Set(xaiClientVersionHeader, xaiClientVersionValue)		r.Header.Set("User-Agent", "xai-grok-workspace/"+xaiClientVersionValue)		r.Header.Set(xaiClientIdentifierHeader, xaiClientIdentifierValue)		r.Header.Set(xaiAuthenticateResponseHeader, xaiAuthenticateResponseValue)	}	applyXAICustomHeaders(r, auth, clientHeaders...)} func xaiResolveComposerSessionID(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, baseModel string) (string, error) {	if sessionID := xaiExecutionSessionID(req, opts); sessionID != "" {		return sessionID, nil	}	if !xaiRequiresIsolatedConversation(baseModel) {		return "", nil	}	cached, ok, errCache := helps.ClaudeCodePromptCache(ctx, baseModel, req.Payload, opts.Headers)	if errCache != nil {		return "", errCache	}	if ok {		return cached.ID, nil	}	return uuid.NewString(), nil} func xaiExecutionSessionID(req cliproxyexecutor.Request, opts cliproxyexecutor.Options) string {	if value := xaiMetadataString(opts.Metadata, cliproxyexecutor.ExecutionSessionMetadataKey); value != "" {		return value	}	if value := xaiMetadataString(req.Metadata, cliproxyexecutor.ExecutionSessionMetadataKey); value != "" {		return value	}	if promptCacheKey := gjson.GetBytes(req.Payload, "prompt_cache_key"); promptCacheKey.Exists() {		if value := strings.TrimSpace(promptCacheKey.String()); value != "" {			return value		}	}	return helps.DerivedSessionUUID("xai", opts.Metadata, req.Metadata)} func xaiRequiresIsolatedConversation(model string) bool {	return strings.HasPrefix(strings.ToLower(strings.TrimSpace(model)), xaiComposerModelPrefix)} func xaiImageEndpointPath(opts cliproxyexecutor.Options) string {	if opts.SourceFormat.String() != xaiImageHandlerType {		return ""	} 	path := xaiMetadataString(opts.Metadata, cliproxyexecutor.RequestPathMetadataKey)	if strings.HasSuffix(path, "/images/edits") {		return xaiImagesEditsPath	}	if strings.HasSuffix(path, "/images/generations") {		return xaiImagesGenerationsPath	}	return xaiDefaultImageEndpointPath} // normalizeXAIImageRefs rewrites OpenAI-style image object fields to the xAI// image API shape before the payload is sent upstream:////	{"image":{"image_url":"https://..."}} → {"image":{"url":"https://..."}}//// Applies to image / images / reference_images anywhere in the JSON tree,// including nested objects and array items. Does not rewrite chat content// parts shaped as {"type":"image_url","image_url":{...}}.func normalizeXAIImageRefs(body []byte) []byte {	if !gjson.ValidBytes(body) {		return body	} 	decoder := json.NewDecoder(bytes.NewReader(body))	decoder.UseNumber()	var payload any	if errDecode := decoder.Decode(&payload); errDecode != nil {		return body	} 	if !normalizeXAIImageRefsValue(payload) {		return body	}	normalized, errMarshal := json.Marshal(payload)	if errMarshal != nil {		return body	}	return normalized} func normalizeXAIImageRefsValue(value any) bool {	changed := false	switch node := value.(type) {	case map[string]any:		for key, child := range node {			switch key {			case "image":				changed = normalizeXAIImageRef(child) || changed			case "images", "reference_images":				if refs, ok := child.([]any); ok {					for _, ref := range refs {						changed = normalizeXAIImageRef(ref) || changed					}				}			}			changed = normalizeXAIImageRefsValue(child) || changed		}	case []any:		for _, child := range node {			changed = normalizeXAIImageRefsValue(child) || changed		}	}	return changed} func normalizeXAIImageRef(value any) bool {	ref, ok := value.(map[string]any)	if !ok {		return false	} 	originalURL, _ := ref["url"].(string)	url := strings.TrimSpace(originalURL)	imageURL, hasImageURL := ref["image_url"]	if url == "" {		switch imageURL := imageURL.(type) {		case string:			url = strings.TrimSpace(imageURL)		case map[string]any:			url, _ = imageURL["url"].(string)			url = strings.TrimSpace(url)		}	}	if url == "" {		return false	}	if url == originalURL && !hasImageURL {		return false	} 	// Always emit the xAI field name and drop the OpenAI alias.	ref["url"] = url	delete(ref, "image_url")	return true} func xaiIsVideoRequest(opts cliproxyexecutor.Options) bool {	return opts.SourceFormat.String() == xaiVideoHandlerType} func xaiVideoEndpointPath(opts cliproxyexecutor.Options) string {	if !xaiIsVideoRequest(opts) {		return ""	}	path := xaiMetadataString(opts.Metadata, cliproxyexecutor.RequestPathMetadataKey)	if strings.HasSuffix(path, "/videos/edits") {		return xaiVideosEditsPath	}	if strings.HasSuffix(path, "/videos/extensions") {		return xaiVideosExtensionsPath	}	if strings.HasSuffix(path, "/videos/generations") {		return xaiVideosGenerationsPath	}	return ""} func xaiMetadataString(meta map[string]any, key string) string {	if len(meta) == 0 || key == "" {

Lines 1–500 of 1644.