907 lines · 29.5 KB · Go
RawDownload
// Package executor provides runtime execution capabilities for various AI service providers.// This file implements the Gemini CLI executor that talks to Cloud Code Assist endpoints// using OAuth credentials from auth metadata.package executor import (	"bufio"	"bytes"	"context"	"encoding/json"	"fmt"	"io"	"net/http"	"regexp"	"strconv"	"strings"	"time" 	"github.com/router-for-me/CLIProxyAPI/v6/internal/config"	"github.com/router-for-me/CLIProxyAPI/v6/internal/misc"	"github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/executor/helps"	"github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/geminicli"	"github.com/router-for-me/CLIProxyAPI/v6/internal/thinking"	"github.com/router-for-me/CLIProxyAPI/v6/internal/util"	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"	sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"	log "github.com/sirupsen/logrus"	"github.com/tidwall/gjson"	"github.com/tidwall/sjson"	"golang.org/x/oauth2"	"golang.org/x/oauth2/google") const (	codeAssistEndpoint      = "https://cloudcode-pa.googleapis.com"	codeAssistVersion       = "v1internal"	geminiOAuthClientID     = "681255809395-oo8ft2oprdrnp9e3aqf6av3hmdib135j.apps.googleusercontent.com"	geminiOAuthClientSecret = "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") var geminiOAuthScopes = []string{	"https://www.googleapis.com/auth/cloud-platform",	"https://www.googleapis.com/auth/userinfo.email",	"https://www.googleapis.com/auth/userinfo.profile",} // GeminiCLIExecutor talks to the Cloud Code Assist endpoint using OAuth credentials from auth metadata.type GeminiCLIExecutor struct {	cfg *config.Config} // NewGeminiCLIExecutor creates a new Gemini CLI executor instance.//// Parameters://   - cfg: The application configuration//// Returns://   - *GeminiCLIExecutor: A new Gemini CLI executor instancefunc NewGeminiCLIExecutor(cfg *config.Config) *GeminiCLIExecutor {	return &GeminiCLIExecutor{cfg: cfg}} // Identifier returns the executor identifier.func (e *GeminiCLIExecutor) Identifier() string { return "gemini-cli" } // PrepareRequest injects Gemini CLI credentials into the outgoing HTTP request.func (e *GeminiCLIExecutor) PrepareRequest(req *http.Request, auth *cliproxyauth.Auth) error {	if req == nil {		return nil	}	tokenSource, _, errSource := prepareGeminiCLITokenSource(req.Context(), e.cfg, auth)	if errSource != nil {		return errSource	}	tok, errTok := tokenSource.Token()	if errTok != nil {		return errTok	}	if strings.TrimSpace(tok.AccessToken) == "" {		return statusErr{code: http.StatusUnauthorized, msg: "missing access token"}	}	req.Header.Set("Authorization", "Bearer "+tok.AccessToken)	applyGeminiCLIHeaders(req, "unknown")	var attrs map[string]string	if auth != nil {		attrs = auth.Attributes	}	util.ApplyCustomHeadersFromAttrs(req, attrs)	return nil} // HttpRequest injects Gemini CLI credentials into the request and executes it.func (e *GeminiCLIExecutor) HttpRequest(ctx context.Context, auth *cliproxyauth.Auth, req *http.Request) (*http.Response, error) {	if req == nil {		return nil, fmt.Errorf("gemini-cli executor: request is nil")	}	if ctx == nil {		ctx = req.Context()	}	httpReq := req.WithContext(ctx)	if err := e.PrepareRequest(httpReq, auth); err != nil {		return nil, err	}	httpClient := newHTTPClient(ctx, e.cfg, auth, 0)	return httpClient.Do(httpReq)} // Execute performs a non-streaming request to the Gemini CLI API.func (e *GeminiCLIExecutor) 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"}	}	baseModel := thinking.ParseSuffix(req.Model).ModelName 	tokenSource, baseTokenData, err := prepareGeminiCLITokenSource(ctx, e.cfg, auth)	if err != nil {		return resp, err	} 	reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)	defer reporter.TrackFailure(ctx, &err) 	from := opts.SourceFormat	to := sdktranslator.FromString("gemini-cli") 	originalPayloadSource := req.Payload	if len(opts.OriginalRequest) > 0 {		originalPayloadSource = opts.OriginalRequest	}	originalPayload := originalPayloadSource	originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, false)	basePayload := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, false) 	basePayload, err = thinking.ApplyThinking(basePayload, req.Model, from.String(), to.String(), e.Identifier())	if err != nil {		return resp, err	} 	basePayload = fixGeminiCLIImageAspectRatio(baseModel, basePayload)	requestedModel := helps.PayloadRequestedModel(opts, req.Model)	basePayload = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "gemini", "request", basePayload, originalTranslated, requestedModel) 	action := "generateContent"	if req.Metadata != nil {		if a, _ := req.Metadata["action"].(string); a == "countTokens" {			action = "countTokens"		}	} 	projectID := resolveGeminiProjectID(auth)	models := cliPreviewFallbackOrder(baseModel)	if len(models) == 0 || models[0] != baseModel {		models = append([]string{baseModel}, models...)	} 	httpClient := newHTTPClient(ctx, e.cfg, auth, 0)	respCtx := context.WithValue(ctx, "alt", opts.Alt) 	var authID, authLabel, authType, authValue string	authID = auth.ID	authLabel = auth.Label	authType, authValue = auth.AccountInfo() 	var lastStatus int	var lastBody []byte 	for idx, attemptModel := range models {		payload := append([]byte(nil), basePayload...)		if action == "countTokens" {			payload = deleteJSONField(payload, "project")			payload = deleteJSONField(payload, "model")		} else {			payload = setJSONField(payload, "project", projectID)			payload = setJSONField(payload, "model", attemptModel)		} 		tok, errTok := tokenSource.Token()		if errTok != nil {			err = errTok			return resp, err		}		updateGeminiCLITokenMetadata(auth, baseTokenData, tok) 		url := fmt.Sprintf("%s/%s:%s", codeAssistEndpoint, codeAssistVersion, action)		if opts.Alt != "" && action != "countTokens" {			url = url + fmt.Sprintf("?$alt=%s", opts.Alt)		} 		reqHTTP, errReq := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload))		if errReq != nil {			err = errReq			return resp, err		}		reqHTTP.Header.Set("Content-Type", "application/json")		reqHTTP.Header.Set("Authorization", "Bearer "+tok.AccessToken)		applyGeminiCLIHeaders(reqHTTP, attemptModel)		reqHTTP.Header.Set("Accept", "application/json")		util.ApplyCustomHeadersFromAttrs(reqHTTP, auth.Attributes)		helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{			URL:       url,			Method:    http.MethodPost,			Headers:   reqHTTP.Header.Clone(),			Body:      payload,			Provider:  e.Identifier(),			AuthID:    authID,			AuthLabel: authLabel,			AuthType:  authType,			AuthValue: authValue,		}) 		httpResp, errDo := httpClient.Do(reqHTTP)		if errDo != nil {			helps.RecordAPIResponseError(ctx, e.cfg, errDo)			err = errDo			return resp, err		} 		data, errRead := io.ReadAll(httpResp.Body)		if errClose := httpResp.Body.Close(); errClose != nil {			log.Errorf("gemini cli executor: close response body error: %v", errClose)		}		helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())		if errRead != nil {			helps.RecordAPIResponseError(ctx, e.cfg, errRead)			err = errRead			return resp, err		}		helps.AppendAPIResponseChunk(ctx, e.cfg, data)		if httpResp.StatusCode >= 200 && httpResp.StatusCode < 300 {			reporter.Publish(ctx, helps.ParseGeminiCLIUsage(data))			var param any			out := sdktranslator.TranslateNonStream(respCtx, to, from, attemptModel, opts.OriginalRequest, payload, data, &param)			resp = cliproxyexecutor.Response{Payload: out, Headers: httpResp.Header.Clone()}			return resp, nil		} 		lastStatus = httpResp.StatusCode		lastBody = append([]byte(nil), data...)		helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data))		if httpResp.StatusCode == 429 {			if idx+1 < len(models) {				log.Debugf("gemini cli executor: rate limited, retrying with next model: %s", models[idx+1])			} else {				log.Debug("gemini cli executor: rate limited, no additional fallback model")			}			continue		} 		err = newGeminiStatusErr(httpResp.StatusCode, data)		return resp, err	} 	if len(lastBody) > 0 {		helps.AppendAPIResponseChunk(ctx, e.cfg, lastBody)	}	if lastStatus == 0 {		lastStatus = 429	}	err = newGeminiStatusErr(lastStatus, lastBody)	return resp, err} // ExecuteStream performs a streaming request to the Gemini CLI API.func (e *GeminiCLIExecutor) 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"}	}	baseModel := thinking.ParseSuffix(req.Model).ModelName 	tokenSource, baseTokenData, err := prepareGeminiCLITokenSource(ctx, e.cfg, auth)	if err != nil {		return nil, err	} 	reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)	defer reporter.TrackFailure(ctx, &err) 	from := opts.SourceFormat	to := sdktranslator.FromString("gemini-cli") 	originalPayloadSource := req.Payload	if len(opts.OriginalRequest) > 0 {		originalPayloadSource = opts.OriginalRequest	}	originalPayload := originalPayloadSource	originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, true)	basePayload := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, true) 	basePayload, err = thinking.ApplyThinking(basePayload, req.Model, from.String(), to.String(), e.Identifier())	if err != nil {		return nil, err	} 	basePayload = fixGeminiCLIImageAspectRatio(baseModel, basePayload)	requestedModel := helps.PayloadRequestedModel(opts, req.Model)	basePayload = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "gemini", "request", basePayload, originalTranslated, requestedModel) 	projectID := resolveGeminiProjectID(auth) 	models := cliPreviewFallbackOrder(baseModel)	if len(models) == 0 || models[0] != baseModel {		models = append([]string{baseModel}, models...)	} 	httpClient := newHTTPClient(ctx, e.cfg, auth, 0)	respCtx := context.WithValue(ctx, "alt", opts.Alt) 	var authID, authLabel, authType, authValue string	authID = auth.ID	authLabel = auth.Label	authType, authValue = auth.AccountInfo() 	var lastStatus int	var lastBody []byte 	for idx, attemptModel := range models {		payload := append([]byte(nil), basePayload...)		payload = setJSONField(payload, "project", projectID)		payload = setJSONField(payload, "model", attemptModel) 		tok, errTok := tokenSource.Token()		if errTok != nil {			err = errTok			return nil, err		}		updateGeminiCLITokenMetadata(auth, baseTokenData, tok) 		url := fmt.Sprintf("%s/%s:%s", codeAssistEndpoint, codeAssistVersion, "streamGenerateContent")		if opts.Alt == "" {			url = url + "?alt=sse"		} else {			url = url + fmt.Sprintf("?$alt=%s", opts.Alt)		} 		reqHTTP, errReq := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload))		if errReq != nil {			err = errReq			return nil, err		}		reqHTTP.Header.Set("Content-Type", "application/json")		reqHTTP.Header.Set("Authorization", "Bearer "+tok.AccessToken)		applyGeminiCLIHeaders(reqHTTP, attemptModel)		reqHTTP.Header.Set("Accept", "text/event-stream")		util.ApplyCustomHeadersFromAttrs(reqHTTP, auth.Attributes)		helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{			URL:       url,			Method:    http.MethodPost,			Headers:   reqHTTP.Header.Clone(),			Body:      payload,			Provider:  e.Identifier(),			AuthID:    authID,			AuthLabel: authLabel,			AuthType:  authType,			AuthValue: authValue,		}) 		httpResp, errDo := httpClient.Do(reqHTTP)		if errDo != nil {			helps.RecordAPIResponseError(ctx, e.cfg, errDo)			err = errDo			return nil, err		}		helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())		if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {			data, errRead := io.ReadAll(httpResp.Body)			if errClose := httpResp.Body.Close(); errClose != nil {				log.Errorf("gemini cli executor: close response body error: %v", errClose)			}			if errRead != nil {				helps.RecordAPIResponseError(ctx, e.cfg, errRead)				err = errRead				return nil, err			}			helps.AppendAPIResponseChunk(ctx, e.cfg, data)			lastStatus = httpResp.StatusCode			lastBody = append([]byte(nil), data...)			helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data))			if httpResp.StatusCode == 429 {				if idx+1 < len(models) {					log.Debugf("gemini cli executor: rate limited, retrying with next model: %s", models[idx+1])				} else {					log.Debug("gemini cli executor: rate limited, no additional fallback model")				}				continue			}			err = newGeminiStatusErr(httpResp.StatusCode, data)			return nil, err		} 		out := make(chan cliproxyexecutor.StreamChunk)		go func(resp *http.Response, reqBody []byte, attemptModel string) {			defer close(out)			defer func() {				if errClose := resp.Body.Close(); errClose != nil {					log.Errorf("gemini cli executor: close response body error: %v", errClose)				}			}()			if opts.Alt == "" {				scanner := bufio.NewScanner(resp.Body)				scanner.Buffer(nil, streamScannerBuffer)				var param any				for scanner.Scan() {					line := scanner.Bytes()					helps.AppendAPIResponseChunk(ctx, e.cfg, line)					if detail, ok := helps.ParseGeminiCLIStreamUsage(line); ok {						reporter.Publish(ctx, detail)					}					if bytes.HasPrefix(line, dataTag) {						segments := sdktranslator.TranslateStream(respCtx, to, from, attemptModel, opts.OriginalRequest, reqBody, bytes.Clone(line), &param)						for i := range segments {							out <- cliproxyexecutor.StreamChunk{Payload: segments[i]}						}					}				} 				segments := sdktranslator.TranslateStream(respCtx, to, from, attemptModel, opts.OriginalRequest, reqBody, []byte("[DONE]"), &param)				for i := range segments {					out <- cliproxyexecutor.StreamChunk{Payload: segments[i]}				}				if errScan := scanner.Err(); errScan != nil {					helps.RecordAPIResponseError(ctx, e.cfg, errScan)					reporter.PublishFailure(ctx)					out <- cliproxyexecutor.StreamChunk{Err: errScan}				}				return			} 			data, errRead := io.ReadAll(resp.Body)			if errRead != nil {				helps.RecordAPIResponseError(ctx, e.cfg, errRead)				reporter.PublishFailure(ctx)				out <- cliproxyexecutor.StreamChunk{Err: errRead}				return			}			helps.AppendAPIResponseChunk(ctx, e.cfg, data)			reporter.Publish(ctx, helps.ParseGeminiCLIUsage(data))			var param any			segments := sdktranslator.TranslateStream(respCtx, to, from, attemptModel, opts.OriginalRequest, reqBody, data, &param)			for i := range segments {				out <- cliproxyexecutor.StreamChunk{Payload: segments[i]}			} 			segments = sdktranslator.TranslateStream(respCtx, to, from, attemptModel, opts.OriginalRequest, reqBody, []byte("[DONE]"), &param)			for i := range segments {				out <- cliproxyexecutor.StreamChunk{Payload: segments[i]}			}		}(httpResp, append([]byte(nil), payload...), attemptModel) 		return &cliproxyexecutor.StreamResult{Headers: httpResp.Header.Clone(), Chunks: out}, nil	} 	if len(lastBody) > 0 {		helps.AppendAPIResponseChunk(ctx, e.cfg, lastBody)	}	if lastStatus == 0 {		lastStatus = 429	}	err = newGeminiStatusErr(lastStatus, lastBody)	return nil, err} // CountTokens counts tokens for the given request using the Gemini CLI API.func (e *GeminiCLIExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {	baseModel := thinking.ParseSuffix(req.Model).ModelName 	tokenSource, baseTokenData, err := prepareGeminiCLITokenSource(ctx, e.cfg, auth)	if err != nil {		return cliproxyexecutor.Response{}, err	} 	from := opts.SourceFormat	to := sdktranslator.FromString("gemini-cli") 	models := cliPreviewFallbackOrder(baseModel)	if len(models) == 0 || models[0] != baseModel {		models = append([]string{baseModel}, models...)	} 	httpClient := newHTTPClient(ctx, e.cfg, auth, 0)	respCtx := context.WithValue(ctx, "alt", opts.Alt) 	var authID, authLabel, authType, authValue string	if auth != nil {		authID = auth.ID		authLabel = auth.Label		authType, authValue = auth.AccountInfo()	} 	var lastStatus int	var lastBody []byte 	// The loop variable attemptModel is only used as the concrete model id sent to the upstream	// Gemini CLI endpoint when iterating fallback variants.	for range models {		payload := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, false) 		payload, err = thinking.ApplyThinking(payload, req.Model, from.String(), to.String(), e.Identifier())		if err != nil {			return cliproxyexecutor.Response{}, err

Lines 1–500 of 907.