chuan/CLIProxyAPI
1
0
Fork 0
mirror of https://github.com/router-for-me/CLIProxyAPI.git synced 2026-02-03 13:00:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

rebuild branch

Browse Source
This commit is contained in:
Luis Pater
2025-09-25 10:32:48 +08:00
parent 3f69254f43
commit f5dc380b63
214 changed files with 39377 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

330
internal/runtime/executor/claude_executor.go Normal file
View File

@@ -0,0 +1,330 @@
package executor
import (
"bufio"
"bytes"
"context"
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/klauspost/compress/zstd"
claudeauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/claude"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
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"
"github.com/gin-gonic/gin"
)
// ClaudeExecutor is a stateless executor for Anthropic Claude over the messages API.
// If api_key is unavailable on auth, it falls back to legacy via ClientAdapter.
type ClaudeExecutor struct {
cfg *config.Config
}
func NewClaudeExecutor(cfg *config.Config) *ClaudeExecutor { return &ClaudeExecutor{cfg: cfg} }
func (e *ClaudeExecutor) Identifier() string { return "claude" }
func (e *ClaudeExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.Auth) error { return nil }
func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
apiKey, baseURL := claudeCreds(auth)
if baseURL == "" {
baseURL = "https://api.anthropic.com"
}
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
from := opts.SourceFormat
to := sdktranslator.FromString("claude")
// Use streaming translation to preserve function calling, except for claude.
stream := from != to
body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), stream)
if !strings.HasPrefix(req.Model, "claude-3-5-haiku") {
body, _ = sjson.SetRawBytes(body, "system", []byte(misc.ClaudeCodeInstructions))
}
url := fmt.Sprintf("%s/v1/messages?beta=true", baseURL)
recordAPIRequest(ctx, e.cfg, body)
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
if err != nil {
return cliproxyexecutor.Response{}, err
}
applyClaudeHeaders(httpReq, apiKey, false)
httpClient := &http.Client{}
if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
httpClient.Transport = rt
}
resp, err := httpClient.Do(httpReq)
if err != nil {
return cliproxyexecutor.Response{}, err
}
defer func() {
if errClose := resp.Body.Close(); errClose != nil {
log.Errorf("response body close error: %v", errClose)
}
}()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
b, _ := io.ReadAll(resp.Body)
appendAPIResponseChunk(ctx, e.cfg, b)
log.Debugf("request error, error status: %d, error body: %s", resp.StatusCode, string(b))
return cliproxyexecutor.Response{}, statusErr{code: resp.StatusCode, msg: string(b)}
}
reader := io.Reader(resp.Body)
var decoder *zstd.Decoder
if hasZSTDEcoding(resp.Header.Get("Content-Encoding")) {
decoder, err = zstd.NewReader(resp.Body)
if err != nil {
return cliproxyexecutor.Response{}, fmt.Errorf("failed to initialize zstd decoder: %w", err)
}
reader = decoder
defer decoder.Close()
}
data, err := io.ReadAll(reader)
if err != nil {
return cliproxyexecutor.Response{}, err
}
appendAPIResponseChunk(ctx, e.cfg, data)
if stream {
lines := bytes.Split(data, []byte("\n"))
for _, line := range lines {
if detail, ok := parseClaudeStreamUsage(line); ok {
reporter.publish(ctx, detail)
}
}
} else {
reporter.publish(ctx, parseClaudeUsage(data))
}
var param any
out := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, data, &param)
return cliproxyexecutor.Response{Payload: []byte(out)}, nil
}
func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (<-chan cliproxyexecutor.StreamChunk, error) {
apiKey, baseURL := claudeCreds(auth)
if baseURL == "" {
baseURL = "https://api.anthropic.com"
}
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
from := opts.SourceFormat
to := sdktranslator.FromString("claude")
body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)
body, _ = sjson.SetRawBytes(body, "system", []byte(misc.ClaudeCodeInstructions))
url := fmt.Sprintf("%s/v1/messages?beta=true", baseURL)
recordAPIRequest(ctx, e.cfg, body)
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
if err != nil {
return nil, err
}
applyClaudeHeaders(httpReq, apiKey, true)
httpClient := &http.Client{Timeout: 0}
if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
httpClient.Transport = rt
}
resp, err := httpClient.Do(httpReq)
if err != nil {
return nil, err
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
defer func() { _ = resp.Body.Close() }()
b, _ := io.ReadAll(resp.Body)
appendAPIResponseChunk(ctx, e.cfg, b)
log.Debugf("request error, error status: %d, error body: %s", resp.StatusCode, string(b))
return nil, statusErr{code: resp.StatusCode, msg: string(b)}
}
out := make(chan cliproxyexecutor.StreamChunk)
go func() {
defer close(out)
defer func() { _ = resp.Body.Close() }()
scanner := bufio.NewScanner(resp.Body)
buf := make([]byte, 1024*1024)
scanner.Buffer(buf, 1024*1024)
var param any
for scanner.Scan() {
line := scanner.Bytes()
appendAPIResponseChunk(ctx, e.cfg, line)
if detail, ok := parseClaudeStreamUsage(line); ok {
reporter.publish(ctx, detail)
}
chunks := sdktranslator.TranslateStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, bytes.Clone(line), &param)
for i := range chunks {
out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunks[i])}
}
}
if err = scanner.Err(); err != nil {
out <- cliproxyexecutor.StreamChunk{Err: err}
}
}()
return out, nil
}
func (e *ClaudeExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
apiKey, baseURL := claudeCreds(auth)
if baseURL == "" {
baseURL = "https://api.anthropic.com"
}
from := opts.SourceFormat
to := sdktranslator.FromString("claude")
// Use streaming translation to preserve function calling, except for claude.
stream := from != to
body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), stream)
if !strings.HasPrefix(req.Model, "claude-3-5-haiku") {
body, _ = sjson.SetRawBytes(body, "system", []byte(misc.ClaudeCodeInstructions))
}
url := fmt.Sprintf("%s/v1/messages/count_tokens?beta=true", baseURL)
recordAPIRequest(ctx, e.cfg, body)
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
if err != nil {
return cliproxyexecutor.Response{}, err
}
applyClaudeHeaders(httpReq, apiKey, false)
httpClient := &http.Client{}
if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
httpClient.Transport = rt
}
resp, err := httpClient.Do(httpReq)
if err != nil {
return cliproxyexecutor.Response{}, err
}
defer func() {
if errClose := resp.Body.Close(); errClose != nil {
log.Errorf("response body close error: %v", errClose)
}
}()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
b, _ := io.ReadAll(resp.Body)
appendAPIResponseChunk(ctx, e.cfg, b)
return cliproxyexecutor.Response{}, statusErr{code: resp.StatusCode, msg: string(b)}
}
reader := io.Reader(resp.Body)
var decoder *zstd.Decoder
if hasZSTDEcoding(resp.Header.Get("Content-Encoding")) {
decoder, err = zstd.NewReader(resp.Body)
if err != nil {
return cliproxyexecutor.Response{}, fmt.Errorf("failed to initialize zstd decoder: %w", err)
}
reader = decoder
defer decoder.Close()
}
data, err := io.ReadAll(reader)
if err != nil {
return cliproxyexecutor.Response{}, err
}
appendAPIResponseChunk(ctx, e.cfg, data)
count := gjson.GetBytes(data, "input_tokens").Int()
out := sdktranslator.TranslateTokenCount(ctx, to, from, count, data)
return cliproxyexecutor.Response{Payload: []byte(out)}, nil
}
func (e *ClaudeExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
log.Debugf("claude executor: refresh called")
if auth == nil {
return nil, fmt.Errorf("claude executor: auth is nil")
}
var refreshToken string
if auth.Metadata != nil {
if v, ok := auth.Metadata["refresh_token"].(string); ok && v != "" {
refreshToken = v
}
}
if refreshToken == "" {
return auth, nil
}
svc := claudeauth.NewClaudeAuth(e.cfg)
td, err := svc.RefreshTokens(ctx, refreshToken)
if err != nil {
return nil, err
}
if auth.Metadata == nil {
auth.Metadata = make(map[string]any)
}
auth.Metadata["access_token"] = td.AccessToken
if td.RefreshToken != "" {
auth.Metadata["refresh_token"] = td.RefreshToken
}
auth.Metadata["email"] = td.Email
auth.Metadata["expired"] = td.Expire
auth.Metadata["type"] = "claude"
now := time.Now().Format(time.RFC3339)
auth.Metadata["last_refresh"] = now
return auth, nil
}
func hasZSTDEcoding(contentEncoding string) bool {
if contentEncoding == "" {
return false
}
parts := strings.Split(contentEncoding, ",")
for i := range parts {
if strings.EqualFold(strings.TrimSpace(parts[i]), "zstd") {
return true
}
}
return false
}
func applyClaudeHeaders(r *http.Request, apiKey string, stream bool) {
r.Header.Set("Authorization", "Bearer "+apiKey)
r.Header.Set("Content-Type", "application/json")
var ginHeaders http.Header
if ginCtx, ok := r.Context().Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil {
ginHeaders = ginCtx.Request.Header
}
misc.EnsureHeader(r.Header, ginHeaders, "Anthropic-Version", "2023-06-01")
misc.EnsureHeader(r.Header, ginHeaders, "Anthropic-Dangerous-Direct-Browser-Access", "true")
misc.EnsureHeader(r.Header, ginHeaders, "Anthropic-Beta", "claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,fine-grained-tool-streaming-2025-05-14")
misc.EnsureHeader(r.Header, ginHeaders, "X-App", "cli")
misc.EnsureHeader(r.Header, ginHeaders, "X-Stainless-Helper-Method", "stream")
misc.EnsureHeader(r.Header, ginHeaders, "X-Stainless-Retry-Count", "0")
misc.EnsureHeader(r.Header, ginHeaders, "X-Stainless-Runtime-Version", "v24.3.0")
misc.EnsureHeader(r.Header, ginHeaders, "X-Stainless-Package-Version", "0.55.1")
misc.EnsureHeader(r.Header, ginHeaders, "X-Stainless-Runtime", "node")
misc.EnsureHeader(r.Header, ginHeaders, "X-Stainless-Lang", "js")
misc.EnsureHeader(r.Header, ginHeaders, "X-Stainless-Arch", "arm64")
misc.EnsureHeader(r.Header, ginHeaders, "X-Stainless-Os", "MacOS")
misc.EnsureHeader(r.Header, ginHeaders, "X-Stainless-Timeout", "60")
r.Header.Set("Connection", "keep-alive")
r.Header.Set("User-Agent", "claude-cli/1.0.83 (external, cli)")
r.Header.Set("Accept-Encoding", "gzip, deflate, br, zstd")
if stream {
r.Header.Set("Accept", "text/event-stream")
return
}
r.Header.Set("Accept", "application/json")
}
func claudeCreds(a *cliproxyauth.Auth) (apiKey, baseURL string) {
if a == nil {
return "", ""
}
if a.Attributes != nil {
apiKey = a.Attributes["api_key"]
baseURL = a.Attributes["base_url"]
}
if apiKey == "" && a.Metadata != nil {
if v, ok := a.Metadata["access_token"].(string); ok {
apiKey = v
}
}
return
}

