Skip to content

Commit 39dc455

Browse files
luispaterluispater
authored
Merge pull request #2412 from sususu98/feat/signature-cache-toggle
feat: configurable signature cache toggle for Antigravity/Claude thinking blocks
2 parents 30e94b6 + 38f0ae5 commit 39dc455

11 files changed

Lines changed: 1534 additions & 62 deletions

‎config.example.yaml‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,16 @@ nonstream-keepalive-interval: 0
114114
# keepalive-seconds: 15 # Default: 0 (disabled). <= 0 disables keep-alives.
115115
# bootstrap-retries: 1 # Default: 0 (disabled). Retries before first byte is sent.
116116

117+
# Signature cache validation for thinking blocks (Antigravity/Claude).
118+
# When true (default), cached signatures are preferred and validated.
119+
# When false, client signatures are used directly after normalization (bypass mode for testing).
120+
# antigravity-signature-cache-enabled: true
121+
122+
# Bypass mode signature validation strictness (only applies when signature cache is disabled).
123+
# When true, validates full Claude protobuf tree (Field 2 -> Field 1 structure).
124+
# When false (default), only checks R/E prefix + base64 + first byte 0x12.
125+
# antigravity-signature-bypass-strict: false
126+
117127
# Gemini API keys
118128
# gemini-api-key:
119129
# - api-key: "AIzaSy...01"

‎internal/api/server.go‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import (
2424
"github.com/router-for-me/CLIProxyAPI/v6/internal/api/middleware"
2525
"github.com/router-for-me/CLIProxyAPI/v6/internal/api/modules"
2626
ampmodule "github.com/router-for-me/CLIProxyAPI/v6/internal/api/modules/amp"
27+
"github.com/router-for-me/CLIProxyAPI/v6/internal/cache"
2728
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
2829
"github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
2930
"github.com/router-for-me/CLIProxyAPI/v6/internal/managementasset"
@@ -261,6 +262,7 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk
261262
}
262263
managementasset.SetCurrentConfig(cfg)
263264
auth.SetQuotaCooldownDisabled(cfg.DisableCooling)
265+
applySignatureCacheConfig(nil, cfg)
264266
// Initialize management handler
265267
s.mgmt = managementHandlers.NewHandler(cfg, configFilePath, authManager)
266268
if optionState.localPassword != "" {
@@ -918,6 +920,8 @@ func (s *Server) UpdateClients(cfg *config.Config) {
918920
auth.SetQuotaCooldownDisabled(cfg.DisableCooling)
919921
}
920922

923+
applySignatureCacheConfig(oldCfg, cfg)
924+
921925
if s.handlers != nil && s.handlers.AuthManager != nil {
922926
s.handlers.AuthManager.SetRetryConfig(cfg.RequestRetry, time.Duration(cfg.MaxRetryInterval)*time.Second, cfg.MaxRetryCredentials)
923927
}
@@ -1056,3 +1060,40 @@ func AuthMiddleware(manager *sdkaccess.Manager) gin.HandlerFunc {
10561060
c.AbortWithStatusJSON(statusCode, gin.H{"error": err.Message})
10571061
}
10581062
}
1063+
1064+
func configuredSignatureCacheEnabled(cfg *config.Config) bool {
1065+
if cfg != nil && cfg.AntigravitySignatureCacheEnabled != nil {
1066+
return *cfg.AntigravitySignatureCacheEnabled
1067+
}
1068+
return true
1069+
}
1070+
1071+
func applySignatureCacheConfig(oldCfg, cfg *config.Config) {
1072+
newVal := configuredSignatureCacheEnabled(cfg)
1073+
newStrict := configuredSignatureBypassStrict(cfg)
1074+
if oldCfg == nil {
1075+
cache.SetSignatureCacheEnabled(newVal)
1076+
cache.SetSignatureBypassStrictMode(newStrict)
1077+
log.Debugf("antigravity_signature_cache_enabled toggled to %t", newVal)
1078+
return
1079+
}
1080+
1081+
oldVal := configuredSignatureCacheEnabled(oldCfg)
1082+
if oldVal != newVal {
1083+
cache.SetSignatureCacheEnabled(newVal)
1084+
log.Debugf("antigravity_signature_cache_enabled updated from %t to %t", oldVal, newVal)
1085+
}
1086+
1087+
oldStrict := configuredSignatureBypassStrict(oldCfg)
1088+
if oldStrict != newStrict {
1089+
cache.SetSignatureBypassStrictMode(newStrict)
1090+
log.Debugf("antigravity_signature_bypass_strict updated from %t to %t", oldStrict, newStrict)
1091+
}
1092+
}
1093+
1094+
func configuredSignatureBypassStrict(cfg *config.Config) bool {
1095+
if cfg != nil && cfg.AntigravitySignatureBypassStrict != nil {
1096+
return *cfg.AntigravitySignatureBypassStrict
1097+
}
1098+
return false
1099+
}

