Merge pull request #296 from router-for-me/antigravity

Antigravity bugfix

docker-compose.yml

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

internal/runtime/executor/aistudio_executor.go

@@ -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 {

internal/runtime/executor/antigravity_executor.go

@@ -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])}

internal/runtime/executor/usage_helpers.go

@@ -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
+}

internal/translator/gemini-cli/openai/chat-completions/gemini-cli_openai_response.go

@@ -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}
 }
 

sdk/auth/antigravity.go

@@ -21,6 +21,7 @@ import (
 const (
 	antigravityClientID     = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"
 	antigravityClientSecret = "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"
+	antigravityCallbackPort = 51121
 )
 
 var antigravityScopes = []string{
@@ -160,7 +161,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
 	}