From 089d14a2e3fce50701d057c43b48882f349b0102 Mon Sep 17 00:00:00 2001 From: erio Date: Mon, 20 Apr 2026 00:47:03 +0800 Subject: [PATCH] fix(signature): isolate pool buckets by platform OAuth/setup-token accounts now use per-platform buckets (oauth:anthropic, oauth:antigravity) instead of a single shared "oauth" bucket. Prevents cross-contamination between platforms whose thinking signatures are incompatible. BucketFor() now takes a platform parameter. All call sites (harvester, rectifier factory, gateway service) updated. --- .../service/antigravity_gateway_service.go | 1 + backend/internal/service/gateway_service.go | 2 + backend/internal/service/signature/pool.go | 33 ++++++-------- .../service/signature/pool_rectifier.go | 4 +- .../internal/service/signature/rectifier.go | 2 + .../service/signature/rectifier_test.go | 44 +++++++++++-------- .../service/signature_rectifier_factory.go | 2 +- 7 files changed, 46 insertions(+), 42 deletions(-) diff --git a/backend/internal/service/antigravity_gateway_service.go b/backend/internal/service/antigravity_gateway_service.go index 1fd0eefd46..87513636c0 100644 --- a/backend/internal/service/antigravity_gateway_service.go +++ b/backend/internal/service/antigravity_gateway_service.go @@ -1486,6 +1486,7 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context, sigIn := signature.AntigravityInput{ AccountType: account.Type, AccountID: account.ID, + Platform: account.Platform, Request: &retryClaudeReq, } applied, proceedAfter, stripErr := rectifier.Apply(ctx, sigIn, stage) diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index a04d862576..f547689c81 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -4212,6 +4212,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A sigIn := signature.ClaudeInput{ AccountType: account.Type, AccountID: account.ID, + Platform: account.Platform, Body: body, } rectifier := s.signatureFactory.ForClaude(ctx, account) @@ -8374,6 +8375,7 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context, sigIn := signature.ClaudeInput{ AccountType: account.Type, AccountID: account.ID, + Platform: account.Platform, Body: body, } rectifier := s.signatureFactory.ForClaude(ctx, account) diff --git a/backend/internal/service/signature/pool.go b/backend/internal/service/signature/pool.go index 26c37e091b..80d45348c5 100644 --- a/backend/internal/service/signature/pool.go +++ b/backend/internal/service/signature/pool.go @@ -15,35 +15,28 @@ const ( accountTypeAPIKey = "apikey" ) -// Bucket keys used by the SignaturePool. -// - BucketOAuthShared: shared pool across all max OAuth and setup-token accounts -// - BucketAPIKey(id): per-account pool for API Key accounts -// -// Bedrock / Upstream / other types currently do not participate in the pool -// (they never emit thinking signatures we can reuse). -const ( - BucketOAuthShared = "oauth" -) - -// BucketAPIKey returns the bucket key for an API Key account. -func BucketAPIKey(accountID int64) string { - return fmt.Sprintf("apikey:%d", accountID) -} - -// BucketFor maps an account type+id to its pool bucket. -// Returns empty string for types that do not participate in the pool. // Account types that do not participate in the signature pool. const ( accountTypeBedrock = "bedrock" accountTypeUpstream = "upstream" + + platformAnthropicDirect = "anthropic" ) -func BucketFor(accountType string, accountID int64) string { +// BucketFor maps (accountType, platform, accountID) to a pool bucket key. +// OAuth/setup-token accounts share a per-platform pool (anthropic signatures +// must not mix with antigravity signatures). API Key accounts get a +// per-account pool regardless of platform. Returns "" for types that do not +// participate (bedrock, upstream). +func BucketFor(accountType, platform string, accountID int64) string { switch accountType { case accountTypeOAuth, accountTypeSetupToken: - return BucketOAuthShared + if platform == "" { + platform = platformAnthropicDirect + } + return fmt.Sprintf("oauth:%s", platform) case accountTypeAPIKey: - return BucketAPIKey(accountID) + return fmt.Sprintf("apikey:%d", accountID) case accountTypeBedrock, accountTypeUpstream: return "" } diff --git a/backend/internal/service/signature/pool_rectifier.go b/backend/internal/service/signature/pool_rectifier.go index be1e6389e6..f5b536fb58 100644 --- a/backend/internal/service/signature/pool_rectifier.go +++ b/backend/internal/service/signature/pool_rectifier.go @@ -26,7 +26,7 @@ func (r *PoolClaudeRectifier) Apply(ctx context.Context, in ClaudeInput, stage S if stage != StageThinkingOnly { return nil, false } - bucket := BucketFor(in.AccountType, in.AccountID) + bucket := BucketFor(in.AccountType, in.Platform, in.AccountID) if bucket == "" { return nil, false } @@ -65,7 +65,7 @@ func (r *PoolAntigravityRectifier) Apply(ctx context.Context, in AntigravityInpu if stage != StageThinkingOnly { return false, false, nil } - bucket := BucketFor(in.AccountType, in.AccountID) + bucket := BucketFor(in.AccountType, in.Platform, in.AccountID) if bucket == "" { return false, false, nil } diff --git a/backend/internal/service/signature/rectifier.go b/backend/internal/service/signature/rectifier.go index b2ba365694..0a3fe497cf 100644 --- a/backend/internal/service/signature/rectifier.go +++ b/backend/internal/service/signature/rectifier.go @@ -47,6 +47,7 @@ func (s Stage) Name() string { type ClaudeInput struct { AccountType string AccountID int64 + Platform string Body []byte LastErrMsg string // error message from the previous stage's 400 response; empty on first stage } @@ -65,6 +66,7 @@ type ClaudeRectifier interface { type AntigravityInput struct { AccountType string AccountID int64 + Platform string Request *antigravity.ClaudeRequest LastErrMsg string } diff --git a/backend/internal/service/signature/rectifier_test.go b/backend/internal/service/signature/rectifier_test.go index 844613a3f6..9d8b147eec 100644 --- a/backend/internal/service/signature/rectifier_test.go +++ b/backend/internal/service/signature/rectifier_test.go @@ -47,16 +47,22 @@ func (p *fakePool) Size(_ context.Context, bucket string) (int64, error) { } func TestBucketFor(t *testing.T) { - if got := BucketFor("oauth", 1); got != "oauth" { - t.Errorf("oauth bucket: got %q, want oauth", got) + if got := BucketFor("oauth", "anthropic", 1); got != "oauth:anthropic" { + t.Errorf("anthropic oauth: got %q, want oauth:anthropic", got) } - if got := BucketFor("setup-token", 1); got != "oauth" { - t.Errorf("setup-token must share oauth bucket, got %q", got) + if got := BucketFor("oauth", "antigravity", 1); got != "oauth:antigravity" { + t.Errorf("antigravity oauth: got %q, want oauth:antigravity", got) } - if got := BucketFor("apikey", 42); got != "apikey:42" { + if got := BucketFor("setup-token", "anthropic", 1); got != "oauth:anthropic" { + t.Errorf("setup-token must share platform oauth bucket, got %q", got) + } + if got := BucketFor("oauth", "", 1); got != "oauth:anthropic" { + t.Errorf("empty platform defaults to anthropic: got %q", got) + } + if got := BucketFor("apikey", "anthropic", 42); got != "apikey:42" { t.Errorf("apikey bucket: got %q, want apikey:42", got) } - if got := BucketFor("bedrock", 1); got != "" { + if got := BucketFor("bedrock", "", 1); got != "" { t.Errorf("bedrock has no pool: got %q, want empty", got) } } @@ -95,7 +101,7 @@ func TestStripClaudeRectifier_Stage2GatedOnToolError(t *testing.T) { func TestPoolClaudeRectifier_EmptyPoolSignalsAbort(t *testing.T) { pool := newFakePool() r := &PoolClaudeRectifier{Pool: pool, Capacity: 10} - in := ClaudeInput{AccountType: "oauth", AccountID: 1, Body: []byte(`{"messages":[]}`)} + in := ClaudeInput{AccountType: "oauth", AccountID: 1, Platform: "anthropic", Body: []byte(`{"messages":[]}`)} out, proceed := r.Apply(context.Background(), in, StageThinkingOnly) if proceed { t.Fatalf("empty pool must signal proceed=false (rule A)") @@ -107,10 +113,10 @@ func TestPoolClaudeRectifier_EmptyPoolSignalsAbort(t *testing.T) { func TestPoolClaudeRectifier_NoThinkingBlocksSignalsAbort(t *testing.T) { pool := newFakePool() - _ = pool.Add(context.Background(), "oauth", "good", time.Now(), 10) + _ = pool.Add(context.Background(), "oauth:anthropic", "good", time.Now(), 10) r := &PoolClaudeRectifier{Pool: pool, Capacity: 10} // Pool has a sig, but the request has no thinking blocks to replace. - in := ClaudeInput{AccountType: "oauth", AccountID: 1, Body: []byte(`{"messages":[{"role":"user","content":"hi"}]}`)} + in := ClaudeInput{AccountType: "oauth", AccountID: 1, Platform: "anthropic", Body: []byte(`{"messages":[{"role":"user","content":"hi"}]}`)} out, proceed := r.Apply(context.Background(), in, StageThinkingOnly) if proceed { t.Fatalf("replaced=0 must signal proceed=false") @@ -122,10 +128,10 @@ func TestPoolClaudeRectifier_NoThinkingBlocksSignalsAbort(t *testing.T) { func TestPoolClaudeRectifier_ReplacesAndStopsAtStage2(t *testing.T) { pool := newFakePool() - _ = pool.Add(context.Background(), "oauth", "good", time.Now(), 10) + _ = pool.Add(context.Background(), "oauth:anthropic", "good", time.Now(), 10) r := &PoolClaudeRectifier{Pool: pool, Capacity: 10} body := []byte(`{"messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"x","signature":"bad"}]}]}`) - in := ClaudeInput{AccountType: "oauth", AccountID: 1, Body: body} + in := ClaudeInput{AccountType: "oauth", AccountID: 1, Platform: "anthropic", Body: body} out, proceed := r.Apply(context.Background(), in, StageThinkingOnly) if !proceed { @@ -158,7 +164,7 @@ func TestPoolAntigravityRectifier_EmptyPool(t *testing.T) { {Role: "assistant", Content: json.RawMessage(`[{"type":"thinking","thinking":"x","signature":"bad"}]`)}, }, } - applied, proceed, err := r.Apply(context.Background(), AntigravityInput{AccountType: "oauth", AccountID: 1, Request: req}, StageThinkingOnly) + applied, proceed, err := r.Apply(context.Background(), AntigravityInput{AccountType: "oauth", AccountID: 1, Platform: "antigravity", Request: req}, StageThinkingOnly) if applied || proceed || err != nil { t.Errorf("empty pool: applied=%v proceed=%v err=%v; want all false/nil", applied, proceed, err) } @@ -166,14 +172,14 @@ func TestPoolAntigravityRectifier_EmptyPool(t *testing.T) { func TestPoolAntigravityRectifier_AppliedAndOneShot(t *testing.T) { pool := newFakePool() - _ = pool.Add(context.Background(), "oauth", "g0", time.Now(), 10) + _ = pool.Add(context.Background(), "oauth:antigravity", "g0", time.Now(), 10) r := &PoolAntigravityRectifier{Pool: pool, Capacity: 10} req := &antigravity.ClaudeRequest{ Messages: []antigravity.ClaudeMessage{ {Role: "assistant", Content: json.RawMessage(`[{"type":"thinking","thinking":"x","signature":"bad"}]`)}, }, } - applied, proceed, err := r.Apply(context.Background(), AntigravityInput{AccountType: "oauth", AccountID: 1, Request: req}, StageThinkingOnly) + applied, proceed, err := r.Apply(context.Background(), AntigravityInput{AccountType: "oauth", AccountID: 1, Platform: "antigravity", Request: req}, StageThinkingOnly) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -199,7 +205,7 @@ func (errPool) Size(context.Context, string) (int64, error) { return 0, errors.N func TestPoolClaudeRectifier_PoolErrorIsTreatedAsEmpty(t *testing.T) { r := &PoolClaudeRectifier{Pool: errPool{}, Capacity: 10} - in := ClaudeInput{AccountType: "oauth", AccountID: 1, Body: []byte(`{"messages":[]}`)} + in := ClaudeInput{AccountType: "oauth", AccountID: 1, Platform: "anthropic", Body: []byte(`{"messages":[]}`)} _, proceed := r.Apply(context.Background(), in, StageThinkingOnly) if proceed { t.Errorf("pool error must behave like empty pool (proceed=false)") @@ -208,9 +214,9 @@ func TestPoolClaudeRectifier_PoolErrorIsTreatedAsEmpty(t *testing.T) { func TestPoolAntigravityRectifier_NilRequestNoop(t *testing.T) { pool := newFakePool() - _ = pool.Add(context.Background(), "oauth", "g", time.Now(), 10) + _ = pool.Add(context.Background(), "oauth:antigravity", "g", time.Now(), 10) r := &PoolAntigravityRectifier{Pool: pool, Capacity: 10} - applied, proceed, err := r.Apply(context.Background(), AntigravityInput{AccountType: "oauth", AccountID: 1, Request: nil}, StageThinkingOnly) + applied, proceed, err := r.Apply(context.Background(), AntigravityInput{AccountType: "oauth", AccountID: 1, Platform: "antigravity", Request: nil}, StageThinkingOnly) if applied || proceed || err != nil { t.Errorf("nil request: applied=%v proceed=%v err=%v; want all false/nil", applied, proceed, err) } @@ -218,7 +224,7 @@ func TestPoolAntigravityRectifier_NilRequestNoop(t *testing.T) { func TestPoolClaudeRectifier_NilReceiverNoPanic(t *testing.T) { var r *PoolClaudeRectifier - _, proceed := r.Apply(context.Background(), ClaudeInput{AccountType: "oauth", AccountID: 1, Body: []byte(`{}`)}, StageThinkingOnly) + _, proceed := r.Apply(context.Background(), ClaudeInput{AccountType: "oauth", AccountID: 1, Platform: "anthropic", Body: []byte(`{}`)}, StageThinkingOnly) if proceed { t.Errorf("nil receiver must return proceed=false") } @@ -235,7 +241,7 @@ func TestStripAntigravityRectifier_BothStages(t *testing.T) { t.Fatalf("expected 2 stages, got %d", len(stages)) } for i, stage := range stages { - in := AntigravityInput{AccountType: "oauth", AccountID: 1, Request: &antigravity.ClaudeRequest{}} + in := AntigravityInput{AccountType: "oauth", AccountID: 1, Platform: "antigravity", Request: &antigravity.ClaudeRequest{}} applied, proceed, err := r.Apply(context.Background(), in, stage) if err != nil || !applied || !proceed { t.Errorf("stage %d: applied=%v proceed=%v err=%v", i, applied, proceed, err) diff --git a/backend/internal/service/signature_rectifier_factory.go b/backend/internal/service/signature_rectifier_factory.go index 07d3fb4a9f..2a6afe7313 100644 --- a/backend/internal/service/signature_rectifier_factory.go +++ b/backend/internal/service/signature_rectifier_factory.go @@ -112,7 +112,7 @@ func (f *signatureRectifierFactory) WrapResponseBody(ctx context.Context, accoun } h := signature.NewHarvester(f.pool, f.poolCapacity(ctx)) return h.Wrap(ctx, body, signature.HarvestOptions{ - Bucket: signature.BucketFor(account.Type, account.ID), + Bucket: signature.BucketFor(account.Type, account.Platform, account.ID), Streaming: streaming, Skip: func() bool { v, _ := ctx.Value(ctxkey.IsSignatureRectifyRetry).(bool)