opensource/sub2api
3
0
Fork 0
mirror of https://github.com/Wei-Shaw/sub2api.git synced 2026-09-24 16:05:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(signature): isolate pool buckets by platform

Browse Source 
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.
This commit is contained in:
erio
2026-04-20 00:47:03 +08:00
parent 53b88213f0
commit 089d14a2e3
7 changed files with 46 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/service/antigravity_gateway_service.go
+1
View File
@@ -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)
backend/internal/service/gateway_service.go
+2
View File
@@ -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)
backend/internal/service/signature/pool.go
+13 -20
View File
@@ -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 ""
}
backend/internal/service/signature/pool_rectifier.go
+2 -2
View File
@@ -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
}
backend/internal/service/signature/rectifier.go
+2
View File
@@ -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
}
backend/internal/service/signature/rectifier_test.go
+25 -19
View File
@@ -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)
backend/internal/service/signature_rectifier_factory.go
+1 -1
View File
@@ -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)