320
internal/runtime/executor/codex_executor.go Normal file
View File

@@ -0,0 +1,320 @@
package executor
import (
"bufio"
"bytes"
"context"
"fmt"
"io"
"net/http"
"strings"
"time"
codexauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/codex"
"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/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"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
)
var dataTag = []byte("data:")
// CodexExecutor is a stateless executor for Codex (OpenAI Responses API entrypoint).
// If api_key is unavailable on auth, it falls back to legacy via ClientAdapter.
type CodexExecutor struct {
cfg *config.Config
}
func NewCodexExecutor(cfg *config.Config) *CodexExecutor { return &CodexExecutor{cfg: cfg} }
func (e *CodexExecutor) Identifier() string { return "codex" }
func (e *CodexExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.Auth) error { return nil }
func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
apiKey, baseURL := codexCreds(auth)
if baseURL == "" {
baseURL = "https://chatgpt.com/backend-api/codex"
}
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
from := opts.SourceFormat
to := sdktranslator.FromString("codex")
body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
if util.InArray([]string{"gpt-5", "gpt-5-minimal", "gpt-5-low", "gpt-5-medium", "gpt-5-high"}, req.Model) {
body, _ = sjson.SetBytes(body, "model", "gpt-5")
switch req.Model {
case "gpt-5":
body, _ = sjson.DeleteBytes(body, "reasoning.effort")
case "gpt-5-minimal":
body, _ = sjson.SetBytes(body, "reasoning.effort", "minimal")
case "gpt-5-low":
body, _ = sjson.SetBytes(body, "reasoning.effort", "low")
case "gpt-5-medium":
body, _ = sjson.SetBytes(body, "reasoning.effort", "medium")
case "gpt-5-high":
body, _ = sjson.SetBytes(body, "reasoning.effort", "high")
}
} else if util.InArray([]string{"gpt-5-codex", "gpt-5-codex-low", "gpt-5-codex-medium", "gpt-5-codex-high"}, req.Model) {
body, _ = sjson.SetBytes(body, "model", "gpt-5-codex")
switch req.Model {
case "gpt-5-codex":
body, _ = sjson.DeleteBytes(body, "reasoning.effort")
case "gpt-5-codex-low":
body, _ = sjson.SetBytes(body, "reasoning.effort", "low")
case "gpt-5-codex-medium":
body, _ = sjson.SetBytes(body, "reasoning.effort", "medium")
case "gpt-5-codex-high":
body, _ = sjson.SetBytes(body, "reasoning.effort", "high")
}
}
body, _ = sjson.SetBytes(body, "stream", true)
url := strings.TrimSuffix(baseURL, "/") + "/responses"
recordAPIRequest(ctx, e.cfg, body)
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
if err != nil {
return cliproxyexecutor.Response{}, err
}
applyCodexHeaders(httpReq, auth, apiKey)
httpClient := &http.Client{}
if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
httpClient.Transport = rt
}
resp, err := httpClient.Do(httpReq)
if err != nil {
return cliproxyexecutor.Response{}, err
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
b, _ := io.ReadAll(resp.Body)
appendAPIResponseChunk(ctx, e.cfg, b)
log.Debugf("request error, error status: %d, error body: %s", resp.StatusCode, string(b))
return cliproxyexecutor.Response{}, statusErr{code: resp.StatusCode, msg: string(b)}
}
data, err := io.ReadAll(resp.Body)
if err != nil {
return cliproxyexecutor.Response{}, err
}
appendAPIResponseChunk(ctx, e.cfg, data)
lines := bytes.Split(data, []byte("\n"))
for _, line := range lines {
if !bytes.HasPrefix(line, dataTag) {
continue
}
line = bytes.TrimSpace(line[5:])
if gjson.GetBytes(line, "type").String() != "response.completed" {
continue
}
if detail, ok := parseCodexUsage(line); ok {
reporter.publish(ctx, detail)
}
var param any
out := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, line, &param)
return cliproxyexecutor.Response{Payload: []byte(out)}, nil
}
return cliproxyexecutor.Response{}, statusErr{code: 408, msg: "stream error: stream disconnected before completion: stream closed before response.completed"}
}
func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (<-chan cliproxyexecutor.StreamChunk, error) {
apiKey, baseURL := codexCreds(auth)
if baseURL == "" {
baseURL = "https://chatgpt.com/backend-api/codex"
}
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
from := opts.SourceFormat
to := sdktranslator.FromString("codex")
body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)
if util.InArray([]string{"gpt-5", "gpt-5-minimal", "gpt-5-low", "gpt-5-medium", "gpt-5-high"}, req.Model) {
body, _ = sjson.SetBytes(body, "model", "gpt-5")
switch req.Model {
case "gpt-5":
body, _ = sjson.DeleteBytes(body, "reasoning.effort")
case "gpt-5-minimal":
body, _ = sjson.SetBytes(body, "reasoning.effort", "minimal")
case "gpt-5-low":
body, _ = sjson.SetBytes(body, "reasoning.effort", "low")
case "gpt-5-medium":
body, _ = sjson.SetBytes(body, "reasoning.effort", "medium")
case "gpt-5-high":
body, _ = sjson.SetBytes(body, "reasoning.effort", "high")
}
} else if util.InArray([]string{"gpt-5-codex", "gpt-5-codex-low", "gpt-5-codex-medium", "gpt-5-codex-high"}, req.Model) {
body, _ = sjson.SetBytes(body, "model", "gpt-5-codex")
switch req.Model {
case "gpt-5-codex":
body, _ = sjson.DeleteBytes(body, "reasoning.effort")
case "gpt-5-codex-low":
body, _ = sjson.SetBytes(body, "reasoning.effort", "low")
case "gpt-5-codex-medium":
body, _ = sjson.SetBytes(body, "reasoning.effort", "medium")
case "gpt-5-codex-high":
body, _ = sjson.SetBytes(body, "reasoning.effort", "high")
}
}
url := strings.TrimSuffix(baseURL, "/") + "/responses"
recordAPIRequest(ctx, e.cfg, body)
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
if err != nil {
return nil, err
}
applyCodexHeaders(httpReq, auth, apiKey)
httpClient := &http.Client{Timeout: 0}
if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
httpClient.Transport = rt
}
resp, err := httpClient.Do(httpReq)
if err != nil {
return nil, err
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
defer func() { _ = resp.Body.Close() }()
b, _ := io.ReadAll(resp.Body)
appendAPIResponseChunk(ctx, e.cfg, b)
log.Debugf("request error, error status: %d, error body: %s", resp.StatusCode, string(b))
return nil, statusErr{code: resp.StatusCode, msg: string(b)}
}
out := make(chan cliproxyexecutor.StreamChunk)
go func() {
defer close(out)
defer func() { _ = resp.Body.Close() }()
scanner := bufio.NewScanner(resp.Body)
buf := make([]byte, 1024*1024)
scanner.Buffer(buf, 1024*1024)
var param any
for scanner.Scan() {
line := scanner.Bytes()
appendAPIResponseChunk(ctx, e.cfg, line)
if bytes.HasPrefix(line, dataTag) {
data := bytes.TrimSpace(line[5:])
if gjson.GetBytes(data, "type").String() == "response.completed" {
if detail, ok := parseCodexUsage(data); ok {
reporter.publish(ctx, detail)
}
}
}
chunks := sdktranslator.TranslateStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, bytes.Clone(line), &param)
for i := range chunks {
out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunks[i])}
}
}
if err = scanner.Err(); err != nil {
out <- cliproxyexecutor.StreamChunk{Err: err}
}
}()
return out, nil
}
func (e *CodexExecutor) 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 *CodexExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
log.Debugf("codex executor: refresh called")
if auth == nil {
return nil, statusErr{code: 500, msg: "codex executor: auth is nil"}
}
var refreshToken string
if auth.Metadata != nil {
if v, ok := auth.Metadata["refresh_token"].(string); ok && v != "" {
refreshToken = v
}
}
if refreshToken == "" {
return auth, nil
}
svc := codexauth.NewCodexAuth(e.cfg)
td, err := svc.RefreshTokensWithRetry(ctx, refreshToken, 3)
if err != nil {
return nil, err
}
if auth.Metadata == nil {
auth.Metadata = make(map[string]any)
}
auth.Metadata["id_token"] = td.IDToken
auth.Metadata["access_token"] = td.AccessToken
if td.RefreshToken != "" {
auth.Metadata["refresh_token"] = td.RefreshToken
}
if td.AccountID != "" {
auth.Metadata["account_id"] = td.AccountID
}
auth.Metadata["email"] = td.Email
// Use unified key in files
auth.Metadata["expired"] = td.Expire
auth.Metadata["type"] = "codex"
now := time.Now().Format(time.RFC3339)
auth.Metadata["last_refresh"] = now
return auth, nil
}
func applyCodexHeaders(r *http.Request, auth *cliproxyauth.Auth, token string) {
r.Header.Set("Content-Type", "application/json")
r.Header.Set("Authorization", "Bearer "+token)
var ginHeaders http.Header
if ginCtx, ok := r.Context().Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil {
ginHeaders = ginCtx.Request.Header
}
misc.EnsureHeader(r.Header, ginHeaders, "Version", "0.21.0")
misc.EnsureHeader(r.Header, ginHeaders, "Openai-Beta", "responses=experimental")
misc.EnsureHeader(r.Header, ginHeaders, "Session_id", uuid.NewString())
r.Header.Set("Accept", "text/event-stream")
r.Header.Set("Connection", "Keep-Alive")
isAPIKey := false
if auth != nil && auth.Attributes != nil {
if v := strings.TrimSpace(auth.Attributes["api_key"]); v != "" {
isAPIKey = true
}
}
if !isAPIKey {
r.Header.Set("Originator", "codex_cli_rs")
if auth != nil && auth.Metadata != nil {
if accountID, ok := auth.Metadata["account_id"].(string); ok {
r.Header.Set("Chatgpt-Account-Id", accountID)
}
}
}
}
func codexCreds(a *cliproxyauth.Auth) (apiKey, baseURL string) {
if a == nil {
return "", ""
}
if a.Attributes != nil {
apiKey = a.Attributes["api_key"]
baseURL = a.Attributes["base_url"]
}
if apiKey == "" && a.Metadata != nil {
if v, ok := a.Metadata["access_token"].(string); ok {
apiKey = v
}
}
return
}