‎internal/cache/signature_cache.go‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,10 @@ import (
55
"encoding/hex"
66
"strings"
77
"sync"
8+
"sync/atomic"
89
"time"
10+
11+
log "github.com/sirupsen/logrus"
912
)
1013

1114
// SignatureEntry holds a cached thinking signature with timestamp
@@ -193,3 +196,39 @@ func GetModelGroup(modelName string) string {
193196
}
194197
return modelName
195198
}
199+
200+
var signatureCacheEnabled atomic.Bool
201+
var signatureBypassStrictMode atomic.Bool
202+
203+
func init() {
204+
signatureCacheEnabled.Store(true)
205+
signatureBypassStrictMode.Store(false)
206+
}
207+
208+
// SetSignatureCacheEnabled switches Antigravity signature handling between cache mode and bypass mode.
209+
func SetSignatureCacheEnabled(enabled bool) {
210+
signatureCacheEnabled.Store(enabled)
211+
if !enabled {
212+
log.Warn("antigravity signature cache DISABLED - bypass mode active, cached signatures will not be used for request translation")
213+
}
214+
}
215+
216+
// SignatureCacheEnabled returns whether signature cache validation is enabled.
217+
func SignatureCacheEnabled() bool {
218+
return signatureCacheEnabled.Load()
219+
}
220+
221+
// SetSignatureBypassStrictMode controls whether bypass mode uses strict protobuf-tree validation.
222+
func SetSignatureBypassStrictMode(strict bool) {
223+
signatureBypassStrictMode.Store(strict)
224+
if strict {
225+
log.Info("antigravity bypass signature validation: strict mode (protobuf tree)")
226+
} else {
227+
log.Info("antigravity bypass signature validation: basic mode (R/E + 0x12)")
228+
}
229+
}
230+
231+
// SignatureBypassStrictMode returns whether bypass mode uses strict protobuf-tree validation.
232+
func SignatureBypassStrictMode() bool {
233+
return signatureBypassStrictMode.Load()
234+
}

‎internal/config/config.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,13 @@ type Config struct {
8585
// WebsocketAuth enables or disables authentication for the WebSocket API.
8686
WebsocketAuth bool `yaml:"ws-auth" json:"ws-auth"`
8787

88+
// AntigravitySignatureCacheEnabled controls whether signature cache validation is enabled for thinking blocks.
89+
// When true (default), cached signatures are preferred and validated.
90+
// When false, client signatures are used directly after normalization (bypass mode).
91+
AntigravitySignatureCacheEnabled *bool `yaml:"antigravity-signature-cache-enabled,omitempty" json:"antigravity-signature-cache-enabled,omitempty"`
92+
93+
AntigravitySignatureBypassStrict *bool `yaml:"antigravity-signature-bypass-strict,omitempty" json:"antigravity-signature-bypass-strict,omitempty"`
94+
8895
// GeminiKey defines Gemini API key configurations with optional routing overrides.
8996
GeminiKey []GeminiKey `yaml:"gemini-api-key" json:"gemini-api-key"`
9097

‎internal/runtime/executor/antigravity_executor.go‎

Lines changed: 60 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,12 @@ import (
2323
"time"
2424

2525
"github.com/google/uuid"
26+
"github.com/router-for-me/CLIProxyAPI/v6/internal/cache"
2627
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
2728
"github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
2829
"github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/executor/helps"
2930
"github.com/router-for-me/CLIProxyAPI/v6/internal/thinking"
31+
antigravityclaude "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/antigravity/claude"
3032
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
3133
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
3234
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
182184
return client
183185
}
184186

187+
func validateAntigravityRequestSignatures(from sdktranslator.Format, rawJSON []byte) error {
188+
if from.String() != "claude" {
189+
return nil
190+
}
191+
if cache.SignatureCacheEnabled() {
192+
return nil
193+
}
194+
if !cache.SignatureBypassStrictMode() {
195+
// Non-strict bypass: let the translator handle invalid signatures
196+
// by dropping unsigned thinking blocks silently (no 400).
197+
return nil
198+
}
199+
if err := antigravityclaude.ValidateClaudeBypassSignatures(rawJSON); err != nil {
200+
return statusErr{code: http.StatusBadRequest, msg: err.Error()}
201+
}
202+
return nil
203+
}
204+
185205
// Identifier returns the executor identifier.
186206
func (e *AntigravityExecutor) Identifier() string { return antigravityAuthType }
187207

@@ -664,14 +684,6 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au
664684
return e.executeClaudeNonStream(ctx, auth, req, opts)
665685
}
666686

