// Package executor provides runtime execution capabilities for various AI service providers.// This file implements the AI Studio executor that routes requests through a websocket-backed// transport for the AI Studio provider.package executor import (	"bytes"	"context"	"encoding/json"	"fmt"	"io"	"net/http"	"net/url"	"strings" 	"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/thinking"	"github.com/router-for-me/CLIProxyAPI/v7/internal/util"	"github.com/router-for-me/CLIProxyAPI/v7/internal/wsrelay"	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"	"github.com/tidwall/gjson"	"github.com/tidwall/sjson") // AIStudioExecutor routes AI Studio requests through a websocket-backed transport.type AIStudioExecutor struct {	provider string	relay    *wsrelay.Manager	cfg      *config.Config} // NewAIStudioExecutor creates a new AI Studio executor instance.//// Parameters://   - cfg: The application configuration//   - provider: The provider name//   - relay: The websocket relay manager//// Returns://   - *AIStudioExecutor: A new AI Studio executor instancefunc NewAIStudioExecutor(cfg *config.Config, provider string, relay *wsrelay.Manager) *AIStudioExecutor {	return &AIStudioExecutor{provider: strings.ToLower(provider), relay: relay, cfg: cfg}} // Identifier returns the executor identifier.func (e *AIStudioExecutor) Identifier() string { return "aistudio" } // PrepareRequest prepares the HTTP request for execution.func (e *AIStudioExecutor) PrepareRequest(req *http.Request, auth *cliproxyauth.Auth) error {	if req == nil {		return nil	}	var attrs map[string]string	if auth != nil {		attrs = auth.Attributes	}	util.ApplyCustomHeadersFromAttrs(req, attrs)	return nil} // HttpRequest forwards an arbitrary HTTP request through the websocket relay.func (e *AIStudioExecutor) HttpRequest(ctx context.Context, auth *cliproxyauth.Auth, req *http.Request) (*http.Response, error) {	if req == nil {		return nil, fmt.Errorf("aistudio executor: request is nil")	}	if ctx == nil {		ctx = req.Context()	}	if e.relay == nil {		return nil, fmt.Errorf("aistudio executor: ws relay is nil")	}	if auth == nil || auth.ID == "" {		return nil, fmt.Errorf("aistudio executor: missing auth")	}	httpReq := req.WithContext(ctx)	if err := e.PrepareRequest(httpReq, auth); err != nil {		return nil, err	}	if httpReq.URL == nil || strings.TrimSpace(httpReq.URL.String()) == "" {		return nil, fmt.Errorf("aistudio executor: request URL is empty")	} 	var body []byte	if httpReq.Body != nil {		b, errRead := io.ReadAll(httpReq.Body)		if errRead != nil {			return nil, errRead		}		body = b		httpReq.Body = io.NopCloser(bytes.NewReader(b))	} 	wsReq := &wsrelay.HTTPRequest{		Method:  httpReq.Method,		URL:     httpReq.URL.String(),		Headers: httpReq.Header.Clone(),		Body:    body,	}	wsResp, errRelay := e.relay.NonStream(ctx, auth.ID, wsReq)	if errRelay != nil {		return nil, errRelay	}	if wsResp == nil {		return nil, fmt.Errorf("aistudio executor: ws response is nil")	} 	statusText := http.StatusText(wsResp.Status)	if statusText == "" {		statusText = "Unknown"	}	resp := &http.Response{		StatusCode:    wsResp.Status,		Status:        fmt.Sprintf("%d %s", wsResp.Status, statusText),		Header:        wsResp.Headers.Clone(),		Body:          io.NopCloser(bytes.NewReader(wsResp.Body)),		ContentLength: int64(len(wsResp.Body)),		Request:       httpReq,	}	return resp, nil} // Execute performs a non-streaming request to the AI Studio API.func (e *AIStudioExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) {	if opts.Alt == "responses/compact" {		return resp, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"}	}	ctx = helps.EnsureSessionContext(ctx, opts, req.Payload)	baseModel := thinking.ParseSuffix(req.Model).ModelName	reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)	defer reporter.TrackFailure(ctx, &err) 	translatedReq, body, err := e.translateRequest(ctx, req, opts, false)	if err != nil {		return resp, err	}	reporter.SetTranslatedReasoningEffort(body.payload, body.toFormat.String()) 	endpoint := e.buildEndpoint(baseModel, body.action, opts.Alt)	wsReq := &wsrelay.HTTPRequest{		Method:  http.MethodPost,		URL:     endpoint,		Headers: http.Header{"Content-Type": []string{"application/json"}},		Body:    body.payload,	}	var attrs map[string]string	if auth != nil {		attrs = auth.Attributes	}	customHeaderReq := &http.Request{Header: wsReq.Headers}	if ctx != nil {		customHeaderReq = customHeaderReq.WithContext(ctx)	}	util.ApplyCustomHeadersFromAttrs(customHeaderReq, attrs, opts.Headers) 	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:       endpoint,		Method:    http.MethodPost,		Headers:   wsReq.Headers.Clone(),		Body:      body.payload,		Provider:  e.Identifier(),		AuthID:    authID,		AuthLabel: authLabel,		AuthType:  authType,		AuthValue: authValue,	}) 	reporter.StartResponseTTFT()	wsResp, err := e.relay.NonStream(ctx, authID, wsReq)	if err != nil {		helps.RecordAPIResponseError(ctx, e.cfg, err)		return resp, err	}	helps.RecordAPIResponseMetadata(ctx, e.cfg, wsResp.Status, wsResp.Headers.Clone())	reporter.StartResponseTTFT()	if len(wsResp.Body) > 0 {		reporter.MarkFirstResponseByte()		helps.AppendAPIResponseChunk(ctx, e.cfg, wsResp.Body)	}	if wsResp.Status < 200 || wsResp.Status >= 300 {		return resp, statusErr{code: wsResp.Status, msg: string(wsResp.Body)}	}	reporter.Publish(ctx, helps.ParseGeminiUsage(wsResp.Body))	responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts)	var param any	out := sdktranslator.TranslateNonStream(ctx, body.toFormat, responseFormat, req.Model, opts.OriginalRequest, translatedReq, wsResp.Body, &param)	if responseFormat == sdktranslator.FormatOpenAIResponse {		out = helps.EnsureResponsesUsageDetails(out)	}	resp = cliproxyexecutor.Response{Payload: ensureColonSpacedJSON(out), Headers: wsResp.Headers.Clone()}	return resp, nil} // ExecuteStream performs a streaming request to the AI Studio API.func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (_ *cliproxyexecutor.StreamResult, err error) {	if opts.Alt == "responses/compact" {		return nil, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"}	}	ctx = helps.EnsureSessionContext(ctx, opts, req.Payload)	baseModel := thinking.ParseSuffix(req.Model).ModelName	reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)	defer reporter.TrackFailure(ctx, &err) 	translatedReq, body, err := e.translateRequest(ctx, req, opts, true)	if err != nil {		return nil, err	}	reporter.SetTranslatedReasoningEffort(body.payload, body.toFormat.String()) 	endpoint := e.buildEndpoint(baseModel, body.action, opts.Alt)	wsReq := &wsrelay.HTTPRequest{		Method:  http.MethodPost,		URL:     endpoint,		Headers: http.Header{"Content-Type": []string{"application/json"}},		Body:    body.payload,	}	var attrs map[string]string	if auth != nil {		attrs = auth.Attributes	}	customHeaderReq := &http.Request{Header: wsReq.Headers}	if ctx != nil {		customHeaderReq = customHeaderReq.WithContext(ctx)	}	util.ApplyCustomHeadersFromAttrs(customHeaderReq, attrs, opts.Headers)	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:       endpoint,		Method:    http.MethodPost,		Headers:   wsReq.Headers.Clone(),		Body:      body.payload,		Provider:  e.Identifier(),		AuthID:    authID,		AuthLabel: authLabel,		AuthType:  authType,		AuthValue: authValue,	})	reporter.StartResponseTTFT()	wsStream, err := e.relay.Stream(ctx, authID, wsReq)	if err != nil {		helps.RecordAPIResponseError(ctx, e.cfg, err)		return nil, err	}	firstEvent, ok := <-wsStream	if !ok {		err = fmt.Errorf("wsrelay: stream closed before start")		helps.RecordAPIResponseError(ctx, e.cfg, err)		return nil, err	}	if firstEvent.Status > 0 && firstEvent.Status != http.StatusOK {		metadataLogged := false		if firstEvent.Status > 0 {			helps.RecordAPIResponseMetadata(ctx, e.cfg, firstEvent.Status, firstEvent.Headers.Clone())			reporter.StartResponseTTFT()			metadataLogged = true		}		var body bytes.Buffer		if len(firstEvent.Payload) > 0 {			reporter.MarkFirstResponseByte()			helps.AppendAPIResponseChunk(ctx, e.cfg, firstEvent.Payload)			body.Write(firstEvent.Payload)		}		if firstEvent.Type == wsrelay.MessageTypeStreamEnd {			return nil, statusErr{code: firstEvent.Status, msg: body.String()}		}		for event := range wsStream {			if event.Err != nil {				helps.RecordAPIResponseError(ctx, e.cfg, event.Err)				if body.Len() == 0 {					body.WriteString(event.Err.Error())				}				break			}			if !metadataLogged && event.Status > 0 {				helps.RecordAPIResponseMetadata(ctx, e.cfg, event.Status, event.Headers.Clone())				reporter.StartResponseTTFT()				metadataLogged = true			}			if len(event.Payload) > 0 {				reporter.MarkFirstResponseByte()				helps.AppendAPIResponseChunk(ctx, e.cfg, event.Payload)				body.Write(event.Payload)			}			if event.Type == wsrelay.MessageTypeStreamEnd {				break			}		}		return nil, statusErr{code: firstEvent.Status, msg: body.String()}	}	out := make(chan cliproxyexecutor.StreamChunk)	go func(first wsrelay.StreamEvent) {		defer close(out)		defer reporter.EnsurePublished(ctx)		responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts)		originalRequest := opts.OriginalRequest		if len(originalRequest) == 0 {			originalRequest = req.Payload		}		claudeInputTokens := helps.NewClaudeInputTokenState(opts.SourceFormat, body.toFormat, responseFormat, originalRequest)		var param any		metadataLogged := false		processEvent := func(event wsrelay.StreamEvent) bool {			if event.Err != nil {				helps.RecordAPIResponseError(ctx, e.cfg, event.Err)				reporter.PublishFailure(ctx, event.Err)				select {				case out <- cliproxyexecutor.StreamChunk{Err: fmt.Errorf("wsrelay: %v", event.Err)}:				case <-ctx.Done():				}				return false			}			switch event.Type {			case wsrelay.MessageTypeStreamStart:				if !metadataLogged && event.Status > 0 {					helps.RecordAPIResponseMetadata(ctx, e.cfg, event.Status, event.Headers.Clone())					reporter.StartResponseTTFT()					metadataLogged = true				}			case wsrelay.MessageTypeStreamChunk:				if len(event.Payload) > 0 {					reporter.MarkFirstResponseByte()					helps.AppendAPIResponseChunk(ctx, e.cfg, event.Payload)					filtered := helps.FilterSSEUsageMetadata(event.Payload)					if detail, ok := helps.ParseGeminiStreamUsage(filtered); ok {						reporter.Publish(ctx, detail)					}					lines := helps.TranslateStreamWithClaudeInputTokens(ctx, body.toFormat, responseFormat, req.Model, opts.OriginalRequest, translatedReq, filtered, &param, claudeInputTokens)					for i := range lines {						select {						case out <- cliproxyexecutor.StreamChunk{Payload: ensureColonSpacedJSON(lines[i])}:						case <-ctx.Done():							return false						}					}					break				}			case wsrelay.MessageTypeStreamEnd:				return false			case wsrelay.MessageTypeHTTPResp:				if !metadataLogged && event.Status > 0 {					helps.RecordAPIResponseMetadata(ctx, e.cfg, event.Status, event.Headers.Clone())					reporter.StartResponseTTFT()					metadataLogged = true				}				if len(event.Payload) > 0 {					reporter.MarkFirstResponseByte()					helps.AppendAPIResponseChunk(ctx, e.cfg, event.Payload)				}				lines := helps.TranslateStreamWithClaudeInputTokens(ctx, body.toFormat, responseFormat, req.Model, opts.OriginalRequest, translatedReq, event.Payload, &param, claudeInputTokens)				for i := range lines {					select {					case out <- cliproxyexecutor.StreamChunk{Payload: ensureColonSpacedJSON(lines[i])}:					case <-ctx.Done():						return false					}				}				reporter.Publish(ctx, helps.ParseGeminiUsage(event.Payload))				return false			case wsrelay.MessageTypeError:				helps.RecordAPIResponseError(ctx, e.cfg, event.Err)				reporter.PublishFailure(ctx, event.Err)				select {				case out <- cliproxyexecutor.StreamChunk{Err: fmt.Errorf("wsrelay: %v", event.Err)}:				case <-ctx.Done():				}				return false			}			return true		}		if !processEvent(first) {			return		}		for event := range wsStream {			if !processEvent(event) {				return			}		}	}(firstEvent)	return &cliproxyexecutor.StreamResult{Headers: firstEvent.Headers.Clone(), Chunks: out}, nil} // CountTokens counts tokens for the given request using the AI Studio API.func (e *AIStudioExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {	baseModel := thinking.ParseSuffix(req.Model).ModelName	countReq := req	countMetadata := make(map[string]any, len(req.Metadata)+1)	for k, v := range req.Metadata {		countMetadata[k] = v	}	countMetadata["action"] = "countTokens"	countReq.Metadata = countMetadata 	_, body, err := e.translateRequest(ctx, countReq, opts, false)	if err != nil {		return cliproxyexecutor.Response{}, err	} 	body.payload, _ = sjson.DeleteBytes(body.payload, "generationConfig")	body.payload, _ = sjson.DeleteBytes(body.payload, "tools")	body.payload, _ = sjson.DeleteBytes(body.payload, "safetySettings") 	endpoint := e.buildEndpoint(baseModel, "countTokens", "")	wsReq := &wsrelay.HTTPRequest{		Method:  http.MethodPost,		URL:     endpoint,		Headers: http.Header{"Content-Type": []string{"application/json"}},		Body:    body.payload,	}	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:       endpoint,		Method:    http.MethodPost,		Headers:   wsReq.Headers.Clone(),		Body:      body.payload,		Provider:  e.Identifier(),		AuthID:    authID,		AuthLabel: authLabel,		AuthType:  authType,		AuthValue: authValue,	})	resp, err := e.relay.NonStream(ctx, authID, wsReq)	if err != nil {		helps.RecordAPIResponseError(ctx, e.cfg, err)		return cliproxyexecutor.Response{}, err	}	helps.RecordAPIResponseMetadata(ctx, e.cfg, resp.Status, resp.Headers.Clone())	if len(resp.Body) > 0 {		helps.AppendAPIResponseChunk(ctx, e.cfg, resp.Body)	}	if resp.Status < 200 || resp.Status >= 300 {		return cliproxyexecutor.Response{}, statusErr{code: resp.Status, msg: string(resp.Body)}	}	totalTokens := gjson.GetBytes(resp.Body, "totalTokens").Int()	if totalTokens <= 0 {		return cliproxyexecutor.Response{}, fmt.Errorf("wsrelay: totalTokens missing in response")	}	responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts)	translated := sdktranslator.TranslateTokenCount(ctx, body.toFormat, responseFormat, totalTokens, resp.Body)	return cliproxyexecutor.Response{Payload: translated}, nil} // Refresh refreshes the authentication credentials (no-op for AI Studio).func (e *AIStudioExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {	if refreshed, handled, err := helps.RefreshAuthViaHome(ctx, e.cfg, auth); handled {		return refreshed, err	}	return auth, nil} type translatedPayload struct {	payload  []byte	action   string	toFormat sdktranslator.Format} func (e *AIStudioExecutor) translateRequest(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, stream bool) ([]byte, translatedPayload, error) {	baseModel := thinking.ParseSuffix(req.Model).ModelName 	from := opts.SourceFormat	to := sdktranslator.FromString("gemini")	originalPayloadSource := req.Payload	if len(opts.OriginalRequest) > 0 {		originalPayloadSource = opts.OriginalRequest	}	originalPayload := originalPayloadSource	originalTranslated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, originalPayload, stream)	payload := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, stream)	payload, err := helps.ApplyThinkingWithSourcePayload(payload, req.Payload, originalPayloadSource, req.Model, from.String(), to.String(), e.Identifier())	if err != nil {		return nil, translatedPayload{}, err	}	payload = fixGeminiImageAspectRatio(baseModel, payload)	requestedModel := helps.PayloadRequestedModel(opts, req.Model)	requestPath := helps.PayloadRequestPath(opts)	payload = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", payload, originalTranslated, requestedModel, requestPath, opts.Headers)	payload, _ = sjson.DeleteBytes(payload, "generationConfig.maxOutputTokens")	payload, _ = sjson.DeleteBytes(payload, "generationConfig.responseMimeType")	payload, _ = sjson.DeleteBytes(payload, "generationConfig.responseJsonSchema")	metadataAction := "generateContent"	if req.Metadata != nil {		if action, _ := req.Metadata["action"].(string); action == "countTokens" {			metadataAction = action

Lines 1–500 of 610. Load more