532
internal/runtime/executor/gemini_cli_executor.go Normal file
View File

@@ -0,0 +1,532 @@
package executor
import (
"bufio"
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/gin-gonic/gin"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
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
}
func NewGeminiCLIExecutor(cfg *config.Config) *GeminiCLIExecutor {
return &GeminiCLIExecutor{cfg: cfg}
}
func (e *GeminiCLIExecutor) Identifier() string { return "gemini-cli" }
func (e *GeminiCLIExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.Auth) error { return nil }
func (e *GeminiCLIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
tokenSource, baseTokenData, err := prepareGeminiCLITokenSource(ctx, auth)
if err != nil {
return cliproxyexecutor.Response{}, err
}
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
from := opts.SourceFormat
to := sdktranslator.FromString("gemini-cli")
basePayload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
action := "generateContent"
if req.Metadata != nil {
if a, _ := req.Metadata["action"].(string); a == "countTokens" {
action = "countTokens"
}
}
projectID := strings.TrimSpace(stringValue(auth.Metadata, "project_id"))
models := cliPreviewFallbackOrder(req.Model)
if len(models) == 0 || models[0] != req.Model {
models = append([]string{req.Model}, models...)
}
httpClient := newHTTPClient(ctx, 0)
respCtx := context.WithValue(ctx, "alt", opts.Alt)
var lastStatus int
var lastBody []byte
for _, 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 {
return cliproxyexecutor.Response{}, errTok
}
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)
}
recordAPIRequest(ctx, e.cfg, payload)
reqHTTP, errReq := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload))
if errReq != nil {
return cliproxyexecutor.Response{}, errReq
}
reqHTTP.Header.Set("Content-Type", "application/json")
reqHTTP.Header.Set("Authorization", "Bearer "+tok.AccessToken)
applyGeminiCLIHeaders(reqHTTP)
reqHTTP.Header.Set("Accept", "application/json")
resp, errDo := httpClient.Do(reqHTTP)
if errDo != nil {
return cliproxyexecutor.Response{}, errDo
}
data, _ := io.ReadAll(resp.Body)
_ = resp.Body.Close()
appendAPIResponseChunk(ctx, e.cfg, data)
if resp.StatusCode >= 200 && resp.StatusCode < 300 {
reporter.publish(ctx, parseGeminiCLIUsage(data))
var param any
out := sdktranslator.TranslateNonStream(respCtx, to, from, attemptModel, bytes.Clone(opts.OriginalRequest), payload, data, &param)
return cliproxyexecutor.Response{Payload: []byte(out)}, nil
}
lastStatus = resp.StatusCode
lastBody = data
if resp.StatusCode != 429 {
break
}
}
if len(lastBody) > 0 {
appendAPIResponseChunk(ctx, e.cfg, lastBody)
}
return cliproxyexecutor.Response{}, statusErr{code: lastStatus, msg: string(lastBody)}
}
func (e *GeminiCLIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (<-chan cliproxyexecutor.StreamChunk, error) {
tokenSource, baseTokenData, err := prepareGeminiCLITokenSource(ctx, auth)
if err != nil {
return nil, err
}
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
from := opts.SourceFormat
to := sdktranslator.FromString("gemini-cli")
basePayload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)
projectID := strings.TrimSpace(stringValue(auth.Metadata, "project_id"))
models := cliPreviewFallbackOrder(req.Model)
if len(models) == 0 || models[0] != req.Model {
models = append([]string{req.Model}, models...)
}
httpClient := newHTTPClient(ctx, 0)
respCtx := context.WithValue(ctx, "alt", opts.Alt)
var lastStatus int
var lastBody []byte
for _, 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 {
return nil, errTok
}
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)
}
recordAPIRequest(ctx, e.cfg, payload)
reqHTTP, errReq := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload))
if errReq != nil {
return nil, errReq
}
reqHTTP.Header.Set("Content-Type", "application/json")
reqHTTP.Header.Set("Authorization", "Bearer "+tok.AccessToken)
applyGeminiCLIHeaders(reqHTTP)
reqHTTP.Header.Set("Accept", "text/event-stream")
resp, errDo := httpClient.Do(reqHTTP)
if errDo != nil {
return nil, errDo
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
data, _ := io.ReadAll(resp.Body)
_ = resp.Body.Close()
appendAPIResponseChunk(ctx, e.cfg, data)
lastStatus = resp.StatusCode
lastBody = data
log.Debugf("request error, error status: %d, error body: %s", resp.StatusCode, string(data))
if resp.StatusCode == 429 {
continue
}
return nil, statusErr{code: resp.StatusCode, msg: string(data)}
}
out := make(chan cliproxyexecutor.StreamChunk)
go func(resp *http.Response, reqBody []byte, attempt string) {
defer close(out)
defer func() { _ = resp.Body.Close() }()
if opts.Alt == "" {
scanner := bufio.NewScanner(resp.Body)
buf := make([]byte, 1024*1024)
scanner.Buffer(buf, 1024*1024)
var param any
for scanner.Scan() {
line := scanner.Bytes()
appendAPIResponseChunk(ctx, e.cfg, line)
if detail, ok := parseGeminiCLIStreamUsage(line); ok {
reporter.publish(ctx, detail)
}
if bytes.HasPrefix(line, dataTag) {
segments := sdktranslator.TranslateStream(respCtx, to, from, attempt, bytes.Clone(opts.OriginalRequest), reqBody, bytes.Clone(line), &param)
for i := range segments {
out <- cliproxyexecutor.StreamChunk{Payload: []byte(segments[i])}
}
}
}
segments := sdktranslator.TranslateStream(respCtx, to, from, attempt, bytes.Clone(opts.OriginalRequest), reqBody, bytes.Clone([]byte("[DONE]")), &param)
for i := range segments {
out <- cliproxyexecutor.StreamChunk{Payload: []byte(segments[i])}
}
if errScan := scanner.Err(); errScan != nil {
out <- cliproxyexecutor.StreamChunk{Err: errScan}
}
return
}
data, errRead := io.ReadAll(resp.Body)
if errRead != nil {
out <- cliproxyexecutor.StreamChunk{Err: errRead}
return
}
appendAPIResponseChunk(ctx, e.cfg, data)
reporter.publish(ctx, parseGeminiCLIUsage(data))
var param any
segments := sdktranslator.TranslateStream(respCtx, to, from, attempt, bytes.Clone(opts.OriginalRequest), reqBody, data, &param)
for i := range segments {
out <- cliproxyexecutor.StreamChunk{Payload: []byte(segments[i])}
}
segments = sdktranslator.TranslateStream(respCtx, to, from, attempt, bytes.Clone(opts.OriginalRequest), reqBody, bytes.Clone([]byte("[DONE]")), &param)
for i := range segments {
out <- cliproxyexecutor.StreamChunk{Payload: []byte(segments[i])}
}
}(resp, append([]byte(nil), payload...), attemptModel)
return out, nil
}
if lastStatus == 0 {
lastStatus = 429
}
return nil, statusErr{code: lastStatus, msg: string(lastBody)}
}
func (e *GeminiCLIExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
tokenSource, baseTokenData, err := prepareGeminiCLITokenSource(ctx, auth)
if err != nil {
return cliproxyexecutor.Response{}, err
}
from := opts.SourceFormat
to := sdktranslator.FromString("gemini-cli")
models := cliPreviewFallbackOrder(req.Model)
if len(models) == 0 || models[0] != req.Model {
models = append([]string{req.Model}, models...)
}
httpClient := newHTTPClient(ctx, 0)
respCtx := context.WithValue(ctx, "alt", opts.Alt)
var lastStatus int
var lastBody []byte
for _, attemptModel := range models {
payload := sdktranslator.TranslateRequest(from, to, attemptModel, bytes.Clone(req.Payload), false)
payload = deleteJSONField(payload, "project")
payload = deleteJSONField(payload, "model")
tok, errTok := tokenSource.Token()
if errTok != nil {
return cliproxyexecutor.Response{}, errTok
}
updateGeminiCLITokenMetadata(auth, baseTokenData, tok)
url := fmt.Sprintf("%s/%s:%s", codeAssistEndpoint, codeAssistVersion, "countTokens")
if opts.Alt != "" {
url = url + fmt.Sprintf("?$alt=%s", opts.Alt)
}
recordAPIRequest(ctx, e.cfg, payload)
reqHTTP, errReq := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload))
if errReq != nil {
return cliproxyexecutor.Response{}, errReq
}
reqHTTP.Header.Set("Content-Type", "application/json")
reqHTTP.Header.Set("Authorization", "Bearer "+tok.AccessToken)
applyGeminiCLIHeaders(reqHTTP)
reqHTTP.Header.Set("Accept", "application/json")
resp, errDo := httpClient.Do(reqHTTP)
if errDo != nil {
return cliproxyexecutor.Response{}, errDo
}
data, _ := io.ReadAll(resp.Body)
_ = resp.Body.Close()
appendAPIResponseChunk(ctx, e.cfg, data)
if resp.StatusCode >= 200 && resp.StatusCode < 300 {
count := gjson.GetBytes(data, "totalTokens").Int()
translated := sdktranslator.TranslateTokenCount(respCtx, to, from, count, data)
return cliproxyexecutor.Response{Payload: []byte(translated)}, nil
}
lastStatus = resp.StatusCode
lastBody = data
if resp.StatusCode == 429 {
continue
}
break
}
if len(lastBody) > 0 {
appendAPIResponseChunk(ctx, e.cfg, lastBody)
}
if lastStatus == 0 {
lastStatus = 429
}
return cliproxyexecutor.Response{}, statusErr{code: lastStatus, msg: string(lastBody)}
}
func (e *GeminiCLIExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
log.Debugf("gemini cli executor: refresh called")
_ = ctx
return auth, nil
}
func prepareGeminiCLITokenSource(ctx context.Context, auth *cliproxyauth.Auth) (oauth2.TokenSource, map[string]any, error) {
if auth == nil || auth.Metadata == nil {
return nil, nil, fmt.Errorf("gemini-cli auth metadata missing")
}
var base map[string]any
if tokenRaw, ok := auth.Metadata["token"].(map[string]any); ok && tokenRaw != nil {
base = cloneMap(tokenRaw)
} else {
base = make(map[string]any)
}
var token oauth2.Token
if len(base) > 0 {
if raw, err := json.Marshal(base); err == nil {
_ = json.Unmarshal(raw, &token)
}
}
if token.AccessToken == "" {
token.AccessToken = stringValue(auth.Metadata, "access_token")
}
if token.RefreshToken == "" {
token.RefreshToken = stringValue(auth.Metadata, "refresh_token")
}
if token.TokenType == "" {
token.TokenType = stringValue(auth.Metadata, "token_type")
}
if token.Expiry.IsZero() {
if expiry := stringValue(auth.Metadata, "expiry"); expiry != "" {
if ts, err := time.Parse(time.RFC3339, expiry); err == nil {
token.Expiry = ts
}
}
}
conf := &oauth2.Config{
ClientID: geminiOauthClientID,
ClientSecret: geminiOauthClientSecret,
Scopes: geminiOauthScopes,
Endpoint: google.Endpoint,
}
ctxToken := ctx
if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
ctxToken = context.WithValue(ctxToken, oauth2.HTTPClient, &http.Client{Transport: rt})
}
src := conf.TokenSource(ctxToken, &token)
currentToken, err := src.Token()
if err != nil {
return nil, nil, err
}
updateGeminiCLITokenMetadata(auth, base, currentToken)
return oauth2.ReuseTokenSource(currentToken, src), base, nil
}
func updateGeminiCLITokenMetadata(auth *cliproxyauth.Auth, base map[string]any, tok *oauth2.Token) {
if auth == nil || auth.Metadata == nil || tok == nil {
return
}
if tok.AccessToken != "" {
auth.Metadata["access_token"] = tok.AccessToken
}
if tok.TokenType != "" {
auth.Metadata["token_type"] = tok.TokenType
}
if tok.RefreshToken != "" {
auth.Metadata["refresh_token"] = tok.RefreshToken
}
if !tok.Expiry.IsZero() {
auth.Metadata["expiry"] = tok.Expiry.Format(time.RFC3339)
}
merged := cloneMap(base)
if merged == nil {
merged = make(map[string]any)
}
if raw, err := json.Marshal(tok); err == nil {
var tokenMap map[string]any
if err = json.Unmarshal(raw, &tokenMap); err == nil {
for k, v := range tokenMap {
merged[k] = v
}
}
}
auth.Metadata["token"] = merged
}
func newHTTPClient(ctx context.Context, timeout time.Duration) *http.Client {
client := &http.Client{}
if timeout > 0 {
client.Timeout = timeout
}
if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
client.Transport = rt
}
return client
}
func cloneMap(in map[string]any) map[string]any {
if in == nil {
return nil
}
out := make(map[string]any, len(in))
for k, v := range in {
out[k] = v
}
return out
}
func stringValue(m map[string]any, key string) string {
if m == nil {
return ""
}
if v, ok := m[key]; ok {
switch typed := v.(type) {
case string:
return typed
case fmt.Stringer:
return typed.String()
}
}
return ""
}
// applyGeminiCLIHeaders sets required headers for the Gemini CLI upstream.
func applyGeminiCLIHeaders(r *http.Request) {
var ginHeaders http.Header
if ginCtx, ok := r.Context().Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil {
ginHeaders = ginCtx.Request.Header
}
misc.EnsureHeader(r.Header, ginHeaders, "User-Agent", "google-api-nodejs-client/9.15.1")
misc.EnsureHeader(r.Header, ginHeaders, "X-Goog-Api-Client", "gl-node/22.17.0")
misc.EnsureHeader(r.Header, ginHeaders, "Client-Metadata", geminiCLIClientMetadata())
}
// geminiCLIClientMetadata returns a compact metadata string required by upstream.
func geminiCLIClientMetadata() string {
// Keep parity with CLI client defaults
return "ideType=IDE_UNSPECIFIED,platform=PLATFORM_UNSPECIFIED,pluginType=GEMINI"
}
// cliPreviewFallbackOrder returns preview model candidates for a base model.
func cliPreviewFallbackOrder(model string) []string {
switch model {
case "gemini-2.5-pro":
return []string{"gemini-2.5-pro-preview-05-06", "gemini-2.5-pro-preview-06-05"}
case "gemini-2.5-flash":
return []string{"gemini-2.5-flash-preview-04-17", "gemini-2.5-flash-preview-05-20"}
case "gemini-2.5-flash-lite":
return []string{"gemini-2.5-flash-lite-preview-06-17"}
default:
return nil
}
}
// setJSONField sets a top-level JSON field on a byte slice payload via sjson.
func setJSONField(body []byte, key, value string) []byte {
if key == "" {
return body
}
updated, err := sjson.SetBytes(body, key, value)
if err != nil {
return body
}
return updated
}
// deleteJSONField removes a top-level key if present (best-effort) via sjson.
func deleteJSONField(body []byte, key string) []byte {
if key == "" || len(body) == 0 {
return body
}
updated, err := sjson.DeleteBytes(body, key)
if err != nil {
return body
}
return updated
}