667-
token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
668-
if errToken != nil {
669-
return resp, errToken
670-
}
671-
if updatedAuth != nil {
672-
auth = updatedAuth
673-
}
674-
675687
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
676688
defer reporter.TrackFailure(ctx, &err)
677689

@@ -683,6 +695,16 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au
683695
originalPayloadSource = opts.OriginalRequest
684696
}
685697
originalPayload := originalPayloadSource
698+
if errValidate := validateAntigravityRequestSignatures(from, originalPayload); errValidate != nil {
699+
return resp, errValidate
700+
}
701+
token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
702+
if errToken != nil {
703+
return resp, errToken
704+
}
705+
if updatedAuth != nil {
706+
auth = updatedAuth
707+
}
686708
originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, false)
687709
translated := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, false)
688710

@@ -874,14 +896,6 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth *
874896
return resp, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d}
875897
}
876898

877-
token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
878-
if errToken != nil {
879-
return resp, errToken
880-
}
881-
if updatedAuth != nil {
882-
auth = updatedAuth
883-
}
884-
885899
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
886900
defer reporter.TrackFailure(ctx, &err)
887901

@@ -893,6 +907,16 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth *
893907
originalPayloadSource = opts.OriginalRequest
894908
}
895909
originalPayload := originalPayloadSource
910+
if errValidate := validateAntigravityRequestSignatures(from, originalPayload); errValidate != nil {
911+
return resp, errValidate
912+
}
913+
token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
914+
if errToken != nil {
915+
return resp, errToken
916+
}
917+
if updatedAuth != nil {
918+
auth = updatedAuth
919+
}
896920
originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, true)
897921
translated := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, true)
898922

@@ -1335,14 +1359,6 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya
13351359
return nil, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d}
13361360
}
13371361

1338-
token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
1339-
if errToken != nil {
1340-
return nil, errToken
1341-
}
1342-
if updatedAuth != nil {
1343-
auth = updatedAuth
1344-
}
1345-
13461362
reporter := helps.NewUsageReporter(ctx, e.Identifier(), baseModel, auth)
13471363
defer reporter.TrackFailure(ctx, &err)
13481364

@@ -1354,6 +1370,16 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya
13541370
originalPayloadSource = opts.OriginalRequest
13551371
}
13561372
originalPayload := originalPayloadSource
1373+
if errValidate := validateAntigravityRequestSignatures(from, originalPayload); errValidate != nil {
1374+
return nil, errValidate
1375+
}
1376+
token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
1377+
if errToken != nil {
1378+
return nil, errToken
1379+
}
1380+
if updatedAuth != nil {
1381+
auth = updatedAuth
1382+
}
13571383
originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, true)
13581384
translated := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, true)
13591385

@@ -1593,6 +1619,16 @@ func (e *AntigravityExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Au
15931619
func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
15941620
baseModel := thinking.ParseSuffix(req.Model).ModelName
15951621

1622+
from := opts.SourceFormat
1623+
to := sdktranslator.FromString("antigravity")
1624+
respCtx := context.WithValue(ctx, "alt", opts.Alt)
1625+
originalPayloadSource := req.Payload
1626+
if len(opts.OriginalRequest) > 0 {
1627+
originalPayloadSource = opts.OriginalRequest
1628+
}
1629+
if errValidate := validateAntigravityRequestSignatures(from, originalPayloadSource); errValidate != nil {
1630+
return cliproxyexecutor.Response{}, errValidate
1631+
}
15961632
token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth)
15971633
if errToken != nil {
15981634
return cliproxyexecutor.Response{}, errToken
@@ -1604,10 +1640,6 @@ func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyaut
16041640
return cliproxyexecutor.Response{}, statusErr{code: http.StatusUnauthorized, msg: "missing access token"}
16051641
}
16061642

1607-
from := opts.SourceFormat
1608-
to := sdktranslator.FromString("antigravity")
1609-
respCtx := context.WithValue(ctx, "alt", opts.Alt)
1610-
16111643
// Prepare payload once (doesn't depend on baseURL)
16121644
payload := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, false)
16131645

0 commit comments

Comments
 (0)