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

feat: configurable signature cache toggle for Antigravity/Claude thinking blocks
This commit is contained in:
Luis Pater
2026-04-09 21:54:47 +08:00
committed by GitHub
11 changed files with 1534 additions and 62 deletions

View File

@@ -114,6 +114,16 @@ nonstream-keepalive-interval: 0
# keepalive-seconds: 15 # Default: 0 (disabled). <= 0 disables keep-alives.
# bootstrap-retries: 1 # Default: 0 (disabled). Retries before first byte is sent.
# Signature cache validation for thinking blocks (Antigravity/Claude).
# When true (default), cached signatures are preferred and validated.
# When false, client signatures are used directly after normalization (bypass mode for testing).
# antigravity-signature-cache-enabled: true
# Bypass mode signature validation strictness (only applies when signature cache is disabled).
# When true, validates full Claude protobuf tree (Field 2 -> Field 1 structure).
# When false (default), only checks R/E prefix + base64 + first byte 0x12.
# antigravity-signature-bypass-strict: false
# Gemini API keys
# gemini-api-key:
# - api-key: "AIzaSy...01"

View File

@@ -24,6 +24,7 @@ import (
"github.com/router-for-me/CLIProxyAPI/v6/internal/api/middleware"
"github.com/router-for-me/CLIProxyAPI/v6/internal/api/modules"
ampmodule "github.com/router-for-me/CLIProxyAPI/v6/internal/api/modules/amp"
"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/logging"
"github.com/router-for-me/CLIProxyAPI/v6/internal/managementasset"
@@ -261,6 +262,7 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk
}
managementasset.SetCurrentConfig(cfg)
auth.SetQuotaCooldownDisabled(cfg.DisableCooling)
applySignatureCacheConfig(nil, cfg)
// Initialize management handler
s.mgmt = managementHandlers.NewHandler(cfg, configFilePath, authManager)
if optionState.localPassword != "" {
@@ -918,6 +920,8 @@ func (s *Server) UpdateClients(cfg *config.Config) {
auth.SetQuotaCooldownDisabled(cfg.DisableCooling)
}
applySignatureCacheConfig(oldCfg, cfg)
if s.handlers != nil && s.handlers.AuthManager != nil {
s.handlers.AuthManager.SetRetryConfig(cfg.RequestRetry, time.Duration(cfg.MaxRetryInterval)*time.Second, cfg.MaxRetryCredentials)
}
@@ -1056,3 +1060,40 @@ func AuthMiddleware(manager *sdkaccess.Manager) gin.HandlerFunc {
c.AbortWithStatusJSON(statusCode, gin.H{"error": err.Message})
}
}
func configuredSignatureCacheEnabled(cfg *config.Config) bool {
if cfg != nil && cfg.AntigravitySignatureCacheEnabled != nil {
return *cfg.AntigravitySignatureCacheEnabled
}
return true
}
func applySignatureCacheConfig(oldCfg, cfg *config.Config) {
newVal := configuredSignatureCacheEnabled(cfg)
newStrict := configuredSignatureBypassStrict(cfg)
if oldCfg == nil {
cache.SetSignatureCacheEnabled(newVal)
cache.SetSignatureBypassStrictMode(newStrict)
log.Debugf("antigravity_signature_cache_enabled toggled to %t", newVal)
return
}
oldVal := configuredSignatureCacheEnabled(oldCfg)
if oldVal != newVal {
cache.SetSignatureCacheEnabled(newVal)
log.Debugf("antigravity_signature_cache_enabled updated from %t to %t", oldVal, newVal)
}
oldStrict := configuredSignatureBypassStrict(oldCfg)
if oldStrict != newStrict {
cache.SetSignatureBypassStrictMode(newStrict)
log.Debugf("antigravity_signature_bypass_strict updated from %t to %t", oldStrict, newStrict)
}
}
func configuredSignatureBypassStrict(cfg *config.Config) bool {
if cfg != nil && cfg.AntigravitySignatureBypassStrict != nil {
return *cfg.AntigravitySignatureBypassStrict
}
return false
}

View File

@@ -5,7 +5,10 @@ import (
"encoding/hex"
"strings"
"sync"
"sync/atomic"
"time"
log "github.com/sirupsen/logrus"
)
// SignatureEntry holds a cached thinking signature with timestamp
@@ -193,3 +196,39 @@ func GetModelGroup(modelName string) string {
}
return modelName
}
var signatureCacheEnabled atomic.Bool
var signatureBypassStrictMode atomic.Bool
func init() {
signatureCacheEnabled.Store(true)
signatureBypassStrictMode.Store(false)
}
// SetSignatureCacheEnabled switches Antigravity signature handling between cache mode and bypass mode.
func SetSignatureCacheEnabled(enabled bool) {
signatureCacheEnabled.Store(enabled)
if !enabled {
log.Warn("antigravity signature cache DISABLED - bypass mode active, cached signatures will not be used for request translation")
}
}
// SignatureCacheEnabled returns whether signature cache validation is enabled.
func SignatureCacheEnabled() bool {
return signatureCacheEnabled.Load()
}
// SetSignatureBypassStrictMode controls whether bypass mode uses strict protobuf-tree validation.
func SetSignatureBypassStrictMode(strict bool) {
signatureBypassStrictMode.Store(strict)
if strict {
log.Info("antigravity bypass signature validation: strict mode (protobuf tree)")
} else {
log.Info("antigravity bypass signature validation: basic mode (R/E + 0x12)")
}
}
// SignatureBypassStrictMode returns whether bypass mode uses strict protobuf-tree validation.
func SignatureBypassStrictMode() bool {
return signatureBypassStrictMode.Load()
}

View File

