package executor import ( "bytes" "context" "fmt" "net/http" "sync" "time" "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/gemini" "github.com/router-for-me/CLIProxyAPI/v6/internal/config" "github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces" geminiwebapi "github.com/router-for-me/CLIProxyAPI/v6/internal/provider/gemini-web" 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" ) type GeminiWebExecutor struct { cfg *config.Config mu sync.Mutex } func NewGeminiWebExecutor(cfg *config.Config) *GeminiWebExecutor { return &GeminiWebExecutor{cfg: cfg} } func (e *GeminiWebExecutor) Identifier() string { return "gemini-web" } func (e *GeminiWebExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.Auth) error { return nil } func (e *GeminiWebExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { state, err := e.stateFor(auth) if err != nil { return cliproxyexecutor.Response{}, err } if err = state.EnsureClient(); err != nil { return cliproxyexecutor.Response{}, err } reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth) mutex := state.GetRequestMutex() if mutex != nil { mutex.Lock() defer mutex.Unlock() } payload := bytes.Clone(req.Payload) resp, errMsg, prep := state.Send(ctx, req.Model, payload, opts) if errMsg != nil { return cliproxyexecutor.Response{}, geminiWebErrorFromMessage(errMsg) } resp = state.ConvertToTarget(ctx, req.Model, prep, resp) reporter.publish(ctx, parseGeminiUsage(resp)) from := opts.SourceFormat to := sdktranslator.FromString("gemini-web") var param any out := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), payload, bytes.Clone(resp), ¶m) return cliproxyexecutor.Response{Payload: []byte(out)}, nil } func (e *GeminiWebExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (<-chan cliproxyexecutor.StreamChunk, error) { state, err := e.stateFor(auth) if err != nil { return nil, err } if err = state.EnsureClient(); err != nil { return nil, err } reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth) mutex := state.GetRequestMutex() if mutex != nil { mutex.Lock() } gemBytes, errMsg, prep := state.Send(ctx, req.Model, bytes.Clone(req.Payload), opts) if errMsg != nil { if mutex != nil { mutex.Unlock() } return nil, geminiWebErrorFromMessage(errMsg) } reporter.publish(ctx, parseGeminiUsage(gemBytes)) from := opts.SourceFormat to := sdktranslator.FromString("gemini-web") var param any lines := state.ConvertStream(ctx, req.Model, prep, gemBytes) done := state.DoneStream(ctx, req.Model, prep) out := make(chan cliproxyexecutor.StreamChunk) go func() { defer close(out) if mutex != nil { defer mutex.Unlock() } for _, line := range lines { lines = sdktranslator.TranslateStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), req.Payload, bytes.Clone([]byte(line)), ¶m) for _, l := range lines { out <- cliproxyexecutor.StreamChunk{Payload: []byte(l)} } } for _, line := range done { lines = sdktranslator.TranslateStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), req.Payload, bytes.Clone([]byte(line)), ¶m) for _, l := range lines { out <- cliproxyexecutor.StreamChunk{Payload: []byte(l)} } } }() return out, nil } func (e *GeminiWebExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { return cliproxyexecutor.Response{Payload: []byte{}}, fmt.Errorf("not implemented") } func (e *GeminiWebExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { log.Debugf("gemini web executor: refresh called") state, err := e.stateFor(auth) if err != nil { return nil, err } if err = state.Refresh(ctx); err != nil { return nil, err } ts := state.TokenSnapshot() if auth.Metadata == nil { auth.Metadata = make(map[string]any) } auth.Metadata["secure_1psid"] = ts.Secure1PSID auth.Metadata["secure_1psidts"] = ts.Secure1PSIDTS auth.Metadata["type"] = "gemini-web" auth.Metadata["last_refresh"] = time.Now().Format(time.RFC3339) return auth, nil } type geminiWebRuntime struct { state *geminiwebapi.GeminiWebState } func (e *GeminiWebExecutor) stateFor(auth *cliproxyauth.Auth) (*geminiwebapi.GeminiWebState, error) { if auth == nil { return nil, fmt.Errorf("gemini-web executor: auth is nil") } if runtime, ok := auth.Runtime.(*geminiWebRuntime); ok && runtime != nil && runtime.state != nil { return runtime.state, nil } e.mu.Lock() defer e.mu.Unlock() if runtime, ok := auth.Runtime.(*geminiWebRuntime); ok && runtime != nil && runtime.state != nil { return runtime.state, nil } ts, err := parseGeminiWebToken(auth) if err != nil { return nil, err } cfg := e.cfg if auth.ProxyURL != "" && cfg != nil { copyCfg := *cfg copyCfg.ProxyURL = auth.ProxyURL cfg = ©Cfg } storagePath := "" if auth.Attributes != nil { if p, ok := auth.Attributes["path"]; ok { storagePath = p } } state := geminiwebapi.NewGeminiWebState(cfg, ts, storagePath) runtime := &geminiWebRuntime{state: state} auth.Runtime = runtime return state, nil } func parseGeminiWebToken(auth *cliproxyauth.Auth) (*gemini.GeminiWebTokenStorage, error) { if auth == nil { return nil, fmt.Errorf("gemini-web executor: auth is nil") } if auth.Metadata == nil { return nil, fmt.Errorf("gemini-web executor: missing metadata") } psid := stringFromMetadata(auth.Metadata, "secure_1psid", "secure_1psid", "__Secure-1PSID") psidts := stringFromMetadata(auth.Metadata, "secure_1psidts", "secure_1psidts", "__Secure-1PSIDTS") if psid == "" || psidts == "" { return nil, fmt.Errorf("gemini-web executor: incomplete cookie metadata") } return &gemini.GeminiWebTokenStorage{Secure1PSID: psid, Secure1PSIDTS: psidts}, nil } func stringFromMetadata(meta map[string]any, keys ...string) string { for _, key := range keys { if val, ok := meta[key]; ok { if s, okStr := val.(string); okStr && s != "" { return s } } } return "" } func geminiWebErrorFromMessage(msg *interfaces.ErrorMessage) error { if msg == nil { return nil } return geminiWebError{message: msg} } type geminiWebError struct { message *interfaces.ErrorMessage } func (e geminiWebError) Error() string { if e.message == nil { return "gemini-web error" } if e.message.Error != nil { return e.message.Error.Error() } return fmt.Sprintf("gemini-web error: status %d", e.message.StatusCode) } func (e geminiWebError) StatusCode() int { if e.message == nil { return 0 } return e.message.StatusCode }