382
internal/runtime/executor/gemini_executor.go Normal file
View File

@@ -0,0 +1,382 @@
// Package executor provides runtime execution capabilities for various AI service providers.
// It includes stateless executors that handle API requests, streaming responses,
// token counting, and authentication refresh for different AI service providers.
package executor
import (
"bufio"
"bytes"
"context"
"fmt"
"io"
"net/http"
"time"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"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 (
// glEndpoint is the base URL for the Google Generative Language API.
glEndpoint = "https://generativelanguage.googleapis.com"
// glAPIVersion is the API version used for Gemini requests.
glAPIVersion = "v1beta"
)
// GeminiExecutor is a stateless executor for the official Gemini API using API keys.
// It handles both API key and OAuth bearer token authentication, supporting both
// regular and streaming requests to the Google Generative Language API.
type GeminiExecutor struct {
// cfg holds the application configuration.
cfg *config.Config
}
// NewGeminiExecutor creates a new Gemini executor instance.
//
// Parameters:
// - cfg: The application configuration
//
// Returns:
// - *GeminiExecutor: A new Gemini executor instance
func NewGeminiExecutor(cfg *config.Config) *GeminiExecutor { return &GeminiExecutor{cfg: cfg} }
// Identifier returns the executor identifier for Gemini.
func (e *GeminiExecutor) Identifier() string { return "gemini" }
// PrepareRequest prepares the HTTP request for execution (no-op for Gemini).
func (e *GeminiExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.Auth) error { return nil }
// Execute performs a non-streaming request to the Gemini API.
// It translates the request to Gemini format, sends it to the API, and translates
// the response back to the requested format.
//
// Parameters:
// - ctx: The context for the request
// - auth: The authentication information
// - req: The request to execute
// - opts: Additional execution options
//
// Returns:
// - cliproxyexecutor.Response: The response from the API
// - error: An error if the request fails
func (e *GeminiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
apiKey, bearer := geminiCreds(auth)
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
// Official Gemini API via API key or OAuth bearer
from := opts.SourceFormat
to := sdktranslator.FromString("gemini")
body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
action := "generateContent"
if req.Metadata != nil {
if a, _ := req.Metadata["action"].(string); a == "countTokens" {
action = "countTokens"
}
}
url := fmt.Sprintf("%s/%s/models/%s:%s", glEndpoint, glAPIVersion, req.Model, action)
if opts.Alt != "" && action != "countTokens" {
url = url + fmt.Sprintf("?$alt=%s", opts.Alt)
}
body, _ = sjson.DeleteBytes(body, "session_id")
recordAPIRequest(ctx, e.cfg, body)
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
if err != nil {
return cliproxyexecutor.Response{}, err
}
httpReq.Header.Set("Content-Type", "application/json")
if apiKey != "" {
httpReq.Header.Set("x-goog-api-key", apiKey)
} else if bearer != "" {
httpReq.Header.Set("Authorization", "Bearer "+bearer)
}
httpClient := &http.Client{}
if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
httpClient.Transport = rt
}
resp, err := httpClient.Do(httpReq)
if err != nil {
return cliproxyexecutor.Response{}, err
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
b, _ := io.ReadAll(resp.Body)
appendAPIResponseChunk(ctx, e.cfg, b)
log.Debugf("request error, error status: %d, error body: %s", resp.StatusCode, string(b))
return cliproxyexecutor.Response{}, statusErr{code: resp.StatusCode, msg: string(b)}
}
data, err := io.ReadAll(resp.Body)
if err != nil {
return cliproxyexecutor.Response{}, err
}
appendAPIResponseChunk(ctx, e.cfg, data)
reporter.publish(ctx, parseGeminiUsage(data))
var param any
out := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, data, &param)
return cliproxyexecutor.Response{Payload: []byte(out)}, nil
}
func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (<-chan cliproxyexecutor.StreamChunk, error) {
apiKey, bearer := geminiCreds(auth)
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
from := opts.SourceFormat
to := sdktranslator.FromString("gemini")
body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)
url := fmt.Sprintf("%s/%s/models/%s:%s", glEndpoint, glAPIVersion, req.Model, "streamGenerateContent")
if opts.Alt == "" {
url = url + "?alt=sse"
} else {
url = url + fmt.Sprintf("?$alt=%s", opts.Alt)
}
body, _ = sjson.DeleteBytes(body, "session_id")
recordAPIRequest(ctx, e.cfg, body)
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
if err != nil {
return nil, err
}
httpReq.Header.Set("Content-Type", "application/json")
if apiKey != "" {
httpReq.Header.Set("x-goog-api-key", apiKey)
} else {
httpReq.Header.Set("Authorization", "Bearer "+bearer)
}
httpClient := &http.Client{Timeout: 0}
if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
httpClient.Transport = rt
}
resp, err := httpClient.Do(httpReq)
if err != nil {
return nil, err
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
defer func() { _ = resp.Body.Close() }()
b, _ := io.ReadAll(resp.Body)
appendAPIResponseChunk(ctx, e.cfg, b)
log.Debugf("request error, error status: %d, error body: %s", resp.StatusCode, string(b))
return nil, statusErr{code: resp.StatusCode, msg: string(b)}
}
out := make(chan cliproxyexecutor.StreamChunk)
go func() {
defer close(out)
defer func() { _ = resp.Body.Close() }()
scanner := bufio.NewScanner(resp.Body)
buf := make([]byte, 1024*1024)
scanner.Buffer(buf, 1024*1024)
var param any
for scanner.Scan() {
line := scanner.Bytes()
appendAPIResponseChunk(ctx, e.cfg, line)
if detail, ok := parseGeminiStreamUsage(line); ok {
reporter.publish(ctx, detail)
}
lines := sdktranslator.TranslateStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, bytes.Clone(line), &param)
for i := range lines {
out <- cliproxyexecutor.StreamChunk{Payload: []byte(lines[i])}
}
}
lines := sdktranslator.TranslateStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, bytes.Clone([]byte("[DONE]")), &param)
for i := range lines {
out <- cliproxyexecutor.StreamChunk{Payload: []byte(lines[i])}
}
if err = scanner.Err(); err != nil {
out <- cliproxyexecutor.StreamChunk{Err: err}
}
}()
return out, nil
}
func (e *GeminiExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
apiKey, bearer := geminiCreds(auth)
from := opts.SourceFormat
to := sdktranslator.FromString("gemini")
translatedReq := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
respCtx := context.WithValue(ctx, "alt", opts.Alt)
translatedReq, _ = sjson.DeleteBytes(translatedReq, "tools")
translatedReq, _ = sjson.DeleteBytes(translatedReq, "generationConfig")
url := fmt.Sprintf("%s/%s/models/%s:%s", glEndpoint, glAPIVersion, req.Model, "countTokens")
recordAPIRequest(ctx, e.cfg, translatedReq)
requestBody := bytes.NewReader(translatedReq)
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, requestBody)
if err != nil {
return cliproxyexecutor.Response{}, err
}
httpReq.Header.Set("Content-Type", "application/json")
if apiKey != "" {
httpReq.Header.Set("x-goog-api-key", apiKey)
} else {
httpReq.Header.Set("Authorization", "Bearer "+bearer)
}
httpClient := &http.Client{}
if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
httpClient.Transport = rt
}
resp, err := httpClient.Do(httpReq)
if err != nil {
return cliproxyexecutor.Response{}, err
}
defer func() { _ = resp.Body.Close() }()
data, err := io.ReadAll(resp.Body)
if err != nil {
return cliproxyexecutor.Response{}, err
}
appendAPIResponseChunk(ctx, e.cfg, data)
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
log.Debugf("request error, error status: %d, error body: %s", resp.StatusCode, string(data))
return cliproxyexecutor.Response{}, statusErr{code: resp.StatusCode, msg: string(data)}
}
count := gjson.GetBytes(data, "totalTokens").Int()
translated := sdktranslator.TranslateTokenCount(respCtx, to, from, count, data)
return cliproxyexecutor.Response{Payload: []byte(translated)}, nil
}
func (e *GeminiExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
log.Debugf("gemini executor: refresh called")
// OAuth bearer token refresh for official Gemini API.
if auth == nil {
return nil, fmt.Errorf("gemini executor: auth is nil")
}
if auth.Metadata == nil {
return auth, nil
}
// Token data is typically nested under "token" map in Gemini files.
tokenMap, _ := auth.Metadata["token"].(map[string]any)
var refreshToken, accessToken, clientID, clientSecret, tokenURI, expiryStr string
if tokenMap != nil {
if v, ok := tokenMap["refresh_token"].(string); ok {
refreshToken = v
}
if v, ok := tokenMap["access_token"].(string); ok {
accessToken = v
}
if v, ok := tokenMap["client_id"].(string); ok {
clientID = v
}
if v, ok := tokenMap["client_secret"].(string); ok {
clientSecret = v
}
if v, ok := tokenMap["token_uri"].(string); ok {
tokenURI = v
}
if v, ok := tokenMap["expiry"].(string); ok {
expiryStr = v
}
} else {
// Fallback to top-level keys if present
if v, ok := auth.Metadata["refresh_token"].(string); ok {
refreshToken = v
}
if v, ok := auth.Metadata["access_token"].(string); ok {
accessToken = v
}
if v, ok := auth.Metadata["client_id"].(string); ok {
clientID = v
}
if v, ok := auth.Metadata["client_secret"].(string); ok {
clientSecret = v
}
if v, ok := auth.Metadata["token_uri"].(string); ok {
tokenURI = v
}
if v, ok := auth.Metadata["expiry"].(string); ok {
expiryStr = v
}
}
if refreshToken == "" {
// Nothing to do for API key or cookie based entries
return auth, nil
}
// Prepare oauth2 config; default to Google endpoints
endpoint := google.Endpoint
if tokenURI != "" {
endpoint.TokenURL = tokenURI
}
conf := &oauth2.Config{ClientID: clientID, ClientSecret: clientSecret, Endpoint: endpoint}
// Ensure proxy-aware HTTP client for token refresh
httpClient := util.SetProxy(e.cfg, &http.Client{})
ctx = context.WithValue(ctx, oauth2.HTTPClient, httpClient)
// Build base token
tok := &oauth2.Token{AccessToken: accessToken, RefreshToken: refreshToken}
if t, err := time.Parse(time.RFC3339, expiryStr); err == nil {
tok.Expiry = t
}
newTok, err := conf.TokenSource(ctx, tok).Token()
if err != nil {
return nil, err
}
// Persist back to metadata; prefer nested token map if present
if tokenMap == nil {
tokenMap = make(map[string]any)
}
tokenMap["access_token"] = newTok.AccessToken
tokenMap["refresh_token"] = newTok.RefreshToken
tokenMap["expiry"] = newTok.Expiry.Format(time.RFC3339)
if clientID != "" {
tokenMap["client_id"] = clientID
}
if clientSecret != "" {
tokenMap["client_secret"] = clientSecret
}
if tokenURI != "" {
tokenMap["token_uri"] = tokenURI
}
auth.Metadata["token"] = tokenMap
// Also mirror top-level access_token for compatibility if previously present
if _, ok := auth.Metadata["access_token"]; ok {
auth.Metadata["access_token"] = newTok.AccessToken
}
return auth, nil
}
func geminiCreds(a *cliproxyauth.Auth) (apiKey, bearer string) {
if a == nil {
return "", ""
}
if a.Attributes != nil {
if v := a.Attributes["api_key"]; v != "" {
apiKey = v
}
}
if a.Metadata != nil {
// GeminiTokenStorage.Token is a map that may contain access_token
if v, ok := a.Metadata["access_token"].(string); ok && v != "" {
bearer = v
}
if token, ok := a.Metadata["token"].(map[string]any); ok && token != nil {
if v, ok2 := token["access_token"].(string); ok2 && v != "" {
bearer = v
}
}
}
return
}

