mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 03:52:26 +08:00
fix(gateway): Claude Code 限制分组的 OpenAI 兼容入口改走降级分组,不再直接 403
分组开启 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 单测覆盖两个入口在有/无降级分组时的行为。
This commit is contained in:
feeeei
@@ -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
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user