diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index a4bcda50fa..e00ca04013 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -610,6 +610,23 @@ func (h *GatewayHandler) Messages(c *gin.Context) { c.Request = c.Request.WithContext(ctx) } + // 会话槽失败释放：failover 链上每个选中的 Anthropic OAuth 账号都注册过同一会话 + // （checkAndRegisterSession）。若请求最终失败（转发失败/客户端中断/换号耗尽）， + // 须立即释放本次会话注册——上游从未真正服务该会话，继续占槽会让 max_sessions + // 受限的账号被失败请求的 session hash 卡满整个空闲窗口，后续新会话全部被拒。 + // 成功请求保持既有空闲超时语义（会话按最后活动时间过期）。 + sessionSlotAccounts := make(map[int64]*service.Account) + upstreamServedSession := false + defer func() { + if upstreamServedSession { + return + } + // 客户端可能已断开、请求 ctx 已取消，用独立 ctx 执行释放 + for _, acc := range sessionSlotAccounts { + h.gatewayService.ReleaseAccountSession(context.Background(), acc, sessionKey) + } + }() + for { fs := NewFailoverState(h.maxAccountSwitches, hasBoundSession) retryWithFallback := false @@ -763,10 +780,15 @@ func (h *GatewayHandler) Messages(c *gin.Context) { h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", profitVetoExhaustedMessage, streamStarted) return } + // 尝试被否决（从未转发），立即释放该账号的会话注册 + h.gatewayService.ReleaseAccountSession(context.Background(), account, sessionKey) + delete(sessionSlotAccounts, account.ID) continue } account = latest selection.Account = latest + // 记录本请求注册过会话槽的账号（profit 准入后账号已定） + sessionSlotAccounts[account.ID] = account // 等待路径保持既有 eager 绑定（无门时 helper 直接绑定）；调度器已 // 抢槽的直达路径无门时由选号内部绑定，这里只在门下补准入后绑定。 if selection.ProfitGateActive() || !selection.Acquired { @@ -983,6 +1005,11 @@ func (h *GatewayHandler) Messages(c *gin.Context) { currentSubscription = nil fallbackUsed = true retryWithFallback = true + // 原分组账号已确定性失败（prompt too long），先释放其会话注册再走兜底分组 + for _, acc := range sessionSlotAccounts { + h.gatewayService.ReleaseAccountSession(context.Background(), acc, sessionKey) + } + sessionSlotAccounts = make(map[int64]*service.Account) break } _ = h.antigravityGatewayService.WriteMappedClaudeError(c, account, promptTooLongErr.StatusCode, promptTooLongErr.RequestID, promptTooLongErr.Body) @@ -998,6 +1025,9 @@ func (h *GatewayHandler) Messages(c *gin.Context) { action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr) switch action { case FailoverContinue: + // 本次尝试已确定性失败，立即释放该账号的会话注册 + h.gatewayService.ReleaseAccountSession(context.Background(), account, sessionKey) + delete(sessionSlotAccounts, account.ID) continue case FailoverExhausted: h.handleFailoverExhausted(c, fs.LastFailoverErr, account.Platform, streamStarted) @@ -1036,6 +1066,8 @@ func (h *GatewayHandler) Messages(c *gin.Context) { // 不会走到这里重复计费。 if result != nil { submitForwardUsage(result) + // 上游已接受并计量本次会话（流中断），会话槽保持既有语义 + upstreamServedSession = true } return } @@ -1061,6 +1093,8 @@ func (h *GatewayHandler) Messages(c *gin.Context) { } submitForwardUsage(result) + // 转发成功，会话槽保持既有空闲超时语义 + upstreamServedSession = true return } if !retryWithFallback { @@ -2163,6 +2197,8 @@ func (h *GatewayHandler) CountTokens(c *gin.Context) { if err := h.gatewayService.ForwardCountTokens(c.Request.Context(), c, account, parsedReq); err != nil { reqLog.Error("gateway.count_tokens_forward_failed", zap.Int64("account_id", account.ID), zap.Error(err)) // 错误响应已在 ForwardCountTokens 中处理 + // 上游未服务该会话，立即释放选号时注册的会话槽（客户端可能已断开，用独立 ctx） + h.gatewayService.ReleaseAccountSession(context.Background(), account, sessionHash) return } } diff --git a/backend/internal/repository/session_limit_cache.go b/backend/internal/repository/session_limit_cache.go index 1711256937..4b38e58f50 100644 --- a/backend/internal/repository/session_limit_cache.go +++ b/backend/internal/repository/session_limit_cache.go @@ -216,6 +216,15 @@ func (c *sessionLimitCache) RegisterSession(ctx context.Context, accountID int64 return result == 1, nil } +// UnregisterSession 立即移除会话注册（不等待空闲超时） +// 请求最终失败时调用：上游从未服务该会话，继续占槽会卡住 max_sessions 受限的账号 +func (c *sessionLimitCache) UnregisterSession(ctx context.Context, accountID int64, sessionUUID string) error { + if sessionUUID == "" { + return nil + } + return c.rdb.ZRem(ctx, sessionLimitKey(accountID), sessionUUID).Err() +} + // RefreshSession 刷新会话时间戳 func (c *sessionLimitCache) RefreshSession(ctx context.Context, accountID int64, sessionUUID string, idleTimeout time.Duration) error { if sessionUUID == "" { diff --git a/backend/internal/repository/session_limit_cache_integration_test.go b/backend/internal/repository/session_limit_cache_integration_test.go new file mode 100644 index 0000000000..7f07300cb3 --- /dev/null +++ b/backend/internal/repository/session_limit_cache_integration_test.go @@ -0,0 +1,92 @@ +//go:build integration + +package repository + +import ( + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/suite" +) + +type SessionLimitCacheSuite struct { + IntegrationRedisSuite + cache service.SessionLimitCache +} + +func (s *SessionLimitCacheSuite) SetupTest() { + s.IntegrationRedisSuite.SetupTest() + s.cache = NewSessionLimitCache(s.rdb, 5) +} + +// TestUnregisterSession_FreesSlotForNewSession 复现并验证修复场景： +// max_sessions=1 的账号上，会话 A 注册后若请求转发失败，会话槽会被 A 占住 +// 整个空闲超时窗口（默认 5 分钟），导致窗口内新会话 B 全部被拒。 +// UnregisterSession 必须立即释放槽位，让 B 无需等待空闲超时即可注册。 +func (s *SessionLimitCacheSuite) TestUnregisterSession_FreesSlotForNewSession() { + const accountID = int64(101) + maxSessions := 1 + idleTimeout := 5 * time.Minute + + // 会话 A 注册成功（请求开始转发） + allowed, err := s.cache.RegisterSession(s.ctx, accountID, "session-A", maxSessions, idleTimeout) + s.RequireNoError(err, "RegisterSession A") + s.True(allowed, "首个会话应被允许") + + // 转发失败后的新会话 B 在空闲超时窗口内被拒（修复前的行为） + allowed, err = s.cache.RegisterSession(s.ctx, accountID, "session-B", maxSessions, idleTimeout) + s.RequireNoError(err, "RegisterSession B before release") + s.False(allowed, "槽位被 A 占用时新会话应被拒绝") + + // 请求最终失败：立即释放 A 的会话注册（不等待空闲超时） + s.RequireNoError(s.cache.UnregisterSession(s.ctx, accountID, "session-A"), "UnregisterSession A") + + // A 释放后 B 立即可注册，无需等待空闲超时 + allowed, err = s.cache.RegisterSession(s.ctx, accountID, "session-B", maxSessions, idleTimeout) + s.RequireNoError(err, "RegisterSession B after release") + s.True(allowed, "释放后新会话应立即可注册") + + count, err := s.cache.GetActiveSessionCount(s.ctx, accountID) + s.RequireNoError(err, "GetActiveSessionCount") + s.Equal(1, count, "释放后活跃会话数应为 1（仅 B）") +} + +// TestUnregisterSession_IdempotentAndMissing 验证释放操作幂等： +// 重复释放与释放不存在的会话均不应报错。 +func (s *SessionLimitCacheSuite) TestUnregisterSession_IdempotentAndMissing() { + const accountID = int64(102) + + s.RequireNoError(s.cache.UnregisterSession(s.ctx, accountID, "missing"), "释放不存在的会话") + s.RequireNoError(s.cache.UnregisterSession(s.ctx, accountID, "missing"), "重复释放") + + s.RequireNoError(s.cache.UnregisterSession(s.ctx, accountID, ""), "空 sessionUUID 应为 no-op") +} + +// TestUnregisterSession_OnlyTargetsSpecifiedSession 验证释放只影响目标会话， +// 不影响同账号上的其他活跃会话。 +func (s *SessionLimitCacheSuite) TestUnregisterSession_OnlyTargetsSpecifiedSession() { + const accountID = int64(103) + maxSessions := 2 + idleTimeout := 5 * time.Minute + + for _, sid := range []string{"session-A", "session-B"} { + allowed, err := s.cache.RegisterSession(s.ctx, accountID, sid, maxSessions, idleTimeout) + s.RequireNoError(err, "RegisterSession "+sid) + s.True(allowed) + } + + s.RequireNoError(s.cache.UnregisterSession(s.ctx, accountID, "session-A"), "UnregisterSession A") + + active, err := s.cache.IsSessionActive(s.ctx, accountID, "session-A") + s.RequireNoError(err, "IsSessionActive A") + s.False(active, "A 应已释放") + + active, err = s.cache.IsSessionActive(s.ctx, accountID, "session-B") + s.RequireNoError(err, "IsSessionActive B") + s.True(active, "B 不应受影响") +} + +func TestSessionLimitCacheSuite(t *testing.T) { + suite.Run(t, new(SessionLimitCacheSuite)) +} diff --git a/backend/internal/service/gateway_scheduling.go b/backend/internal/service/gateway_scheduling.go index 977a28a99d..2cc0c8ba75 100644 --- a/backend/internal/service/gateway_scheduling.go +++ b/backend/internal/service/gateway_scheduling.go @@ -1435,6 +1435,28 @@ func (s *GatewayService) checkAndRegisterSession(ctx context.Context, account *A return allowed } +// ReleaseAccountSession 立即释放会话槽（不等待空闲超时） +// 供 handler 在请求最终失败（选号成功但转发失败/客户端中断）时调用： +// 上游从未真正服务该会话，若继续占槽，max_sessions 受限的账号会被失败请求的 +// session hash 卡满整个空闲窗口，后续新会话全部被拒。 +// 适用条件与 checkAndRegisterSession 对齐；不适用账号为 no-op，幂等可安全重复调用。 +func (s *GatewayService) ReleaseAccountSession(ctx context.Context, account *Account, sessionID string) { + if s == nil || s.sessionLimitCache == nil || account == nil || sessionID == "" { + return + } + if !account.IsAnthropicOAuthOrSetupToken() { + return + } + if account.GetMaxSessions() <= 0 { + return + } + if err := s.sessionLimitCache.UnregisterSession(ctx, account.ID, sessionID); err != nil { + slog.Debug("session_limit.release_failed", + "account_id", account.ID, + "error", err) + } +} + func (s *GatewayService) getSchedulableAccount(ctx context.Context, accountID int64) (*Account, error) { var ( account *Account diff --git a/backend/internal/service/session_limit_cache.go b/backend/internal/service/session_limit_cache.go index 5482d61063..41f8fd5b0b 100644 --- a/backend/internal/service/session_limit_cache.go +++ b/backend/internal/service/session_limit_cache.go @@ -33,6 +33,12 @@ type SessionLimitCache interface { // 用于活跃会话保持活动状态 RefreshSession(ctx context.Context, accountID int64, sessionUUID string, idleTimeout time.Duration) error + // UnregisterSession 立即移除会话注册（不等待空闲超时） + // 用于请求最终失败的场景：上游从未真正服务该会话，不应继续占用会话槽， + // 否则 max_sessions 受限的账号会被失败请求的 session hash 卡满整个空闲窗口， + // 导致后续新会话在超时窗口内全部被拒 + UnregisterSession(ctx context.Context, accountID int64, sessionUUID string) error + // GetActiveSessionCount 获取当前活跃会话数 // 返回未过期的会话数量 GetActiveSessionCount(ctx context.Context, accountID int64) (int, error) diff --git a/backend/internal/service/session_limit_release_test.go b/backend/internal/service/session_limit_release_test.go new file mode 100644 index 0000000000..ba29ccec27 --- /dev/null +++ b/backend/internal/service/session_limit_release_test.go @@ -0,0 +1,121 @@ +package service + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/require" +) + +// sessionLimitReleaseCacheStub 记录 UnregisterSession 调用，用于验证释放逻辑 +type sessionLimitReleaseCacheStub struct { + SessionLimitCache + + unregistered map[int64][]string + err error +} + +func newSessionLimitReleaseCacheStub() *sessionLimitReleaseCacheStub { + return &sessionLimitReleaseCacheStub{ + unregistered: make(map[int64][]string), + } +} + +func (s *sessionLimitReleaseCacheStub) UnregisterSession(_ context.Context, accountID int64, sessionUUID string) error { + if s.err != nil { + return s.err + } + s.unregistered[accountID] = append(s.unregistered[accountID], sessionUUID) + return nil +} + +func newSessionLimitTestAccount() *Account { + return &Account{ + ID: 42, + Platform: PlatformAnthropic, + Type: AccountTypeOAuth, + Extra: map[string]any{"max_sessions": 1}, + } +} + +// TestReleaseAccountSession_ReleasesRegisteredSlot 验证： +// 对启用会话限制的 Anthropic OAuth 账号，ReleaseAccountSession 必须立即移除 +// 该账号上注册的会话（不等待空闲超时）。 +func TestReleaseAccountSession_ReleasesRegisteredSlot(t *testing.T) { + cache := newSessionLimitReleaseCacheStub() + svc := &GatewayService{sessionLimitCache: cache} + acc := newSessionLimitTestAccount() + + svc.ReleaseAccountSession(context.Background(), acc, "session-hash-1") + + require.Equal(t, []string{"session-hash-1"}, cache.unregistered[42], + "应立即移除注册的会话槽") +} + +// TestReleaseAccountSession_NoOpForInapplicableAccounts 验证： +// - 非 Anthropic OAuth/SetupToken 账号 +// - 未启用 max_sessions 的账号 +// - 空 sessionID +// 以上场景均为 no-op，不得触发 UnregisterSession。 +func TestReleaseAccountSession_NoOpForInapplicableAccounts(t *testing.T) { + apiKeyAcc := &Account{ + ID: 43, + Platform: PlatformAnthropic, + Type: AccountTypeAPIKey, + Extra: map[string]any{"max_sessions": 1}, + } + noLimitAcc := &Account{ + ID: 44, + Platform: PlatformAnthropic, + Type: AccountTypeOAuth, + } + enabledAcc := newSessionLimitTestAccount() + + cases := []struct { + name string + account *Account + sessionID string + }{ + {"api_key_account", apiKeyAcc, "session-hash"}, + {"max_sessions_disabled", noLimitAcc, "session-hash"}, + {"empty_session_id", enabledAcc, ""}, + {"nil_account", nil, "session-hash"}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + cache := newSessionLimitReleaseCacheStub() + svc := &GatewayService{sessionLimitCache: cache} + svc.ReleaseAccountSession(context.Background(), tc.account, tc.sessionID) + require.Empty(t, cache.unregistered, "不适用账号不应触发释放") + }) + } +} + +// TestReleaseAccountSession_NilCacheAndErrorTolerance 验证： +// sessionLimitCache 不可用时 no-op；UnregisterSession 返回错误时不 panic（仅记录日志）。 +func TestReleaseAccountSession_NilCacheAndErrorTolerance(t *testing.T) { + // nil cache：no-op + svc := &GatewayService{} + svc.ReleaseAccountSession(context.Background(), newSessionLimitTestAccount(), "session-hash") + + // 底层错误：不 panic + cache := &sessionLimitReleaseCacheStub{err: errors.New("redis down")} + svc = &GatewayService{sessionLimitCache: cache} + svc.ReleaseAccountSession(context.Background(), newSessionLimitTestAccount(), "session-hash") +} + +// TestReleaseAccountSession_Idempotent 验证释放操作幂等，可安全重复调用 +// （failover 链上按次释放 + defer 兜底可能对同一账号重复释放）。 +func TestReleaseAccountSession_Idempotent(t *testing.T) { + cache := newSessionLimitReleaseCacheStub() + svc := &GatewayService{sessionLimitCache: cache} + acc := newSessionLimitTestAccount() + + svc.ReleaseAccountSession(context.Background(), acc, "session-hash") + svc.ReleaseAccountSession(context.Background(), acc, "session-hash") + + // 两次调用都透传到缓存层（Redis ZREM 本身幂等，重复移除无副作用） + require.Len(t, cache.unregistered[42], 2) +} diff --git a/backend/internal/testutil/stubs.go b/backend/internal/testutil/stubs.go index 339d57d867..c9b37efb9f 100644 --- a/backend/internal/testutil/stubs.go +++ b/backend/internal/testutil/stubs.go @@ -139,6 +139,9 @@ func (c StubSessionLimitCache) RegisterSession(_ context.Context, _ int64, _ str func (c StubSessionLimitCache) RefreshSession(_ context.Context, _ int64, _ string, _ time.Duration) error { return nil } +func (c StubSessionLimitCache) UnregisterSession(_ context.Context, _ int64, _ string) error { + return nil +} func (c StubSessionLimitCache) GetActiveSessionCount(_ context.Context, _ int64) (int, error) { return 0, nil }