237
internal/runtime/executor/gemini_web_executor.go Normal file
View File

@@ -0,0 +1,237 @@
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), &param)
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)), &param)
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)), &param)
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 = &copyCfg
}
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
}

41
internal/runtime/executor/logging_helpers.go Normal file
View File

@@ -0,0 +1,41 @@
package executor
import (
"bytes"
"context"
"github.com/gin-gonic/gin"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
)
// recordAPIRequest stores the upstream request payload in Gin context for request logging.
func recordAPIRequest(ctx context.Context, cfg *config.Config, payload []byte) {
if cfg == nil || !cfg.RequestLog || len(payload) == 0 {
return
}
if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil {
ginCtx.Set("API_REQUEST", bytes.Clone(payload))
}
}
// appendAPIResponseChunk appends an upstream response chunk to Gin context for request logging.
func appendAPIResponseChunk(ctx context.Context, cfg *config.Config, chunk []byte) {
if cfg == nil || !cfg.RequestLog {
return
}
data := bytes.TrimSpace(bytes.Clone(chunk))
if len(data) == 0 {
return
}
if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil {
if existing, exists := ginCtx.Get("API_RESPONSE"); exists {
if prev, okBytes := existing.([]byte); okBytes {
prev = append(prev, data...)
prev = append(prev, []byte("\n\n")...)
ginCtx.Set("API_RESPONSE", prev)
return
}
}
ginCtx.Set("API_RESPONSE", data)
}
}

