417 lines · 16.0 KB · Go
RawDownload
package executor import (	"bytes"	"context"	"fmt"	"io"	"net/http"	"strings"	"time" 	"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"	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") func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) {	ctx = helps.EnsureSessionContext(ctx, opts, req.Payload)	if opts.Alt == "responses/compact" {		return e.executeCompact(ctx, auth, req, opts)	}	if endpointPath := xaiImageEndpointPath(opts); endpointPath != "" {		return e.executeImages(ctx, auth, req, opts, endpointPath)	}	if xaiIsVideoRequest(opts) {		return e.executeVideos(ctx, auth, req, opts)	} 	token, _ := xaiCreds(auth)	baseURL := xaiChatBaseURL(auth)	logXAIResolvedBaseURL(ctx, baseURL) 	prepared, err := e.prepareResponsesRequest(ctx, req, opts, true)	if err != nil {		return resp, err	} 	reporter := helps.NewExecutorUsageReporter(ctx, e, prepared.baseModel, auth)	defer reporter.TrackFailure(ctx, &err)	reporter.SetTranslatedReasoningEffort(prepared.body, e.Identifier()) 	url := strings.TrimSuffix(baseURL, "/") + "/responses"	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(prepared.body))	if err != nil {		return resp, err	}	applyXAIChatHeaders(httpReq, auth, token, true, prepared.sessionID, opts.Headers)	e.recordXAIRequest(ctx, auth, url, httpReq.Header.Clone(), prepared.body) 	httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0)	httpClient = reporter.TrackHTTPClient(httpClient)	httpResp, err := httpClient.Do(httpReq)	if err != nil {		helps.RecordAPIResponseError(ctx, e.cfg, err)		return resp, err	}	defer func() {		if errClose := httpResp.Body.Close(); errClose != nil {			log.Errorf("xai executor: close response body error: %v", errClose)		}	}()	helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())	if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {		data, errRead := io.ReadAll(httpResp.Body)		if errRead != nil {			helps.RecordAPIResponseError(ctx, e.cfg, errRead)			return resp, errRead		}		helps.AppendAPIResponseChunk(ctx, e.cfg, data)		helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data))		return resp, xaiStatusErr(httpResp.StatusCode, data)	} 	data, err := io.ReadAll(httpResp.Body)	if err != nil {		helps.RecordAPIResponseError(ctx, e.cfg, err)		return resp, err	}	helps.AppendAPIResponseChunk(ctx, e.cfg, data) 	outputItemsByIndex := make(map[int64][]byte)	var outputItemsFallback [][]byte	responseFilter := newXAIInternalXSearchResponseFilter(prepared.filterInternalXSearch, prepared.clientDeclaredTools)	namespaceRestorer := newXAINamespaceRestorer(prepared.namespaceTools)	for _, line := range bytes.Split(data, []byte("\n")) {		if !bytes.HasPrefix(line, xaiDataTag) {			continue		}		eventData := xaiNormalizeReasoningSummaryData(bytes.TrimSpace(line[len(xaiDataTag):]))		eventData = namespaceRestorer.restore(eventData)		eventData = responseFilter.apply(eventData)		if len(eventData) == 0 {			continue		}		eventType := gjson.GetBytes(eventData, "type").String()		switch eventType {		case "response.output_item.done":			xaiCollectOutputItemDone(eventData, outputItemsByIndex, &outputItemsFallback)		case "response.completed", "response.incomplete":			if detail, ok := helps.ParseCodexUsage(eventData); ok {				reporter.Publish(ctx, detail)			}			completedData := xaiPatchCompletedOutput(eventData, outputItemsByIndex, outputItemsFallback)			completedData = xaiNormalizeReasoningSummaryData(completedData)			if eventType == "response.completed" {				// A truncated turn carries no replayable terminal state, so only a				// completed response may refresh the reasoning replay cache.				cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, completedData)			}			var param any			out := sdktranslator.TranslateNonStream(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, completedData, &param)			if prepared.responseFormat == sdktranslator.FormatOpenAIResponse {				out = helps.EnsureResponsesUsageDetails(out)			}			return cliproxyexecutor.Response{Payload: out, Headers: httpResp.Header.Clone()}, nil		}	} 	return resp, statusErr{code: http.StatusRequestTimeout, msg: "xai stream error: stream disconnected before response.completed or response.incomplete"}} func (e *XAIExecutor) executeCompact(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) {	prepared, data, headers, errCompact := e.executeCompactRequest(ctx, auth, req, opts)	if errCompact != nil {		return resp, errCompact	} 	var param any	out := sdktranslator.TranslateNonStream(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, data, &param)	if prepared.responseFormat == sdktranslator.FormatOpenAIResponse {		out = helps.EnsureResponsesUsageDetails(out)	}	return cliproxyexecutor.Response{Payload: out, Headers: headers}, nil} func (e *XAIExecutor) executeCompactRequest(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*xaiPreparedRequest, []byte, http.Header, error) {	token, _ := xaiCreds(auth)	// Compact must not use xaiChatBaseURL: CLI chat-proxy returns 404 for	// /responses/compact and a 404 cools down the whole xAI auth pool.	baseURL := xaiCompactBaseURL(auth)	logXAIResolvedBaseURL(ctx, baseURL) 	prepared, err := e.prepareResponsesRequestTo(ctx, req, opts, false, sdktranslator.FormatOpenAIResponse)	if err != nil {		return nil, nil, nil, err	}	prepared.body, _ = sjson.DeleteBytes(prepared.body, "stream")	prepared.body, _ = sjson.DeleteBytes(prepared.body, "tools")	// Compact deletes tools after prepareResponsesRequestTo, which can now keep	// image_generation and rewrite its forced choice to "required" on grok-4.6+.	// Drop the leftover selection so compact does not send tool_choice without tools.	prepared.body = normalizeXAIToolChoiceForTools(prepared.body)	for _, field := range []string{"max_output_tokens", "temperature", "top_p", "top_k", "stop"} {		prepared.body, _ = sjson.DeleteBytes(prepared.body, field)	}	prepared.body = xaiRemoveInputItemsByType(prepared.body, "compaction_trigger")	if previousResponseID := strings.TrimSpace(gjson.GetBytes(req.Payload, "previous_response_id").String()); previousResponseID != "" {		prepared.body, _ = sjson.SetBytes(prepared.body, "previous_response_id", previousResponseID)	} 	reporter := helps.NewExecutorUsageReporter(ctx, e, prepared.baseModel, auth)	defer reporter.TrackFailure(ctx, &err)	reporter.SetTranslatedReasoningEffort(prepared.body, e.Identifier()) 	requestURL := strings.TrimSuffix(baseURL, "/") + "/responses/compact"	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, requestURL, bytes.NewReader(prepared.body))	if err != nil {		return nil, nil, nil, err	}	// Official API / custom compact endpoints use standard API headers, not CLI	// chat-proxy identity headers (which applyXAIChatHeaders may still attach for OAuth chat).	applyXAIHeaders(httpReq, auth, token, false, prepared.sessionID, opts.Headers)	e.recordXAIRequest(ctx, auth, requestURL, httpReq.Header.Clone(), prepared.body) 	httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0)	httpClient = reporter.TrackHTTPClient(httpClient)	httpResp, err := httpClient.Do(httpReq)	if err != nil {		helps.RecordAPIResponseError(ctx, e.cfg, err)		return nil, nil, nil, err	}	defer func() {		if errClose := httpResp.Body.Close(); errClose != nil {			log.Errorf("xai executor: close response body error: %v", errClose)		}	}()	helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) 	data, err := io.ReadAll(httpResp.Body)	if err != nil {		helps.RecordAPIResponseError(ctx, e.cfg, err)		return nil, nil, nil, err	}	helps.AppendAPIResponseChunk(ctx, e.cfg, data) 	if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {		helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data))		err = xaiStatusErr(httpResp.StatusCode, data)		return nil, nil, nil, err	} 	reporter.Publish(ctx, helps.ParseOpenAIUsage(data))	reporter.EnsurePublished(ctx)	clearXAIReasoningReplayAfterCompaction(ctx, prepared.replayScope)	return prepared, data, httpResp.Header.Clone(), nil} func (e *XAIExecutor) executeCompactionTriggerStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {	prepared, data, headers, err := e.executeCompactRequest(ctx, auth, req, opts)	if err != nil {		return nil, err	} 	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 xaiInputHasItemType(body []byte, itemType string) bool {	input := gjson.GetBytes(body, "input")	if !input.IsArray() {		return false	}	for _, item := range input.Array() {		if item.Get("type").String() == itemType {			return true		}	}	return false} func xaiRemoveInputItemsByType(body []byte, itemType string) []byte {	input := gjson.GetBytes(body, "input")	if !input.IsArray() {		return body	} 	var buf bytes.Buffer	buf.WriteByte('[')	kept := 0	for _, item := range input.Array() {		if item.Get("type").String() == itemType {			continue		}		if kept > 0 {			buf.WriteByte(',')		}		buf.WriteString(item.Raw)		kept++	}	buf.WriteByte(']') 	updated, err := sjson.SetRawBytes(body, "input", buf.Bytes())	if err != nil {		return body	}	return updated} func xaiBuildCompactionTriggerStreamChunks(prepared *xaiPreparedRequest, compactData []byte) [][]byte {	responseID := xaiCompactionResponseID(compactData)	now := time.Now().Unix()	createdAt := gjson.GetBytes(compactData, "created_at").Int()	if createdAt == 0 {		createdAt = now	}	completedAt := gjson.GetBytes(compactData, "completed_at").Int()	if completedAt == 0 {		completedAt = now	} 	item := xaiCompactionOutputItem(compactData, responseID)	output := make([]byte, 0, len(item)+2)	output = append(output, '[')	output = append(output, item...)	output = append(output, ']') 	createdResponse := xaiBuildCompactionBaseResponse(prepared, compactData, responseID, createdAt, "in_progress")	inProgressResponse := xaiBuildCompactionBaseResponse(prepared, compactData, responseID, createdAt, "in_progress")	completedResponse := xaiBuildCompactionBaseResponse(prepared, compactData, responseID, createdAt, "completed")	requestModelName := ""	if prepared != nil {		requestModelName = gjson.GetBytes(prepared.originalPayload, "model").String()		if requestModelName == "" {			requestModelName = prepared.baseModel		}	}	if requestModelName == "" {		requestModelName = gjson.GetBytes(compactData, "model").String()	}	if requestModelName != "" {		createdResponse, _ = sjson.SetBytes(createdResponse, "model", requestModelName)		inProgressResponse, _ = sjson.SetBytes(inProgressResponse, "model", requestModelName)	}	completedResponse, _ = sjson.SetBytes(completedResponse, "completed_at", completedAt)	completedResponse, _ = sjson.SetRawBytes(completedResponse, "output", output)	if usage := gjson.GetBytes(compactData, "usage"); usage.Exists() {		completedResponse, _ = sjson.SetRawBytes(completedResponse, "usage", []byte(usage.Raw))	} 	createdPayload := []byte(`{"type":"response.created","sequence_number":0}`)	createdPayload, _ = sjson.SetRawBytes(createdPayload, "response", createdResponse)	inProgressPayload := []byte(`{"type":"response.in_progress","sequence_number":1}`)	inProgressPayload, _ = sjson.SetRawBytes(inProgressPayload, "response", inProgressResponse)	addedPayload := []byte(`{"type":"response.output_item.added","sequence_number":2,"output_index":0}`)	addedPayload, _ = sjson.SetRawBytes(addedPayload, "item", item)	keepalivePayload := []byte(`{"type":"keepalive","sequence_number":3}`)	donePayload := []byte(`{"type":"response.output_item.done","sequence_number":4,"output_index":0}`)	donePayload, _ = sjson.SetRawBytes(donePayload, "item", item)	completedPayload := []byte(`{"type":"response.completed","sequence_number":5}`)	completedPayload, _ = sjson.SetRawBytes(completedPayload, "response", completedResponse)	completedPayload = helps.EnsureResponsesUsageDetails(completedPayload) 	return [][]byte{		xaiBuildSSEFrame("response.created", createdPayload),		xaiBuildSSEFrame("response.in_progress", inProgressPayload),		xaiBuildSSEFrame("response.output_item.added", addedPayload),		xaiBuildSSEFrame("keepalive", keepalivePayload),		xaiBuildSSEFrame("response.output_item.done", donePayload),		xaiBuildSSEFrame("response.completed", completedPayload),	}} func xaiBuildCompactionBaseResponse(prepared *xaiPreparedRequest, compactData []byte, responseID string, createdAt int64, status string) []byte {	response := []byte(`{"id":"","object":"response","created_at":0,"status":"","background":false,"error":null,"incomplete_details":null,"output":[]}`)	response, _ = sjson.SetBytes(response, "id", responseID)	response, _ = sjson.SetBytes(response, "created_at", createdAt)	response, _ = sjson.SetBytes(response, "status", status)	if model := gjson.GetBytes(compactData, "model").String(); model != "" {		response, _ = sjson.SetBytes(response, "model", model)	} else if prepared != nil && prepared.baseModel != "" {		response, _ = sjson.SetBytes(response, "model", prepared.baseModel)	} 	if prepared == nil {		return response	}	for _, field := range []string{		"instructions",		"max_output_tokens",		"max_tool_calls",		"parallel_tool_calls",		"previous_response_id",		"prompt_cache_key",		"reasoning",		"text",		"tool_choice",		"tools",		"top_logprobs",		"top_p",		"truncation",		"user",		"metadata",	} {		if value := gjson.GetBytes(prepared.body, field); value.Exists() {			response, _ = sjson.SetRawBytes(response, field, []byte(value.Raw))		}	}	return response} func xaiCompactionOutputItem(compactData []byte, responseID string) []byte {	itemResult := gjson.GetBytes(compactData, "output.0")	item := []byte(`{"type":"compaction"}`)	if itemResult.Exists() && itemResult.Type == gjson.JSON {		item = []byte(itemResult.Raw)	}	if !gjson.GetBytes(item, "type").Exists() {		item, _ = sjson.SetBytes(item, "type", "compaction")	}	if !gjson.GetBytes(item, "id").Exists() {		item, _ = sjson.SetBytes(item, "id", xaiCompactionItemID(responseID))	}	return item} func xaiCompactionResponseID(compactData []byte) string {	if responseID := strings.TrimSpace(gjson.GetBytes(compactData, "id").String()); responseID != "" {		if strings.HasPrefix(responseID, "resp_") {			return responseID		}		return "resp_" + strings.TrimPrefix(responseID, "cmp_")	}	return fmt.Sprintf("resp_xai_compaction_%d", time.Now().UnixNano())} func xaiCompactionItemID(responseID string) string {	if suffix := strings.TrimPrefix(responseID, "resp_"); suffix != "" && suffix != responseID {		return "cmp_" + suffix	}	return "cmp_" + responseID} func xaiBuildSSEFrame(eventName string, data []byte) []byte {	out := make([]byte, 0, len(eventName)+len(data)+16)	out = append(out, "event: "...)	out = append(out, eventName...)	out = append(out, '\n')	out = append(out, "data: "...)	out = append(out, data...)	out = append(out, '\n', '\n')	return out}