From b31b435091709920c87c2b5d4d66bf9e8d25134b Mon Sep 17 00:00:00 2001 From: feeeei Date: Thu, 24 Sep 2026 16:20:16 +0800 Subject: [PATCH] =?UTF-8?q?fix(gateway):=20Claude=20Code=20=E9=99=90?= =?UTF-8?q?=E5=88=B6=E5=88=86=E7=BB=84=E7=9A=84=20OpenAI=20=E5=85=BC?= =?UTF-8?q?=E5=AE=B9=E5=85=A5=E5=8F=A3=E6=94=B9=E8=B5=B0=E9=99=8D=E7=BA=A7?= =?UTF-8?q?=E5=88=86=E7=BB=84,=E4=B8=8D=E5=86=8D=E7=9B=B4=E6=8E=A5=20403?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 分组开启 claude_code_only 并配置降级分组时,/v1/messages 会在选号入口的 checkClaudeCodeRestriction 中把非 Claude Code 请求解析到降级分组;但 GatewayHandler.ChatCompletions / Responses 在 handler 层只要分组开了 claude_code_only 就直接返回 403,不看降级分组,请求到不了调度层:只会发 Chat Completions / Responses 的第三方客户端拿这类分组的 Key 一律 403。 - handler 层提前拒绝只在未配置降级分组时生效,拒绝的状态码与文案不变 - 配置了降级分组时请求继续,由 SelectAccountWithLoadAwareness 解析到降级分组 选号,再经已有的 Chat Completions / Responses ⇄ Anthropic Messages 转换转发; service 层无改动 - 与 /v1/messages 的降级语义一致:选号与渠道限制预检按降级分组,计费与用量 记录仍按 API Key 所属分组 新增 handler 单测覆盖两个入口在有/无降级分组时的行为。 --- .../gateway_handler_chat_completions.go | 7 +- ...teway_handler_claude_code_fallback_test.go | 181 ++++++++++++++++++ .../handler/gateway_handler_responses.go | 12 +- 3 files changed, 191 insertions(+), 9 deletions(-) create mode 100644 backend/internal/handler/gateway_handler_claude_code_fallback_test.go diff --git a/backend/internal/handler/gateway_handler_chat_completions.go b/backend/internal/handler/gateway_handler_chat_completions.go index 7a98504596..6a5fd0160f 100644 --- a/backend/internal/handler/gateway_handler_chat_completions.go +++ b/backend/internal/handler/gateway_handler_chat_completions.go @@ -101,8 +101,11 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) { // 解析渠道级模型映射 channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, reqModel) - // Claude Code only restriction - if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly { + // Claude Code only restriction: /v1/chat/completions is never a Claude Code + // endpoint. With a fallback group the request continues and account selection + // (checkClaudeCodeRestriction) schedules it in the fallback group; without one + // it is rejected here. + if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly && apiKey.Group.FallbackGroupID == nil { h.chatCompletionsErrorResponse(c, http.StatusForbidden, "permission_error", "This group is restricted to Claude Code clients (/v1/messages only)") return diff --git a/backend/internal/handler/gateway_handler_claude_code_fallback_test.go b/backend/internal/handler/gateway_handler_claude_code_fallback_test.go new file mode 100644 index 0000000000..ae15cd887d --- /dev/null +++ b/backend/internal/handler/gateway_handler_claude_code_fallback_test.go @@ -0,0 +1,181 @@ +//go:build unit + +package handler + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + "sync" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" + middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +// groupScopedSchedulerCache 按桶的分组只返回该分组的成员账号，并记录被查询过的分组。 +type groupScopedSchedulerCache struct { + *fakeSchedulerCache + mu sync.Mutex + groupIDs []int64 +} + +func (c *groupScopedSchedulerCache) GetSnapshot(_ context.Context, bucket service.SchedulerBucket) ([]*service.Account, bool, error) { + c.mu.Lock() + c.groupIDs = append(c.groupIDs, bucket.GroupID) + c.mu.Unlock() + var members []*service.Account + for _, account := range c.accounts { + for _, ag := range account.AccountGroups { + if ag.GroupID == bucket.GroupID { + members = append(members, account) + break + } + } + } + return members, true, nil +} + +func (c *groupScopedSchedulerCache) queriedGroupIDs() []int64 { + c.mu.Lock() + defer c.mu.Unlock() + return append([]int64(nil), c.groupIDs...) +} + +type groupMapRepo struct { + *fakeGroupRepo + groups map[int64]*service.Group +} + +func (r *groupMapRepo) GetByID(_ context.Context, id int64) (*service.Group, error) { + if group, ok := r.groups[id]; ok { + return group, nil + } + return nil, service.ErrGroupNotFound +} + +func (r *groupMapRepo) GetByIDLite(ctx context.Context, id int64) (*service.Group, error) { + return r.GetByID(ctx, id) +} + +func TestGatewayOpenAICompatibleHandlersClaudeCodeOnlyFallback(t *testing.T) { + gin.SetMode(gin.TestMode) + + const ( + primaryGroupID = int64(9300) + fallbackGroupID = int64(9301) + primaryAccountID = int64(9310) + fallbackAccountID = int64(9311) + ) + + // 账号不带 api_key：转发在取令牌时失败，请求不会触达上游，断言只看选号结果。 + newAccount := func(id, groupID int64) *service.Account { + return &service.Account{ + ID: id, Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey, + Status: service.StatusActive, Schedulable: true, Concurrency: 1, + AccountGroups: []service.AccountGroup{{AccountID: id, GroupID: groupID}}, + } + } + + endpoints := []struct { + name string + path string + body string + call func(*GatewayHandler, *gin.Context) + }{ + { + name: "responses", path: "/v1/responses", + body: `{"model":"claude-sonnet-4-5","input":"hello","stream":false}`, + call: (*GatewayHandler).Responses, + }, + { + name: "chat completions", path: "/v1/chat/completions", + body: `{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hello"}],"stream":false}`, + call: (*GatewayHandler).ChatCompletions, + }, + } + + for _, ep := range endpoints { + for _, tc := range []struct { + name string + hasFallback bool + }{ + {name: "with fallback group", hasFallback: true}, + {name: "without fallback group", hasFallback: false}, + } { + t.Run(ep.name+"/"+tc.name, func(t *testing.T) { + primary := &service.Group{ + ID: primaryGroupID, Hydrated: true, Platform: service.PlatformAnthropic, + Status: service.StatusActive, ClaudeCodeOnly: true, + } + if tc.hasFallback { + fallbackID := fallbackGroupID + primary.FallbackGroupID = &fallbackID + } + fallback := &service.Group{ + ID: fallbackGroupID, Hydrated: true, Platform: service.PlatformAnthropic, + Status: service.StatusActive, + } + + schedulerCache := &groupScopedSchedulerCache{fakeSchedulerCache: &fakeSchedulerCache{accounts: []*service.Account{ + newAccount(primaryAccountID, primaryGroupID), + newAccount(fallbackAccountID, fallbackGroupID), + }}} + gatewayService := service.NewGatewayService( + nil, &groupMapRepo{fakeGroupRepo: &fakeGroupRepo{}, groups: map[int64]*service.Group{ + primaryGroupID: primary, + fallbackGroupID: fallback, + }}, nil, nil, nil, nil, nil, nil, nil, + service.NewSchedulerSnapshotService(schedulerCache, nil, nil, nil, nil), + nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, + ) + cfg := &config.Config{RunMode: config.RunModeSimple} + billingCacheService := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(billingCacheService.Stop) + h := &GatewayHandler{ + gatewayService: gatewayService, + billingCacheService: billingCacheService, + concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(&fakeConcurrencyCache{}), SSEPingFormatClaude, 0), + maxAccountSwitches: 1, + cfg: cfg, + } + + primaryGroupIDRef := primaryGroupID + apiKey := &service.APIKey{ + ID: 9320, UserID: 9330, GroupID: &primaryGroupIDRef, Group: primary, Status: service.StatusActive, + User: &service.User{ID: 9330, Concurrency: 10, Balance: 100}, + } + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + ctx := context.WithValue(context.Background(), ctxkey.Group, primary) + req := httptest.NewRequest(http.MethodPost, ep.path, bytes.NewBufferString(ep.body)).WithContext(ctx) + req.Header.Set("Content-Type", "application/json") + c.Request = req + c.Set(string(middleware.ContextKeyAPIKey), apiKey) + c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.UserID, Concurrency: 10}) + + ep.call(h, c) + + selected, reachedSelection := c.Get(opsAccountIDKey) + if !tc.hasFallback { + require.Equal(t, http.StatusForbidden, recorder.Code) + require.Contains(t, recorder.Body.String(), "This group is restricted to Claude Code clients") + require.False(t, reachedSelection, "a Claude Code only group without fallback must be rejected before account selection") + require.Empty(t, schedulerCache.queriedGroupIDs()) + return + } + require.NotContains(t, recorder.Body.String(), "restricted to Claude Code clients") + require.True(t, reachedSelection, "a Claude Code only group with fallback must reach account selection") + require.Equal(t, fallbackAccountID, selected) + queried := schedulerCache.queriedGroupIDs() + require.Contains(t, queried, fallbackGroupID) + require.NotContains(t, queried, primaryGroupID) + }) + } + } +} diff --git a/backend/internal/handler/gateway_handler_responses.go b/backend/internal/handler/gateway_handler_responses.go index a170877315..1cfb122cc0 100644 --- a/backend/internal/handler/gateway_handler_responses.go +++ b/backend/internal/handler/gateway_handler_responses.go @@ -105,13 +105,11 @@ func (h *GatewayHandler) Responses(c *gin.Context) { // 解析渠道级模型映射 channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(requestCtx, apiKey.GroupID, reqModel) - // Claude Code only restriction: - // /v1/responses is never a Claude Code endpoint. - // When claude_code_only is enabled, this endpoint is rejected. - // The existing service-layer checkClaudeCodeRestriction handles degradation - // to fallback groups when the Forward path calls SelectAccountForModelWithExclusions. - // Here we just reject at handler level since /v1/responses clients can't be Claude Code. - if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly { + // Claude Code only restriction: /v1/responses is never a Claude Code + // endpoint. With a fallback group the request continues and account selection + // (checkClaudeCodeRestriction) schedules it in the fallback group; without one + // it is rejected here. + if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly && apiKey.Group.FallbackGroupID == nil { h.responsesErrorResponse(c, http.StatusForbidden, "permission_error", "This group is restricted to Claude Code clients (/v1/messages only)") return