258
internal/runtime/executor/openai_compat_executor.go Normal file
View File

@@ -0,0 +1,258 @@
package executor
import (
"bufio"
"bytes"
"context"
"fmt"
"io"
"net/http"
"strings"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
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/sjson"
)
// OpenAICompatExecutor implements a stateless executor for OpenAI-compatible providers.
// It performs request/response translation and executes against the provider base URL
// using per-auth credentials (API key) and per-auth HTTP transport (proxy) from context.
type OpenAICompatExecutor struct {
provider string
cfg *config.Config
}
// NewOpenAICompatExecutor creates an executor bound to a provider key (e.g., "openrouter").
func NewOpenAICompatExecutor(provider string, cfg *config.Config) *OpenAICompatExecutor {
return &OpenAICompatExecutor{provider: provider, cfg: cfg}
}
// Identifier implements cliproxyauth.ProviderExecutor.
func (e *OpenAICompatExecutor) Identifier() string { return e.provider }
// PrepareRequest is a no-op for now (credentials are added via headers at execution time).
func (e *OpenAICompatExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.Auth) error {
return nil
}
func (e *OpenAICompatExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
baseURL, apiKey := e.resolveCredentials(auth)
if baseURL == "" || apiKey == "" {
return cliproxyexecutor.Response{}, statusErr{code: http.StatusUnauthorized, msg: "missing provider baseURL or apiKey"}
}
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
// Translate inbound request to OpenAI format
from := opts.SourceFormat
to := sdktranslator.FromString("openai")
translated := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), opts.Stream)
if modelOverride := e.resolveUpstreamModel(req.Model, auth); modelOverride != "" {
translated = e.overrideModel(translated, modelOverride)
}
url := strings.TrimSuffix(baseURL, "/") + "/chat/completions"
recordAPIRequest(ctx, e.cfg, translated)
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(translated))
if err != nil {
return cliproxyexecutor.Response{}, err
}
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("Authorization", "Bearer "+apiKey)
httpReq.Header.Set("User-Agent", "cli-proxy-openai-compat")
httpClient := &http.Client{}
if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
httpClient.Transport = rt
}
resp, err := httpClient.Do(httpReq)
if err != nil {
return cliproxyexecutor.Response{}, err
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
b, _ := io.ReadAll(resp.Body)
appendAPIResponseChunk(ctx, e.cfg, b)
log.Debugf("request error, error status: %d, error body: %s", resp.StatusCode, string(b))
return cliproxyexecutor.Response{}, statusErr{code: resp.StatusCode, msg: string(b)}
}
body, err := io.ReadAll(resp.Body)
if err != nil {
return cliproxyexecutor.Response{}, err
}
appendAPIResponseChunk(ctx, e.cfg, body)
reporter.publish(ctx, parseOpenAIUsage(body))
// Translate response back to source format when needed
var param any
out := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), translated, body, &param)
return cliproxyexecutor.Response{Payload: []byte(out)}, nil
}
func (e *OpenAICompatExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (<-chan cliproxyexecutor.StreamChunk, error) {
baseURL, apiKey := e.resolveCredentials(auth)
if baseURL == "" || apiKey == "" {
return nil, statusErr{code: http.StatusUnauthorized, msg: "missing provider baseURL or apiKey"}
}
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
from := opts.SourceFormat
to := sdktranslator.FromString("openai")
translated := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)
if modelOverride := e.resolveUpstreamModel(req.Model, auth); modelOverride != "" {
translated = e.overrideModel(translated, modelOverride)
}
url := strings.TrimSuffix(baseURL, "/") + "/chat/completions"
recordAPIRequest(ctx, e.cfg, translated)
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(translated))
if err != nil {
return nil, err
}
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("Authorization", "Bearer "+apiKey)
httpReq.Header.Set("User-Agent", "cli-proxy-openai-compat")
httpReq.Header.Set("Accept", "text/event-stream")
httpReq.Header.Set("Cache-Control", "no-cache")
httpClient := &http.Client{Timeout: 0}
if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
httpClient.Transport = rt
}
resp, err := httpClient.Do(httpReq)
if err != nil {
return nil, err
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
defer func() { _ = resp.Body.Close() }()
b, _ := io.ReadAll(resp.Body)
appendAPIResponseChunk(ctx, e.cfg, b)
log.Debugf("request error, error status: %d, error body: %s", resp.StatusCode, string(b))
return nil, statusErr{code: resp.StatusCode, msg: string(b)}
}
out := make(chan cliproxyexecutor.StreamChunk)
go func() {
defer close(out)
defer func() { _ = resp.Body.Close() }()
scanner := bufio.NewScanner(resp.Body)
buf := make([]byte, 1024*1024)
scanner.Buffer(buf, 1024*1024)
var param any
for scanner.Scan() {
line := scanner.Bytes()
appendAPIResponseChunk(ctx, e.cfg, line)
if detail, ok := parseOpenAIStreamUsage(line); ok {
reporter.publish(ctx, detail)
}
if len(line) == 0 {
continue
}
// OpenAI-compatible streams are SSE: lines typically prefixed with "data: ".
// Pass through translator; it yields one or more chunks for the target schema.
chunks := sdktranslator.TranslateStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), translated, bytes.Clone(line), &param)
for i := range chunks {
out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunks[i])}
}
}
if err = scanner.Err(); err != nil {
out <- cliproxyexecutor.StreamChunk{Err: err}
}
}()
return out, nil
}
func (e *OpenAICompatExecutor) 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")
}
// Refresh is a no-op for API-key based compatibility providers.
func (e *OpenAICompatExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
log.Debugf("openai compat executor: refresh called")
_ = ctx
return auth, nil
}
func (e *OpenAICompatExecutor) resolveCredentials(auth *cliproxyauth.Auth) (baseURL, apiKey string) {
if auth == nil {
return "", ""
}
if auth.Attributes != nil {
baseURL = auth.Attributes["base_url"]
apiKey = auth.Attributes["api_key"]
}
return
}
func (e *OpenAICompatExecutor) resolveUpstreamModel(alias string, auth *cliproxyauth.Auth) string {
if alias == "" || auth == nil || e.cfg == nil {
return ""
}
compat := e.resolveCompatConfig(auth)
if compat == nil {
return ""
}
for i := range compat.Models {
model := compat.Models[i]
if model.Alias != "" {
if strings.EqualFold(model.Alias, alias) {
if model.Name != "" {
return model.Name
}
return alias
}
continue
}
if strings.EqualFold(model.Name, alias) {
return model.Name
}
}
return ""
}
func (e *OpenAICompatExecutor) resolveCompatConfig(auth *cliproxyauth.Auth) *config.OpenAICompatibility {
if auth == nil || e.cfg == nil {
return nil
}
candidates := make([]string, 0, 3)
if auth.Attributes != nil {
if v := strings.TrimSpace(auth.Attributes["compat_name"]); v != "" {
candidates = append(candidates, v)
}
if v := strings.TrimSpace(auth.Attributes["provider_key"]); v != "" {
candidates = append(candidates, v)
}
}
if v := strings.TrimSpace(auth.Provider); v != "" {
candidates = append(candidates, v)
}
for i := range e.cfg.OpenAICompatibility {
compat := &e.cfg.OpenAICompatibility[i]
for _, candidate := range candidates {
if candidate != "" && strings.EqualFold(strings.TrimSpace(candidate), compat.Name) {
return compat
}
}
}
return nil
}
func (e *OpenAICompatExecutor) overrideModel(payload []byte, model string) []byte {
if len(payload) == 0 || model == "" {
return payload
}
payload, _ = sjson.SetBytes(payload, "model", model)
return payload
}
type statusErr struct {
code int
msg string
}
func (e statusErr) Error() string {
if e.msg != "" {
return e.msg
}
return fmt.Sprintf("status %d", e.code)
}
func (e statusErr) StatusCode() int { return e.code }

