feat(executor): add CountTokens support across all executors

- Introduced `CountTokens` method to Codex, Claude, Gemini, Qwen, OpenAI-compatible, and other executors.
- Implemented `ExecuteCount` in `AuthManager` for token counting via provider round-robin.
- Updated handlers to leverage `ExecuteCountWithAuthManager` for streamlined token counting.
- Added fallback and error handling logic for token counting requests.
This commit is contained in:
Luis Pater
2025-09-23 02:27:51 +08:00
parent d32fc0400e
commit ac59023abb
11 changed files with 210 additions and 28 deletions

View File

@@ -142,6 +142,10 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
return out, nil
}
func (e *ClaudeExecutor) 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 *ClaudeExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
log.Debugf("claude executor: refresh called")
if auth == nil {

View File

@@ -100,6 +100,10 @@ func (a *ClientAdapter) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au
return out, nil
}
func (e *ClientAdapter) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
return cliproxyexecutor.Response{Payload: []byte{}}, nil
}
// Refresh delegates to the legacy client's refresh logic when available.
func (a *ClientAdapter) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
client, _, err := resolveLegacyClient(auth)

View File

@@ -4,6 +4,7 @@ import (
"bufio"
"bytes"
"context"
"fmt"
"io"
"net/http"
"strings"
@@ -192,6 +193,10 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au
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 {

View File

@@ -260,6 +260,81 @@ func (e *GeminiCLIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyaut
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 {
var param any
translated := sdktranslator.TranslateNonStream(respCtx, to, from, attemptModel, bytes.Clone(opts.OriginalRequest), payload, data, &param)
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

View File

@@ -167,6 +167,10 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
return out, nil
}
func (e *GeminiExecutor) 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 *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.

View File

@@ -106,6 +106,10 @@ func (e *GeminiWebExecutor) ExecuteStream(ctx context.Context, auth *cliproxyaut
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)

View File

@@ -155,6 +155,10 @@ func (e *OpenAICompatExecutor) ExecuteStream(ctx context.Context, auth *cliproxy
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")

View File

@@ -150,6 +150,10 @@ func (e *QwenExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut
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 {