Merge pull request #2412 from sususu98/feat/signature-cache-toggle

feat: configurable signature cache toggle for Antigravity/Claude thinking blocks

internal/runtime/executor/antigravity_executor.go

@@ -23,10 +23,12 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"github.com/router-for-me/CLIProxyAPI/v6/internal/cache"
 	"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/runtime/executor/helps"
 	"github.com/router-for-me/CLIProxyAPI/v6/internal/thinking"
+	antigravityclaude "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/antigravity/claude"
 	"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
 	sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
 	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
@@ -182,6 +184,24 @@ func newAntigravityHTTPClient(ctx context.Context, cfg *config.Config, auth *cli
 	return client
 }
 
+func validateAntigravityRequestSignatures(from sdktranslator.Format, rawJSON []byte) error {
+	if from.String() != "claude" {
+		return nil
+	}
+	if cache.SignatureCacheEnabled() {
+		return nil
+	}
+	if !cache.SignatureBypassStrictMode() {
+		// Non-strict bypass: let the translator handle invalid signatures
+		// by dropping unsigned thinking blocks silently (no 400).
+		return nil
+	}
+	if err := antigravityclaude.ValidateClaudeBypassSignatures(rawJSON); err != nil {
+		return statusErr{code: http.StatusBadRequest, msg: err.Error()}
+	}
+	return nil
+}
+
 // Identifier returns the executor identifier.
 func (e *AntigravityExecutor) Identifier() string { return antigravityAuthType }
 
@@ -664,14 +684,6 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au
 		return e.executeClaudeNonStream(ctx, auth, req, opts)
 	}
 
-	token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
-	if errToken != nil {
-		return resp, errToken
-	}
-	if updatedAuth != nil {
-		auth = updatedAuth
-	}
-
 	reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
 	defer reporter.TrackFailure(ctx, &err)
 
@@ -683,6 +695,16 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au
 		originalPayloadSource = opts.OriginalRequest
 	}
 	originalPayload := originalPayloadSource
+	if errValidate := validateAntigravityRequestSignatures(from, originalPayload); errValidate != nil {
+		return resp, errValidate
+	}
+	token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
+	if errToken != nil {
+		return resp, errToken
+	}
+	if updatedAuth != nil {
+		auth = updatedAuth
+	}
 	originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, false)
 	translated := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, false)
 
@@ -874,14 +896,6 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth *
 		return resp, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d}
 	}
 
-	token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
-	if errToken != nil {
-		return resp, errToken
-	}
-	if updatedAuth != nil {
-		auth = updatedAuth
-	}
-
 	reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
 	defer reporter.TrackFailure(ctx, &err)
 
@@ -893,6 +907,16 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth *
 		originalPayloadSource = opts.OriginalRequest
 	}
 	originalPayload := originalPayloadSource
+	if errValidate := validateAntigravityRequestSignatures(from, originalPayload); errValidate != nil {
+		return resp, errValidate
+	}
+	token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
+	if errToken != nil {
+		return resp, errToken
+	}
+	if updatedAuth != nil {
+		auth = updatedAuth
+	}
 	originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, true)
 	translated := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, true)
 
@@ -1335,14 +1359,6 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya
 		return nil, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d}
 	}
 
-	token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
-	if errToken != nil {
-		return nil, errToken
-	}
-	if updatedAuth != nil {
-		auth = updatedAuth
-	}
-
 	reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
 	defer reporter.TrackFailure(ctx, &err)
 
@@ -1354,6 +1370,16 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya
 		originalPayloadSource = opts.OriginalRequest
 	}
 	originalPayload := originalPayloadSource
+	if errValidate := validateAntigravityRequestSignatures(from, originalPayload); errValidate != nil {
+		return nil, errValidate
+	}
+	token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
+	if errToken != nil {
+		return nil, errToken
+	}
+	if updatedAuth != nil {
+		auth = updatedAuth
+	}
 	originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, true)
 	translated := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, true)
 
@@ -1593,6 +1619,16 @@ func (e *AntigravityExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Au
 func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
 	baseModel := thinking.ParseSuffix(req.Model).ModelName
 
+	from := opts.SourceFormat
+	to := sdktranslator.FromString("antigravity")
+	respCtx := context.WithValue(ctx, "alt", opts.Alt)
+	originalPayloadSource := req.Payload
+	if len(opts.OriginalRequest) > 0 {
+		originalPayloadSource = opts.OriginalRequest
+	}
+	if errValidate := validateAntigravityRequestSignatures(from, originalPayloadSource); errValidate != nil {
+		return cliproxyexecutor.Response{}, errValidate
+	}
 	token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
 	if errToken != nil {
 		return cliproxyexecutor.Response{}, errToken
@@ -1604,10 +1640,6 @@ func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyaut
 		return cliproxyexecutor.Response{}, statusErr{code: http.StatusUnauthorized, msg: "missing access token"}
 	}
 
-	from := opts.SourceFormat
-	to := sdktranslator.FromString("antigravity")
-	respCtx := context.WithValue(ctx, "alt", opts.Alt)
-
 	// Prepare payload once (doesn't depend on baseURL)
 	payload := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, false)
 