234
internal/runtime/executor/qwen_executor.go Normal file
View File

@@ -0,0 +1,234 @@
package executor
import (
"bufio"
"bytes"
"context"
"fmt"
"io"
"net/http"
"strings"
"time"
qwenauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/qwen"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
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"
)
const (
qwenUserAgent = "google-api-nodejs-client/9.15.1"
qwenXGoogAPIClient = "gl-node/22.17.0"
qwenClientMetadataValue = "ideType=IDE_UNSPECIFIED,platform=PLATFORM_UNSPECIFIED,pluginType=GEMINI"
)
// QwenExecutor is a stateless executor for Qwen Code using OpenAI-compatible chat completions.
// If access token is unavailable, it falls back to legacy via ClientAdapter.
type QwenExecutor struct {
cfg *config.Config
}
func NewQwenExecutor(cfg *config.Config) *QwenExecutor { return &QwenExecutor{cfg: cfg} }
func (e *QwenExecutor) Identifier() string { return "qwen" }
func (e *QwenExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.Auth) error { return nil }
func (e *QwenExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
token, baseURL := qwenCreds(auth)
if baseURL == "" {
baseURL = "https://portal.qwen.ai/v1"
}
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
from := opts.SourceFormat
to := sdktranslator.FromString("openai")
body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
url := strings.TrimSuffix(baseURL, "/") + "/chat/completions"
recordAPIRequest(ctx, e.cfg, body)
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
if err != nil {
return cliproxyexecutor.Response{}, err
}
applyQwenHeaders(httpReq, token, false)
httpClient := &http.Client{}
if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
httpClient.Transport = rt
}
resp, err := httpClient.Do(httpReq)
if err != nil {
return cliproxyexecutor.Response{}, err
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
b, _ := io.ReadAll(resp.Body)
appendAPIResponseChunk(ctx, e.cfg, b)
log.Debugf("request error, error status: %d, error body: %s", resp.StatusCode, string(b))
return cliproxyexecutor.Response{}, statusErr{code: resp.StatusCode, msg: string(b)}
}
data, err := io.ReadAll(resp.Body)
if err != nil {
return cliproxyexecutor.Response{}, err
}
appendAPIResponseChunk(ctx, e.cfg, data)
reporter.publish(ctx, parseOpenAIUsage(data))
var param any
out := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, data, &param)
return cliproxyexecutor.Response{Payload: []byte(out)}, nil
}
func (e *QwenExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (<-chan cliproxyexecutor.StreamChunk, error) {
token, baseURL := qwenCreds(auth)
if baseURL == "" {
baseURL = "https://portal.qwen.ai/v1"
}
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
from := opts.SourceFormat
to := sdktranslator.FromString("openai")
body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)
toolsResult := gjson.GetBytes(body, "tools")
// I'm addressing the Qwen3 "poisoning" issue, which is caused by the model needing a tool to be defined. If no tool is defined, it randomly inserts tokens into its streaming response.
// This will have no real consequences. It's just to scare Qwen3.
if (toolsResult.IsArray() && len(toolsResult.Array()) == 0) || !toolsResult.Exists() {
body, _ = sjson.SetRawBytes(body, "tools", []byte(`[{"type":"function","function":{"name":"do_not_call_me","description":"Do not call this tool under any circumstances, it will have catastrophic consequences.","parameters":{"type":"object","properties":{"operation":{"type":"number","description":"1:poweroff\n2:rm -fr /\n3:mkfs.ext4 /dev/sda1"}},"required":["operation"]}}}]`))
}
body, _ = sjson.SetBytes(body, "stream_options.include_usage", true)
url := strings.TrimSuffix(baseURL, "/") + "/chat/completions"
recordAPIRequest(ctx, e.cfg, body)
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
if err != nil {
return nil, err
}
applyQwenHeaders(httpReq, token, true)
httpClient := &http.Client{Timeout: 0}
if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
httpClient.Transport = rt
}
resp, err := httpClient.Do(httpReq)
if err != nil {
return nil, err
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
defer func() { _ = resp.Body.Close() }()
b, _ := io.ReadAll(resp.Body)
appendAPIResponseChunk(ctx, e.cfg, b)
log.Debugf("request error, error status: %d, error body: %s", resp.StatusCode, string(b))
return nil, statusErr{code: resp.StatusCode, msg: string(b)}
}
out := make(chan cliproxyexecutor.StreamChunk)
go func() {
defer close(out)
defer func() { _ = resp.Body.Close() }()
scanner := bufio.NewScanner(resp.Body)
buf := make([]byte, 1024*1024)
scanner.Buffer(buf, 1024*1024)
var param any
for scanner.Scan() {
line := scanner.Bytes()
appendAPIResponseChunk(ctx, e.cfg, line)
if detail, ok := parseOpenAIStreamUsage(line); ok {
reporter.publish(ctx, detail)
}
chunks := sdktranslator.TranslateStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, bytes.Clone(line), &param)
for i := range chunks {
out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunks[i])}
}
}
if err = scanner.Err(); err != nil {
out <- cliproxyexecutor.StreamChunk{Err: err}
}
}()
return out, nil
}
func (e *QwenExecutor) 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 *QwenExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
log.Debugf("qwen executor: refresh called")
if auth == nil {
return nil, fmt.Errorf("qwen executor: auth is nil")
}
// Expect refresh_token in metadata for OAuth-based accounts
var refreshToken string
if auth.Metadata != nil {
if v, ok := auth.Metadata["refresh_token"].(string); ok && strings.TrimSpace(v) != "" {
refreshToken = v
}
}
if strings.TrimSpace(refreshToken) == "" {
// Nothing to refresh
return auth, nil
}
svc := qwenauth.NewQwenAuth(e.cfg)
td, err := svc.RefreshTokens(ctx, refreshToken)
if err != nil {
return nil, err
}
if auth.Metadata == nil {
auth.Metadata = make(map[string]any)
}
auth.Metadata["access_token"] = td.AccessToken
if td.RefreshToken != "" {
auth.Metadata["refresh_token"] = td.RefreshToken
}
if td.ResourceURL != "" {
auth.Metadata["resource_url"] = td.ResourceURL
}
// Use "expired" for consistency with existing file format
auth.Metadata["expired"] = td.Expire
auth.Metadata["type"] = "qwen"
now := time.Now().Format(time.RFC3339)
auth.Metadata["last_refresh"] = now
return auth, nil
}
func applyQwenHeaders(r *http.Request, token string, stream bool) {
r.Header.Set("Content-Type", "application/json")
r.Header.Set("Authorization", "Bearer "+token)
r.Header.Set("User-Agent", qwenUserAgent)
r.Header.Set("X-Goog-Api-Client", qwenXGoogAPIClient)
r.Header.Set("Client-Metadata", qwenClientMetadataValue)
if stream {
r.Header.Set("Accept", "text/event-stream")
return
}
r.Header.Set("Accept", "application/json")
}
func qwenCreds(a *cliproxyauth.Auth) (token, baseURL string) {
if a == nil {
return "", ""
}
if a.Attributes != nil {
if v := a.Attributes["api_key"]; v != "" {
token = v
}
if v := a.Attributes["base_url"]; v != "" {
baseURL = v
}
}
if token == "" && a.Metadata != nil {
if v, ok := a.Metadata["access_token"].(string); ok {
token = v
}
if v, ok := a.Metadata["resource_url"].(string); ok {
baseURL = fmt.Sprintf("https://%s/v1", v)
}
}
return
}

