Compare commits

...

5 Commits

Author SHA1 Message Date
Luis Pater
327cc7039e **refactor(auth): use customizable HTTP client for Antigravity requests**
Some checks failed
docker-image / docker (push) Has been cancelled
goreleaser / goreleaser (push) Has been cancelled
- Replaced `http.DefaultClient` with a configurable `http.Client` instance for Antigravity OAuth flow methods.
- Updated `exchangeAntigravityCode` and `fetchAntigravityUserInfo` to accept `httpClient` as a parameter.
- Added `util.SetProxy` usage to initialize the `httpClient` with proxy support.
2025-11-21 20:54:56 +08:00
Luis Pater
b4d15ace91 Merge pull request #296 from router-for-me/antigravity
Some checks failed
docker-image / docker (push) Has been cancelled
goreleaser / goreleaser (push) Has been cancelled
Antigravity bugfix
2025-11-21 17:32:36 +08:00
hkfires
abc2465b29 fix(gemini-cli): ignore thoughtSignature and empty parts 2025-11-21 17:12:56 +08:00
hkfires
4ba5b43d82 feat(executor): share SSE usage filtering across streams 2025-11-21 16:51:05 +08:00
hkfires
27faf718a3 fix(auth): use fixed antigravity callback port 51121 2025-11-21 13:56:33 +08:00
6 changed files with 133 additions and 69 deletions

View File

@@ -19,6 +19,7 @@ services:
- "8085:8085"
- "1455:1455"
- "54545:54545"
- "51121:51121"
- "11451:11451"
volumes:
- ./config.yaml:/CLIProxyAPI/config.yaml

View File