@@ -85,6 +85,13 @@ type Config struct {
// WebsocketAuth enables or disables authentication for the WebSocket API.
WebsocketAuth bool `yaml:"ws-auth" json:"ws-auth"`
// AntigravitySignatureCacheEnabled controls whether signature cache validation is enabled for thinking blocks.
// When true (default), cached signatures are preferred and validated.
// When false, client signatures are used directly after normalization (bypass mode).
AntigravitySignatureCacheEnabled *bool `yaml:"antigravity-signature-cache-enabled,omitempty" json:"antigravity-signature-cache-enabled,omitempty"`
AntigravitySignatureBypassStrict *bool `yaml:"antigravity-signature-bypass-strict,omitempty" json:"antigravity-signature-bypass-strict,omitempty"`
// GeminiKey defines Gemini API key configurations with optional routing overrides.
GeminiKey []GeminiKey `yaml:"gemini-api-key" json:"gemini-api-key"`

View File

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

View File

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

View File

@@ -17,6 +17,56 @@ import (
"github.com/tidwall/sjson"
)
func resolveThinkingSignature(modelName, thinkingText, rawSignature string) string {
if cache.SignatureCacheEnabled() {
return resolveCacheModeSignature(modelName, thinkingText, rawSignature)
}
return resolveBypassModeSignature(rawSignature)
}
func resolveCacheModeSignature(modelName, thinkingText, rawSignature string) string {
if thinkingText != "" {
if cachedSig := cache.GetCachedSignature(modelName, thinkingText); cachedSig != "" {
return cachedSig
}
}
if rawSignature == "" {
return ""
}
clientSignature := ""
arrayClientSignatures := strings.SplitN(rawSignature, "#", 2)
if len(arrayClientSignatures) == 2 {
if cache.GetModelGroup(modelName) == arrayClientSignatures[0] {
clientSignature = arrayClientSignatures[1]
}
}
if cache.HasValidSignature(modelName, clientSignature) {
return clientSignature
}
return ""
}
func resolveBypassModeSignature(rawSignature string) string {
if rawSignature == "" {
return ""
}
normalized, err := normalizeClaudeBypassSignature(rawSignature)
if err != nil {
return ""
}
return normalized
}
func hasResolvedThinkingSignature(modelName, signature string) bool {
if cache.SignatureCacheEnabled() {
return cache.HasValidSignature(modelName, signature)
}
return signature != ""
}
// ConvertClaudeRequestToAntigravity parses and transforms a Claude Code API request into Gemini CLI API format.
// It extracts the model name, system instruction, message contents, and tool declarations
// from the raw JSON request and returns them in the format expected by the Gemini CLI API.
@@ -101,42 +151,15 @@ func ConvertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _
if contentTypeResult.Type == gjson.String && contentTypeResult.String() == "thinking" {
// Use GetThinkingText to handle wrapped thinking objects
thinkingText := thinking.GetThinkingText(contentResult)
// Always try cached signature first (more reliable than client-provided)
// Client may send stale or invalid signatures from different sessions
signature := ""
if thinkingText != "" {
if cachedSig := cache.GetCachedSignature(modelName, thinkingText); cachedSig != "" {
signature = cachedSig
// log.Debugf("Using cached signature for thinking block")
}
}
// Fallback to client signature only if cache miss and client signature is valid
if signature == "" {
signatureResult := contentResult.Get("signature")
clientSignature := ""
if signatureResult.Exists() && signatureResult.String() != "" {
arrayClientSignatures := strings.SplitN(signatureResult.String(), "#", 2)
if len(arrayClientSignatures) == 2 {
if cache.GetModelGroup(modelName) == arrayClientSignatures[0] {
clientSignature = arrayClientSignatures[1]
}
}
}
if cache.HasValidSignature(modelName, clientSignature) {
signature = clientSignature
}
// log.Debugf("Using client-provided signature for thinking block")
}
signature := resolveThinkingSignature(modelName, thinkingText, contentResult.Get("signature").String())
// Store for subsequent tool_use in the same message
if cache.HasValidSignature(modelName, signature) {
if hasResolvedThinkingSignature(modelName, signature) {
currentMessageThinkingSignature = signature
}
// Skip trailing unsigned thinking blocks on last assistant message
isUnsigned := !cache.HasValidSignature(modelName, signature)
// Skip unsigned thinking blocks instead of converting them to text.
isUnsigned := !hasResolvedThinkingSignature(modelName, signature)
// If unsigned, skip entirely (don't convert to text)
// Claude requires assistant messages to start with thinking blocks when thinking is enabled
@@ -198,7 +221,7 @@ func ConvertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _
// This is the approach used in opencode-google-antigravity-auth for Gemini
// and also works for Claude through Antigravity API
const skipSentinel = "skip_thought_signature_validator"
if cache.HasValidSignature(modelName, currentMessageThinkingSignature) {
if hasResolvedThinkingSignature(modelName, currentMessageThinkingSignature) {
partJSON, _ = sjson.SetBytes(partJSON, "thoughtSignature", currentMessageThinkingSignature)
} else {
// No valid signature - use skip sentinel to bypass validation

View File

@@ -1,13 +1,97 @@
package claude
import (
"bytes"
"encoding/base64"
"strings"
"testing"
"github.com/router-for-me/CLIProxyAPI/v6/internal/cache"
"github.com/tidwall/gjson"
"google.golang.org/protobuf/encoding/protowire"
)
func testAnthropicNativeSignature(t *testing.T) string {
t.Helper()
payload := buildClaudeSignaturePayload(t, 12, uint64Ptr(2), "claude-sonnet-4-6", true)
signature := base64.StdEncoding.EncodeToString(payload)
if len(signature) < cache.MinValidSignatureLen {
t.Fatalf("test signature too short: %d", len(signature))
}
return signature
}
func testMinimalAnthropicSignature(t *testing.T) string {
t.Helper()
payload := buildClaudeSignaturePayload(t, 12, nil, "", false)
return base64.StdEncoding.EncodeToString(payload)
}
func buildClaudeSignaturePayload(t *testing.T, channelID uint64, field2 *uint64, modelText string, includeField7 bool) []byte {
t.Helper()
channelBlock := []byte{}
channelBlock = protowire.AppendTag(channelBlock, 1, protowire.VarintType)
channelBlock = protowire.AppendVarint(channelBlock, channelID)
if field2 != nil {
channelBlock = protowire.AppendTag(channelBlock, 2, protowire.VarintType)
channelBlock = protowire.AppendVarint(channelBlock, *field2)
}
if modelText != "" {
channelBlock = protowire.AppendTag(channelBlock, 6, protowire.BytesType)
channelBlock = protowire.AppendString(channelBlock, modelText)
}
if includeField7 {
channelBlock = protowire.AppendTag(channelBlock, 7, protowire.VarintType)
channelBlock = protowire.AppendVarint(channelBlock, 0)
}
container := []byte{}
container = protowire.AppendTag(container, 1, protowire.BytesType)
container = protowire.AppendBytes(container, channelBlock)
container = protowire.AppendTag(container, 2, protowire.BytesType)
container = protowire.AppendBytes(container, bytes.Repeat([]byte{0x11}, 12))
container = protowire.AppendTag(container, 3, protowire.BytesType)
container = protowire.AppendBytes(container, bytes.Repeat([]byte{0x22}, 12))
container = protowire.AppendTag(container, 4, protowire.BytesType)
container = protowire.AppendBytes(container, bytes.Repeat([]byte{0x33}, 48))
payload := []byte{}
payload = protowire.AppendTag(payload, 2, protowire.BytesType)
payload = protowire.AppendBytes(payload, container)
payload = protowire.AppendTag(payload, 3, protowire.VarintType)
payload = protowire.AppendVarint(payload, 1)
return payload
}
func uint64Ptr(v uint64) *uint64 {
return &v
}
func testNonAnthropicRawSignature(t *testing.T) string {
t.Helper()
payload := bytes.Repeat([]byte{0x34}, 48)
signature := base64.StdEncoding.EncodeToString(payload)
if len(signature) < cache.MinValidSignatureLen {
t.Fatalf("test signature too short: %d", len(signature))
}
return signature
}
func testGeminiRawSignature(t *testing.T) string {
t.Helper()
payload := append([]byte{0x0A}, bytes.Repeat([]byte{0x56}, 48)...)
signature := base64.StdEncoding.EncodeToString(payload)
if len(signature) < cache.MinValidSignatureLen {
t.Fatalf("test signature too short: %d", len(signature))
}
return signature
}
func TestConvertClaudeRequestToAntigravity_BasicStructure(t *testing.T) {
inputJSON := []byte(`{
"model": "claude-3-5-sonnet-20240620",
@@ -116,6 +200,545 @@ func TestConvertClaudeRequestToAntigravity_ThinkingBlocks(t *testing.T) {
}
}
func TestValidateBypassMode_AcceptsClaudeSingleAndDoubleLayer(t *testing.T) {
rawSignature := testAnthropicNativeSignature(t)
doubleEncoded := base64.StdEncoding.EncodeToString([]byte(rawSignature))
inputJSON := []byte(`{
"messages": [
{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "one", "signature": "` + rawSignature + `"},
{"type": "thinking", "thinking": "two", "signature": "claude#` + doubleEncoded + `"}
]
}
]
}`)
if err := ValidateClaudeBypassSignatures(inputJSON); err != nil {
t.Fatalf("ValidateBypassModeSignatures returned error: %v", err)
}
}
func TestValidateBypassMode_RejectsGeminiSignature(t *testing.T) {
inputJSON := []byte(`{
"messages": [
{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "one", "signature": "` + testGeminiRawSignature(t) + `"}
]
}
]
}`)
err := ValidateClaudeBypassSignatures(inputJSON)
if err == nil {
t.Fatal("expected Gemini signature to be rejected")
}
}
func TestValidateBypassMode_RejectsMissingSignature(t *testing.T) {
inputJSON := []byte(`{
"messages": [
{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "one"}
]
}
]
}`)
err := ValidateClaudeBypassSignatures(inputJSON)
if err == nil {
t.Fatal("expected missing signature to be rejected")
}
if !strings.Contains(err.Error(), "missing thinking signature") {
t.Fatalf("expected missing signature message, got: %v", err)
}
}
func TestValidateBypassMode_RejectsNonREPrefix(t *testing.T) {
inputJSON := []byte(`{
"messages": [
{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "one", "signature": "` + testNonAnthropicRawSignature(t) + `"}
]
}
]
}`)
err := ValidateClaudeBypassSignatures(inputJSON)
if err == nil {
t.Fatal("expected non-R/E signature to be rejected")
}
}
func TestValidateBypassMode_RejectsEPrefixWrongFirstByte(t *testing.T) {
t.Parallel()
payload := append([]byte{0x10}, bytes.Repeat([]byte{0x34}, 48)...)
sig := base64.StdEncoding.EncodeToString(payload)
if sig[0] != 'E' {
t.Fatalf("test setup: expected E prefix, got %c", sig[0])
}
inputJSON := []byte(`{
"messages": [{"role": "assistant", "content": [
{"type": "thinking", "thinking": "t", "signature": "` + sig + `"}
]}]
}`)
err := ValidateClaudeBypassSignatures(inputJSON)
if err == nil {
t.Fatal("expected E-prefix with wrong first byte (0x10) to be rejected")
}
if !strings.Contains(err.Error(), "0x10") {
t.Fatalf("expected error to mention 0x10, got: %v", err)
}
}
func TestValidateBypassMode_RejectsTopLevel12WithoutClaudeTree(t *testing.T) {
previous := cache.SignatureBypassStrictMode()
cache.SetSignatureBypassStrictMode(true)
t.Cleanup(func() {
cache.SetSignatureBypassStrictMode(previous)
})
payload := append([]byte{0x12}, bytes.Repeat([]byte{0x34}, 48)...)
sig := base64.StdEncoding.EncodeToString(payload)
inputJSON := []byte(`{
"messages": [{"role": "assistant", "content": [
{"type": "thinking", "thinking": "t", "signature": "` + sig + `"}
]}]
}`)
err := ValidateClaudeBypassSignatures(inputJSON)
if err == nil {
t.Fatal("expected non-Claude protobuf tree to be rejected in strict mode")
}
if !strings.Contains(err.Error(), "malformed protobuf") && !strings.Contains(err.Error(), "Field 2") {
t.Fatalf("expected protobuf tree error, got: %v", err)
}
}
func TestValidateBypassMode_NonStrictAccepts12WithoutClaudeTree(t *testing.T) {
previous := cache.SignatureBypassStrictMode()
cache.SetSignatureBypassStrictMode(false)
t.Cleanup(func() {
cache.SetSignatureBypassStrictMode(previous)
})
payload := append([]byte{0x12}, bytes.Repeat([]byte{0x34}, 48)...)
sig := base64.StdEncoding.EncodeToString(payload)
inputJSON := []byte(`{
"messages": [{"role": "assistant", "content": [
{"type": "thinking", "thinking": "t", "signature": "` + sig + `"}
]}]
}`)
err := ValidateClaudeBypassSignatures(inputJSON)
if err != nil {
t.Fatalf("non-strict mode should accept 0x12 without protobuf tree, got: %v", err)
}
}
func TestValidateBypassMode_RejectsRPrefixInnerNotE(t *testing.T) {
t.Parallel()
inner := "F" + strings.Repeat("a", 60)
outer := base64.StdEncoding.EncodeToString([]byte(inner))
if outer[0] != 'R' {
t.Fatalf("test setup: expected R prefix, got %c", outer[0])
}
inputJSON := []byte(`{
"messages": [{"role": "assistant", "content": [
{"type": "thinking", "thinking": "t", "signature": "` + outer + `"}
]}]
}`)
err := ValidateClaudeBypassSignatures(inputJSON)
if err == nil {
t.Fatal("expected R-prefix with non-E inner to be rejected")
}
}
func TestValidateBypassMode_RejectsInvalidBase64(t *testing.T) {
t.Parallel()
tests := []struct {
name string
sig string
}{
{"E invalid", "E!!!invalid!!!"},
{"R invalid", "R$$$invalid$$$"},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
inputJSON := []byte(`{
"messages": [{"role": "assistant", "content": [
{"type": "thinking", "thinking": "t", "signature": "` + tt.sig + `"}
]}]
}`)
err := ValidateClaudeBypassSignatures(inputJSON)
if err == nil {
t.Fatal("expected invalid base64 to be rejected")
}
if !strings.Contains(err.Error(), "base64") {
t.Fatalf("expected base64 error, got: %v", err)
}
})
}
}
func TestValidateBypassMode_RejectsPrefixStrippedToEmpty(t *testing.T) {
t.Parallel()
tests := []struct {
name string
sig string
}{
{"prefix only", "claude#"},
{"prefix with spaces", "claude# "},
{"hash only", "#"},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
inputJSON := []byte(`{
"messages": [{"role": "assistant", "content": [
{"type": "thinking", "thinking": "t", "signature": "` + tt.sig + `"}
]}]
}`)
err := ValidateClaudeBypassSignatures(inputJSON)
if err == nil {
t.Fatal("expected prefix-only signature to be rejected")
}
})
}
}
func TestValidateBypassMode_HandlesMultipleHashMarks(t *testing.T) {
t.Parallel()
rawSignature := testAnthropicNativeSignature(t)
sig := "claude#" + rawSignature + "#extra"
inputJSON := []byte(`{
"messages": [{"role": "assistant", "content": [
{"type": "thinking", "thinking": "t", "signature": "` + sig + `"}
]}]
}`)
err := ValidateClaudeBypassSignatures(inputJSON)
if err == nil {
t.Fatal("expected signature with trailing # to be rejected (invalid base64)")
}
}
func TestValidateBypassMode_HandlesWhitespace(t *testing.T) {
t.Parallel()
rawSignature := testAnthropicNativeSignature(t)
tests := []struct {
name string
sig string
}{
{"leading space", " " + rawSignature},
{"trailing space", rawSignature + " "},
{"both spaces", " " + rawSignature + " "},
{"leading tab", "\t" + rawSignature},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
inputJSON := []byte(`{
"messages": [{"role": "assistant", "content": [
{"type": "thinking", "thinking": "t", "signature": "` + tt.sig + `"}
]}]
}`)
if err := ValidateClaudeBypassSignatures(inputJSON); err != nil {
t.Fatalf("expected whitespace-padded signature to be accepted, got: %v", err)
}
})
}
}
func TestValidateBypassMode_RejectsOversizedSignature(t *testing.T) {
t.Parallel()
payload := append([]byte{0x12}, bytes.Repeat([]byte{0x34}, maxBypassSignatureLen)...)
sig := base64.StdEncoding.EncodeToString(payload)
if len(sig) <= maxBypassSignatureLen {
t.Fatalf("test setup: signature should exceed max length, got %d", len(sig))
}
inputJSON := []byte(`{
"messages": [{"role": "assistant", "content": [
{"type": "thinking", "thinking": "t", "signature": "` + sig + `"}
]}]
}`)
err := ValidateClaudeBypassSignatures(inputJSON)
if err == nil {
t.Fatal("expected oversized signature to be rejected")
}
if !strings.Contains(err.Error(), "maximum length") {
t.Fatalf("expected length error, got: %v", err)
}
}
func TestResolveBypassModeSignature_TrimsWhitespace(t *testing.T) {
previous := cache.SignatureCacheEnabled()
cache.SetSignatureCacheEnabled(false)
t.Cleanup(func() {
cache.SetSignatureCacheEnabled(previous)
})
rawSignature := testAnthropicNativeSignature(t)
expected := resolveBypassModeSignature(rawSignature)
if expected == "" {
t.Fatal("test setup: expected non-empty normalized signature")
}
got := resolveBypassModeSignature(rawSignature + " ")
if got != expected {
t.Fatalf("expected trailing whitespace to be trimmed:\n got: %q\n want: %q", got, expected)
}
}
func TestConvertClaudeRequestToAntigravity_BypassModeNormalizesESignature(t *testing.T) {
cache.ClearSignatureCache("")
previous := cache.SignatureCacheEnabled()
cache.SetSignatureCacheEnabled(false)
t.Cleanup(func() {
cache.SetSignatureCacheEnabled(previous)
cache.ClearSignatureCache("")
})
thinkingText := "Let me think..."
cachedSignature := "cachedSignature1234567890123456789012345678901234567890123"
rawSignature := testAnthropicNativeSignature(t)
expectedSignature := base64.StdEncoding.EncodeToString([]byte(rawSignature))
cache.CacheSignature("claude-sonnet-4-5-thinking", thinkingText, cachedSignature)
inputJSON := []byte(`{
"model": "claude-sonnet-4-5-thinking",
"messages": [
{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "` + thinkingText + `", "signature": "` + rawSignature + `"},
{"type": "text", "text": "Answer"}
]
}
]
}`)
output := ConvertClaudeRequestToAntigravity("claude-sonnet-4-5-thinking", inputJSON, false)
outputStr := string(output)
part := gjson.Get(outputStr, "request.contents.0.parts.0")
if part.Get("thoughtSignature").String() != expectedSignature {
t.Fatalf("Expected bypass-mode signature '%s', got '%s'", expectedSignature, part.Get("thoughtSignature").String())
}
if part.Get("thoughtSignature").String() == cachedSignature {
t.Fatal("Bypass mode should not reuse cached signature")
}
}
func TestConvertClaudeRequestToAntigravity_BypassModePreservesShortValidSignature(t *testing.T) {
cache.ClearSignatureCache("")
previous := cache.SignatureCacheEnabled()
cache.SetSignatureCacheEnabled(false)
t.Cleanup(func() {
cache.SetSignatureCacheEnabled(previous)
cache.ClearSignatureCache("")
})
rawSignature := testMinimalAnthropicSignature(t)
expectedSignature := base64.StdEncoding.EncodeToString([]byte(rawSignature))
inputJSON := []byte(`{
"model": "claude-sonnet-4-5-thinking",
"messages": [
{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "tiny", "signature": "` + rawSignature + `"},
{"type": "text", "text": "Answer"}
]
}
]
}`)
output := ConvertClaudeRequestToAntigravity("claude-sonnet-4-5-thinking", inputJSON, false)
parts := gjson.GetBytes(output, "request.contents.0.parts").Array()
if len(parts) != 2 {
t.Fatalf("expected thinking part to be preserved in bypass mode, got %d parts", len(parts))
}
if parts[0].Get("thoughtSignature").String() != expectedSignature {
t.Fatalf("expected normalized short signature %q, got %q", expectedSignature, parts[0].Get("thoughtSignature").String())
}
if !parts[0].Get("thought").Bool() {
t.Fatalf("expected first part to remain a thought block, got %s", parts[0].Raw)
}
if parts[1].Get("text").String() != "Answer" {
t.Fatalf("expected trailing text part, got %s", parts[1].Raw)
}
if thoughtSig := gjson.GetBytes(output, "request.contents.0.parts.1.thoughtSignature").String(); thoughtSig != "" {
t.Fatalf("expected plain text part to have no thought signature, got %q", thoughtSig)
}
if functionSig := gjson.GetBytes(output, "request.contents.0.parts.0.functionCall.thoughtSignature").String(); functionSig != "" {
t.Fatalf("unexpected functionCall payload in thinking part: %q", functionSig)
}
}
func TestInspectClaudeSignaturePayload_ExtractsSpecTree(t *testing.T) {
t.Parallel()
payload := buildClaudeSignaturePayload(t, 12, uint64Ptr(2), "claude-sonnet-4-6", true)
tree, err := inspectClaudeSignaturePayload(payload, 1)
if err != nil {
t.Fatalf("expected structured Claude payload to parse, got: %v", err)
}
if tree.RoutingClass != "routing_class_12" {
t.Fatalf("routing_class = %q, want routing_class_12", tree.RoutingClass)
}
if tree.InfrastructureClass != "infra_google" {
t.Fatalf("infrastructure_class = %q, want infra_google", tree.InfrastructureClass)
}
if tree.SchemaFeatures != "extended_model_tagged_schema" {
t.Fatalf("schema_features = %q, want extended_model_tagged_schema", tree.SchemaFeatures)
}
if tree.ModelText != "claude-sonnet-4-6" {
t.Fatalf("model_text = %q, want claude-sonnet-4-6", tree.ModelText)
}
}
func TestInspectDoubleLayerSignature_TracksEncodingLayers(t *testing.T) {
t.Parallel()
inner := base64.StdEncoding.EncodeToString(buildClaudeSignaturePayload(t, 11, uint64Ptr(2), "", false))
outer := base64.StdEncoding.EncodeToString([]byte(inner))
tree, err := inspectDoubleLayerSignature(outer)
if err != nil {
t.Fatalf("expected double-layer Claude signature to parse, got: %v", err)
}
if tree.EncodingLayers != 2 {
t.Fatalf("encoding_layers = %d, want 2", tree.EncodingLayers)
}
if tree.LegacyRouteHint != "legacy_vertex_direct" {
t.Fatalf("legacy_route_hint = %q, want legacy_vertex_direct", tree.LegacyRouteHint)
}
}
func TestConvertClaudeRequestToAntigravity_CacheModeDropsRawSignature(t *testing.T) {
cache.ClearSignatureCache("")
previous := cache.SignatureCacheEnabled()
cache.SetSignatureCacheEnabled(true)
t.Cleanup(func() {
cache.SetSignatureCacheEnabled(previous)
cache.ClearSignatureCache("")
})
rawSignature := testAnthropicNativeSignature(t)
inputJSON := []byte(`{
"model": "claude-sonnet-4-5-thinking",
"messages": [
{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "Let me think...", "signature": "` + rawSignature + `"},
{"type": "text", "text": "Answer"}
]
}
]
}`)
output := ConvertClaudeRequestToAntigravity("claude-sonnet-4-5-thinking", inputJSON, false)
parts := gjson.GetBytes(output, "request.contents.0.parts").Array()
if len(parts) != 1 {
t.Fatalf("Expected raw signature thinking block to be dropped in cache mode, got %d parts", len(parts))
}
if parts[0].Get("text").String() != "Answer" {
t.Fatalf("Expected remaining text part, got %s", parts[0].Raw)
}
}
func TestConvertClaudeRequestToAntigravity_BypassModeDropsInvalidSignature(t *testing.T) {
cache.ClearSignatureCache("")
previous := cache.SignatureCacheEnabled()
cache.SetSignatureCacheEnabled(false)
t.Cleanup(func() {
cache.SetSignatureCacheEnabled(previous)
cache.ClearSignatureCache("")
})
invalidRawSignature := testNonAnthropicRawSignature(t)
inputJSON := []byte(`{
"model": "claude-sonnet-4-5-thinking",
"messages": [
{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "Let me think...", "signature": "` + invalidRawSignature + `"},
{"type": "text", "text": "Answer"}
]
}
]
}`)
output := ConvertClaudeRequestToAntigravity("claude-sonnet-4-5-thinking", inputJSON, false)
outputStr := string(output)
parts := gjson.Get(outputStr, "request.contents.0.parts").Array()
if len(parts) != 1 {
t.Fatalf("Expected invalid thinking block to be removed, got %d parts", len(parts))
}
if parts[0].Get("text").String() != "Answer" {
t.Fatalf("Expected remaining text part, got %s", parts[0].Raw)
}
if parts[0].Get("thought").Bool() {
t.Fatal("Invalid raw signature should not preserve thinking block")
}
}
func TestConvertClaudeRequestToAntigravity_BypassModeDropsGeminiSignature(t *testing.T) {
cache.ClearSignatureCache("")
previous := cache.SignatureCacheEnabled()
cache.SetSignatureCacheEnabled(false)
t.Cleanup(func() {
cache.SetSignatureCacheEnabled(previous)
cache.ClearSignatureCache("")
})
geminiPayload := append([]byte{0x0A}, bytes.Repeat([]byte{0x56}, 48)...)
geminiSig := base64.StdEncoding.EncodeToString(geminiPayload)
inputJSON := []byte(`{
"model": "claude-sonnet-4-5-thinking",
"messages": [
{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "hmm", "signature": "` + geminiSig + `"},
{"type": "text", "text": "Answer"}
]
}
]
}`)
output := ConvertClaudeRequestToAntigravity("claude-sonnet-4-5-thinking", inputJSON, false)
parts := gjson.GetBytes(output, "request.contents.0.parts").Array()
if len(parts) != 1 {
t.Fatalf("expected Gemini-signed thinking block to be dropped, got %d parts", len(parts))
}
if parts[0].Get("text").String() != "Answer" {
t.Fatalf("expected remaining text part, got %s", parts[0].Raw)
}
}
func TestConvertClaudeRequestToAntigravity_ThinkingBlockWithoutSignature(t *testing.T) {
cache.ClearSignatureCache("")

View File

@@ -9,6 +9,7 @@ package claude
import (
"bytes"
"context"
"encoding/base64"
"fmt"
"strings"
"sync/atomic"
@@ -23,6 +24,33 @@ import (
"github.com/tidwall/sjson"
)
// decodeSignature decodes R... (2-layer Base64) to E... (1-layer Base64, Anthropic format).
// Returns empty string if decoding fails (skip invalid signatures).
func decodeSignature(signature string) string {
if signature == "" {
return signature
}
if strings.HasPrefix(signature, "R") {
decoded, err := base64.StdEncoding.DecodeString(signature)
if err != nil {
log.Warnf("antigravity claude response: failed to decode signature, skipping")
return ""
}
return string(decoded)
}
return signature
}
func formatClaudeSignatureValue(modelName, signature string) string {
if cache.SignatureCacheEnabled() {
return fmt.Sprintf("%s#%s", cache.GetModelGroup(modelName), signature)
}
if cache.GetModelGroup(modelName) == "claude" {
return decodeSignature(signature)
}
return signature
}
// Params holds parameters for response conversion and maintains state across streaming chunks.
// This structure tracks the current state of the response translation process to ensure
// proper sequencing of SSE events and transitions between different content types.
@@ -144,13 +172,30 @@ func ConvertAntigravityResponseToClaude(_ context.Context, _ string, originalReq
if thoughtSignature := partResult.Get("thoughtSignature"); thoughtSignature.Exists() && thoughtSignature.String() != "" {
// log.Debug("Branch: signature_delta")
// Flush co-located text before emitting the signature
if partText := partTextResult.String(); partText != "" {
if params.ResponseType != 2 {
if params.ResponseType != 0 {
appendEvent("content_block_stop", fmt.Sprintf(`{"type":"content_block_stop","index":%d}`, params.ResponseIndex))
params.ResponseIndex++
}
appendEvent("content_block_start", fmt.Sprintf(`{"type":"content_block_start","index":%d,"content_block":{"type":"thinking","thinking":""}}`, params.ResponseIndex))
params.ResponseType = 2
params.CurrentThinkingText.Reset()
}
params.CurrentThinkingText.WriteString(partText)
data, _ := sjson.SetBytes([]byte(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"thinking_delta","thinking":""}}`, params.ResponseIndex)), "delta.thinking", partText)
appendEvent("content_block_delta", string(data))
}
if params.CurrentThinkingText.Len() > 0 {
cache.CacheSignature(modelName, params.CurrentThinkingText.String(), thoughtSignature.String())
// log.Debugf("Cached signature for thinking block (textLen=%d)", params.CurrentThinkingText.Len())
params.CurrentThinkingText.Reset()
}
data, _ := sjson.SetBytes([]byte(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"signature_delta","signature":""}}`, params.ResponseIndex)), "delta.signature", fmt.Sprintf("%s#%s", cache.GetModelGroup(modelName), thoughtSignature.String()))
sigValue := formatClaudeSignatureValue(modelName, thoughtSignature.String())
data, _ := sjson.SetBytes([]byte(fmt.Sprintf(`{"type":"content_block_delta","index":%d,"delta":{"type":"signature_delta","signature":""}}`, params.ResponseIndex)), "delta.signature", sigValue)
appendEvent("content_block_delta", string(data))
params.HasContent = true
} else if params.ResponseType == 2 { // Continue existing thinking block if already in thinking state
@@ -419,7 +464,8 @@ func ConvertAntigravityResponseToClaudeNonStream(_ context.Context, _ string, or
block := []byte(`{"type":"thinking","thinking":""}`)
block, _ = sjson.SetBytes(block, "thinking", thinkingBuilder.String())
if thinkingSignature != "" {
block, _ = sjson.SetBytes(block, "signature", fmt.Sprintf("%s#%s", cache.GetModelGroup(modelName), thinkingSignature))
sigValue := formatClaudeSignatureValue(modelName, thinkingSignature)
block, _ = sjson.SetBytes(block, "signature", sigValue)
}
responseJSON, _ = sjson.SetRawBytes(responseJSON, "content.-1", block)
thinkingBuilder.Reset()

View File

@@ -1,6 +1,7 @@
package claude
import (
"bytes"
"context"
"strings"
"testing"
@@ -244,3 +245,105 @@ func TestConvertAntigravityResponseToClaude_MultipleThinkingBlocks(t *testing.T)
t.Error("Second thinking block signature should be cached")
}
}
func TestConvertAntigravityResponseToClaude_TextAndSignatureInSameChunk(t *testing.T) {
cache.ClearSignatureCache("")
requestJSON := []byte(`{
"model": "claude-sonnet-4-5-thinking",
"messages": [{"role": "user", "content": [{"type": "text", "text": "Test"}]}]
}`)
validSignature := "RtestSig1234567890123456789012345678901234567890123456789"
// Chunk 1: thinking text only (no signature)
chunk1 := []byte(`{
"response": {
"candidates": [{
"content": {
"parts": [{"text": "First part.", "thought": true}]
}
}]
}
}`)
// Chunk 2: thinking text AND signature in the same part
chunk2 := []byte(`{
"response": {
"candidates": [{
"content": {
"parts": [{"text": " Second part.", "thought": true, "thoughtSignature": "` + validSignature + `"}]
}
}]
}
}`)
var param any
ctx := context.Background()
result1 := ConvertAntigravityResponseToClaude(ctx, "claude-sonnet-4-5-thinking", requestJSON, requestJSON, chunk1, &param)
result2 := ConvertAntigravityResponseToClaude(ctx, "claude-sonnet-4-5-thinking", requestJSON, requestJSON, chunk2, &param)
allOutput := string(bytes.Join(result1, nil)) + string(bytes.Join(result2, nil))
// The text " Second part." must appear as a thinking_delta, not be silently dropped
if !strings.Contains(allOutput, "Second part.") {
t.Error("Text co-located with signature must be emitted as thinking_delta before the signature")
}
// The signature must also be emitted
if !strings.Contains(allOutput, "signature_delta") {
t.Error("Signature delta must still be emitted")
}
// Verify the cached signature covers the FULL text (both parts)
fullText := "First part. Second part."
cachedSig := cache.GetCachedSignature("claude-sonnet-4-5-thinking", fullText)
if cachedSig != validSignature {
t.Errorf("Cached signature should cover full text %q, got sig=%q", fullText, cachedSig)
}
}
func TestConvertAntigravityResponseToClaude_SignatureOnlyChunk(t *testing.T) {
cache.ClearSignatureCache("")
requestJSON := []byte(`{
"model": "claude-sonnet-4-5-thinking",
"messages": [{"role": "user", "content": [{"type": "text", "text": "Test"}]}]
}`)
validSignature := "RtestSig1234567890123456789012345678901234567890123456789"
// Chunk 1: thinking text
chunk1 := []byte(`{
"response": {
"candidates": [{
"content": {
"parts": [{"text": "Full thinking text.", "thought": true}]
}
}]
}
}`)
// Chunk 2: signature only (empty text) — the normal case
chunk2 := []byte(`{
"response": {
"candidates": [{
"content": {
"parts": [{"text": "", "thought": true, "thoughtSignature": "` + validSignature + `"}]
}
}]
}
}`)
var param any
ctx := context.Background()
ConvertAntigravityResponseToClaude(ctx, "claude-sonnet-4-5-thinking", requestJSON, requestJSON, chunk1, &param)
ConvertAntigravityResponseToClaude(ctx, "claude-sonnet-4-5-thinking", requestJSON, requestJSON, chunk2, &param)
cachedSig := cache.GetCachedSignature("claude-sonnet-4-5-thinking", "Full thinking text.")
if cachedSig != validSignature {
t.Errorf("Signature-only chunk should still cache correctly, got %q", cachedSig)
}
}

View File

@@ -0,0 +1,391 @@
// Claude thinking signature validation for Antigravity bypass mode.
//
// Spec reference: SIGNATURE-CHANNEL-SPEC.md
//
// # Encoding Detection (Spec §3)
//
// Claude signatures use base64 encoding in one or two layers. The raw string's
// first character determines the encoding depth — this is mathematically equivalent
// to the spec's "decode first, check byte" approach:
//
// - 'E' prefix → single-layer: payload[0]==0x12, first 6 bits = 000100 = base64 index 4 = 'E'
// - 'R' prefix → double-layer: inner[0]=='E' (0x45), first 6 bits = 010001 = base64 index 17 = 'R'
//
// All valid signatures are normalized to R-form (double-layer base64) before
// sending to the Antigravity backend.
//
// # Protobuf Structure (Spec §4.1, §4.2) — strict mode only
//
// After base64 decoding to raw bytes (first byte must be 0x12):
//
// Top-level protobuf
// ├── Field 2 (bytes): container ← extractBytesField(payload, 2)
// │ ├── Field 1 (bytes): channel block ← extractBytesField(container, 1)
// │ │ ├── Field 1 (varint): channel_id [required] → routing_class (11 | 12)
// │ │ ├── Field 2 (varint): infra [optional] → infrastructure_class (aws=1 | google=2)
// │ │ ├── Field 3 (varint): version=2 [skipped]
// │ │ ├── Field 5 (bytes): ECDSA sig [skipped, per Spec §11]
// │ │ ├── Field 6 (bytes): model_text [optional] → schema_features
// │ │ └── Field 7 (varint): unknown [optional] → schema_features
// │ ├── Field 2 (bytes): nonce 12B [skipped]
// │ ├── Field 3 (bytes): session 12B [skipped]
// │ ├── Field 4 (bytes): SHA-384 48B [skipped]
// │ └── Field 5 (bytes): metadata [skipped, per Spec §11]
// └── Field 3 (varint): =1 [skipped]
//
// # Output Dimensions (Spec §8)
//
// routing_class: routing_class_11 | routing_class_12 | unknown
// infrastructure_class: infra_default (absent) | infra_aws (1) | infra_google (2) | infra_unknown
// schema_features: compact_schema (len 70-72, no f6/f7) | extended_model_tagged_schema (f6 exists) | unknown
// legacy_route_hint: only for ch=11 — legacy_default_group | legacy_aws_group | legacy_vertex_direct/proxy
//
// # Compatibility
//
// Verified against all confirmed spec samples (Anthropic Max 20x, Azure, Vertex,
// Bedrock) and legacy ch=11 signatures. Both single-layer (E) and double-layer (R)
// encodings are supported. Historical cache-mode 'modelGroup#' prefixes are stripped.
package claude
import (
"encoding/base64"
"fmt"
"strings"
"unicode/utf8"
"github.com/router-for-me/CLIProxyAPI/v6/internal/cache"
"github.com/tidwall/gjson"
"google.golang.org/protobuf/encoding/protowire"
)
const maxBypassSignatureLen = 8192
type claudeSignatureTree struct {
EncodingLayers int
ChannelID uint64
Field2 *uint64
RoutingClass string
InfrastructureClass string
SchemaFeatures string
ModelText string
LegacyRouteHint string
HasField7 bool
}
func ValidateClaudeBypassSignatures(inputRawJSON []byte) error {
messages := gjson.GetBytes(inputRawJSON, "messages")
if !messages.IsArray() {
return nil
}
messageResults := messages.Array()
for i := 0; i < len(messageResults); i++ {
contentResults := messageResults[i].Get("content")
if !contentResults.IsArray() {
continue
}
parts := contentResults.Array()
for j := 0; j < len(parts); j++ {
part := parts[j]
if part.Get("type").String() != "thinking" {
continue
}
rawSignature := strings.TrimSpace(part.Get("signature").String())
if rawSignature == "" {
return fmt.Errorf("messages[%d].content[%d]: missing thinking signature", i, j)
}
if _, err := normalizeClaudeBypassSignature(rawSignature); err != nil {
return fmt.Errorf("messages[%d].content[%d]: %w", i, j, err)
}
}
}
return nil
}
func normalizeClaudeBypassSignature(rawSignature string) (string, error) {
sig := strings.TrimSpace(rawSignature)
if sig == "" {
return "", fmt.Errorf("empty signature")
}
if idx := strings.IndexByte(sig, '#'); idx >= 0 {
sig = strings.TrimSpace(sig[idx+1:])
}
if sig == "" {
return "", fmt.Errorf("empty signature after stripping prefix")
}
if len(sig) > maxBypassSignatureLen {
return "", fmt.Errorf("signature exceeds maximum length (%d bytes)", maxBypassSignatureLen)
}
switch sig[0] {
case 'R':
if err := validateDoubleLayerSignature(sig); err != nil {
return "", err
}
return sig, nil
case 'E':
if err := validateSingleLayerSignature(sig); err != nil {
return "", err
}
return base64.StdEncoding.EncodeToString([]byte(sig)), nil
default:
return "", fmt.Errorf("invalid signature: expected 'E' or 'R' prefix, got %q", string(sig[0]))
}
}
func validateDoubleLayerSignature(sig string) error {
decoded, err := base64.StdEncoding.DecodeString(sig)
if err != nil {
return fmt.Errorf("invalid double-layer signature: base64 decode failed: %w", err)
}
if len(decoded) == 0 {
return fmt.Errorf("invalid double-layer signature: empty after decode")
}
if decoded[0] != 'E' {
return fmt.Errorf("invalid double-layer signature: inner does not start with 'E', got 0x%02x", decoded[0])
}
return validateSingleLayerSignatureContent(string(decoded), 2)
}
func validateSingleLayerSignature(sig string) error {
return validateSingleLayerSignatureContent(sig, 1)
}
func validateSingleLayerSignatureContent(sig string, encodingLayers int) error {
decoded, err := base64.StdEncoding.DecodeString(sig)
if err != nil {
return fmt.Errorf("invalid single-layer signature: base64 decode failed: %w", err)
}
if len(decoded) == 0 {
return fmt.Errorf("invalid single-layer signature: empty after decode")
}
if decoded[0] != 0x12 {
return fmt.Errorf("invalid Claude signature: expected first byte 0x12, got 0x%02x", decoded[0])
}
if !cache.SignatureBypassStrictMode() {
return nil
}
_, err = inspectClaudeSignaturePayload(decoded, encodingLayers)
return err
}
func inspectDoubleLayerSignature(sig string) (*claudeSignatureTree, error) {
decoded, err := base64.StdEncoding.DecodeString(sig)
if err != nil {
return nil, fmt.Errorf("invalid double-layer signature: base64 decode failed: %w", err)
}
if len(decoded) == 0 {
return nil, fmt.Errorf("invalid double-layer signature: empty after decode")
}
if decoded[0] != 'E' {
return nil, fmt.Errorf("invalid double-layer signature: inner does not start with 'E', got 0x%02x", decoded[0])
}
return inspectSingleLayerSignatureWithLayers(string(decoded), 2)
}
func inspectSingleLayerSignature(sig string) (*claudeSignatureTree, error) {
return inspectSingleLayerSignatureWithLayers(sig, 1)
}
func inspectSingleLayerSignatureWithLayers(sig string, encodingLayers int) (*claudeSignatureTree, error) {
decoded, err := base64.StdEncoding.DecodeString(sig)
if err != nil {
return nil, fmt.Errorf("invalid single-layer signature: base64 decode failed: %w", err)
}
if len(decoded) == 0 {
return nil, fmt.Errorf("invalid single-layer signature: empty after decode")
}
return inspectClaudeSignaturePayload(decoded, encodingLayers)
}
func inspectClaudeSignaturePayload(payload []byte, encodingLayers int) (*claudeSignatureTree, error) {
if len(payload) == 0 {
return nil, fmt.Errorf("invalid Claude signature: empty payload")
}
if payload[0] != 0x12 {
return nil, fmt.Errorf("invalid Claude signature: expected first byte 0x12, got 0x%02x", payload[0])
}
container, err := extractBytesField(payload, 2, "top-level protobuf")
if err != nil {
return nil, err
}
channelBlock, err := extractBytesField(container, 1, "Claude Field 2 container")
if err != nil {
return nil, err
}
return inspectClaudeChannelBlock(channelBlock, encodingLayers)
}
func inspectClaudeChannelBlock(channelBlock []byte, encodingLayers int) (*claudeSignatureTree, error) {
tree := &claudeSignatureTree{
EncodingLayers: encodingLayers,
RoutingClass: "unknown",
InfrastructureClass: "infra_unknown",
SchemaFeatures: "unknown_schema_features",
}
haveChannelID := false
hasField6 := false
hasField7 := false
err := walkProtobufFields(channelBlock, func(num protowire.Number, typ protowire.Type, raw []byte) error {
switch num {
case 1:
if typ != protowire.VarintType {
return fmt.Errorf("invalid Claude signature: Field 2.1.1 channel_id must be varint")
}
channelID, err := decodeVarintField(raw, "Field 2.1.1 channel_id")
if err != nil {
return err
}
tree.ChannelID = channelID
haveChannelID = true
case 2:
if typ != protowire.VarintType {
return fmt.Errorf("invalid Claude signature: Field 2.1.2 field2 must be varint")
}
field2, err := decodeVarintField(raw, "Field 2.1.2 field2")
if err != nil {
return err
}
tree.Field2 = &field2
case 6:
if typ != protowire.BytesType {
return fmt.Errorf("invalid Claude signature: Field 2.1.6 model_text must be bytes")
}
modelBytes, err := decodeBytesField(raw, "Field 2.1.6 model_text")
if err != nil {
return err
}
if !utf8.Valid(modelBytes) {
return fmt.Errorf("invalid Claude signature: Field 2.1.6 model_text is not valid UTF-8")
}
tree.ModelText = string(modelBytes)
hasField6 = true
case 7:
if typ != protowire.VarintType {
return fmt.Errorf("invalid Claude signature: Field 2.1.7 must be varint")
}
if _, err := decodeVarintField(raw, "Field 2.1.7"); err != nil {
return err
}
hasField7 = true
tree.HasField7 = true
}
return nil
})
if err != nil {
return nil, err
}
if !haveChannelID {
return nil, fmt.Errorf("invalid Claude signature: missing Field 2.1.1 channel_id")
}
switch tree.ChannelID {
case 11:
tree.RoutingClass = "routing_class_11"
case 12:
tree.RoutingClass = "routing_class_12"
}
if tree.Field2 == nil {
tree.InfrastructureClass = "infra_default"
} else {
switch *tree.Field2 {
case 1:
tree.InfrastructureClass = "infra_aws"
case 2:
tree.InfrastructureClass = "infra_google"
default:
tree.InfrastructureClass = "infra_unknown"
}
}
switch {
case hasField6:
tree.SchemaFeatures = "extended_model_tagged_schema"
case !hasField6 && !hasField7 && len(channelBlock) >= 70 && len(channelBlock) <= 72:
tree.SchemaFeatures = "compact_schema"
}
if tree.ChannelID == 11 {
switch {
case tree.Field2 == nil:
tree.LegacyRouteHint = "legacy_default_group"
case *tree.Field2 == 1:
tree.LegacyRouteHint = "legacy_aws_group"
case *tree.Field2 == 2 && tree.EncodingLayers == 2:
tree.LegacyRouteHint = "legacy_vertex_direct"
case *tree.Field2 == 2 && tree.EncodingLayers == 1:
tree.LegacyRouteHint = "legacy_vertex_proxy"
}
}
return tree, nil
}
func extractBytesField(msg []byte, fieldNum protowire.Number, scope string) ([]byte, error) {
var value []byte
err := walkProtobufFields(msg, func(num protowire.Number, typ protowire.Type, raw []byte) error {
if num != fieldNum {
return nil
}
if typ != protowire.BytesType {
return fmt.Errorf("invalid Claude signature: %s field %d must be bytes", scope, fieldNum)
}
bytesValue, err := decodeBytesField(raw, fmt.Sprintf("%s field %d", scope, fieldNum))
if err != nil {
return err
}
value = bytesValue
return nil
})
if err != nil {
return nil, err
}
if value == nil {
return nil, fmt.Errorf("invalid Claude signature: missing %s field %d", scope, fieldNum)
}
return value, nil
}
func walkProtobufFields(msg []byte, visit func(num protowire.Number, typ protowire.Type, raw []byte) error) error {
for offset := 0; offset < len(msg); {
num, typ, n := protowire.ConsumeTag(msg[offset:])
if n < 0 {
return fmt.Errorf("invalid Claude signature: malformed protobuf tag: %w", protowire.ParseError(n))
}
offset += n
valueLen := protowire.ConsumeFieldValue(num, typ, msg[offset:])
if valueLen < 0 {
return fmt.Errorf("invalid Claude signature: malformed protobuf field %d: %w", num, protowire.ParseError(valueLen))
}
fieldRaw := msg[offset : offset+valueLen]
if err := visit(num, typ, fieldRaw); err != nil {
return err
}
offset += valueLen
}
return nil
}
func decodeVarintField(raw []byte, label string) (uint64, error) {
value, n := protowire.ConsumeVarint(raw)
if n < 0 {
return 0, fmt.Errorf("invalid Claude signature: failed to decode %s: %w", label, protowire.ParseError(n))
}
return value, nil
}
func decodeBytesField(raw []byte, label string) ([]byte, error) {
value, n := protowire.ConsumeBytes(raw)
if n < 0 {
return nil, fmt.Errorf("invalid Claude signature: failed to decode %s: %w", label, protowire.ParseError(n))
}
return value, nil
}