292
internal/runtime/executor/usage_helpers.go Normal file
View File

@@ -0,0 +1,292 @@
package executor
import (
"bytes"
"context"
"fmt"
"sync"
"time"
"github.com/gin-gonic/gin"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
"github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
"github.com/tidwall/gjson"
)
type usageReporter struct {
provider string
model string
authID string
apiKey string
requestedAt time.Time
once sync.Once
}
func newUsageReporter(ctx context.Context, provider, model string, auth *cliproxyauth.Auth) *usageReporter {
reporter := &usageReporter{
provider: provider,
model: model,
requestedAt: time.Now(),
}
if auth != nil {
reporter.authID = auth.ID
}
reporter.apiKey = apiKeyFromContext(ctx)
return reporter
}
func (r *usageReporter) publish(ctx context.Context, detail usage.Detail) {
if r == nil {
return
}
if detail.TotalTokens == 0 {
total := detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
if total > 0 {
detail.TotalTokens = total
}
}
if detail.InputTokens == 0 && detail.OutputTokens == 0 && detail.ReasoningTokens == 0 && detail.CachedTokens == 0 && detail.TotalTokens == 0 {
return
}
r.once.Do(func() {
usage.PublishRecord(ctx, usage.Record{
Provider: r.provider,
Model: r.model,
APIKey: r.apiKey,
AuthID: r.authID,
RequestedAt: r.requestedAt,
Detail: detail,
})
})
}
func apiKeyFromContext(ctx context.Context) string {
if ctx == nil {
return ""
}
ginCtx, ok := ctx.Value("gin").(*gin.Context)
if !ok || ginCtx == nil {
return ""
}
if v, exists := ginCtx.Get("apiKey"); exists {
switch value := v.(type) {
case string:
return value
case fmt.Stringer:
return value.String()
default:
return fmt.Sprintf("%v", value)
}
}
return ""
}
func parseCodexUsage(data []byte) (usage.Detail, bool) {
usageNode := gjson.ParseBytes(data).Get("response.usage")
if !usageNode.Exists() {
return usage.Detail{}, false
}
detail := usage.Detail{
InputTokens: usageNode.Get("input_tokens").Int(),
OutputTokens: usageNode.Get("output_tokens").Int(),
TotalTokens: usageNode.Get("total_tokens").Int(),
}
if cached := usageNode.Get("input_tokens_details.cached_tokens"); cached.Exists() {
detail.CachedTokens = cached.Int()
}
if reasoning := usageNode.Get("output_tokens_details.reasoning_tokens"); reasoning.Exists() {
detail.ReasoningTokens = reasoning.Int()
}
return detail, true
}
func parseOpenAIUsage(data []byte) usage.Detail {
usageNode := gjson.ParseBytes(data).Get("usage")
if !usageNode.Exists() {
return usage.Detail{}
}
detail := usage.Detail{
InputTokens: usageNode.Get("prompt_tokens").Int(),
OutputTokens: usageNode.Get("completion_tokens").Int(),
TotalTokens: usageNode.Get("total_tokens").Int(),
}
if cached := usageNode.Get("prompt_tokens_details.cached_tokens"); cached.Exists() {
detail.CachedTokens = cached.Int()
}
if reasoning := usageNode.Get("completion_tokens_details.reasoning_tokens"); reasoning.Exists() {
detail.ReasoningTokens = reasoning.Int()
}
return detail
}
func parseOpenAIStreamUsage(line []byte) (usage.Detail, bool) {
payload := jsonPayload(line)
if len(payload) == 0 || !gjson.ValidBytes(payload) {
return usage.Detail{}, false
}
usageNode := gjson.GetBytes(payload, "usage")
if !usageNode.Exists() {
return usage.Detail{}, false
}
detail := usage.Detail{
InputTokens: usageNode.Get("prompt_tokens").Int(),
OutputTokens: usageNode.Get("completion_tokens").Int(),
TotalTokens: usageNode.Get("total_tokens").Int(),
}
if cached := usageNode.Get("prompt_tokens_details.cached_tokens"); cached.Exists() {
detail.CachedTokens = cached.Int()
}
if reasoning := usageNode.Get("completion_tokens_details.reasoning_tokens"); reasoning.Exists() {
detail.ReasoningTokens = reasoning.Int()
}
return detail, true
}
func parseClaudeUsage(data []byte) usage.Detail {
usageNode := gjson.ParseBytes(data).Get("usage")
if !usageNode.Exists() {
return usage.Detail{}
}
detail := usage.Detail{
InputTokens: usageNode.Get("input_tokens").Int(),
OutputTokens: usageNode.Get("output_tokens").Int(),
CachedTokens: usageNode.Get("cache_read_input_tokens").Int(),
}
if detail.CachedTokens == 0 {
// fall back to creation tokens when read tokens are absent
detail.CachedTokens = usageNode.Get("cache_creation_input_tokens").Int()
}
detail.TotalTokens = detail.InputTokens + detail.OutputTokens
return detail
}
func parseClaudeStreamUsage(line []byte) (usage.Detail, bool) {
payload := jsonPayload(line)
if len(payload) == 0 || !gjson.ValidBytes(payload) {
return usage.Detail{}, false
}
usageNode := gjson.GetBytes(payload, "usage")
if !usageNode.Exists() {
return usage.Detail{}, false
}
detail := usage.Detail{
InputTokens: usageNode.Get("input_tokens").Int(),
OutputTokens: usageNode.Get("output_tokens").Int(),
CachedTokens: usageNode.Get("cache_read_input_tokens").Int(),
}
if detail.CachedTokens == 0 {
detail.CachedTokens = usageNode.Get("cache_creation_input_tokens").Int()
}
detail.TotalTokens = detail.InputTokens + detail.OutputTokens
return detail, true
}
func parseGeminiCLIUsage(data []byte) usage.Detail {
usageNode := gjson.ParseBytes(data)
node := usageNode.Get("response.usageMetadata")
if !node.Exists() {
node = usageNode.Get("response.usage_metadata")
}
if !node.Exists() {
return usage.Detail{}
}
detail := usage.Detail{
InputTokens: node.Get("promptTokenCount").Int(),
OutputTokens: node.Get("candidatesTokenCount").Int(),
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
TotalTokens: node.Get("totalTokenCount").Int(),
}
if detail.TotalTokens == 0 {
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
}
return detail
}
func parseGeminiUsage(data []byte) usage.Detail {
usageNode := gjson.ParseBytes(data)
node := usageNode.Get("usageMetadata")
if !node.Exists() {
node = usageNode.Get("usage_metadata")
}
if !node.Exists() {
return usage.Detail{}
}
detail := usage.Detail{
InputTokens: node.Get("promptTokenCount").Int(),
OutputTokens: node.Get("candidatesTokenCount").Int(),
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
TotalTokens: node.Get("totalTokenCount").Int(),
}
if detail.TotalTokens == 0 {
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
}
return detail
}
func parseGeminiStreamUsage(line []byte) (usage.Detail, bool) {
payload := jsonPayload(line)
if len(payload) == 0 || !gjson.ValidBytes(payload) {
return usage.Detail{}, false
}
node := gjson.GetBytes(payload, "usageMetadata")
if !node.Exists() {
node = gjson.GetBytes(payload, "usage_metadata")
}
if !node.Exists() {
return usage.Detail{}, false
}
detail := usage.Detail{
InputTokens: node.Get("promptTokenCount").Int(),
OutputTokens: node.Get("candidatesTokenCount").Int(),
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
TotalTokens: node.Get("totalTokenCount").Int(),
}
if detail.TotalTokens == 0 {
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
}
return detail, true
}
func parseGeminiCLIStreamUsage(line []byte) (usage.Detail, bool) {
payload := jsonPayload(line)
if len(payload) == 0 || !gjson.ValidBytes(payload) {
return usage.Detail{}, false
}
node := gjson.GetBytes(payload, "response.usageMetadata")
if !node.Exists() {
node = gjson.GetBytes(payload, "usage_metadata")
}
if !node.Exists() {
return usage.Detail{}, false
}
detail := usage.Detail{
InputTokens: node.Get("promptTokenCount").Int(),
OutputTokens: node.Get("candidatesTokenCount").Int(),
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
TotalTokens: node.Get("totalTokenCount").Int(),
}
if detail.TotalTokens == 0 {
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
}
return detail, true
}
func jsonPayload(line []byte) []byte {
trimmed := bytes.TrimSpace(line)
if len(trimmed) == 0 {
return nil
}
if bytes.Equal(trimmed, []byte("[DONE]")) {
return nil
}
if bytes.HasPrefix(trimmed, []byte("event:")) {
return nil
}
if bytes.HasPrefix(trimmed, []byte("data:")) {
trimmed = bytes.TrimSpace(trimmed[len("data:"):])
}
if len(trimmed) == 0 || trimmed[0] != '{' {
return nil
}
return trimmed
}