@@ -151,7 +151,7 @@ func (e *AIStudioExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth
case wsrelay.MessageTypeStreamChunk:
if len(event.Payload) > 0 {
appendAPIResponseChunk(ctx, e.cfg, bytes.Clone(event.Payload))
filtered := filterAIStudioUsageMetadata(event.Payload)
filtered := FilterSSEUsageMetadata(event.Payload)
if detail, ok := parseGeminiStreamUsage(filtered); ok {
reporter.publish(ctx, detail)
}
@@ -296,65 +296,6 @@ func (e *AIStudioExecutor) buildEndpoint(model, action, alt string) string {
return base
}
// filterAIStudioUsageMetadata removes usageMetadata from intermediate SSE events so that
// only the terminal chunk retains token statistics.
func filterAIStudioUsageMetadata(payload []byte) []byte {
if len(payload) == 0 {
return payload
}
lines := bytes.Split(payload, []byte("\n"))
modified := false
for idx, line := range lines {
trimmed := bytes.TrimSpace(line)
if len(trimmed) == 0 || !bytes.HasPrefix(trimmed, []byte("data:")) {
continue
}
dataIdx := bytes.Index(line, []byte("data:"))
if dataIdx < 0 {
continue
}
rawJSON := bytes.TrimSpace(line[dataIdx+5:])
cleaned, changed := stripUsageMetadataFromJSON(rawJSON)
if !changed {
continue
}
var rebuilt []byte
rebuilt = append(rebuilt, line[:dataIdx]...)
rebuilt = append(rebuilt, []byte("data:")...)
if len(cleaned) > 0 {
rebuilt = append(rebuilt, ' ')
rebuilt = append(rebuilt, cleaned...)
}
lines[idx] = rebuilt
modified = true
}
if !modified {
return payload
}
return bytes.Join(lines, []byte("\n"))
}
// stripUsageMetadataFromJSON drops usageMetadata when no finishReason is present.
func stripUsageMetadataFromJSON(rawJSON []byte) ([]byte, bool) {
jsonBytes := bytes.TrimSpace(rawJSON)
if len(jsonBytes) == 0 || !gjson.ValidBytes(jsonBytes) {
return rawJSON, false
}
finishReason := gjson.GetBytes(jsonBytes, "candidates.0.finishReason")
if finishReason.Exists() && finishReason.String() != "" {
return rawJSON, false
}
if !gjson.GetBytes(jsonBytes, "usageMetadata").Exists() {
return rawJSON, false
}
cleaned, err := sjson.DeleteBytes(jsonBytes, "usageMetadata")
if err != nil {
return rawJSON, false
}
return cleaned, true
}
// ensureColonSpacedJSON normalizes JSON objects so that colons are followed by a single space while
// keeping the payload otherwise compact. Non-JSON inputs are returned unchanged.
func ensureColonSpacedJSON(payload []byte) []byte {

View File

@@ -167,6 +167,11 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya
for scanner.Scan() {
line := scanner.Bytes()
appendAPIResponseChunk(ctx, e.cfg, line)
// Filter usage metadata for all models
// Only retain usage statistics in the terminal chunk
line = FilterSSEUsageMetadata(line)
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])}

View File

@@ -12,6 +12,7 @@ import (
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"
"github.com/tidwall/sjson"
)
type usageReporter struct {
@@ -383,3 +384,94 @@ func jsonPayload(line []byte) []byte {
}
return trimmed
}
// FilterSSEUsageMetadata removes usageMetadata from intermediate SSE events so that
// only the terminal chunk retains token statistics.
// This function is shared between aistudio and antigravity executors.
func FilterSSEUsageMetadata(payload []byte) []byte {
if len(payload) == 0 {
return payload
}
lines := bytes.Split(payload, []byte("\n"))
modified := false
for idx, line := range lines {
trimmed := bytes.TrimSpace(line)
if len(trimmed) == 0 || !bytes.HasPrefix(trimmed, []byte("data:")) {
continue
}
dataIdx := bytes.Index(line, []byte("data:"))
if dataIdx < 0 {
continue
}
rawJSON := bytes.TrimSpace(line[dataIdx+5:])
cleaned, changed := StripUsageMetadataFromJSON(rawJSON)
if !changed {
continue
}
var rebuilt []byte
rebuilt = append(rebuilt, line[:dataIdx]...)
rebuilt = append(rebuilt, []byte("data:")...)
if len(cleaned) > 0 {
rebuilt = append(rebuilt, ' ')
rebuilt = append(rebuilt, cleaned...)
}
lines[idx] = rebuilt
modified = true
}
if !modified {
return payload
}
return bytes.Join(lines, []byte("\n"))
}
// StripUsageMetadataFromJSON drops usageMetadata when no finishReason is present.
// This function is shared between aistudio and antigravity executors.
// It handles both formats:
// - Aistudio: candidates.0.finishReason
// - Antigravity: response.candidates.0.finishReason
func StripUsageMetadataFromJSON(rawJSON []byte) ([]byte, bool) {
jsonBytes := bytes.TrimSpace(rawJSON)
if len(jsonBytes) == 0 || !gjson.ValidBytes(jsonBytes) {
return rawJSON, false
}
// Check for finishReason in both aistudio and antigravity formats
finishReason := gjson.GetBytes(jsonBytes, "candidates.0.finishReason")
if !finishReason.Exists() {
finishReason = gjson.GetBytes(jsonBytes, "response.candidates.0.finishReason")
}
// If finishReason exists and is not empty, keep the usageMetadata
if finishReason.Exists() && finishReason.String() != "" {
return rawJSON, false
}
// Check for usageMetadata in both possible locations
usageMetadata := gjson.GetBytes(jsonBytes, "usageMetadata")
if !usageMetadata.Exists() {
usageMetadata = gjson.GetBytes(jsonBytes, "response.usageMetadata")
}
if !usageMetadata.Exists() {
return rawJSON, false
}
// Remove usageMetadata from both possible locations
cleaned := jsonBytes
var changed bool
// Try to remove usageMetadata from root level
if gjson.GetBytes(cleaned, "usageMetadata").Exists() {
cleaned, _ = sjson.DeleteBytes(cleaned, "usageMetadata")
changed = true
}
// Try to remove usageMetadata from response level
if gjson.GetBytes(cleaned, "response.usageMetadata").Exists() {
cleaned, _ = sjson.DeleteBytes(cleaned, "response.usageMetadata")
changed = true
}
return cleaned, changed
}

View File

@@ -98,25 +98,40 @@ func ConvertCliResponseToOpenAI(_ context.Context, _ string, originalRequestRawJ
// Process the main content part of the response.
partsResult := gjson.GetBytes(rawJSON, "response.candidates.0.content.parts")
hasFunctionCall := false
hasValidContent := false
if partsResult.IsArray() {
partResults := partsResult.Array()
for i := 0; i < len(partResults); i++ {
partResult := partResults[i]
partTextResult := partResult.Get("text")
functionCallResult := partResult.Get("functionCall")
thoughtSignatureResult := partResult.Get("thoughtSignature")
inlineDataResult := partResult.Get("inlineData")
if !inlineDataResult.Exists() {
inlineDataResult = partResult.Get("inline_data")
}
// Handle thoughtSignature - this is encrypted reasoning content that should not be exposed to the client
if thoughtSignatureResult.Exists() && thoughtSignatureResult.String() != "" {
// Skip thoughtSignature processing - it's internal encrypted data
continue
}
if partTextResult.Exists() {
textContent := partTextResult.String()
// Skip empty text content to avoid generating unnecessary chunks
if textContent == "" {
continue
}
// Handle text content, distinguishing between regular content and reasoning/thoughts.
if partResult.Get("thought").Bool() {
template, _ = sjson.Set(template, "choices.0.delta.reasoning_content", partTextResult.String())
template, _ = sjson.Set(template, "choices.0.delta.reasoning_content", textContent)
} else {
template, _ = sjson.Set(template, "choices.0.delta.content", partTextResult.String())
template, _ = sjson.Set(template, "choices.0.delta.content", textContent)
}
template, _ = sjson.Set(template, "choices.0.delta.role", "assistant")
hasValidContent = true
} else if functionCallResult.Exists() {
// Handle function call content.
hasFunctionCall = true
@@ -176,6 +191,12 @@ func ConvertCliResponseToOpenAI(_ context.Context, _ string, originalRequestRawJ
template, _ = sjson.Set(template, "choices.0.native_finish_reason", "tool_calls")
}
// Only return a chunk if there's actual content or a finish reason
finishReason := gjson.GetBytes(rawJSON, "response.candidates.0.finishReason")
if !hasValidContent && !finishReason.Exists() {
return []string{}
}
return []string{template}
}

View File

@@ -21,6 +21,7 @@ import (
const (
antigravityClientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"
antigravityClientSecret = "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"
antigravityCallbackPort = 51121
)
var antigravityScopes = []string{
@@ -58,6 +59,8 @@ func (AntigravityAuthenticator) Login(ctx context.Context, cfg *config.Config, o
opts = &LoginOptions{}
}
httpClient := util.SetProxy(&cfg.SDKConfig, &http.Client{})
state, err := misc.GenerateRandomState()
if err != nil {
return nil, fmt.Errorf("antigravity: failed to generate state: %w", err)
@@ -112,14 +115,14 @@ func (AntigravityAuthenticator) Login(ctx context.Context, cfg *config.Config, o
return nil, fmt.Errorf("antigravity: missing authorization code")
}
tokenResp, errToken := exchangeAntigravityCode(ctx, cbRes.Code, redirectURI)
tokenResp, errToken := exchangeAntigravityCode(ctx, cbRes.Code, redirectURI, httpClient)
if errToken != nil {
return nil, fmt.Errorf("antigravity: token exchange failed: %w", errToken)
}
email := ""
if tokenResp.AccessToken != "" {
if info, errInfo := fetchAntigravityUserInfo(ctx, tokenResp.AccessToken); errInfo == nil && strings.TrimSpace(info.Email) != "" {
if info, errInfo := fetchAntigravityUserInfo(ctx, tokenResp.AccessToken, httpClient); errInfo == nil && strings.TrimSpace(info.Email) != "" {
email = strings.TrimSpace(info.Email)
}
}
@@ -160,7 +163,8 @@ type callbackResult struct {
}
func startAntigravityCallbackServer() (*http.Server, int, <-chan callbackResult, error) {
listener, err := net.Listen("tcp", "127.0.0.1:0")
addr := fmt.Sprintf(":%d", antigravityCallbackPort)
listener, err := net.Listen("tcp", addr)
if err != nil {
return nil, 0, nil, err
}
@@ -200,7 +204,7 @@ type antigravityTokenResponse struct {
TokenType string `json:"token_type"`
}
func exchangeAntigravityCode(ctx context.Context, code, redirectURI string) (*antigravityTokenResponse, error) {
func exchangeAntigravityCode(ctx context.Context, code, redirectURI string, httpClient *http.Client) (*antigravityTokenResponse, error) {
data := url.Values{}
data.Set("code", code)
data.Set("client_id", antigravityClientID)
@@ -214,7 +218,7 @@ func exchangeAntigravityCode(ctx context.Context, code, redirectURI string) (*an
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
resp, errDo := http.DefaultClient.Do(req)
resp, errDo := httpClient.Do(req)
if errDo != nil {
return nil, errDo
}
@@ -238,7 +242,7 @@ type antigravityUserInfo struct {
Email string `json:"email"`
}
func fetchAntigravityUserInfo(ctx context.Context, accessToken string) (*antigravityUserInfo, error) {
func fetchAntigravityUserInfo(ctx context.Context, accessToken string, httpClient *http.Client) (*antigravityUserInfo, error) {
if strings.TrimSpace(accessToken) == "" {
return &antigravityUserInfo{}, nil
}
@@ -248,7 +252,7 @@ func fetchAntigravityUserInfo(ctx context.Context, accessToken string) (*antigra
}
req.Header.Set("Authorization", "Bearer "+accessToken)
resp, errDo := http.DefaultClient.Do(req)
resp, errDo := httpClient.Do(req)
if errDo != nil {
return nil, errDo
}