internal/runtime/executor/antigravity_executor_signature_test.go

@@ -0,0 +1,157 @@
+package executor
+
+import (
+	"bytes"
+	"context"
+	"encoding/base64"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/router-for-me/CLIProxyAPI/v6/internal/cache"
+	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"
+)
+
+func testGeminiSignaturePayload() string {
+	payload := append([]byte{0x0A}, bytes.Repeat([]byte{0x56}, 48)...)
+	return base64.StdEncoding.EncodeToString(payload)
+}
+
+func testAntigravityAuth(baseURL string) *cliproxyauth.Auth {
+	return &cliproxyauth.Auth{
+		Attributes: map[string]string{
+			"base_url": baseURL,
+		},
+		Metadata: map[string]any{
+			"access_token": "token-123",
+			"expired":      time.Now().Add(24 * time.Hour).Format(time.RFC3339),
+		},
+	}
+}
+
+func invalidClaudeThinkingPayload() []byte {
+	return []byte(`{
+		"model": "claude-sonnet-4-5-thinking",
+		"messages": [
+			{
+				"role": "assistant",
+				"content": [
+					{"type": "thinking", "thinking": "bad", "signature": "` + testGeminiSignaturePayload() + `"},
+					{"type": "text", "text": "hello"}
+				]
+			}
+		]
+	}`)
+}
+
+func TestAntigravityExecutor_StrictBypassRejectsInvalidSignature(t *testing.T) {
+	previousCache := cache.SignatureCacheEnabled()
+	previousStrict := cache.SignatureBypassStrictMode()
+	cache.SetSignatureCacheEnabled(false)
+	cache.SetSignatureBypassStrictMode(true)
+	t.Cleanup(func() {
+		cache.SetSignatureCacheEnabled(previousCache)
+		cache.SetSignatureBypassStrictMode(previousStrict)
+	})
+
+	var hits atomic.Int32
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		hits.Add(1)
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte(`{"response":{"candidates":[{"content":{"parts":[{"text":"ok"}]}}]}}`))
+	}))
+	defer server.Close()
+
+	executor := NewAntigravityExecutor(nil)
+	auth := testAntigravityAuth(server.URL)
+	payload := invalidClaudeThinkingPayload()
+	opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("claude"), OriginalRequest: payload}
+	req := cliproxyexecutor.Request{Model: "claude-sonnet-4-5-thinking", Payload: payload}
+
+	tests := []struct {
+		name   string
+		invoke func() error
+	}{
+		{
+			name: "execute",
+			invoke: func() error {
+				_, err := executor.Execute(context.Background(), auth, req, opts)
+				return err
+			},
+		},
+		{
+			name: "stream",
+			invoke: func() error {
+				_, err := executor.ExecuteStream(context.Background(), auth, req, cliproxyexecutor.Options{SourceFormat: opts.SourceFormat, OriginalRequest: payload, Stream: true})
+				return err
+			},
+		},
+		{
+			name: "count tokens",
+			invoke: func() error {
+				_, err := executor.CountTokens(context.Background(), auth, req, opts)
+				return err
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			err := tt.invoke()
+			if err == nil {
+				t.Fatal("expected invalid signature to return an error")
+			}
+			statusProvider, ok := err.(interface{ StatusCode() int })
+			if !ok {
+				t.Fatalf("expected status error, got %T: %v", err, err)
+			}
+			if statusProvider.StatusCode() != http.StatusBadRequest {
+				t.Fatalf("status = %d, want %d", statusProvider.StatusCode(), http.StatusBadRequest)
+			}
+		})
+	}
+
+	if got := hits.Load(); got != 0 {
+		t.Fatalf("expected invalid signature to be rejected before upstream request, got %d upstream hits", got)
+	}
+}
+
+func TestAntigravityExecutor_NonStrictBypassSkipsPrecheck(t *testing.T) {
+	previousCache := cache.SignatureCacheEnabled()
+	previousStrict := cache.SignatureBypassStrictMode()
+	cache.SetSignatureCacheEnabled(false)
+	cache.SetSignatureBypassStrictMode(false)
+	t.Cleanup(func() {
+		cache.SetSignatureCacheEnabled(previousCache)
+		cache.SetSignatureBypassStrictMode(previousStrict)
+	})
+
+	payload := invalidClaudeThinkingPayload()
+	from := sdktranslator.FromString("claude")
+
+	err := validateAntigravityRequestSignatures(from, payload)
+	if err != nil {
+		t.Fatalf("non-strict bypass should skip precheck, got: %v", err)
+	}
+}
+
+func TestAntigravityExecutor_CacheModeSkipsPrecheck(t *testing.T) {
+	previous := cache.SignatureCacheEnabled()
+	cache.SetSignatureCacheEnabled(true)
+	t.Cleanup(func() {
+		cache.SetSignatureCacheEnabled(previous)
+	})
+
+	payload := invalidClaudeThinkingPayload()
+	from := sdktranslator.FromString("claude")
+
+	err := validateAntigravityRequestSignatures(from, payload)
+	if err != nil {
+		t.Fatalf("cache mode should skip precheck, got: %v", err)
+	}
+}